Add Lazy handling to AsyncDataloader

Author: rmosolgo

lib/graphql/dataloader/async_dataloader.rb

@@ -29,16 +29,7 @@ def run(trace_query_lazy: nil)
             first_pass = false
             fiber_vars = get_fiber_variables
 
-            while (f = (job_fibers.shift || (((job_fibers.size + next_job_fibers.size + source_tasks.size) < jobs_fiber_limit) && spawn_job_fiber(trace))))
-              if f.alive?
-                finished = run_fiber(f)
-                if !finished
-                  next_job_fibers << f
-                end
-              end
-            end
-            job_fibers.concat(next_job_fibers)
-            next_job_fibers.clear
+            run_pending_steps(job_fibers, next_job_fibers, source_tasks, jobs_fiber_limit, trace)
 
             Sync do |root_task|
               set_fiber_variables(fiber_vars)
@@ -54,6 +45,18 @@ def run(trace_query_lazy: nil)
                 next_source_tasks.clear
               end
             end
+
+            if !@lazies_at_depth.empty?
+              with_trace_query_lazy(trace_query_lazy) do
+                run_next_pending_lazies(job_fibers, trace)
+                run_pending_steps(job_fibers, next_job_fibers, source_tasks, jobs_fiber_limit, trace)
+              end
+            elsif !@steps_to_rerun_after_lazy.empty?
+              @pending_jobs.concat(@steps_to_rerun_after_lazy)
+              f = spawn_job_fiber(trace)
+              job_fibers << f
+              @steps_to_rerun_after_lazy.clear
+            end
           end
           trace&.end_dataloader(self)
         end
@@ -69,6 +72,19 @@ def run(trace_query_lazy: nil)
 
       private
 
+      def run_pending_steps(job_fibers, next_job_fibers, source_tasks, jobs_fiber_limit, trace)
+        while (f = (job_fibers.shift || (((job_fibers.size + next_job_fibers.size + source_tasks.size) < jobs_fiber_limit) && spawn_job_fiber(trace))))
+          if f.alive?
+            finished = run_fiber(f)
+            if !finished
+              next_job_fibers << f
+            end
+          end
+        end
+        job_fibers.concat(next_job_fibers)
+        next_job_fibers.clear
+      end
+
       def spawn_source_task(parent_task, condition, trace)
         pending_sources = nil
         @source_cache.each_value do |source_by_batch_params|

spec/graphql/dataloader_spec.rb

@@ -940,6 +940,7 @@ def self.included(child_class)
           query_str = "{ cookbooks { featuredRecipe { name } } }"
           context = { batched_calls_counter: BatchedCallsCounter.new }
           result = schema.execute(query_str, context: context)
+          assert_equal ["Cornbread", "Grits"], result["data"]["cookbooks"].map { |c| c["featuredRecipe"]["name"] }
           refute result.key?("errors")
           assert_equal 1, context[:batched_calls_counter].count
         end