From e3ebf2fddf8dcf2182cdab9cdcf5095c79dedf0b Mon Sep 17 00:00:00 2001 From: Robert Mosolgo Date: Tue, 28 Jul 2026 05:40:45 -0400 Subject: [PATCH] AsyncDataloader: refactor job accounting --- lib/graphql/dataloader/async_dataloader.rb | 97 +++++++++++++------ .../dataloader/async_dataloader_spec.rb | 85 ++++++++++++++++ 2 files changed, 150 insertions(+), 32 deletions(-) diff --git a/lib/graphql/dataloader/async_dataloader.rb b/lib/graphql/dataloader/async_dataloader.rb index 30ebca3999..a1ddf0ff3f 100644 --- a/lib/graphql/dataloader/async_dataloader.rb +++ b/lib/graphql/dataloader/async_dataloader.rb @@ -4,6 +4,11 @@ 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 @@ -11,7 +16,6 @@ def self.use(...) 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 diff --git a/spec/graphql/dataloader/async_dataloader_spec.rb b/spec/graphql/dataloader/async_dataloader_spec.rb index 28c8758472..9eaad318a9 100644 --- a/spec/graphql/dataloader/async_dataloader_spec.rb +++ b/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