From 3bf2b0863a798ae47fb50be479ce5ede3b690065 Mon Sep 17 00:00:00 2001 From: Greg MacWilliam Date: Fri, 17 Jul 2026 16:31:55 -0400 Subject: [PATCH] add async limiter. --- README.md | 3 +- lib/graphql/breadth.rb | 1 + lib/graphql/breadth/executor/lazy_async.rb | 61 ++++++++++------ lib/graphql/breadth/lazy_loader.rb | 20 +++--- sorbet/rbi/shims/async.rbi | 9 +++ test/graphql/breadth/executor/loaders_test.rb | 72 +++++++++++++++++++ 6 files changed, 135 insertions(+), 31 deletions(-) diff --git a/README.md b/README.md index db7b9ea..a47f7ff 100644 --- a/README.md +++ b/README.md @@ -441,7 +441,8 @@ All concurrency settings are optional: * `resource: Symbol`, coordinates limits across related loaders hitting the same upstream resource. Defaults to a unique resource identity per loader class. * `limit: Integer`, number of concurrent operations allowed while hitting the resource. Defaults to 8. -* `timeout: Integer`, number of seconds to wait on load operations before timing out. Defaults to no limit. +* `timeout: Integer`, number of seconds to wait on load operations before timing out. Defaults to none. +* `throttle`, an [async-limiter](https://github.com/socketry/async-limiter) that throttles rate. Defaults to none. #### Async fan-out diff --git a/lib/graphql/breadth.rb b/lib/graphql/breadth.rb index f16b0d2..91d43f3 100644 --- a/lib/graphql/breadth.rb +++ b/lib/graphql/breadth.rb @@ -53,6 +53,7 @@ class ExecutionField; end #: type graphql_result = Hash[String, untyped] #: type variables_hash = Hash[String, untyped] #: type selection_node = GraphQL::Language::Nodes::Field | GraphQL::Language::Nodes::InlineFragment | GraphQL::Language::Nodes::FragmentSpread + #: type async_throttle = ::Async::Limiter::Generic end end diff --git a/lib/graphql/breadth/executor/lazy_async.rb b/lib/graphql/breadth/executor/lazy_async.rb index 34efc48..709997f 100644 --- a/lib/graphql/breadth/executor/lazy_async.rb +++ b/lib/graphql/breadth/executor/lazy_async.rb @@ -6,6 +6,32 @@ module Breadth class Executor # @requires_ancestor: Executor module LazyAsync + #: ( + #| ::Async::Barrier | ::Async::Semaphore, + #| ?owner: ::Async::Barrier?, + #| ?throttle: async_throttle?, + #| ?timeout: Numeric?, + #| ) { (::Async::Task) -> untyped } -> ::Async::Task + def self.fork(parent, owner: nil, throttle: nil, timeout: nil, &block) + handler = if timeout + -> (task) { task.with_timeout(timeout, ::Async::TimeoutError) { block.call(task) } } + else + block + end + + if throttle && owner + throttle.async(parent: owner) do |throttle_task| + parent.async(parent: throttle_task, &handler).wait + end + elsif throttle + throttle.async(parent: parent, &handler) + elsif owner + parent.async(parent: owner, &handler) + else + parent.async(&handler) + end + end + class ExecutorContext #: Executor attr_reader :executor @@ -63,19 +89,20 @@ def initialize(loader, async_context, parent_task) @parent_semaphore = nil end - #: (?resource: Symbol?, ?limit: Integer?, ?timeout: Numeric?) ?{ -> untyped } -> Future - def async(resource: nil, limit: nil, timeout: nil, &block) + #: (?resource: Symbol?, ?limit: Integer?, ?timeout: Numeric?, ?throttle: async_throttle?) ?{ -> untyped } -> Future + def async(resource: nil, limit: nil, timeout: nil, throttle: nil, &block) Kernel.raise ArgumentError, "LazyLoader#async requires a block" unless block settings = @loader.async_settings resource = settings.resource if resource.nil? limit = settings.limit if limit.nil? timeout = settings.timeout if timeout.nil? + throttle = settings.throttle if throttle.nil? Kernel.raise ArgumentError, "Lazy async limit must be positive" unless limit.nil? || limit > 0 Kernel.raise ArgumentError, "Lazy async timeout must be positive" unless timeout.nil? || timeout > 0 - # child fan-out shares the parent resource budget, so release the parent slot before waiting on child slots. + # Child fan-out shares the parent resource budget, so release the parent slot before waiting on child slots. release_parent_permit if @parent_semaphore && @loader.async_settings.resource == resource parent = if limit @@ -85,13 +112,10 @@ def async(resource: nil, limit: nil, timeout: nil, &block) @barrier end - task = parent.async(parent: @barrier) do |async_task| + owner = @barrier if limit + task = LazyAsync.fork(parent, owner: owner, throttle: throttle, timeout: timeout) do begin - if timeout - async_task.with_timeout(timeout, ::Async::TimeoutError, &block) - else - block.call - end + block.call rescue StandardError => e @async_context.executor.handle_or_reraise(e) end @@ -101,13 +125,13 @@ def async(resource: nil, limit: nil, timeout: nil, &block) @futures.last end - #: [T, U] (Enumerable[T], ?resource: Symbol?, ?limit: Integer?, ?timeout: Numeric?) ?{ (T) -> U } -> Array[U] - def async_map(collection, resource: nil, limit: nil, timeout: nil, &block) + #: [T, U] (Enumerable[T], ?resource: Symbol?, ?limit: Integer?, ?timeout: Numeric?, ?throttle: async_throttle?) ?{ (T) -> U } -> Array[U] + def async_map(collection, resource: nil, limit: nil, timeout: nil, throttle: nil, &block) Kernel.raise ArgumentError, "LazyLoader#async_map requires a block" unless block completed = false futures = collection.map do |item| - async(resource: resource, limit: limit, timeout: timeout) { block.call(item) } + async(resource: resource, limit: limit, timeout: timeout, throttle: throttle) { block.call(item) } end results = futures.map(&:wait) @@ -224,19 +248,14 @@ def execute_async_lazy_batches(sync_batches, async_batches) def schedule_async_lazy_batch(async_context, batch) async_context.active_loaders.add(batch.loader) - async_context.barrier.async do |async_task| - settings = batch.loader.async_settings + settings = batch.loader.async_settings + + LazyAsync.fork(async_context.barrier, throttle: settings.throttle, timeout: settings.timeout) do |async_task| loader_context = LoaderContext.new(batch.loader, async_context, async_task) begin begin - if settings.timeout - async_task.with_timeout(settings.timeout) do - loader_context.run { execute_lazy_batch(batch, async_context: loader_context) } - end - else - loader_context.run { execute_lazy_batch(batch, async_context: loader_context) } - end + loader_context.run { execute_lazy_batch(batch, async_context: loader_context) } rescue StandardError => e apply_lazy_error(batch, handle_or_reraise(e)) end diff --git a/lib/graphql/breadth/lazy_loader.rb b/lib/graphql/breadth/lazy_loader.rb index b00a66f..3b6432d 100644 --- a/lib/graphql/breadth/lazy_loader.rb +++ b/lib/graphql/breadth/lazy_loader.rb @@ -5,7 +5,7 @@ module GraphQL module Breadth #: [ContextType < GraphQL::Query::Context] class LazyLoader - AsyncSettings = Data.define(:enabled, :limit, :resource, :timeout) do + AsyncSettings = Data.define(:enabled, :limit, :resource, :timeout, :throttle) do alias_method :enabled?, :enabled end @@ -70,11 +70,12 @@ def resolve(results) limit: nil, resource: nil, timeout: nil, + throttle: nil, ).freeze class << self - #: (?limit: Integer?, ?resource: Symbol?, ?timeout: Numeric?) -> void - def async(limit: 8, resource: nil, timeout: nil) + #: (?limit: Integer?, ?resource: Symbol?, ?timeout: Numeric?, ?throttle: async_throttle?) -> void + def async(limit: 8, resource: nil, timeout: nil, throttle: nil) unless GraphQL::Breadth.async_enabled? Kernel.raise ImplementationError, "Async lazy loaders require `GraphQL::Breadth.enable_async!` during initialization." end @@ -87,6 +88,7 @@ def async(limit: 8, resource: nil, timeout: nil) limit: limit, resource: resource || name&.to_sym || :"lazy_loader_#{object_id}", timeout: timeout, + throttle: throttle, ).freeze end @@ -133,22 +135,22 @@ def async_settings self.class.async_settings end - #: (?resource: Symbol?, ?limit: Integer?, ?timeout: Numeric?) ?{ -> untyped } -> Executor::LazyAsync::Future - def async(resource: nil, limit: nil, timeout: nil, &block) + #: (?resource: Symbol?, ?limit: Integer?, ?timeout: Numeric?, ?throttle: async_throttle?) ?{ -> untyped } -> Executor::LazyAsync::Future + def async(resource: nil, limit: nil, timeout: nil, throttle: nil, &block) unless @async_context raise ImplementationError, "LazyLoader#async requires the loader class to opt into async features via `async(...)`." end - @async_context.async(resource: resource, limit: limit, timeout: timeout, &block) + @async_context.async(resource: resource, limit: limit, timeout: timeout, throttle: throttle, &block) end - #: [T, U] (Enumerable[T], ?resource: Symbol?, ?limit: Integer?, ?timeout: Numeric?) ?{ (T) -> U } -> Array[U] - def async_map(collection, resource: nil, limit: nil, timeout: nil, &block) + #: [T, U] (Enumerable[T], ?resource: Symbol?, ?limit: Integer?, ?timeout: Numeric?, ?throttle: async_throttle?) ?{ (T) -> U } -> Array[U] + def async_map(collection, resource: nil, limit: nil, timeout: nil, throttle: nil, &block) unless @async_context raise ImplementationError, "LazyLoader#async_map requires the loader class to opt into async features via `async(...)`." end - @async_context.async_map(collection, resource: resource, limit: limit, timeout: timeout, &block) + @async_context.async_map(collection, resource: resource, limit: limit, timeout: timeout, throttle: throttle, &block) end #: (Array[untyped], ContextType) -> void diff --git a/sorbet/rbi/shims/async.rbi b/sorbet/rbi/shims/async.rbi index 9dd0ac3..3db65ea 100644 --- a/sorbet/rbi/shims/async.rbi +++ b/sorbet/rbi/shims/async.rbi @@ -1,6 +1,15 @@ # typed: true module Async + module Limiter + class Generic + def async(parent: T.unsafe(nil), **options, &block); end + end + + class Limited < Generic; end + class Queued < Generic; end + end + class TimeoutError < StandardError; end class Barrier diff --git a/test/graphql/breadth/executor/loaders_test.rb b/test/graphql/breadth/executor/loaders_test.rb index d49ca9b..5f126d2 100644 --- a/test/graphql/breadth/executor/loaders_test.rb +++ b/test/graphql/breadth/executor/loaders_test.rb @@ -302,6 +302,40 @@ def perform_map(keys, _ctx) end end + class TestAsyncThrottle + attr_accessor :fork_count + + def initialize + @fork_count = 0 + end + + def async(parent:, &block) + @fork_count += 1 + parent.async(&block) + end + end + + class ProvidedThrottleFanoutLoader < GraphQL::Breadth::LazyLoader + THROTTLE = TestAsyncThrottle.new + + async limit: 1, throttle: THROTTLE + + def map? + true + end + + def perform_map(keys, _ctx) + async_map(keys) do |key| + label = [:provided_throttle, key] + AsyncFlowTracker.enter(label) + sleep(key == "Apple" ? 0.01 : 0.001) + "#{key}-provided-throttle" + ensure + AsyncFlowTracker.leave(label) + end + end + end + class ConflictingLimitOneLoader < GraphQL::Breadth::LazyLoader async limit: 1, resource: :conflicting_async_resource @@ -733,6 +767,7 @@ def setup FanoutMapLoader.reset_tracking! SharedResourceTracker.reset_tracking! AsyncFlowTracker.reset_tracking! + ProvidedThrottleFanoutLoader::THROTTLE.fork_count = 0 AbortedPendingLoader.reset_tracking! end @@ -1206,6 +1241,43 @@ def test_async_limit_nil_runs_child_work_without_a_semaphore assert_equal 2, AsyncFlowTracker.max_active end + def test_async_map_uses_a_loader_provided_throttle_for_parent_and_child_forks + document = GraphQL.parse(%|{ + widgets { + first + } + }|) + + source = { + "widgets" => [ + { "first" => "Apple" }, + { "first" => "Banana" }, + ], + } + + resolvers = LOADER_RESOLVERS.merge( + "Widget" => LOADER_RESOLVERS["Widget"].merge( + "first" => LoaderResolver.new(loader_class: ProvidedThrottleFanoutLoader, key: "first"), + ), + ) + + expected = { + "data" => { + "widgets" => [ + { "first" => "Apple-provided-throttle" }, + { "first" => "Banana-provided-throttle" }, + ], + }, + } + + executor = GraphQL::Breadth::Executor.new(LOADER_SCHEMA, document, resolvers: resolvers, root_object: source) + assert_equal expected, Timeout.timeout(1) { executor.result } + assert_equal 1, ProvidedThrottleFanoutLoader.async_settings.limit + assert_same ProvidedThrottleFanoutLoader::THROTTLE, ProvidedThrottleFanoutLoader.async_settings.throttle + assert_equal 3, ProvidedThrottleFanoutLoader::THROTTLE.fork_count + assert_equal 1, AsyncFlowTracker.max_active + end + def test_async_scheduler_waits_for_requeued_work_targeting_active_loader document = GraphQL.parse(%|{ widget {