Expose poller automation

Author: Sushisource

temporalio/ext/src/worker.rs

@@ -6,23 +6,23 @@ use std::{
 };
 
 use crate::{
+    ROOT_MOD,
     client::Client,
     enter_sync, error, id, new_error,
     runtime::{AsyncCommand, Runtime, RuntimeHandle},
     util::{AsyncCallback, Struct},
-    ROOT_MOD,
 };
+use futures::{StreamExt, stream::BoxStream};
 use futures::{future, stream};
-use futures::{stream::BoxStream, StreamExt};
 use magnus::{
-    class, function, method, prelude::*, typed_data, DataTypeFunctions, Error, IntoValue, RArray,
-    RString, RTypedData, Ruby, TypedData, Value,
+    DataTypeFunctions, Error, IntoValue, RArray, RString, RTypedData, Ruby, TypedData, Value,
+    class, function, method, prelude::*, typed_data,
 };
 use prost::Message;
 use temporal_sdk_core::{
-    replay::{HistoryForReplay, ReplayWorkerInput},
     ResourceBasedSlotsOptions, ResourceBasedSlotsOptionsBuilder, ResourceSlotOptions,
     SlotSupplierOptions, TunerHolder, TunerHolderOptionsBuilder, WorkerConfig, WorkerConfigBuilder,
+    replay::{HistoryForReplay, ReplayWorkerInput},
 };
 use temporal_sdk_core_api::{
     errors::{PollError, WorkflowErrorType},
@@ -34,7 +34,7 @@ use temporal_sdk_core_api::{
 use temporal_sdk_core_protos::coresdk::workflow_completion::WorkflowActivationCompletion;
 use temporal_sdk_core_protos::coresdk::{ActivityHeartbeat, ActivityTaskCompletion};
 use temporal_sdk_core_protos::temporal::api::history::v1::History;
-use tokio::sync::mpsc::{channel, Sender};
+use tokio::sync::mpsc::{Sender, channel};
 use tokio_stream::wrappers::ReceiverStream;
 
 pub fn init(ruby: &Ruby) -> Result<(), Error> {
@@ -480,15 +480,21 @@ fn build_config(options: Struct) -> Result<WorkerConfig, Error> {
         })
         .client_identity_override(options.member::<Option<String>>(id!("identity_override"))?)
         .max_cached_workflows(options.member::<usize>(id!("max_cached_workflows"))?)
-        .workflow_task_poller_behavior(PollerBehavior::SimpleMaximum(
-            options.member::<usize>(id!("max_concurrent_workflow_task_polls"))?,
-        ))
+        .workflow_task_poller_behavior({
+            let poller_behavior = options
+                .child(id!("workflow_task_poller_behavior"))?
+                .ok_or_else(|| error!("Worker options must have workflow_task_poller_behavior"))?;
+            extract_poller_behavior(poller_behavior)?
+        })
         .nonsticky_to_sticky_poll_ratio(
             options.member::<f32>(id!("nonsticky_to_sticky_poll_ratio"))?,
         )
-        .activity_task_poller_behavior(PollerBehavior::SimpleMaximum(
-            options.member::<usize>(id!("max_concurrent_activity_task_polls"))?,
-        ))
+        .activity_task_poller_behavior({
+            let poller_behavior = options
+                .child(id!("activity_task_poller_behavior"))?
+                .ok_or_else(|| error!("Worker options must have activity_task_poller_behavior"))?;
+            extract_poller_behavior(poller_behavior)?
+        })
         .no_remote_activities(options.member::<bool>(id!("no_remote_activities"))?)
         .sticky_queue_schedule_to_start_timeout(Duration::from_secs_f64(
             options.member(id!("sticky_queue_schedule_to_start_timeout"))?,
@@ -605,3 +611,15 @@ fn build_tuner_resource_options<SK: SlotKind>(
         Some(slots_options),
     ))
 }
+
+fn extract_poller_behavior(poller_behavior: Struct) -> Result<PollerBehavior, Error> {
+    Ok(if poller_behavior.member::<usize>(id!("initial")).is_ok() {
+        PollerBehavior::Autoscaling {
+            minimum: poller_behavior.member::<usize>(id!("minimum"))?,
+            maximum: poller_behavior.member::<usize>(id!("maximum"))?,
+            initial: poller_behavior.member::<usize>(id!("initial"))?,
+        }
+    } else {
+        PollerBehavior::SimpleMaximum(poller_behavior.member::<usize>(id!("maximum"))?)
+    })
+}

temporalio/lib/temporalio/internal/bridge/worker.rb

@@ -14,9 +14,9 @@ class Worker
           :tuner,
           :identity_override,
           :max_cached_workflows,
-          :max_concurrent_workflow_task_polls,
+          :workflow_task_poller_behavior,
           :nonsticky_to_sticky_poll_ratio,
-          :max_concurrent_activity_task_polls,
+          :activity_task_poller_behavior,
           :no_remote_activities,
           :sticky_queue_schedule_to_start_timeout,
           :max_heartbeat_throttle_interval,
@@ -65,6 +65,18 @@ class Worker
           keyword_init: true
         )
 
+        PollerBehaviorSimpleMaximum = Struct.new(
+          :simple_maximum,
+          keyword_init: true
+        )
+
+        PollerBehaviorAutoscaling = Struct.new(
+          :minimum,
+          :maximum,
+          :initial,
+          keyword_init: true
+        )
+
         def self.finalize_shutdown_all(workers)
           queue = Queue.new
           async_finalize_all(workers, queue)

temporalio/lib/temporalio/worker.rb

@@ -13,6 +13,7 @@
 require 'temporalio/internal/worker/workflow_worker'
 require 'temporalio/worker/activity_executor'
 require 'temporalio/worker/interceptor'
+require 'temporalio/worker/poller_behavior'
 require 'temporalio/worker/thread_pool'
 require 'temporalio/worker/tuner'
 require 'temporalio/worker/workflow_executor'
@@ -52,7 +53,9 @@ class Worker
       :workflow_payload_codec_thread_pool,
       :unsafe_workflow_io_enabled,
       :deployment_options,
-      :debug_mode
+      :debug_mode,
+      :workflow_task_poller_behavior,
+      :activity_task_poller_behavior
     )
 
     # Options as returned from {options} for `**to_h` splat use in {initialize}. See {initialize} for details.
@@ -385,6 +388,8 @@ def initialize(
       workflow_payload_codec_thread_pool: nil,
       unsafe_workflow_io_enabled: false,
       deployment_options: Worker.default_deployment_options,
+      workflow_task_poller_behavior: nil,
+      activity_task_poller_behavior: nil,
       debug_mode: %w[true 1].include?(ENV['TEMPORAL_DEBUG'].to_s.downcase)
     )
       raise ArgumentError, 'Must have at least one activity or workflow' if activities.empty? && workflows.empty?
@@ -419,7 +424,9 @@ def initialize(
         workflow_payload_codec_thread_pool:,
         unsafe_workflow_io_enabled:,
         deployment_options:,
-        debug_mode:
+        debug_mode:,
+        workflow_task_poller_behavior:,
+        activity_task_poller_behavior:
       ).freeze
 
       should_enforce_versioning_behavior =
@@ -435,6 +442,10 @@ def initialize(
           workflow_failure_exception_types:, workflow_definitions:
         )
 
+      # Convert deprecated max concurrent polls to poller behaviors if not specified
+      workflow_task_poller_behavior ||= SimpleMaximumPollerBehavior.new(maximum: max_concurrent_workflow_task_polls)
+      activity_task_poller_behavior ||= SimpleMaximumPollerBehavior.new(maximum: max_concurrent_activity_task_polls)
+
       # Create the bridge worker
       @bridge_worker = Internal::Bridge::Worker.new(
         client.connection._core_client,
@@ -446,9 +457,9 @@ def initialize(
           tuner: tuner._to_bridge_options,
           identity_override: identity,
           max_cached_workflows:,
-          max_concurrent_workflow_task_polls:,
+          workflow_task_poller_behavior: workflow_task_poller_behavior._to_bridge_options,
           nonsticky_to_sticky_poll_ratio:,
-          max_concurrent_activity_task_polls:,
+          activity_task_poller_behavior: activity_task_poller_behavior._to_bridge_options,
           # For shutdown to work properly, we must disable remote activities
           # ourselves if there are no activities
           no_remote_activities: no_remote_activities || activities.empty?,

temporalio/lib/temporalio/worker/poller_behavior.rb

@@ -0,0 +1,85 @@
+# frozen_string_literal: true
+
+module Temporalio
+  class Worker
+    # Base class for poller behaviors that control how polling scales.
+    #
+    # Use factory methods {.simple_maximum} or {.autoscaling} to create instances.
+    class PollerBehavior
+      # Creates a simple maximum poller behavior
+      # The poller will attempt to poll as long as a slot is available, up to the
+      # provided maximum. Cannot be less than two for workflow tasks, or one for other tasks.
+      #
+      # @param maximum [Integer] Maximum number of concurrent poll requests.
+      # @return [SimpleMaximumPollerBehavior] A simple maximum poller behavior
+      def self.simple_maximum(maximum = 5)
+        SimpleMaximumPollerBehavior.new(maximum: maximum)
+      end
+
+      # Creates an autoscaling poller behavior
+      # The poller will automatically scale the number of pollers based on feedback
+      # from the server. A slot must be available before beginning polling.
+      #
+      # @param minimum [Integer] Minimum number of poll calls (assuming slots are available).
+      # @param maximum [Integer] Maximum number of poll calls that will ever be open at once.
+      # @param initial [Integer] Number of polls attempted initially before scaling kicks in.
+      # @return [AutoscalingPollerBehavior] An autoscaling poller behavior
+      def self.autoscaling(minimum: 1, maximum: 100, initial: 5)
+        AutoscalingPollerBehavior.new(minimum: minimum, maximum: maximum, initial: initial)
+      end
+
+      # @!visibility private
+      def _to_bridge_options
+        raise NotImplementedError, 'Subclasses must implement this method'
+      end
+    end
+
+    # A poller behavior that attempts to poll as long as a slot is available, up to the
+    # provided maximum. Cannot be less than two for workflow tasks, or one for other tasks.
+    class SimpleMaximumPollerBehavior < PollerBehavior
+      # @return [Integer] Maximum number of concurrent poll requests.
+      attr_reader :maximum
+
+      # @param maximum [Integer] Maximum number of concurrent poll requests.
+      def initialize(maximum: 5)
+        super()
+        @maximum = maximum
+      end
+
+      # @!visibility private
+      def _to_bridge_options
+        Internal::Bridge::Worker::PollerBehaviorSimpleMaximum.new(simple_maximum: maximum)
+      end
+    end
+
+    # A poller behavior that automatically scales the number of pollers based on feedback
+    # from the server. A slot must be available before beginning polling.
+    class AutoscalingPollerBehavior < PollerBehavior
+      # @return [Integer] Minimum number of poll calls (assuming slots are available).
+      attr_reader :minimum
+      # @return [Integer] Maximum number of poll calls that will ever be open at once.
+      attr_reader :maximum
+      # @return [Integer] Number of polls attempted initially before scaling kicks in.
+      attr_reader :initial
+
+      # @param minimum [Integer] Minimum number of poll calls (assuming slots are available).
+      # @param maximum [Integer] Maximum number of poll calls that will ever be open at once.
+      # @param initial [Integer] Number of polls attempted initially before scaling kicks in.
+      def initialize(minimum: 1, maximum: 100, initial: 5)
+        super()
+        @minimum = minimum
+        @maximum = maximum
+        @initial = initial
+      end
+
+      # @!visibility private
+      def _to_bridge_options
+        Internal::Bridge::Worker::PollerBehaviorAutoscaling.new(
+          minimum: minimum,
+          maximum: maximum,
+          initial: initial
+        )
+      end
+    end
+  end
+end

temporalio/lib/temporalio/worker/workflow_replayer.rb

@@ -7,6 +7,7 @@
 require 'temporalio/internal/worker/multi_runner'
 require 'temporalio/internal/worker/workflow_worker'
 require 'temporalio/worker/interceptor'
+require 'temporalio/worker/poller_behavior'
 require 'temporalio/worker/thread_pool'
 require 'temporalio/worker/tuner'
 require 'temporalio/worker/workflow_executor'
@@ -205,9 +206,9 @@ def initialize(
               )._to_bridge_options,
               identity_override: options.identity,
               max_cached_workflows: 2,
-              max_concurrent_workflow_task_polls: 2,
+              workflow_task_poller_behavior: Temporalio::Worker::PollerBehavior.simple_maximum(2)._to_bridge_options,
               nonsticky_to_sticky_poll_ratio: 1.0,
-              max_concurrent_activity_task_polls: 1,
+              activity_task_poller_behavior: Temporalio::Worker::PollerBehavior.simple_maximum(1)._to_bridge_options,
               no_remote_activities: true,
               sticky_queue_schedule_to_start_timeout: 1.0,
               max_heartbeat_throttle_interval: 1.0,

temporalio/sig/temporalio/internal/bridge/worker.rbs

@@ -10,9 +10,9 @@ module Temporalio
           attr_accessor tuner: TunerOptions
           attr_accessor identity_override: String?
           attr_accessor max_cached_workflows: Integer
-          attr_accessor max_concurrent_workflow_task_polls: Integer
+          attr_accessor workflow_task_poller_behavior: PollerBehaviorSimpleMaximum | PollerBehaviorAutoscaling
           attr_accessor nonsticky_to_sticky_poll_ratio: Float
-          attr_accessor max_concurrent_activity_task_polls: Integer
+          attr_accessor activity_task_poller_behavior: PollerBehaviorSimpleMaximum | PollerBehaviorAutoscaling
           attr_accessor no_remote_activities: bool
           attr_accessor sticky_queue_schedule_to_start_timeout: Float
           attr_accessor max_heartbeat_throttle_interval: Float
@@ -32,9 +32,9 @@ module Temporalio
             tuner: TunerOptions,
             identity_override: String?,
             max_cached_workflows: Integer,
-            max_concurrent_workflow_task_polls: Integer,
+            workflow_task_poller_behavior: PollerBehaviorSimpleMaximum | PollerBehaviorAutoscaling,
             nonsticky_to_sticky_poll_ratio: Float,
-            max_concurrent_activity_task_polls: Integer,
+            activity_task_poller_behavior: PollerBehaviorSimpleMaximum | PollerBehaviorAutoscaling,
             no_remote_activities: bool,
             sticky_queue_schedule_to_start_timeout: Float,
             max_heartbeat_throttle_interval: Float,
@@ -142,6 +142,20 @@ module Temporalio
           ) -> void
         end
 
+        class PollerBehaviorSimpleMaximum
+          attr_accessor simple_maximum: Integer
+
+          def initialize: (simple_maximum: Integer) -> void
+        end
+
+        class PollerBehaviorAutoscaling
+          attr_accessor minimum: Integer
+          attr_accessor maximum: Integer
+          attr_accessor initial: Integer
+
+          def initialize: (minimum: Integer, maximum: Integer, initial: Integer) -> void
+        end
+
         class WorkflowReplayer
           def self.new: (Runtime runtime, Options options) -> [WorkflowReplayer, Worker]
 

temporalio/sig/temporalio/worker.rbs

@@ -1,5 +1,6 @@
 module Temporalio
   class Worker
+
     class Options
       attr_reader client: Client
       attr_reader task_queue: String
@@ -27,6 +28,8 @@ module Temporalio
       attr_reader workflow_failure_exception_types: Array[singleton(Exception)]
       attr_reader workflow_payload_codec_thread_pool: ThreadPool?
       attr_reader unsafe_workflow_io_enabled: bool
+      attr_reader workflow_task_poller_behavior: PollerBehavior
+      attr_reader activity_task_poller_behavior: PollerBehavior
       attr_reader debug_mode: bool
 
       def initialize: (
@@ -56,6 +59,8 @@ module Temporalio
         workflow_failure_exception_types: Array[singleton(Exception)],
         workflow_payload_codec_thread_pool: ThreadPool?,
         unsafe_workflow_io_enabled: bool,
+        workflow_task_poller_behavior: PollerBehavior,
+        activity_task_poller_behavior: PollerBehavior,
         deployment_options: Worker::DeploymentOptions,
         debug_mode: bool
       ) -> void
@@ -106,6 +111,8 @@ module Temporalio
       ?workflow_failure_exception_types: Array[singleton(Exception)],
       ?workflow_payload_codec_thread_pool: ThreadPool?,
       ?unsafe_workflow_io_enabled: bool,
+      ?workflow_task_poller_behavior: SimpleMaximumPollerBehavior | AutoscalingPollerBehavior,
+      ?activity_task_poller_behavior: SimpleMaximumPollerBehavior | AutoscalingPollerBehavior,
       ?deployment_options: Worker::DeploymentOptions,
       ?debug_mode: bool
     ) -> void

temporalio/sig/temporalio/worker/poller_behavior.rbs

@@ -0,0 +1,25 @@
+module Temporalio
+  class Worker
+    class PollerBehavior
+      def self.simple_maximum: (Integer) -> SimpleMaximumPollerBehavior
+      def self.autoscaling: (?minimum: Integer, ?maximum: Integer, ?initial: Integer) -> AutoscalingPollerBehavior
+      def _to_bridge_options: -> untyped
+    end
+
+    class SimpleMaximumPollerBehavior < PollerBehavior
+      attr_reader maximum: Integer
+
+      def initialize: (?maximum: Integer) -> void
+      def _to_bridge_options: -> Internal::Bridge::Worker::PollerBehaviorSimpleMaximum
+    end
+
+    class AutoscalingPollerBehavior < PollerBehavior
+      attr_reader minimum: Integer
+      attr_reader maximum: Integer
+      attr_reader initial: Integer
+
+      def initialize: (?minimum: Integer, ?maximum: Integer, ?initial: Integer) -> void
+      def _to_bridge_options: -> Internal::Bridge::Worker::PollerBehaviorAutoscaling
+    end
+  end
+end