From 063a76f6629ee5611acfb4b870f7142dd381c863 Mon Sep 17 00:00:00 2001 From: ydah Date: Fri, 10 Jul 2026 07:50:48 +0900 Subject: [PATCH 1/2] Queue pending Dataloader sources Co-authored-by: Robert Mosolgo --- lib/graphql/dataloader.rb | 50 +++++++++++----- lib/graphql/dataloader/async_dataloader.rb | 69 +++++++++------------- lib/graphql/dataloader/source.rb | 23 +++++--- spec/graphql/dataloader_spec.rb | 53 +++++++++++++++++ 4 files changed, 130 insertions(+), 65 deletions(-) diff --git a/lib/graphql/dataloader.rb b/lib/graphql/dataloader.rb index 4fe37cabee..d0f7f91ef9 100644 --- a/lib/graphql/dataloader.rb +++ b/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 = {}.compare_by_identity @pending_jobs = [] if !nonblocking.nil? @nonblocking = nonblocking @@ -148,6 +150,15 @@ def append_job(callable = nil, &job) nil end + # @api private + def queue_pending_source(source) + return nil if @pending_source_set.key?(source) + + @pending_source_set[source] = true + @pending_sources << source + nil + end + # Clear any already-loaded objects from {Source} caches # @return [void] def clear_cache @@ -187,9 +198,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 +317,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 +368,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 diff --git a/lib/graphql/dataloader/async_dataloader.rb b/lib/graphql/dataloader/async_dataloader.rb index 09765eabff..374dc7057f 100644 --- a/lib/graphql/dataloader/async_dataloader.rb +++ b/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 diff --git a/lib/graphql/dataloader/source.rb b/lib/graphql/dataloader/source.rb index 44d82d94a5..d3f9490385 100644 --- a/lib/graphql/dataloader/source.rb +++ b/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`. diff --git a/spec/graphql/dataloader_spec.rb b/spec/graphql/dataloader_spec.rb index b90df16038..874961183c 100644 --- a/spec/graphql/dataloader_spec.rb +++ b/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 = {} From 4e6ccf0f9bfbb92455a9a05bce043ebbe30af426 Mon Sep 17 00:00:00 2001 From: ydah Date: Sat, 11 Jul 2026 09:00:08 +0900 Subject: [PATCH 2/2] Use Set for pending Dataloader sources Co-authored-by: Robert Mosolgo --- lib/graphql/dataloader.rb | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/lib/graphql/dataloader.rb b/lib/graphql/dataloader.rb index d0f7f91ef9..c5fff349a2 100644 --- a/lib/graphql/dataloader.rb +++ b/lib/graphql/dataloader.rb @@ -60,7 +60,7 @@ 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 = {}.compare_by_identity + @pending_source_set = Set.new.compare_by_identity @pending_jobs = [] if !nonblocking.nil? @nonblocking = nonblocking @@ -152,9 +152,8 @@ def append_job(callable = nil, &job) # @api private def queue_pending_source(source) - return nil if @pending_source_set.key?(source) + return nil if !@pending_source_set.add?(source) - @pending_source_set[source] = true @pending_sources << source nil end