Add extstore runner

temporalio/lib/temporalio/converters/data_converter.rb

@@ -1,9 +1,11 @@
 # frozen_string_literal: true
 
 require 'temporalio/api'
+require 'temporalio/converters/external_storage'
 require 'temporalio/converters/failure_converter'
 require 'temporalio/converters/payload_codec'
 require 'temporalio/converters/payload_converter'
+require 'temporalio/internal/external_storage_runner'
 
 module Temporalio
   module Converters
@@ -18,6 +20,13 @@ class DataConverter
       # @return [PayloadCodec, nil] Optional codec for encoding/decoding payload bytes such as for encryption.
       attr_reader :payload_codec
 
+      # @return [ExternalStorage, nil] Optional storage for offloading large payloads.
+      #
+      # @note WARNING: This API is experimental and may change in the future.
+      #
+      # @!visibility private
+      attr_reader :external_storage
+
       # @return [DataConverter] Default data converter.
       def self.default
         @default ||= DataConverter.new
@@ -28,14 +37,22 @@ def self.default
       # @param payload_converter [PayloadConverter] Payload converter to use.
       # @param failure_converter [FailureConverter] Failure converter to use.
       # @param payload_codec [PayloadCodec, nil] Payload codec to use.
+      # @param external_storage [ExternalStorage, nil] External storage to use. WARNING: This is experimental.
       def initialize(
         payload_converter: PayloadConverter.default,
         failure_converter: FailureConverter.default,
-        payload_codec: nil
+        payload_codec: nil,
+        external_storage: nil
       )
         @payload_converter = payload_converter
         @failure_converter = failure_converter
         @payload_codec = payload_codec
+        @external_storage = external_storage
+        # Storage is applied in the payload methods below so every client and activity call gets it without changes at
+        # each call site. This is temporary. It will move to the call sites as a pass over each whole request, which
+        # batches all of a request's payloads into one driver call and can tell drivers which execution they belong to.
+        # Retrieving runs even without storage, so a reference is reported rather than decoded as an ordinary value.
+        @external_storage_runner = Internal::ExternalStorageRunner.new(external_storage)
       end
 
       # Convert a Ruby value to a payload and encode it.
@@ -46,6 +63,7 @@ def initialize(
       def to_payload(value, hint: nil)
         payload = payload_converter.to_payload(value, hint:)
         payload = payload_codec.encode([payload]).first if payload_codec
+        payload = @external_storage_runner.store([payload]).first if external_storage
         payload
       end
 
@@ -60,6 +78,9 @@ def to_payload(value, hint: nil)
       def to_payloads(values, hints: nil)
         payloads = payload_converter.to_payloads(values, hints:)
         payloads.payloads.replace(payload_codec.encode(payloads.payloads)) if payload_codec && !payloads.payloads.empty?
+        if external_storage && !payloads.payloads.empty?
+          payloads.payloads.replace(@external_storage_runner.store(payloads.payloads.to_a))
+        end
         payloads
       end
 
@@ -69,6 +90,7 @@ def to_payloads(values, hints: nil)
       # @param hint [Object, nil] Hint, if any, to assist conversion.
       # @return [Object] Decoded and converted Ruby value.
       def from_payload(payload, hint: nil)
+        payload = @external_storage_runner.retrieve([payload]).first
         payload = payload_codec.decode([payload]).first if payload_codec
         payload_converter.from_payload(payload, hint:)
       end
@@ -84,6 +106,7 @@ def from_payload(payload, hint: nil)
       def from_payloads(payloads, hints: nil)
         return [] unless payloads && !payloads.payloads.empty?
 
+        payloads = Api::Common::V1::Payloads.new(payloads: @external_storage_runner.retrieve(payloads.payloads.to_a))
         if payload_codec && !payloads.payloads.empty?
           payloads = Api::Common::V1::Payloads.new(payloads: payload_codec.decode(payloads.payloads))
         end

temporalio/lib/temporalio/converters/payload_converter/composite.rb

@@ -4,6 +4,7 @@
 require 'temporalio/converters/payload_converter'
 require 'temporalio/converters/raw_value'
 require 'temporalio/error'
+require 'temporalio/internal/external_storage_references'
 
 module Temporalio
   module Converters
@@ -56,6 +57,9 @@ def to_payload(value, hint: nil)
         # @raise [ConverterNotFound] If no converter found for the encoding.
         def from_payload(payload, hint: nil)
           return nil unless payload
+          # A reference that reaches conversion was never retrieved, and would otherwise be converted into the reference
+          # message itself
+          raise Error::ExternalStorageNotConfiguredError if Internal::ExternalStorageReferences.reference?(payload)
 
           encoding = payload.metadata['encoding']
           raise EncodingNotSet, 'Missing payload encoding' unless encoding

temporalio/lib/temporalio/internal/external_storage_runner.rb

@@ -0,0 +1,129 @@
+# frozen_string_literal: true
+
+require 'temporalio/api'
+require 'temporalio/cancellation'
+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'
+require 'temporalio/error'
+require 'temporalio/internal/external_storage_references'
+
+module Temporalio
+  module Internal
+    # Stores large payloads in external storage and fetches them back.
+    #
+    # Storing replaces each payload at or above the size threshold with a reference to what a driver stored. Retrieving
+    # swaps references back for the payloads they point at. Payloads are expected to already be codec-encoded when
+    # stored, and are returned still encoded when retrieved.
+    #
+    # Each driver gets one batched request per call, and requests to different drivers run one after another. The
+    # configured concurrency limits are not enforced and the driver target is always nil.
+    #
+    # @!visibility private
+    class ExternalStorageRunner
+      # @param storage [Converters::ExternalStorage, nil] Configured storage. When nil, storing does nothing and
+      #   retrieving raises if it finds a reference, so a reference is never mistaken for an ordinary payload.
+      def initialize(storage)
+        @storage = storage
+      end
+
+      # @param payloads [Array<Api::Common::V1::Payload>] Codec-encoded payloads. Not mutated.
+      # @return [Array<Api::Common::V1::Payload>] Payloads with every eligible one replaced by a reference.
+      def store(payloads)
+        storage = @storage
+        return payloads unless storage
+
+        select_context = Converters::StorageDriverSelectContext.new(target: nil, cancellation: Cancellation.new)
+
+        # Driver => [[index, serialized payload], ...] for each payload that driver stores
+        payloads_by_driver = Hash.new { |hash, driver| hash[driver] = [] }
+        payloads.each_with_index do |payload, index|
+          serialized = payload.to_proto
+          next if serialized.bytesize < storage.payload_size_threshold
+
+          driver = storage.driver_selector.call(select_context, payload)
+          next unless driver
+
+          # Checked by identity, so a driver built to look like a registered one is rejected.
+          unless storage.driver(driver.name).equal?(driver)
+            raise Error, "External storage driver selector returned a driver that is not registered: #{driver.inspect}"
+          end
+
+          payloads_by_driver[driver] << [index, serialized]
+        end
+
+        result = payloads.dup
+        payloads_by_driver.each do |driver, indexes_and_serialized|
+          context = Converters::StorageDriverStoreContext.new(
+            target: nil, cancellation: Cancellation.new, limiter: UNBOUNDED_LIMITER
+          )
+          # Drivers get copies, since a driver may keep a payload after returning and the SDK modifies some payloads in
+          # place
+          copies = indexes_and_serialized.map { |_index, serialized| Api::Common::V1::Payload.decode(serialized) }
+          claims = driver.store(context, copies)
+          check_result_count(driver, 'store', indexes_and_serialized.size, claims)
+
+          indexes_and_serialized.zip(claims) do |(index, serialized), claim|
+            result[index] = ExternalStorageReferences.create_reference_payload(
+              driver_name: driver.name, claim_data: claim.claim_data, original_size_bytes: serialized.bytesize
+            )
+          end
+        end
+        result
+      end
+
+      # @param payloads [Array<Api::Common::V1::Payload>] Payloads that may contain references. Not mutated.
+      # @return [Array<Api::Common::V1::Payload>] Payloads with every reference replaced by what it points at.
+      # @raise [Error::ExternalStorageNotConfiguredError] If there is a reference but no storage is configured.
+      def retrieve(payloads)
+        storage = @storage
+        # Driver => [[index, claim], ...] for each reference that driver stored. Every driver is looked up here, so an
+        # unregistered name fails before any request is made.
+        claims_by_driver = Hash.new { |hash, driver| hash[driver] = [] }
+        payloads.each_with_index do |payload, index|
+          reference = ExternalStorageReferences.parse_reference(payload)
+          next unless reference
+          raise Error::ExternalStorageNotConfiguredError unless storage
+
+          driver = storage.driver(reference.driver_name)
+          raise Error, "No external storage driver registered with name '#{reference.driver_name}'" unless driver
+
+          claims_by_driver[driver] << [index, Converters::StorageDriverClaim.new(claim_data: reference.claim_data.to_h)]
+        end
+
+        result = payloads.dup
+        claims_by_driver.each do |driver, indexes_and_claims|
+          context = Converters::StorageDriverRetrieveContext.new(
+            cancellation: Cancellation.new, limiter: UNBOUNDED_LIMITER
+          )
+          retrieved = driver.retrieve(context, indexes_and_claims.map { |_index, claim| claim })
+          check_result_count(driver, 'retrieve', indexes_and_claims.size, retrieved)
+
+          indexes_and_claims.zip(retrieved) { |(index, _claim), payload| result[index] = payload }
+        end
+        result
+      end
+
+      private
+
+      def check_result_count(driver, operation, requested, returned)
+        return if returned.is_a?(Array) && returned.size == requested
+
+        raise Error, "External storage driver '#{driver.name}' #{operation} was given #{requested} item(s) but " \
+                     "returned #{returned.is_a?(Array) ? returned.size : returned.inspect}"
+      end
+
+      # Grants every permit immediately, since concurrency limits are not enforced yet.
+      class UnboundedLimiter < Converters::StorageDriverLimiter
+        def permit(_item)
+          yield
+        end
+      end
+
+      UNBOUNDED_LIMITER = UnboundedLimiter.new
+      private_constant :UnboundedLimiter, :UNBOUNDED_LIMITER
+    end
+  end
+end

temporalio/lib/temporalio/internal/worker/activity_worker.rb

@@ -5,6 +5,7 @@
 require 'temporalio/cancellation'
 require 'temporalio/converters/raw_value'
 require 'temporalio/internal/bridge/api'
+require 'temporalio/internal/external_storage_runner'
 require 'temporalio/internal/proto_utils'
 require 'temporalio/scoped_logger'
 require 'temporalio/worker/interceptor'
@@ -191,7 +192,9 @@ def execute_activity(task_token, defn, start)
             priority: Priority._from_proto(start.priority),
             raw_heartbeat_details: begin
               payloads = start.heartbeat_details.to_ary
-              codec = @worker.options.client.data_converter.payload_codec
+              data_converter = @worker.options.client.data_converter
+              payloads = ExternalStorageRunner.new(data_converter.external_storage).retrieve(payloads)
+              codec = data_converter.payload_codec
               payloads = codec.decode(payloads) if codec
               payloads.map { |p| Temporalio::Converters::RawValue.new(p) }
             end,
@@ -214,7 +217,9 @@ def execute_activity(task_token, defn, start)
             # If the activity wants raw_args, we only decode we don't convert
             args: if defn.raw_args
                     payloads = start.input.to_ary
-                    codec = @worker.options.client.data_converter.payload_codec
+                    data_converter = @worker.options.client.data_converter
+                    payloads = ExternalStorageRunner.new(data_converter.external_storage).retrieve(payloads)
+                    codec = data_converter.payload_codec
                     payloads = codec.decode(payloads) if codec
                     payloads.map { |p| Temporalio::Converters::RawValue.new(p) }
                   else

temporalio/lib/temporalio/internal/worker/workflow_worker.rb

@@ -2,6 +2,7 @@
 
 require 'temporalio/api/payload_visitor'
 require 'temporalio/error'
+require 'temporalio/internal/external_storage_runner'
 require 'temporalio/internal/worker/workflow_instance'
 require 'temporalio/scoped_logger'
 require 'temporalio/workflow'
@@ -81,18 +82,30 @@ def initialize(
           @executor = workflow_executor
 
           payload_codec = data_converter.payload_codec
+          external_storage = data_converter.external_storage
           @workflow_payload_codec_thread_pool = workflow_payload_codec_thread_pool
-          if !Fiber.current_scheduler && payload_codec && !@workflow_payload_codec_thread_pool
-            raise ArgumentError, 'Must have workflow payload codec thread pool if providing codec and not using fibers'
+          if !Fiber.current_scheduler && (payload_codec || external_storage) && !@workflow_payload_codec_thread_pool
+            raise ArgumentError, 'Must have workflow payload codec thread pool if providing codec or external ' \
+                                 'storage and not using fibers'
           end
 
-          # If there is a payload codec, we need to build encoding and decoding visitors
-          if payload_codec
+          # If there is a payload codec or external storage, we need to build encoding and decoding visitors. Storage
+          # comes after the codec when encoding and before it when decoding.
+          if payload_codec || external_storage
+            storage_runner = ExternalStorageRunner.new(external_storage) if external_storage
             @payload_encoding_visitor = Api::PayloadVisitor.new(skip_search_attributes: true) do |payload_or_payloads|
-              apply_codec_on_payload_visit(payload_or_payloads) { |payloads| payload_codec.encode(payloads) }
+              replace_payloads_on_visit(payload_or_payloads) do |payloads|
+                payloads = payload_codec.encode(payloads) if payload_codec
+                payloads = storage_runner.store(payloads.to_a) if storage_runner
+                payloads
+              end
             end
             @payload_decoding_visitor = Api::PayloadVisitor.new(skip_search_attributes: true) do |payload_or_payloads|
-              apply_codec_on_payload_visit(payload_or_payloads) { |payloads| payload_codec.decode(payloads) }
+              replace_payloads_on_visit(payload_or_payloads) do |payloads|
+                payloads = storage_runner.retrieve(payloads.to_a) if storage_runner
+                payloads = payload_codec.decode(payloads) if payload_codec
+                payloads
+              end
             end
           end
 
@@ -173,22 +186,41 @@ def on_shutdown_complete
 
         private
 
+        # Decoding and encoding run in a thread pool that only warns about errors, so a failure must complete the
+        # activation as failed. Otherwise the run waits forever for a completion that never comes.
+
         def decode_activation(runner, activation)
           @payload_decoding_visitor.run(activation)
           runner.apply_workflow_activation_decoded(workflow_worker: self, activation:)
+        rescue Exception => e # rubocop:disable Lint/RescueException
+          fail_activation(runner, activation.run_id, e)
         end
 
         def encode_activation_completion(runner, activation_completion)
           @payload_encoding_visitor.run(activation_completion)
           runner.apply_workflow_activation_complete(workflow_worker: self, activation_completion:, encoded: true)
+        rescue Exception => e # rubocop:disable Lint/RescueException
+          fail_activation(runner, activation_completion.run_id, e)
+        end
+
+        def fail_activation(runner, run_id, error)
+          activation_completion = WorkflowInstance.new_completion_with_failure(
+            run_id:,
+            error:,
+            failure_converter: @state.data_converter.failure_converter,
+            payload_converter: @state.data_converter.payload_converter
+          )
+          runner.apply_workflow_activation_complete(workflow_worker: self, activation_completion:, encoded: true)
         end
 
-        def apply_codec_on_payload_visit(payload_or_payloads, &)
+        def replace_payloads_on_visit(payload_or_payloads, &)
           case payload_or_payloads
           when Temporalio::Api::Common::V1::Payload
             new_payloads = yield [payload_or_payloads]
             payload_or_payloads.metadata = new_payloads.first.metadata
             payload_or_payloads.data = new_payloads.first.data
+            # External storage records the original size of a stored payload here
+            payload_or_payloads.external_payloads.replace(new_payloads.first.external_payloads.to_a)
           when Enumerable
             payload_or_payloads.replace(yield payload_or_payloads) # steep:ignore
           else

temporalio/lib/temporalio/worker.rb

@@ -440,9 +440,9 @@ def self._validate_plugins!(plugins)
     #   instead of suspending the workflow via task failure. These are applied in addition to the
     #   `workflow_failure_exception_type` on the workflow definition class itself. If {::Exception} is set, it
     #   effectively will fail a workflow/update in all user exception cases.
-    # @param workflow_payload_codec_thread_pool [ThreadPool, nil] Thread pool to run payload codec encode/decode within.
-    #   This is required if a payload codec exists and the worker is not fiber based. Codecs can potentially block
-    #   execution which is why they need to be run in the background.
+    # @param workflow_payload_codec_thread_pool [ThreadPool, nil] Thread pool to run payload codec encode/decode and
+    #   external storage within. This is required if a payload codec or external storage exists and the worker is not
+    #   fiber based. Both can potentially block execution which is why they need to be run in the background.
     # @param unsafe_workflow_io_enabled [Boolean] If false, the default, workflow code that invokes io_wait on the fiber
     #   scheduler will fail. Instead of setting this to true, users are encouraged to use {Workflow::Unsafe.io_enabled}
     #   with a block for narrower enabling of IO.

temporalio/lib/temporalio/worker/workflow_replayer.rb

@@ -68,9 +68,9 @@ class Options; end # rubocop:disable Lint/EmptyClass
       #   instead of suspending the workflow via task failure. These are applied in addition to the
       #   `workflow_failure_exception_type` on the workflow definition class itself. If {::Exception} is set, it
       #   effectively will fail a workflow/update in all user exception cases.
-      # @param workflow_payload_codec_thread_pool [ThreadPool, nil] Thread pool to run payload codec encode/decode
-      #   within. This is required if a payload codec exists and the worker is not fiber based. Codecs can potentially
-      #   block execution which is why they need to be run in the background.
+      # @param workflow_payload_codec_thread_pool [ThreadPool, nil] Thread pool to run payload codec encode/decode and
+      #   external storage within. This is required if a payload codec or external storage exists and the worker is not
+      #   fiber based. Both can potentially block execution which is why they need to be run in the background.
       # @param unsafe_workflow_io_enabled [Boolean] If false, the default, workflow code that invokes io_wait on the
       #   fiber scheduler will fail. Instead of setting this to true, users are encouraged to use
       #   {Workflow::Unsafe.io_enabled} with a block for narrower enabling of IO.