Queue pending Dataloader sources

lib/graphql/dataloader.rb

@@ -59,6 +59,8 @@ def self.with_dataloading(&block)
 
     def initialize(nonblocking: self.class.default_nonblocking, fiber_limit: self.class.default_fiber_limit)
       @source_cache = Hash.new { |h, k| h[k] = {} }.compare_by_identity
+      @pending_sources = []
+      @pending_source_set = Set.new.compare_by_identity
       @pending_jobs = []
       if !nonblocking.nil?
         @nonblocking = nonblocking
@@ -148,6 +150,14 @@ def append_job(callable = nil, &job)
       nil
     end
 
+    # @api private
+    def queue_pending_source(source)
+      return nil if !@pending_source_set.add?(source)
+
+      @pending_sources << source
+      nil
+    end
+
     # Clear any already-loaded objects from {Source} caches
     # @return [void]
     def clear_cache
@@ -187,9 +197,10 @@ def run_isolated
       @lazies_at_depth = prev_lazies_at_depth
       prev_pending_keys.each do |source_instance, pending|
         pending.each do |key, value|
-          if !source_instance.results.key?(key)
-            source_instance.pending[key] = value
-          end
+          next if source_instance.results.key?(key)
+
+          queue_pending_source(source_instance) if source_instance.pending.empty?
+          source_instance.pending[key] = value
         end
       end
     end
@@ -305,7 +316,7 @@ def run_pending_steps(trace, job_fibers, next_job_fibers, jobs_fiber_limit, sour
       end
       join_queues(job_fibers, next_job_fibers)
 
-      while (!source_fibers.empty? || @source_cache.each_value.any? { |group_sources| group_sources.each_value.any?(&:pending?) })
+      while (!source_fibers.empty? || !@pending_sources.empty?)
         while (f = source_fibers.shift || (((job_fibers.size + source_fibers.size + next_source_fibers.size + next_job_fibers.size) < total_fiber_limit) && spawn_source_fiber(trace)))
           if f.alive?
             finished = run_fiber(f)
@@ -356,21 +367,29 @@ def spawn_job_fiber(trace)
       end
     end
 
-    def spawn_source_fiber(trace)
-      pending_sources = nil
-      @source_cache.each_value do |source_by_batch_params|
-        source_by_batch_params.each_value do |source|
-          if source.pending?
-            pending_sources ||= []
-            pending_sources << source
-          end
-        end
+    def drain_pending_sources
+      pending_sources = @pending_sources
+      @pending_sources = []
+      @pending_source_set.clear
+
+      pending_sources.select!(&:pending?)
+      pending_sources.empty? ? nil : pending_sources
+    end
+
+    def dequeue_pending_source
+      while (source = @pending_sources.shift)
+        @pending_source_set.delete(source)
+        return source if source.pending?
       end
+    end
+
+    def spawn_source_fiber(trace)
+      pending_sources = @pending_sources.dup
 
-      if pending_sources
+      if !pending_sources.empty?
         spawn_fiber do
           trace&.dataloader_spawn_source_fiber(pending_sources)
-          pending_sources.each do |source|
+          while (source = dequeue_pending_source)
             Fiber[:__graphql_current_dataloader_source] = source
             trace&.begin_dataloader_source(source)
             source.run_pending_keys

lib/graphql/dataloader/async_dataloader.rb

@@ -177,9 +177,10 @@ def run_isolated
         end
         prev_pending_keys.each do |source_instance, pending|
           pending.each do |key, value|
-            if !source_instance.results.key?(key)
-              source_instance.pending[key] = value
-            end
+            next if source_instance.results.key?(key)
+
+            queue_pending_source(source_instance) if source_instance.pending.empty?
+            source_instance.pending[key] = value
           end
         end
       end
@@ -280,11 +281,9 @@ def run_sources(run)
         end
 
         allowed_tasks = run.allowed_sources_tasks
-        while (has_pending = @source_cache.each_value.any? { |group_sources| group_sources.each_value.any?(&:pending?) } ) || unsnoozed
+        while (pending_sources = drain_pending_sources) || unsnoozed
           unsnoozed = false
-          if has_pending
-            spawn_source_task(run, allowed_tasks)
-          end
+          spawn_source_task(run, allowed_tasks, pending_sources) if pending_sources
           run.wait_for_queues
         end
       ensure
@@ -318,43 +317,31 @@ def run_next_pending_lazies(run)
         end
       end
 
-      def spawn_source_task(run, num_tasks)
-        pending_sources = nil
-        @source_cache.each_value do |source_by_batch_params|
-          source_by_batch_params.each_value do |source|
-            if source.pending?
-              pending_sources ||= []
-              pending_sources << source
-            end
-          end
+      def spawn_source_task(run, num_tasks, pending_sources)
+        if num_tasks == Float::INFINITY
+          num_tasks = pending_sources.size
         end
-
-        if pending_sources
-          if num_tasks == Float::INFINITY
-            num_tasks = pending_sources.size
-          end
-          fiber_vars = get_fiber_variables
-          trace = run.trace
-          num_tasks.times do
-            new_task = Async::Task.new(run.root_task) do |task|
-              task.graphql_async_dataloader_run = run
-              task.graphql_async_dataloader_condition = run.snoozed_sources_condition
-              trace&.dataloader_spawn_source_fiber(pending_sources)
-              set_fiber_variables(fiber_vars)
-              while (source = pending_sources.shift)
-                trace&.begin_dataloader_source(source)
-                source.run_pending_keys
-                trace&.end_dataloader_source(source)
-              end
-              nil
-            ensure
-              run.finished_tasks.push($! || task)
-              cleanup_fiber
-              trace&.dataloader_fiber_exit
+        fiber_vars = get_fiber_variables
+        trace = run.trace
+        num_tasks.times do
+          new_task = Async::Task.new(run.root_task) do |task|
+            task.graphql_async_dataloader_run = run
+            task.graphql_async_dataloader_condition = run.snoozed_sources_condition
+            trace&.dataloader_spawn_source_fiber(pending_sources)
+            set_fiber_variables(fiber_vars)
+            while (source = pending_sources.shift)
+              trace&.begin_dataloader_source(source)
+              source.run_pending_keys
+              trace&.end_dataloader_source(source)
             end
-            run.started_tasks.push(new_task)
-            new_task
+            nil
+          ensure
+            run.finished_tasks.push($! || task)
+            cleanup_fiber
+            trace&.dataloader_fiber_exit
           end
+          run.started_tasks.push(new_task)
+          new_task
         end
       end
     end

lib/graphql/dataloader/source.rb

@@ -20,9 +20,7 @@ def setup(dataloader)
       # @return [Dataloader::Request] a pending request for a value from `key`. Call `.load` on that object to wait for the result.
       def request(value)
         res_key = result_key_for(value)
-        if !@results.key?(res_key)
-          @pending[res_key] ||= normalize_fetch_key(value)
-        end
+        add_pending_key(res_key, value)
         Dataloader::Request.new(self, value)
       end
 
@@ -51,9 +49,7 @@ def normalize_fetch_key(value)
       def request_all(values)
         values.each do |v|
           res_key = result_key_for(v)
-          if !@results.key?(res_key)
-            @pending[res_key] ||= normalize_fetch_key(v)
-          end
+          add_pending_key(res_key, v)
         end
         Dataloader::RequestAll.new(self, values)
       end
@@ -65,7 +61,7 @@ def load(value)
         if @results.key?(result_key)
           result_for(result_key)
         else
-          @pending[result_key] ||= normalize_fetch_key(value)
+          add_pending_key(result_key, value)
           sync([result_key])
           result_for(result_key)
         end
@@ -79,8 +75,7 @@ def load_all(values)
         values.each { |v|
           k = result_key_for(v)
           result_keys << k
-          if !@results.key?(k)
-            @pending[k] ||= normalize_fetch_key(v)
+          if add_pending_key(k, v)
             pending_keys << k
           end
         }
@@ -106,6 +101,7 @@ def fetch(keys)
       # Then run the batch and update the cache.
       # @return [void]
       def sync(pending_result_keys)
+        @dataloader.queue_pending_source(self) if pending?
         @dataloader.yield(self)
         iterations = 0
         while pending_result_keys.any? { |key| !@results.key?(key) }
@@ -193,6 +189,15 @@ def clear_cache
 
       private
 
+      def add_pending_key(result_key, value)
+        return false if @results.key?(result_key)
+
+        was_empty = @pending.empty?
+        @pending[result_key] ||= normalize_fetch_key(value)
+        @dataloader.queue_pending_source(self) if was_empty
+        true
+      end
+
       # Reads and returns the result for the key from the internal cache, or raises an error if the result was an error
       # @param key [Object] key passed to {#load} or {#load_all}
       # @return [Object] The result from {#fetch} for `key`.

spec/graphql/dataloader_spec.rb

@@ -68,6 +68,27 @@ def fetch(keys)
       end
     end
 
+    class PendingCheckSource < GraphQL::Dataloader::Source
+      class << self
+        attr_accessor :pending_checks
+      end
+
+      self.pending_checks = 0
+
+      def initialize(batch_key)
+        @batch_key = batch_key
+      end
+
+      def fetch(keys)
+        keys
+      end
+
+      def pending?
+        self.class.pending_checks += 1
+        super
+      end
+    end
+
     class NestedDataObject < GraphQL::Dataloader::Source
       def fetch(ids)
         @dataloader.with(DataObject).load_all(ids)
@@ -1496,6 +1517,21 @@ def request_all
     assert :world, value
   end
 
+  it "tracks pending sources without scanning the entire source cache" do
+    dataloader = GraphQL::Dataloader.new
+    100.times do |idx|
+      dataloader.with(FiberSchema::PendingCheckSource, idx)
+    end
+
+    FiberSchema::PendingCheckSource.pending_checks = 0
+    dataloader.append_job do
+      dataloader.with(FiberSchema::PendingCheckSource, 0).load(1)
+    end
+    dataloader.run
+
+    assert_operator FiberSchema::PendingCheckSource.pending_checks, :<, 100
+  end
+
   class CanaryDataloader < GraphQL::Dataloader::NullDataloader
   end
 
@@ -1559,6 +1595,23 @@ def fetch(ids)
       assert_equal({ a: 1, b: 2, c: 3, d: 4, e: 3 }, result)
     end
 
+    it "restores pending sources from the outer queue" do
+      dl = GraphQL::Dataloader.new
+      result = {}
+      outer_request = dl.with(RunIsolated::CountSource).request(1)
+
+      dl.run_isolated {
+        result[:isolated] = dl.with(RunIsolated::CountSource).load(2)
+      }
+
+      dl.append_job {
+        result[:outer] = outer_request.load
+      }
+      dl.run
+
+      assert_equal({ isolated: 1, outer: 2 }, result)
+    end
+
     it "shares a cache" do
       dl = GraphQL::Dataloader.new
       result = {}