Use .yield instead of async { ... }.wait to give tasks a chance to run

Author: rmosolgo

lib/graphql/dataloader/async_dataloader.rb

@@ -19,29 +19,27 @@ def run
           while first_pass || job_tasks.any?
             first_pass = false
 
-            root_task.async do |jobs_task|
-              while (task = job_tasks.shift || spawn_job_task(jobs_task, jobs_condition))
-                if task.alive?
-                  next_job_tasks << task
-                elsif task.failed?
-                  # re-raise a raised error -
-                  # this also covers errors from sources since
-                  # these jobs wait for sources as needed.
-                  task.wait
-                end
+            while (task = job_tasks.shift || spawn_job_task(root_task, jobs_condition))
+              if task.alive?
+                next_job_tasks << task
+              elsif task.failed?
+                # re-raise a raised error -
+                # this also covers errors from sources since
+                # these jobs wait for sources as needed.
+                task.wait
               end
-            end.wait
+            end
+            root_task.yield # give job tasks a chance to run
             job_tasks.concat(next_job_tasks)
             next_job_tasks.clear
 
             while source_tasks.any? || @source_cache.each_value.any? { |group_sources| group_sources.each_value.any?(&:pending?) }
-              root_task.async do |sources_loop_task|
-                while (task = source_tasks.shift || spawn_source_task(sources_loop_task, sources_condition))
-                  if task.alive?
-                    next_source_tasks << task
-                  end
+              while (task = source_tasks.shift || spawn_source_task(root_task, sources_condition))
+                if task.alive?
+                  next_source_tasks << task
                 end
-              end.wait
+              end
+              root_task.yield # give source tasks a chance to run
               sources_condition.signal
               source_tasks.concat(next_source_tasks)
               next_source_tasks.clear
@@ -58,7 +56,7 @@ def run
       def spawn_job_task(parent_task, condition)
         if @pending_jobs.any?
           fiber_vars = get_fiber_variables
-          parent_task.async do |t|
+          parent_task.async do
             set_fiber_variables(fiber_vars)
             Thread.current[:graphql_dataloader_next_tick] = condition
             while job = @pending_jobs.shift

spec/graphql/dataloader_spec.rb

@@ -2,6 +2,10 @@
 require "spec_helper"
 require "fiber"
 
+if defined?(Console) && defined?(Async)
+  Console.logger.disable(Async::Task)
+end
+
 describe GraphQL::Dataloader do
   class BatchedCallsCounter
     def initialize
@@ -1005,27 +1009,30 @@ def self.included(child_class)
 
         it "works with very very large queries" do
           query_str = "{".dup
-          1100.times do |i|
+          fields = 1100
+          fields.times do |i|
             query_str << "\n  field#{i}: lookaheadIngredient(input: { id: 1, batchKey: \"key-#{i}\"}) { name }"
           end
           query_str << "\n}"
           GC.start
+          GC.disable
           res = schema.execute(query_str)
-          assert_equal 1100, res["data"].keys.size
+          assert_equal fields, res["data"].keys.size
           all_fibers = []
           ObjectSpace.each_object(Fiber) do |f|
             all_fibers << f
           end
           all_fibers.delete(Fiber.current)
-          if schema.dataloader_class == GraphQL::Dataloader::AsyncDataloader
-            skip <<~ERR
-              TODO: AsyncDataloader leaves orphan suspended fibers :'(
+          if all_fibers.any?(&:alive?)
+              puts <<~ERR
+            Alive fibers:
 
-                - #{all_fibers.select(&:alive?).join("\n  -")}
+              - #{all_fibers.select(&:alive?).join("\n  - ")}
             ERR
-          else
-            assert_equal [false], all_fibers.map(&:alive?).uniq
           end
+          assert_equal [false], all_fibers.map(&:alive?).uniq
+        ensure
+          GC.enable
         end
 
         it "doesn't perform duplicate source fetches" do