Add header propagation and workflow middleware. (#226)

Co-authored-by: Mark Chua <[email protected]>

Author: markchua

examples/bin/trigger

@@ -2,6 +2,7 @@
 require_relative '../init'
 
 Dir[File.expand_path('../workflows/*.rb', __dir__)].each { |f| require f }
+Dir[File.expand_path('../middleware/*.rb', __dir__)].each { |f| require f }
 
 workflow_class_name, *args = ARGV
 workflow_class = Object.const_get(workflow_class_name)
@@ -10,5 +11,9 @@ workflow_id = SecureRandom.uuid
 # Convert integer strings to integers
 input = args.map { |arg| Integer(arg) rescue arg }
 
+Temporal.configure do |config|
+  config.add_header_propagator(SamplePropagator)
+end
+
 run_id = Temporal.start_workflow(workflow_class, *input, options: { workflow_id: workflow_id })
 Temporal.logger.info "Started workflow", { workflow_id: workflow_id, run_id: run_id }

examples/bin/worker

@@ -24,6 +24,10 @@ if !ENV['USE_ERROR_SERIALIZATION_V2'].nil?
   end
 end
 
+Temporal.configure do |config|
+  config.add_header_propagator(SamplePropagator)
+end
+
 worker = Temporal::Worker.new(binary_checksum: `git show HEAD -s --format=%H`.strip)
 
 worker.register_workflow(AsyncActivityWorkflow)
@@ -94,5 +98,7 @@ worker.register_dynamic_activity(DelegatorActivity)
 
 worker.add_workflow_task_middleware(LoggingMiddleware, 'EXAMPLE')
 worker.add_activity_middleware(LoggingMiddleware, 'EXAMPLE')
+worker.add_activity_middleware(SamplePropagator)
+worker.add_workflow_middleware(SamplePropagator)
 
 worker.start

examples/middleware/sample_propagator.rb

@@ -0,0 +1,10 @@
+class SamplePropagator
+  def inject!(headers)
+    headers['test-header'] = 'test'
+  end
+
+  def call(metadata)
+    Temporal.logger.info("Got headers!", headers: metadata.headers.to_h)
+    yield
+  end
+end

lib/temporal/client.rb

@@ -61,7 +61,7 @@ def start_workflow(workflow, *input, options: {}, **args)
           run_timeout: compute_run_timeout(execution_options),
           task_timeout: execution_options.timeouts[:task],
           workflow_id_reuse_policy: options[:workflow_id_reuse_policy],
-          headers: execution_options.headers,
+          headers: config.header_propagator_chain.inject(execution_options.headers),
           memo: execution_options.memo,
           search_attributes: Workflow::Context::Helpers.process_search_attributes(execution_options.search_attributes),
         )
@@ -78,7 +78,7 @@ def start_workflow(workflow, *input, options: {}, **args)
           run_timeout: compute_run_timeout(execution_options),
           task_timeout: execution_options.timeouts[:task],
           workflow_id_reuse_policy: options[:workflow_id_reuse_policy],
-          headers: execution_options.headers,
+          headers: config.header_propagator_chain.inject(execution_options.headers),
           memo: execution_options.memo,
           search_attributes: Workflow::Context::Helpers.process_search_attributes(execution_options.search_attributes),
           signal_name: signal_name,
@@ -127,7 +127,7 @@ def schedule_workflow(workflow, cron_schedule, *input, options: {}, **args)
         run_timeout: compute_run_timeout(execution_options),
         task_timeout: execution_options.timeouts[:task],
         workflow_id_reuse_policy: options[:workflow_id_reuse_policy],
-        headers: execution_options.headers,
+        headers: config.header_propagator_chain.inject(execution_options.headers),
         cron_schedule: cron_schedule,
         memo: execution_options.memo,
         search_attributes: Workflow::Context::Helpers.process_search_attributes(execution_options.search_attributes),

lib/temporal/configuration.rb

@@ -1,5 +1,7 @@
 require 'temporal/logger'
 require 'temporal/metrics_adapters/null'
+require 'temporal/middleware/header_propagator_chain'
+require 'temporal/middleware/entry'
 require 'temporal/connection/converter/payload/nil'
 require 'temporal/connection/converter/payload/bytes'
 require 'temporal/connection/converter/payload/json'
@@ -12,7 +14,7 @@ class Configuration
     Execution = Struct.new(:namespace, :task_queue, :timeouts, :headers, :search_attributes, keyword_init: true)
 
     attr_reader :timeouts, :error_handlers
-    attr_accessor :connection_type, :converter, :use_error_serialization_v2, :host, :port, :credentials, :identity, :logger, :metrics_adapter, :namespace, :task_queue, :headers, :search_attributes
+    attr_accessor :connection_type, :converter, :use_error_serialization_v2, :host, :port, :credentials, :identity, :logger, :metrics_adapter, :namespace, :task_queue, :headers, :search_attributes, :header_propagators
 
     # See https://docs.temporal.io/blog/activity-timeouts/ for general docs.
     # We want an infinite execution timeout for cron schedules and other perpetual workflows.
@@ -58,6 +60,7 @@ def initialize
       @credentials = :this_channel_is_insecure
       @identity = nil
       @search_attributes = {}
+      @header_propagators = []
     end
 
     def on_error(&block)
@@ -96,6 +99,15 @@ def default_execution_options
       ).freeze
     end
 
+    def add_header_propagator(propagator_class, *args)
+      raise 'header propagator must implement `def inject!(headers)`' unless propagator_class.method_defined? :inject!
+      @header_propagators << Middleware::Entry.new(propagator_class, args)
+    end
+
+    def header_propagator_chain
+      Middleware::HeaderPropagatorChain.new(header_propagators)
+    end
+
     private
 
     def default_identity

lib/temporal/connection/serializer/schedule_activity.rb

@@ -32,7 +32,7 @@ def to_proto
         def serialize_headers(headers)
           return unless headers
 
-          Temporalio::Api::Common::V1::Header.new(fields: object.headers)
+          Temporalio::Api::Common::V1::Header.new(fields: to_payload_map(headers))
         end
       end
     end

lib/temporal/middleware/header_propagator_chain.rb

@@ -0,0 +1,22 @@
+module Temporal
+  module Middleware
+    class HeaderPropagatorChain
+      def initialize(entries = [])
+        @propagators = entries.map(&:init_middleware)
+      end
+
+      def inject(headers)
+        return headers if propagators.empty?
+        h = headers.dup
+        for propagator in propagators
+          propagator.inject!(h)
+        end
+        h
+      end
+
+      private
+
+      attr_reader :propagators
+    end
+  end
+end

lib/temporal/worker.rb

@@ -31,6 +31,7 @@ def initialize(
       @activities = Hash.new { |hash, key| hash[key] = ExecutableLookup.new }
       @pollers = []
       @workflow_task_middleware = []
+      @workflow_middleware = []
       @activity_middleware = []
       @shutting_down = false
       @activity_poller_options = {
@@ -90,6 +91,10 @@ def add_workflow_task_middleware(middleware_class, *args)
       @workflow_task_middleware << Middleware::Entry.new(middleware_class, args)
     end
 
+    def add_workflow_middleware(middleware_class, *args)
+      @workflow_middleware << Middleware::Entry.new(middleware_class, args)
+    end
+
     def add_activity_middleware(middleware_class, *args)
       @activity_middleware << Middleware::Entry.new(middleware_class, args)
     end
@@ -128,14 +133,14 @@ def stop
 
     attr_reader :config, :activity_poller_options, :workflow_poller_options,
                 :activities, :workflows, :pollers,
-                :workflow_task_middleware, :activity_middleware
+                :workflow_task_middleware, :workflow_middleware, :activity_middleware
 
     def shutting_down?
       @shutting_down
     end
 
     def workflow_poller_for(namespace, task_queue, lookup)
-      Workflow::Poller.new(namespace, task_queue, lookup.freeze, config, workflow_task_middleware, workflow_poller_options)
+      Workflow::Poller.new(namespace, task_queue, lookup.freeze, config, workflow_task_middleware, workflow_middleware, workflow_poller_options)
     end
 
     def activity_poller_for(namespace, task_queue, lookup)

lib/temporal/workflow/context.rb

@@ -79,7 +79,7 @@ def execute_activity(activity_class, *input, **args)
           task_queue: execution_options.task_queue,
           retry_policy: execution_options.retry_policy,
           timeouts: execution_options.timeouts,
-          headers: execution_options.headers
+          headers: config.header_propagator_chain.inject(execution_options.headers)
         )
 
         target, cancelation_id = schedule_command(command)
@@ -136,7 +136,7 @@ def execute_workflow(workflow_class, *input, **args)
           retry_policy: execution_options.retry_policy,
           parent_close_policy: parent_close_policy,
           timeouts: execution_options.timeouts,
-          headers: execution_options.headers,
+          headers: config.header_propagator_chain.inject(execution_options.headers),
           cron_schedule: cron_schedule,
           memo: execution_options.memo,
           workflow_id_reuse_policy: workflow_id_reuse_policy,
@@ -261,7 +261,7 @@ def continue_as_new(*input, **args)
           input: input,
           timeouts: execution_options.timeouts,
           retry_policy: execution_options.retry_policy,
-          headers: execution_options.headers,
+          headers: config.header_propagator_chain.inject(execution_options.headers),
           memo: execution_options.memo,
           search_attributes: Helpers.process_search_attributes(execution_options.search_attributes)
         )

lib/temporal/workflow/executor.rb

@@ -16,7 +16,7 @@ class Executor
       # @param task_metadata [Metadata::WorkflowTask]
       # @param config [Configuration]
       # @param track_stack_trace [Boolean]
-      def initialize(workflow_class, history, task_metadata, config, track_stack_trace)
+      def initialize(workflow_class, history, task_metadata, config, track_stack_trace, middleware_chain)
         @workflow_class = workflow_class
         @dispatcher = Dispatcher.new
         @query_registry = QueryRegistry.new
@@ -25,6 +25,7 @@ def initialize(workflow_class, history, task_metadata, config, track_stack_trace
         @task_metadata = task_metadata
         @config = config
         @track_stack_trace = track_stack_trace
+        @middleware_chain = middleware_chain
       end
 
       def run
@@ -55,7 +56,8 @@ def process_queries(queries)
 
       private
 
-      attr_reader :workflow_class, :dispatcher, :query_registry, :state_manager, :task_metadata, :history, :config, :track_stack_trace
+      attr_reader :workflow_class, :dispatcher, :query_registry, :state_manager,
+                  :task_metadata, :history, :config, :track_stack_trace, :middleware_chain
 
       def process_query(query)
         result = query_registry.handle(query.query_type, query.query_args)
@@ -70,7 +72,9 @@ def execute_workflow(input, workflow_started_event)
         context = Workflow::Context.new(state_manager, dispatcher, workflow_class, metadata, config, query_registry, track_stack_trace)
 
         Fiber.new do
-          workflow_class.execute_in_context(context, input)
+          middleware_chain.invoke(metadata) do
+            workflow_class.execute_in_context(context, input)
+          end
         end.resume
       end
     end

lib/temporal/workflow/poller.rb

@@ -14,12 +14,13 @@ class Poller
         binary_checksum: nil
       }.freeze
 
-      def initialize(namespace, task_queue, workflow_lookup, config, middleware = [], options = {})
+      def initialize(namespace, task_queue, workflow_lookup, config, middleware = [], workflow_middleware = [], options = {})
         @namespace = namespace
         @task_queue = task_queue
         @workflow_lookup = workflow_lookup
         @config = config
         @middleware = middleware
+        @workflow_middleware = workflow_middleware
         @shutting_down = false
         @options = DEFAULT_OPTIONS.merge(options)
       end
@@ -45,7 +46,7 @@ def wait
 
       private
 
-      attr_reader :namespace, :task_queue, :connection, :workflow_lookup, :config, :middleware, :options, :thread
+      attr_reader :namespace, :task_queue, :connection, :workflow_lookup, :config, :middleware, :workflow_middleware, :options, :thread
 
       def connection
         @connection ||= Temporal::Connection.generate(config.for_connection)
@@ -96,8 +97,9 @@ def poll_for_task
 
       def process(task)
         middleware_chain = Middleware::Chain.new(middleware)
+        workflow_middleware_chain = Middleware::Chain.new(workflow_middleware)
 
-        TaskProcessor.new(task, namespace, workflow_lookup, middleware_chain, config, binary_checksum).process
+        TaskProcessor.new(task, namespace, workflow_lookup, middleware_chain, workflow_middleware_chain, config, binary_checksum).process
       end
 
       def thread_pool

lib/temporal/workflow/task_processor.rb

@@ -24,14 +24,15 @@ def query_args
       MAX_FAILED_ATTEMPTS = 1
       LEGACY_QUERY_KEY = :legacy_query
 
-      def initialize(task, namespace, workflow_lookup, middleware_chain, config, binary_checksum)
+      def initialize(task, namespace, workflow_lookup, middleware_chain, workflow_middleware_chain, config, binary_checksum)
         @task = task
         @namespace = namespace
         @metadata = Metadata.generate_workflow_task_metadata(task, namespace)
         @task_token = task.task_token
         @workflow_name = task.workflow_type.name
         @workflow_class = workflow_lookup.find(workflow_name)
         @middleware_chain = middleware_chain
+        @workflow_middleware_chain = workflow_middleware_chain
         @config = config
         @binary_checksum = binary_checksum
       end
@@ -53,7 +54,7 @@ def process
         track_stack_trace = queries.values.map(&:query_type).include?(StackTraceTracker::STACK_TRACE_QUERY_NAME)
 
         # TODO: For sticky workflows we need to cache the Executor instance
-        executor = Workflow::Executor.new(workflow_class, history, metadata, config, track_stack_trace)
+        executor = Workflow::Executor.new(workflow_class, history, metadata, config, track_stack_trace, workflow_middleware_chain)
 
         commands = middleware_chain.invoke(metadata) do
           executor.run
@@ -79,7 +80,7 @@ def process
       private
 
       attr_reader :task, :namespace, :task_token, :workflow_name, :workflow_class,
-        :middleware_chain, :metadata, :config, :binary_checksum
+        :middleware_chain, :workflow_middleware_chain, :metadata, :config, :binary_checksum
 
       def connection
         @connection ||= Temporal::Connection.generate(config.for_connection)

spec/unit/lib/temporal/client_spec.rb

@@ -39,6 +39,35 @@ class TestStartWorkflow < Temporal::Workflow
 
     before { allow(connection).to receive(:start_workflow_execution).and_return(temporal_response) }
 
+    context 'with header propagator' do
+      class TestHeaderPropagator
+        def inject!(header)
+          header['test'] = 'asdf'
+        end
+      end
+
+      it 'updates the header' do
+        config.add_header_propagator(TestHeaderPropagator)
+        subject.start_workflow(TestStartWorkflow, 42)
+        expect(connection)
+          .to have_received(:start_workflow_execution)
+                .with(
+                  namespace: 'default-test-namespace',
+                  workflow_id: an_instance_of(String),
+                  workflow_name: 'TestStartWorkflow',
+                  task_queue: 'default-test-task-queue',
+                  input: [42],
+                  task_timeout: Temporal.configuration.timeouts[:task],
+                  run_timeout: Temporal.configuration.timeouts[:run],
+                  execution_timeout: Temporal.configuration.timeouts[:execution],
+                  workflow_id_reuse_policy: nil,
+                  headers: { 'test' => 'asdf' },
+                  memo: {},
+                  search_attributes: {},
+                )
+      end
+    end
+
     context 'using a workflow class' do
       it 'returns run_id' do
         result = subject.start_workflow(TestStartWorkflow, 42)

spec/unit/lib/temporal/configuration_spec.rb

@@ -1,6 +1,10 @@
 require 'temporal/configuration'
 
 describe Temporal::Configuration do
+  class TestHeaderPropagator
+    def inject!(_); end
+  end
+
   describe '#initialize' do
     it 'initializes proper default workflow timeouts' do 
       timeouts = subject.timeouts
@@ -27,6 +31,25 @@
     end
   end
 
+  describe '#add_header_propagator' do
+    let(:header_propagators) { subject.send(:header_propagators) }
+
+    it 'adds middleware entry to the list of middlewares' do
+      subject.add_header_propagator(TestHeaderPropagator)
+      subject.add_header_propagator(TestHeaderPropagator, 'arg1', 'arg2')
+
+      expect(header_propagators.size).to eq(2)
+
+      expect(header_propagators[0]).to be_an_instance_of(Temporal::Middleware::Entry)
+      expect(header_propagators[0].klass).to eq(TestHeaderPropagator)
+      expect(header_propagators[0].args).to eq([])
+
+      expect(header_propagators[1]).to be_an_instance_of(Temporal::Middleware::Entry)
+      expect(header_propagators[1].klass).to eq(TestHeaderPropagator)
+      expect(header_propagators[1].args).to eq(['arg1', 'arg2'])
+    end
+  end
+
   describe '#for_connection' do
     let (:new_identity) { 'new_identity' }