fix: call resolve_references once per __typename (#244)

The previous implementation called `resolve_references` once per
contiguous `__typename` to ensure the correct ordering of the response.
This could have resulted in wildly varying runtime efficiencies.
A request calling for N unique types across M entities could have
as little as N queries or as many as M queries, depending on if the
entities happened to be in large groups of the same type or very
interleaved.

By keeping track of the requested representations and their original
indices, we can rebuild the response in the correct order even when
things get shuffled during computation.

Author: zachahn

lib/apollo-federation/entities_field.rb

@@ -27,9 +27,17 @@ def define_entities_field(possible_entities)
     end
 
     def _entities(representations:)
-      chunked_references = representations.chunk { |r| r[:__typename] }
+      final_result = Array.new(representations.size)
+      grouped_references_with_indices =
+        representations
+        .map
+        .with_index { |r, i| [r, i] }
+        .group_by { |(r, _i)| r[:__typename] }
+
+      grouped_references_with_indices.each do |typename, references_with_indices|
+        references = references_with_indices.map(&:first)
+        indices = references_with_indices.map(&:last)
 
-      chunked_references.flat_map do |typename, references|
         # TODO: Use warden or schema?
         type = context.warden.get_type(typename)
         if type.nil? || type.kind != GraphQL::TypeKinds::OBJECT
@@ -49,8 +57,10 @@ def _entities(representations:)
           results = references
         end
 
-        results.map do |result|
-          context.schema.after_lazy(result) do |resolved_value|
+        results_with_indices = results.zip(indices)
+
+        results_with_indices.each do |result, i|
+          final_result[i] = 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,
@@ -60,6 +70,8 @@ def _entities(representations:)
           end
         end
       end
+
+      final_result
     end
 
     private

spec/apollo-federation/entities_field_spec.rb

@@ -327,6 +327,13 @@ def self.resolve_references(references, _context)
                       ],
                     )
                   end
+
+                  it 'calls resolve_references once per __typename' do
+                    allow(type_with_key).to receive(:resolve_references).and_call_original
+                    allow(another_type_with_key).to receive(:resolve_references).and_call_original
+                    subject
+                    expect([type_with_key, another_type_with_key]).to all have_received(:resolve_references).once
+                  end
                 end
               end