feat: add external storage concurrency limiter APIs

temporalio/lib/temporalio/converters.rb

@@ -2,9 +2,11 @@
 
 require 'temporalio/converters/data_converter'
 require 'temporalio/converters/external_storage'
+require 'temporalio/converters/external_storage_concurrency'
 require 'temporalio/converters/storage_driver'
 require 'temporalio/converters/storage_driver_activity_info'
 require 'temporalio/converters/storage_driver_claim'
+require 'temporalio/converters/storage_driver_limiter'
 require 'temporalio/converters/storage_driver_retrieve_context'
 require 'temporalio/converters/storage_driver_select_context'
 require 'temporalio/converters/storage_driver_store_context'

temporalio/lib/temporalio/converters/external_storage.rb

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 
+require 'temporalio/converters/external_storage_concurrency'
 require 'temporalio/converters/storage_driver'
 
 module Temporalio
@@ -24,6 +25,9 @@ class ExternalStorage
       #   Only called for payloads that meet {#payload_size_threshold}, and must return one of {#drivers} or nil.
       attr_reader :driver_selector
 
+      # @return [ExternalStorageConcurrency] Limits on how many storage requests may be in flight at once.
+      attr_reader :concurrency
+
       # @return [Integer] Minimum encoded payload size, in bytes, that is offloaded. Payloads at or above this size are
       #   offloaded; smaller ones are left inline. Zero offloads every payload.
       attr_reader :payload_size_threshold
@@ -38,7 +42,13 @@ class ExternalStorage
       #   payload.
       # @param payload_size_threshold [Integer] Minimum encoded payload size, in bytes, to offload. Must be a
       #   non-negative Integer.
-      def initialize(drivers:, driver_selector: nil, payload_size_threshold: DEFAULT_PAYLOAD_SIZE_THRESHOLD)
+      # @param concurrency [ExternalStorageConcurrency] Limits on how many storage requests may be in flight at once.
+      def initialize(
+        drivers:,
+        driver_selector: nil,
+        payload_size_threshold: DEFAULT_PAYLOAD_SIZE_THRESHOLD,
+        concurrency: ExternalStorageConcurrency.new
+      )
         raise ArgumentError, 'At least one driver is required' if drivers.empty?
 
         unless payload_size_threshold.is_a?(Integer) && payload_size_threshold >= 0
@@ -65,6 +75,7 @@ def initialize(drivers:, driver_selector: nil, payload_size_threshold: DEFAULT_P
         # Normalized so callers never have to distinguish the single-driver case.
         @driver_selector = driver_selector || ->(_context, _payload) { @drivers.first }
         @payload_size_threshold = payload_size_threshold
+        @concurrency = concurrency
       end
 
       # Get the driver registered under the given name.

temporalio/lib/temporalio/converters/external_storage_concurrency.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+module Temporalio
+  module Converters
+    ExternalStorageConcurrency = Data.define(
+      :max_driver_operations,
+      :max_operations_per_message
+    )
+
+    # Limits on how many external storage requests may be in flight at once. Set on
+    # {ExternalStorage#concurrency}.
+    #
+    # @note WARNING: This API is experimental and may change in the future.
+    #
+    # @!attribute [r] max_driver_operations
+    #   @return [Integer] Maximum requests in flight at once across every driver registered on the owning
+    #     {ExternalStorage}. Each worker and each client gets its own budget of this size.
+    # @!attribute [r] max_operations_per_message
+    #   @return [Integer] Maximum requests in flight at once on behalf of a single message, such as one workflow task
+    #     activation, activity task, Nexus operation, or client request. Every message gets its own budget of this
+    #     size, which stops a message carrying many large payloads from starving the rest.
+    #
+    # @!visibility private
+    class ExternalStorageConcurrency
+      # Shared with every other SDK.
+      DEFAULT_MAX_DRIVER_OPERATIONS = 64
+
+      # Shared with every other SDK.
+      DEFAULT_MAX_OPERATIONS_PER_MESSAGE = 8
+
+      # Create concurrency limits.
+      #
+      # @param max_driver_operations [Integer] Maximum requests in flight across the configuration.
+      # @param max_operations_per_message [Integer] Maximum requests in flight for one message.
+      def initialize(
+        max_driver_operations: DEFAULT_MAX_DRIVER_OPERATIONS,
+        max_operations_per_message: DEFAULT_MAX_OPERATIONS_PER_MESSAGE
+      )
+        # A budget of zero would stop every storage operation rather than disabling the limit, so it is rejected
+        # here instead of deadlocking the first offload.
+        unless max_driver_operations.is_a?(Integer) && max_driver_operations >= 1
+          raise ArgumentError, 'max_driver_operations must be an Integer of at least 1'
+        end
+        unless max_operations_per_message.is_a?(Integer) && max_operations_per_message >= 1
+          raise ArgumentError, 'max_operations_per_message must be an Integer of at least 1'
+        end
+
+        super
+      end
+    end
+  end
+end

temporalio/lib/temporalio/converters/storage_driver_limiter.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module Temporalio
+  module Converters
+    # Bounds how many external storage requests are in flight at once. The SDK supplies one on
+    # {StorageDriverStoreContext} and {StorageDriverRetrieveContext}; a driver calls it rather than
+    # implementing it.
+    #
+    # A driver is handed a whole batch and decides for itself how many requests to make of it, so the SDK cannot
+    # bound that fan out on the driver's behalf. The limit is cooperative: a driver that never calls {#permit} makes
+    # requests that no budget counts, and the SDK warns when it detects that.
+    #
+    # @note WARNING: This API is experimental and may change in the future.
+    #
+    # @!visibility private
+    class StorageDriverLimiter
+      # Hold a permit for the duration of the given block, waiting for one to be available before the block runs
+      # and releasing it when the block returns or raises.
+      #
+      # Take one permit per request, not one around a whole batch: wrapping a whole fan out in a single permit
+      # satisfies this signature while quietly defeating the limit. Taking a permit while already holding one can
+      # deadlock once the nesting reaches the configured limit, so a driver must not nest these.
+      #
+      # Waiting for a permit observes the cancellation on the context this limiter came from, so a driver that is
+      # waiting when the SDK abandons the operation stops waiting.
+      #
+      # @param item [Api::Common::V1::Payload, StorageDriverClaim] What this one request covers: a payload when
+      #   storing, a claim when retrieving.
+      # @yield The request to make while holding the permit.
+      # @return [Object] Whatever the block returned.
+      # @raise [Error::CanceledError] If cancelled while waiting for a permit.
+      def permit(item, &)
+        raise NotImplementedError
+      end
+    end
+  end
+end

temporalio/lib/temporalio/converters/storage_driver_retrieve_context.rb

@@ -1,10 +1,14 @@
 # frozen_string_literal: true
 
 require 'temporalio/cancellation'
+require 'temporalio/converters/storage_driver_limiter'
 
 module Temporalio
   module Converters
-    StorageDriverRetrieveContext = Data.define(:cancellation)
+    StorageDriverRetrieveContext = Data.define(
+      :cancellation,
+      :limiter
+    )
 
     # Context given to {StorageDriver#retrieve}.
     #
@@ -16,6 +20,9 @@ module Converters
     #
     # @!attribute [r] cancellation
     #   @return [Cancellation] Cancelled when the SDK abandons this retrieve operation.
+    # @!attribute [r] limiter
+    #   @return [StorageDriverLimiter] Bounds how many retrieve requests this driver may have in flight. Take one
+    #     permit per request.
     #
     # @!visibility private
     class StorageDriverRetrieveContext; end # rubocop:disable Lint/EmptyClass

temporalio/lib/temporalio/converters/storage_driver_store_context.rb

@@ -1,12 +1,14 @@
 # frozen_string_literal: true
 
 require 'temporalio/cancellation'
+require 'temporalio/converters/storage_driver_limiter'
 
 module Temporalio
   module Converters
     StorageDriverStoreContext = Data.define(
       :target,
-      :cancellation
+      :cancellation,
+      :limiter
     )
 
     # Context given to {StorageDriver#store}.
@@ -19,6 +21,9 @@ module Converters
     #     on behalf of, or nil when there is no associated execution.
     # @!attribute [r] cancellation
     #   @return [Cancellation] Cancelled when the SDK abandons this store operation.
+    # @!attribute [r] limiter
+    #   @return [StorageDriverLimiter] Bounds how many store requests this driver may have in flight. Take one permit
+    #     per request.
     #
     # @!visibility private
     class StorageDriverStoreContext; end # rubocop:disable Lint/EmptyClass

temporalio/rbi/temporalio/converters/external_storage.rbi

@@ -21,6 +21,9 @@ class Temporalio::Converters::ExternalStorage
   sig { returns(Integer) }
   attr_reader :payload_size_threshold
 
+  sig { returns(Temporalio::Converters::ExternalStorageConcurrency) }
+  attr_reader :concurrency
+
   sig do
     params(
       drivers: T::Array[Temporalio::Converters::StorageDriver],
@@ -30,10 +33,16 @@ class Temporalio::Converters::ExternalStorage
           payload: Temporalio::Api::Common::V1::Payload
         ).returns(T.nilable(Temporalio::Converters::StorageDriver))
       ),
-      payload_size_threshold: Integer
+      payload_size_threshold: Integer,
+      concurrency: Temporalio::Converters::ExternalStorageConcurrency
     ).void
   end
-  def initialize(drivers:, driver_selector: T.unsafe(nil), payload_size_threshold: T.unsafe(nil)); end
+  def initialize(
+    drivers:,
+    driver_selector: T.unsafe(nil),
+    payload_size_threshold: T.unsafe(nil),
+    concurrency: T.unsafe(nil)
+  ); end
 
   sig { params(name: String).returns(T.nilable(Temporalio::Converters::StorageDriver)) }
   def driver(name); end

temporalio/rbi/temporalio/converters/external_storage_concurrency.rbi

@@ -0,0 +1,26 @@
+# typed: true
+
+class Temporalio::Converters::ExternalStorageConcurrency < ::Data
+  DEFAULT_MAX_DRIVER_OPERATIONS = T.let(T.unsafe(nil), Integer)
+  DEFAULT_MAX_OPERATIONS_PER_MESSAGE = T.let(T.unsafe(nil), Integer)
+
+  sig { params(max_driver_operations: Integer, max_operations_per_message: Integer).void }
+  def initialize(max_driver_operations: 64, max_operations_per_message: 8); end
+
+  sig { returns(Integer) }
+  def max_driver_operations; end
+
+  sig { returns(Integer) }
+  def max_operations_per_message; end
+
+  class << self
+    sig { params(args: T.untyped).returns(Temporalio::Converters::ExternalStorageConcurrency) }
+    def [](*args); end
+
+    sig { returns(T::Array[Symbol]) }
+    def members; end
+
+    sig { params(args: T.untyped).returns(Temporalio::Converters::ExternalStorageConcurrency) }
+    def new(*args); end
+  end
+end

temporalio/rbi/temporalio/converters/storage_driver_limiter.rbi

@@ -0,0 +1,13 @@
+# typed: true
+
+class Temporalio::Converters::StorageDriverLimiter
+  sig do
+    type_parameters(:T)
+      .params(
+        item: T.any(Temporalio::Api::Common::V1::Payload, Temporalio::Converters::StorageDriverClaim),
+        blk: T.proc.returns(T.type_parameter(:T))
+      )
+      .returns(T.type_parameter(:T))
+  end
+  def permit(item, &blk); end
+end

temporalio/rbi/temporalio/converters/storage_driver_retrieve_context.rbi

@@ -1,12 +1,15 @@
 # typed: true
 
 class Temporalio::Converters::StorageDriverRetrieveContext < ::Data
-  sig { params(cancellation: Temporalio::Cancellation).void }
-  def initialize(cancellation:); end
+  sig { params(cancellation: Temporalio::Cancellation, limiter: Temporalio::Converters::StorageDriverLimiter).void }
+  def initialize(cancellation:, limiter:); end
 
   sig { returns(Temporalio::Cancellation) }
   def cancellation; end
 
+  sig { returns(Temporalio::Converters::StorageDriverLimiter) }
+  def limiter; end
+
   class << self
     sig { params(args: T.untyped).returns(Temporalio::Converters::StorageDriverRetrieveContext) }
     def [](*args); end

temporalio/rbi/temporalio/converters/storage_driver_store_context.rbi

@@ -9,10 +9,11 @@ class Temporalio::Converters::StorageDriverStoreContext < ::Data
           Temporalio::Converters::StorageDriverActivityInfo
         )
       ),
-      cancellation: Temporalio::Cancellation
+      cancellation: Temporalio::Cancellation,
+      limiter: Temporalio::Converters::StorageDriverLimiter
     ).void
   end
-  def initialize(target:, cancellation:); end
+  def initialize(target:, cancellation:, limiter:); end
 
   sig do
     returns(
@@ -29,6 +30,9 @@ class Temporalio::Converters::StorageDriverStoreContext < ::Data
   sig { returns(Temporalio::Cancellation) }
   def cancellation; end
 
+  sig { returns(Temporalio::Converters::StorageDriverLimiter) }
+  def limiter; end
+
   class << self
     sig { params(args: T.untyped).returns(Temporalio::Converters::StorageDriverStoreContext) }
     def [](*args); end

temporalio/sig/temporalio/converters/external_storage.rbs

@@ -8,13 +8,15 @@ module Temporalio
       attr_reader drivers: Array[StorageDriver]
       attr_reader driver_selector: driverSelector
       attr_reader payload_size_threshold: Integer
+      attr_reader concurrency: ExternalStorageConcurrency
 
       @drivers_by_name: Hash[String, StorageDriver]
 
       def initialize: (
         drivers: Array[StorageDriver],
         ?driver_selector: driverSelector?,
-        ?payload_size_threshold: Integer
+        ?payload_size_threshold: Integer,
+        ?concurrency: ExternalStorageConcurrency
       ) -> void
 
       def driver: (String name) -> StorageDriver?

temporalio/sig/temporalio/converters/external_storage_concurrency.rbs

@@ -0,0 +1,16 @@
+module Temporalio
+  module Converters
+    class ExternalStorageConcurrency
+      DEFAULT_MAX_DRIVER_OPERATIONS: Integer
+      DEFAULT_MAX_OPERATIONS_PER_MESSAGE: Integer
+
+      attr_reader max_driver_operations: Integer
+      attr_reader max_operations_per_message: Integer
+
+      def initialize: (
+        ?max_driver_operations: Integer,
+        ?max_operations_per_message: Integer
+      ) -> void
+    end
+  end
+end

temporalio/sig/temporalio/converters/storage_driver_limiter.rbs

@@ -0,0 +1,7 @@
+module Temporalio
+  module Converters
+    class StorageDriverLimiter
+      def permit: [T] (Api::Common::V1::Payload | StorageDriverClaim item) { () -> T } -> T
+    end
+  end
+end

temporalio/sig/temporalio/converters/storage_driver_retrieve_context.rbs

@@ -2,8 +2,12 @@ module Temporalio
   module Converters
     class StorageDriverRetrieveContext
       attr_reader cancellation: Cancellation
+      attr_reader limiter: StorageDriverLimiter
 
-      def initialize: (cancellation: Cancellation) -> void
+      def initialize: (
+        cancellation: Cancellation,
+        limiter: StorageDriverLimiter
+      ) -> void
     end
   end
 end

temporalio/sig/temporalio/converters/storage_driver_store_context.rbs

@@ -3,10 +3,12 @@ module Temporalio
     class StorageDriverStoreContext
       attr_reader target: storageDriverTargetInfo?
       attr_reader cancellation: Cancellation
+      attr_reader limiter: StorageDriverLimiter
 
       def initialize: (
         target: storageDriverTargetInfo?,
-        cancellation: Cancellation
+        cancellation: Cancellation,
+        limiter: StorageDriverLimiter
       ) -> void
     end
   end