Add async limiter

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
 

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
 

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

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

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

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 {