fix(tracing): Properly handle tracing fields that resolve an array of lazy values (#87)

* test: Add failing test

* fix(tracing): Properly handle arrays of lazy values

* refactor: Move EntityTypeResolutionExtension.after_resolve behavior directly into the _entities resolver

* fix: Remove broken require;

* style: Fix linter errors

* test: Add more test cases

* 🎨 Cleanup

* Clarify todo

* style: Fix linter error

Author: rylanc

.gitignore

@@ -7,3 +7,5 @@ node_modules/
 
 # built CircleCI configs
 .circleci/processed-config.yml
+
+.DS_Store

lib/apollo-federation/entities_field.rb

@@ -2,7 +2,6 @@
 
 require 'graphql'
 require 'apollo-federation/any'
-require 'apollo-federation/entity_type_resolution_extension'
 
 module ApolloFederation
   module EntitiesField
@@ -23,7 +22,6 @@ def define_entities_field(possible_entities)
 
         field(:_entities, [entity_type, null: true], null: false) do
           argument :representations, [Any], required: true
-          extension(EntityTypeResolutionExtension)
         end
       end
     end
@@ -47,7 +45,14 @@ def _entities(representations:)
           result = reference
         end
 
-        [type, result]
+        context.schema.after_lazy(result) do |resolved_value|
+          # TODO: This isn't 100% correct: if (for some reason) 2 different resolve_reference calls
+          # return the same object, it might not have the right type
+          # Right now, apollo-federation just adds a __typename property to the result,
+          # but I don't really like the idea of modifying the resolved object
+          context[resolved_value] = type
+          resolved_value
+        end
       end
     end
   end

lib/apollo-federation/entity_type_resolution_extension.rb

@@ -27,6 +27,8 @@ def attach_trace_to_result(result)
       return result unless result.context[:tracing_enabled]
 
       trace = result.context.namespace(KEY)
+      # TODO: If there are errors during query validation, that could also cause a missing
+      # start_time
       raise NotInstalledError unless trace[:start_time]
 
       result['errors']&.each do |error|

lib/apollo-federation/tracing.rb

@@ -159,17 +159,27 @@ def self.execute_field_lazy(data, &block)
 
         end_time_nanos = Process.clock_gettime(Process::CLOCK_MONOTONIC, :nanosecond)
 
-        # interpreter runtime
+        # legacy runtime
         if data.include?(:context)
-          context = data.fetch(:context)
           path = context.path
-        else # legacy runtime
-          context = data.fetch(:query).context
+          field = context.field
+        else # interpreter runtime
           path = data.fetch(:path)
+          field = data.fetch(:field)
         end
 
         trace = context.namespace(ApolloFederation::Tracing::KEY)
 
+        # When a field is resolved with an array of lazy values, the interpreter fires an
+        # `execute_field` for the resolution of the field and then a `execute_field_lazy` event for
+        # each lazy value in the array. Since the path here will contain an index (indicating which
+        # lazy value we're executing: e.g. ['arrayOfLazies', 0]), we won't have a node for the path.
+        # We only care about the end of the parent field (e.g. ['arrayOfLazies']), so we get the
+        # node for that path. What ends up happening is we update the end_time for the parent node
+        # for each of the lazy values. The last one that's executed becomes the final end time.
+        if field.type.list? && path.last.is_a?(Integer)
+          path = path[0...-1]
+        end
         node = trace[:node_map].node_for_path(path)
         node.end_time = end_time_nanos - trace[:start_time_nanos]