AsyncDataloader: refactor job accounting

lib/graphql/dataloader/async_dataloader.rb

@@ -4,14 +4,18 @@ module GraphQL
   class Dataloader
     class AsyncDataloader < Dataloader
       def self.use(...)
+        install_graphql_methods
+        super
+      end
+
+      def self.install_graphql_methods
         if !Async::Task.method_defined?(:cancel)
           Async::Task.alias_method(:cancel, :stop)
         end
         if !Async::Task.method_defined?(:graphql_async_dataloader_run)
           Async::Task.attr_accessor(:graphql_async_dataloader_run)
           Async::Task.attr_accessor(:graphql_async_dataloader_condition)
         end
-        super
       end
 
       def initialize(...)
@@ -54,7 +58,10 @@ def initialize(dataloader, total_fiber_limit, jobs_fiber_limit)
           @running_tasks = nil
           @tasks_channel = nil
           @tasks_channel_task = nil
-          @finished_all_tasks = nil
+          @activity = nil
+          @task_error = nil
+          @expected_resumes = 0
+          @mode = nil
 
           @snoozed_jobs_condition = Async::Condition.new
           @snoozed_sources_condition = Async::Condition.new
@@ -77,20 +84,41 @@ def close_queues
           @tasks_channel_task.cancel
         end
 
-        def wait_for_queues
-          @finished_all_tasks.wait
-          @finished_all_tasks = Async::Promise.new
+        def wait_for_activity
+          @activity.wait
+        end
+
+        def quiesced?
+          @running_tasks.empty? && @tasks_channel.empty? && @expected_resumes == 0
+        end
+
+        def has_pending_work?
+          @mode == :jobs ? @jobs.any? : @dataloader.pending_sources.any?(&:pending?) # rubocop:disable Development/NoneWithoutBlockCop
+        end
+
+        def has_bandwidth?
+          @mode == :jobs ? jobs_bandwidth? : sources_bandwidth?
+        end
+
+        # Signalled tasks don't appear in any accounting until their first slice
+        # pushes `:resumed_task`, so they have to be counted at signal time:
+        def expect_resumes(count)
+          @expected_resumes = count
         end
 
-        def wait_for_no_running_tasks
-          @no_running_tasks.wait
-          @no_running_tasks = Async::Promise.new
+        def check_error!
+          if (err = @task_error)
+            @task_error = nil
+            raise err
+          end
         end
 
         def new_queues(mode)
+          @mode = mode
           @tasks_channel = Async::Queue.new(parent: @root_task)
-          @no_running_tasks = Async::Promise.new
-          @finished_all_tasks = Async::Promise.new
+          @activity = Async::Condition.new
+          @task_error = nil
+          @expected_resumes = 0
           @running_tasks = []
           @tasks_channel_task = @root_task.async do |_t|
             while ((msg, data) = @tasks_channel.wait)
@@ -99,23 +127,18 @@ def new_queues(mode)
                 @running_tasks.push(data)
                 data.run
               when :resumed_task
+                if @expected_resumes > 0
+                  @expected_resumes -= 1
+                end
                 @running_tasks.push(data)
               when :finished_task, :paused_task
                 @running_tasks.delete(data)
-                has_pending_work = mode == :jobs ? @jobs.any? : @dataloader.pending_sources.any?(&:pending?) # rubocop:disable Development/NoneWithoutBlockCop
-                if @running_tasks.empty?
-                  @no_running_tasks.resolve(true)
-                  has_bandwidth = mode == :jobs ? jobs_bandwidth? : sources_bandwidth?
-                  if (!has_pending_work) || (!has_bandwidth)
-                    @finished_all_tasks.resolve(true)
-                  end
-                end
               when :task_error
-                @no_running_tasks.resolve(true)
-                @finished_all_tasks.reject(data)
+                @task_error ||= data
               else
                 raise ArgumentError, "Unknown tasks_channel action: #{msg.inspect}"
               end
+              @activity.signal
             end
           end
         end
@@ -240,19 +263,20 @@ def run(trace_query_lazy: nil)
       private
 
       def run_queue(run, condition, mode)
-        should_wait_for_all_tasks = false
+        opened_queues = false
 
-        if (unsnoozed = condition.waiting?)
-          should_wait_for_all_tasks = true
+        if condition.waiting?
+          opened_queues = true
           run.new_queues(mode)
+          run.expect_resumes(condition.instance_variable_get(:@ready).num_waiting)
           condition.signal
         end
 
-        while (pending_work = (mode == :jobs) ? (!run.jobs.empty? && run.jobs_bandwidth? ? run.jobs : nil) : (drain_pending_sources)) || unsnoozed
-          unsnoozed = false
+        loop do
+          pending_work = (mode == :jobs) ? (!run.jobs.empty? && run.jobs_bandwidth? ? run.jobs : nil) : (drain_pending_sources)
           if pending_work
-            if should_wait_for_all_tasks == false
-              should_wait_for_all_tasks = true
+            if opened_queues == false
+              opened_queues = true
               run.new_queues(mode)
             end
             num_tasks = mode == :sources ? run.current_sources_fiber_limit : 1
@@ -262,14 +286,23 @@ def run_queue(run, condition, mode)
             spawn_tasks(run, mode, condition, pending_work, num_tasks)
           end
 
-          run.wait_for_no_running_tasks
-        end
+          if !opened_queues
+            break
+          end
+
+          run.check_error!
 
-        if should_wait_for_all_tasks
-          run.wait_for_queues
+          if run.quiesced?
+            if !run.has_pending_work? || !run.has_bandwidth?
+              break
+            end
+            # Quiesced, but more work appeared - loop around to drain it.
+          else
+            run.wait_for_activity
+          end
         end
       ensure
-        if should_wait_for_all_tasks
+        if opened_queues
           run.close_queues
         end
       end

spec/graphql/dataloader/async_dataloader_spec.rb

@@ -464,5 +464,90 @@ def author(name:)
         end
       end
     end
+
+    describe "stress test" do
+      RNG = Random.new(20260724)
+      MAX_ITERATIONS = 2000
+      WEDGE_AFTER = 10
+      class SlowSource < GraphQL::Dataloader::Source
+        def initialize(delay)
+          @delay = delay
+        end
+
+        def fetch(keys)
+          sleep(@delay)
+          keys.map { |k| "v#{k}" }
+        end
+      end
+
+      it "works" do
+        # Standalone reproduction of graphql-ruby#5671 (AsyncDataloader hang / ClosedQueueError)
+        $stdout.sync = true
+
+        GraphQL::Dataloader::AsyncDataloader.install_graphql_methods
+
+
+        iter_started_at = nil
+        iter = nil
+        iter_finished = false
+        main = Thread.current
+
+        watchdog = Thread.new do
+          loop do
+            sleep 1
+            started = iter_started_at
+            break if iter_finished
+            next unless started
+            if Process.clock_gettime(Process::CLOCK_MONOTONIC) - started > WEDGE_AFTER
+              puts "\n=== WEDGE: iteration #{iter} hung for >#{WEDGE_AFTER}s ==="
+              puts "--- main thread backtrace ---"
+              puts((main.backtrace || []).first(15))
+              assert false, "Failed"
+            end
+          end
+        end
+
+        MAX_ITERATIONS.times do |i|
+          iter = i
+          iter_started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          n_jobs = 4 + RNG.rand(12)
+          seed_delay = RNG.rand(0.002)
+
+          begin
+            dl = GraphQL::Dataloader::AsyncDataloader.new
+            n_jobs.times do |j|
+              dl.append_job do
+                # Everyone snoozes on the shared source first:
+                dl.with(SlowSource, seed_delay).load(j % 3)
+                case j % 3
+                when 0
+                  # finishes in its first resumed slice - pure CPU, no yields.
+                when 1
+                  # stays "running" across the generation boundary via non-dataloader IO,
+                  # then needs another source round:
+                  sleep(RNG.rand(0.002))
+                  dl.with(SlowSource, RNG.rand(0.001)).load(100 + j)
+                when 2
+                  # a second wave of snoozers to force more generations:
+                  dl.with(SlowSource, RNG.rand(0.001)).load(200 + (j % 2))
+                  sleep(RNG.rand(0.001))
+                  dl.with(SlowSource, RNG.rand(0.001)).load(300 + j)
+                end
+              end
+            end
+            dl.run
+            print "," if (i + 1) % 50 == 0
+          rescue => e
+            puts "\n=== ERROR at iteration #{i}: #{e.class}: #{e.message} ==="
+            puts e.backtrace.first(12)
+            exit!(1)
+          end
+          iter_started_at = nil
+        end
+        iter_finished = true
+        assert "completed #{MAX_ITERATIONS} iterations cleanly"
+        assert watchdog.join
+      end
+    end
   end
 end