From 70c33943c57465dda15b39092e7b086ca5e8ad55 Mon Sep 17 00:00:00 2001 From: Veeral Patel Date: Fri, 24 Jul 2026 12:22:50 -0700 Subject: [PATCH 1/5] feat(sdk): auto-enroll default pollers into autoscaling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When a namespace advertises the `poller_autoscaling_auto_enroll` capability, workers automatically switch workflow, activity, and nexus pollers to autoscaling — but only for poller types the user left at their default (neither a fixed poller count nor an explicit behavior). Explicitly configured pollers are left unchanged. The effective `PollerBehavior` is resolved once, after the namespace capability is known. Poller construction is deferred out of the synchronous `Worker::new_with_pollers` into a lazily-built step that runs after `Worker::validate()` populates capabilities (and on first poll for paths that skip validate, e.g. replay). This keeps `poll_buffer.rs` / `PollScaler` free of any runtime mode-switching — the behavior is simply passed in already resolved. Eligibility (user left the poller at its default) is signalled from the layers above core: the Rust SDK exposes the poller options as `Option`, and the c-bridge treats unset (both FFI pointers null) as eligible. Ports temporalio/sdk-go#2442. --- CHANGELOG.md | 4 + crates/sdk-core-c-bridge/src/worker.rs | 94 ++- crates/sdk-core/src/core_tests/mod.rs | 2 +- crates/sdk-core/src/core_tests/workers.rs | 58 ++ .../sdk-core/src/core_tests/workflow_tasks.rs | 6 +- crates/sdk-core/src/replay/mod.rs | 2 +- crates/sdk-core/src/worker/mod.rs | 696 ++++++++++++------ .../src/worker/workflow/wft_poller.rs | 9 +- crates/sdk-core/tests/heavy_tests.rs | 18 +- .../tests/integ_tests/metrics_tests.rs | 4 +- .../tests/integ_tests/polling_tests.rs | 8 +- .../integ_tests/worker_heartbeat_tests.rs | 8 +- .../tests/integ_tests/worker_tests.rs | 2 +- .../tests/integ_tests/workflow_tests.rs | 3 +- .../integ_tests/workflow_tests/activities.rs | 8 +- .../integ_tests/workflow_tests/stickyness.rs | 2 +- crates/sdk-core/tests/manual_tests.rs | 20 +- crates/sdk/src/lib.rs | 27 +- 18 files changed, 652 insertions(+), 319 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2ebcfb625..3b21c70bb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,10 @@ to docs, or any other relevant information. ## [0.5.0] ### Added +* Workers are now automatically enrolled into poller autoscaling when the namespace advertises the + `poller_autoscaling_auto_enroll` capability. This only applies to poller types left at their + default (the worker set neither a fixed poller count nor a poller behavior); explicitly + configured pollers are left unchanged. * `client()` and `workflow_handle()` helpers to `ActivityContext` for easily obtaining a Temporal client * Exposed `backoff_start_interval` when continuing as new, which will delay the first task of the continued workflow by the configured interval. diff --git a/crates/sdk-core-c-bridge/src/worker.rs b/crates/sdk-core-c-bridge/src/worker.rs index a2916af43..050be7f68 100644 --- a/crates/sdk-core-c-bridge/src/worker.rs +++ b/crates/sdk-core-c-bridge/src/worker.rs @@ -101,24 +101,27 @@ pub struct PollerBehavior { pub autoscaling: *const PollerBehaviorAutoscaling, } -impl TryFrom<&PollerBehavior> for temporalio_sdk_core::PollerBehavior { - type Error = anyhow::Error; - fn try_from(value: &PollerBehavior) -> Result { - if !value.simple_maximum.is_null() && !value.autoscaling.is_null() { +impl PollerBehavior { + /// Converts an FFI poller behavior into an optional core poller behavior. Both fields null + /// means the poller was left unset by lang: core applies the default behavior and treats the + /// poller as eligible for automatic enrollment into poller autoscaling. + fn to_core(&self) -> anyhow::Result> { + if !self.simple_maximum.is_null() && !self.autoscaling.is_null() { bail!("simple_maximum and autoscaling cannot both be non-null values"); } - if let Some(value) = unsafe { value.simple_maximum.as_ref() } { - return Ok(temporalio_sdk_core::PollerBehavior::SimpleMaximum( - value.simple_maximum, - )); - } else if let Some(value) = unsafe { value.autoscaling.as_ref() } { - return Ok(temporalio_sdk_core::PollerBehavior::Autoscaling { - minimum: value.minimum, - maximum: value.maximum, - initial: value.initial, - }); + if let Some(sm) = unsafe { self.simple_maximum.as_ref() } { + Ok(Some(temporalio_sdk_core::PollerBehavior::SimpleMaximum( + sm.simple_maximum, + ))) + } else if let Some(a) = unsafe { self.autoscaling.as_ref() } { + Ok(Some(temporalio_sdk_core::PollerBehavior::Autoscaling { + minimum: a.minimum, + maximum: a.maximum, + initial: a.initial, + })) + } else { + Ok(None) } - bail!("simple_maximum and autoscaling cannot both be null values"); } } @@ -1245,16 +1248,10 @@ impl TryFrom<&WorkerOptions> for temporalio_sdk_core::WorkerConfig { // auto-cancel-activity behavior or shutdown will not occur, so we // always set it even if 0. .graceful_shutdown_period(Duration::from_millis(opt.graceful_shutdown_period_millis)) - .workflow_task_poller_behavior(temporalio_sdk_core::PollerBehavior::try_from( - &opt.workflow_task_poller_behavior, - )?) + .maybe_workflow_task_poller_behavior(opt.workflow_task_poller_behavior.to_core()?) .nonsticky_to_sticky_poll_ratio(opt.nonsticky_to_sticky_poll_ratio) - .activity_task_poller_behavior(temporalio_sdk_core::PollerBehavior::try_from( - &opt.activity_task_poller_behavior, - )?) - .nexus_task_poller_behavior(temporalio_sdk_core::PollerBehavior::try_from( - &opt.nexus_task_poller_behavior, - )?) + .maybe_activity_task_poller_behavior(opt.activity_task_poller_behavior.to_core()?) + .maybe_nexus_task_poller_behavior(opt.nexus_task_poller_behavior.to_core()?) .workflow_failure_errors(if opt.nondeterminism_as_workflow_fail { HashSet::from([WorkflowErrorType::Nondeterminism]) } else { @@ -1418,6 +1415,55 @@ mod tests { } } + #[test] + fn ffi_poller_behavior_opt_maps_unset_to_none() { + let unset = PollerBehavior { + simple_maximum: std::ptr::null(), + autoscaling: std::ptr::null(), + }; + assert!(unset.to_core().unwrap().is_none()); + } + + #[test] + fn ffi_poller_behavior_opt_maps_configured_behaviors() { + assert_eq!( + simple_poller_behavior(3).to_core().unwrap(), + Some(temporalio_sdk_core::PollerBehavior::SimpleMaximum(3)), + ); + let autoscaling = Box::leak(Box::new(PollerBehaviorAutoscaling { + minimum: 2, + maximum: 20, + initial: 4, + })); + let pb = PollerBehavior { + simple_maximum: std::ptr::null(), + autoscaling: autoscaling as *const _, + }; + assert_eq!( + pb.to_core().unwrap(), + Some(temporalio_sdk_core::PollerBehavior::Autoscaling { + minimum: 2, + maximum: 20, + initial: 4, + }), + ); + } + + #[test] + fn ffi_poller_behavior_opt_rejects_both_set() { + let simple = Box::leak(Box::new(PollerBehaviorSimpleMaximum { simple_maximum: 1 })); + let autoscaling = Box::leak(Box::new(PollerBehaviorAutoscaling { + minimum: 1, + maximum: 2, + initial: 1, + })); + let pb = PollerBehavior { + simple_maximum: simple as *const _, + autoscaling: autoscaling as *const _, + }; + assert!(pb.to_core().is_err()); + } + fn base_worker_options( namespace: &str, task_queue: &str, diff --git a/crates/sdk-core/src/core_tests/mod.rs b/crates/sdk-core/src/core_tests/mod.rs index 29f5df27c..17ac93868 100644 --- a/crates/sdk-core/src/core_tests/mod.rs +++ b/crates/sdk-core/src/core_tests/mod.rs @@ -89,7 +89,7 @@ async fn shutdown_interrupts_both_polls() { .activity_task_poller_behavior(PollerBehavior::SimpleMaximum(1_usize)) .build() .unwrap(); - cfg.workflow_task_poller_behavior = PollerBehavior::SimpleMaximum(1_usize); + cfg.workflow_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(1_usize)); cfg }, mock_client, diff --git a/crates/sdk-core/src/core_tests/workers.rs b/crates/sdk-core/src/core_tests/workers.rs index f227bd38c..efb2b207c 100644 --- a/crates/sdk-core/src/core_tests/workers.rs +++ b/crates/sdk-core/src/core_tests/workers.rs @@ -1328,3 +1328,61 @@ async fn graceful_shutdown_sends_shutdown_worker_rpc_during_initiate() { worker.finalize_shutdown().await; } + +#[tokio::test] +async fn validate_enables_auto_enroll_capability() { + let mut mock = mock_worker_client(); + mock.expect_describe_namespace().returning(|| { + Ok(DescribeNamespaceResponse { + namespace_info: Some(NamespaceInfo { + capabilities: Some(Capabilities { + poller_autoscaling_auto_enroll: true, + ..Capabilities::default() + }), + ..NamespaceInfo::default() + }), + ..DescribeNamespaceResponse::default() + }) + }); + let t = canned_histories::single_timer("1"); + let mut mh = MockPollCfg::from_resp_batches("fakeid", t, [1], mock); + mh.enforce_correct_number_of_polls = false; + let worker = mock_worker(build_mock_pollers(mh)); + + worker.validate().await.unwrap(); + let caps = worker.get_namespace_capabilities(); + assert!(caps.poller_autoscaling_auto_enroll()); + // The two capabilities are independent: advertising auto-enroll alone does not imply the + // `poller_autoscaling` capability (the server advertises each on its own). + assert!(!caps.poller_autoscaling()); + + worker.drain_pollers_and_shutdown().await; +} + +#[tokio::test] +async fn validate_without_auto_enroll_leaves_capabilities_off() { + let mut mock = mock_worker_client(); + mock.expect_describe_namespace().returning(|| { + Ok(DescribeNamespaceResponse { + namespace_info: Some(NamespaceInfo { + capabilities: Some(Capabilities { + poller_autoscaling_auto_enroll: false, + ..Capabilities::default() + }), + ..NamespaceInfo::default() + }), + ..DescribeNamespaceResponse::default() + }) + }); + let t = canned_histories::single_timer("1"); + let mut mh = MockPollCfg::from_resp_batches("fakeid", t, [1], mock); + mh.enforce_correct_number_of_polls = false; + let worker = mock_worker(build_mock_pollers(mh)); + + worker.validate().await.unwrap(); + let caps = worker.get_namespace_capabilities(); + assert!(!caps.poller_autoscaling_auto_enroll()); + assert!(!caps.poller_autoscaling()); + + worker.drain_pollers_and_shutdown().await; +} diff --git a/crates/sdk-core/src/core_tests/workflow_tasks.rs b/crates/sdk-core/src/core_tests/workflow_tasks.rs index ea56169ba..4eebc17df 100644 --- a/crates/sdk-core/src/core_tests/workflow_tasks.rs +++ b/crates/sdk-core/src/core_tests/workflow_tasks.rs @@ -2805,7 +2805,7 @@ async fn poller_wont_run_ahead_of_task_slots() { let mut cfg = test_worker_cfg().build().unwrap(); cfg.max_cached_workflows = 10_usize; cfg.max_outstanding_workflow_tasks = Some(10_usize); - cfg.workflow_task_poller_behavior = PollerBehavior::SimpleMaximum(10_usize); + cfg.workflow_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(10_usize)); cfg.task_types = WorkerTaskTypes::workflow_only(); cfg }, @@ -3012,7 +3012,7 @@ async fn slot_provider_cant_hand_out_more_permits_than_cache_size() { .workflow_slot_supplier(Arc::new(EndlessSupplier {})) .build(), )); - cfg.workflow_task_poller_behavior = PollerBehavior::SimpleMaximum(10_usize); + cfg.workflow_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(10_usize)); cfg.task_types = WorkerTaskTypes::workflow_only(); cfg }, @@ -3160,7 +3160,7 @@ async fn both_normal_and_sticky_pollers_poll_concurrently() { let mut cfg = test_worker_cfg().build().unwrap(); cfg.max_cached_workflows = 500_usize; // We need cache, but don't want to deal with evictions cfg.max_outstanding_workflow_tasks = Some(2_usize); - cfg.workflow_task_poller_behavior = PollerBehavior::SimpleMaximum(2_usize); + cfg.workflow_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(2_usize)); cfg.nonsticky_to_sticky_poll_ratio = 0.2; cfg.task_types = WorkerTaskTypes::workflow_only(); cfg diff --git a/crates/sdk-core/src/replay/mod.rs b/crates/sdk-core/src/replay/mod.rs index 218ccbbce..c1cc6ef29 100644 --- a/crates/sdk-core/src/replay/mod.rs +++ b/crates/sdk-core/src/replay/mod.rs @@ -75,7 +75,7 @@ where pub(crate) fn into_core_worker(mut self) -> Result { self.config.max_cached_workflows = 1; - self.config.workflow_task_poller_behavior = PollerBehavior::SimpleMaximum(1); + self.config.workflow_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(1)); self.config.task_types = WorkerTaskTypes::workflow_only(); self.config.skip_client_worker_set_check = true; let historator = Historator::new(self.history_stream); diff --git a/crates/sdk-core/src/worker/mod.rs b/crates/sdk-core/src/worker/mod.rs index fab7c939b..b21b246d6 100644 --- a/crates/sdk-core/src/worker/mod.rs +++ b/crates/sdk-core/src/worker/mod.rs @@ -76,7 +76,7 @@ use std::{ convert::TryInto, future, sync::{ - Arc, + Arc, LazyLock, OnceLock, atomic::{AtomicBool, AtomicU64, Ordering}, }, time::{Duration, SystemTime}, @@ -149,8 +149,12 @@ pub struct WorkerConfig { /// Maximum number of concurrent poll workflow task requests we will perform at a time on this /// worker's task queue. See also [WorkerConfig::nonsticky_to_sticky_poll_ratio]. /// If using SimpleMaximum, Must be at least 2 when `max_cached_workflows` > 0, or is an error. - #[builder(default = PollerBehavior::SimpleMaximum(5))] - pub workflow_task_poller_behavior: PollerBehavior, + /// + /// When left unset (`None`), the effective behavior is resolved at worker start: pollers are + /// automatically enrolled into poller autoscaling when the namespace advertises the + /// `poller_autoscaling_auto_enroll` capability, otherwise the default + /// [PollerBehavior::SimpleMaximum] is used. + pub workflow_task_poller_behavior: Option, /// Only applies when using [PollerBehavior::SimpleMaximum] /// /// (max workflow task polls * this number) = the number of max pollers that will be allowed for @@ -161,13 +165,13 @@ pub struct WorkerConfig { #[builder(default = 0.2)] pub nonsticky_to_sticky_poll_ratio: f32, /// Maximum number of concurrent poll activity task requests we will perform at a time on this - /// worker's task queue - #[builder(default = PollerBehavior::SimpleMaximum(5))] - pub activity_task_poller_behavior: PollerBehavior, + /// worker's task queue. See [WorkerConfig::workflow_task_poller_behavior] for the meaning of + /// leaving this unset (`None`). + pub activity_task_poller_behavior: Option, /// Maximum number of concurrent poll nexus task requests we will perform at a time on this - /// worker's task queue - #[builder(default = PollerBehavior::SimpleMaximum(5))] - pub nexus_task_poller_behavior: PollerBehavior, + /// worker's task queue. See [WorkerConfig::workflow_task_poller_behavior] for the meaning of + /// leaving this unset (`None`). + pub nexus_task_poller_behavior: Option, /// Specifies which task types this worker will poll for. /// /// Note: At least one task type must be specified or the worker will fail validation. @@ -339,9 +343,15 @@ impl WorkerConfigBuilder { ); } - config.workflow_task_poller_behavior.validate()?; - config.activity_task_poller_behavior.validate()?; - config.nexus_task_poller_behavior.validate()?; + if let Some(b) = config.workflow_task_poller_behavior { + b.validate()?; + } + if let Some(b) = config.activity_task_poller_behavior { + b.validate()?; + } + if let Some(b) = config.nexus_task_poller_behavior { + b.validate()?; + } if let Some(ref x) = config.max_worker_activities_per_second && (!x.is_normal() || x.is_sign_negative()) @@ -373,7 +383,7 @@ impl WorkerConfigBuilder { .to_string(), ); } - if matches!(config.workflow_task_poller_behavior, PollerBehavior::SimpleMaximum(u) if u < 2) + if matches!(config.workflow_task_poller_behavior, Some(PollerBehavior::SimpleMaximum(u)) if u < 2) { return Err("`max_cached_workflows` > 0 requires `workflow_task_poller_behavior` to be at least 2".to_string()); } @@ -419,14 +429,12 @@ pub struct Worker { client: Arc, /// Worker instance key, unique identifier for this worker worker_instance_key: Uuid, - /// Manages all workflows and WFT processing. None if workflow polling is disabled - workflows: Option, - /// Manages activity tasks for this worker/task queue - at_task_mgr: Option, + /// The task subsystems (workflows, activity task manager, nexus manager). Built lazily once + /// namespace capabilities are known so effective poller behavior can be resolved. + #[allow(clippy::type_complexity)] + task_subsystems: LazyLock TaskSubsystems + Send>>, /// Manages local activities. None if workflow polling is disabled (local activities require workflows) local_act_mgr: Option>, - /// Manages Nexus tasks - nexus_mgr: Option, /// Has shutdown been called? shutdown_token: CancellationToken, /// Will be called at the end of each activation completion @@ -454,6 +462,7 @@ pub struct Worker { pub struct NamespaceCapabilities { pub(crate) graceful_poll_shutdown: AtomicBool, pub(crate) poller_autoscaling: AtomicBool, + pub(crate) poller_autoscaling_auto_enroll: AtomicBool, pub(crate) worker_commands: AtomicBool, } @@ -470,12 +479,45 @@ impl NamespaceCapabilities { self.poller_autoscaling.load(Ordering::Relaxed) } + /// Returns true if the namespace opts workers into poller autoscaling by default. Poller types + /// left at their default are automatically enrolled into autoscaling when this is set. + pub fn poller_autoscaling_auto_enroll(&self) -> bool { + self.poller_autoscaling_auto_enroll.load(Ordering::Relaxed) + } + /// Returns true if worker commands are supported in this namespace. pub fn worker_commands(&self) -> bool { self.worker_commands.load(Ordering::Relaxed) } } +/// Resolve the effective poller behavior. When no behavior was configured (`None`), pollers are +/// automatically enrolled into autoscaling if the namespace advertises the +/// `poller_autoscaling_auto_enroll` capability; otherwise the default [PollerBehavior::SimpleMaximum] +/// is used. A configured behavior is always used as-is. +pub(crate) fn resolve_effective_behavior( + configured: Option, + capabilities: &NamespaceCapabilities, +) -> PollerBehavior { + match configured { + Some(b) => b, + None if capabilities.poller_autoscaling_auto_enroll() => PollerBehavior::Autoscaling { + minimum: 1, + maximum: 100, + initial: 5, + }, + None => PollerBehavior::SimpleMaximum(5), + } +} + +/// The task subsystems, constructed lazily after namespace capabilities are known so that the +/// effective poller behavior can be resolved once (see [resolve_effective_behavior]). +struct TaskSubsystems { + workflows: Option, + at_task_mgr: Option, + nexus_mgr: Option, +} + struct AllPermitsTracker { wft_permits: watch::Receiver, act_permits: watch::Receiver, @@ -569,12 +611,20 @@ impl Worker { .poller_autoscaling .store(true, Ordering::Relaxed); } + if caps.poller_autoscaling_auto_enroll { + self.capabilities + .poller_autoscaling_auto_enroll + .store(true, Ordering::Relaxed); + } if caps.worker_commands { self.capabilities .worker_commands .store(true, Ordering::Relaxed); } } + // Now that capabilities are known, eagerly build the pollers so effective poller + // behavior is resolved during the normal check-in path. + LazyLock::force(&self.task_subsystems); Ok(NamespaceInfo { limits }) } Err(e) if e.code() == tonic::Code::Unimplemented => { @@ -697,6 +747,7 @@ impl Worker { let capabilities = Arc::new(NamespaceCapabilities { graceful_poll_shutdown: AtomicBool::new(false), poller_autoscaling: AtomicBool::new(false), + poller_autoscaling_auto_enroll: AtomicBool::new(false), worker_commands: AtomicBool::new(false), }); @@ -707,121 +758,6 @@ impl Worker { slot_context_data.clone(), meter.clone(), ); - let (wft_stream, act_poller, nexus_poller) = match task_pollers { - TaskPollers::Real => { - let wft_stream = if config.task_types.enable_workflows { - let stream = make_wft_poller( - &config, - &sticky_queue_name, - &client, - &metrics, - &shutdown_token, - &wft_slots, - wf_last_suc_poll_time.clone(), - wf_sticky_last_suc_poll_time.clone(), - capabilities.clone(), - ) - .boxed(); - let stream = if !client.is_mock() { - // Some replay tests combine a mock client with real pollers, - // and they don't need to use the external stream - stream::select(stream, UnboundedReceiverStream::new(external_wft_rx)) - .left_stream() - } else { - stream.right_stream() - }; - Some(stream) - } else { - None - }; - - let act_poll_buffer = if config.task_types.enable_remote_activities { - let act_metrics = metrics.with_new_attrs([activity_poller()]); - let ap = LongPollBuffer::new_activity_task( - client.clone(), - config.task_queue.clone(), - config.activity_task_poller_behavior, - act_slots.clone(), - shutdown_token.child_token(), - Some(move |np| act_metrics.record_num_pollers(np)), - ActivityTaskOptions { - max_worker_acts_per_second: config.max_worker_activities_per_second, - max_tps: config.max_task_queue_activities_per_second, - }, - act_last_suc_poll_time.clone(), - capabilities.clone(), - ); - Some(Box::from(ap) as BoxedActPoller) - } else { - None - }; - - let nexus_poll_buffer = if config.task_types.enable_nexus { - let np_metrics = metrics.with_new_attrs([nexus_poller()]); - Some(Box::new(LongPollBuffer::new_nexus_task( - client.clone(), - config.task_queue.clone(), - config.nexus_task_poller_behavior, - nexus_slots.clone(), - shutdown_token.child_token(), - Some(move |np| np_metrics.record_num_pollers(np)), - nexus_last_suc_poll_time.clone(), - capabilities.clone(), - shared_namespace_worker, - )) as BoxedNexusPoller) - } else { - None - }; - - #[cfg(any(feature = "test-utilities", test))] - let wft_stream = wft_stream.map(|s| s.left_stream()); - (wft_stream, act_poll_buffer, nexus_poll_buffer) - } - #[cfg(any(feature = "test-utilities", test))] - TaskPollers::Mocked { - wft_stream, - act_poller, - nexus_poller, - } => { - let wft_stream = config - .task_types - .enable_workflows - .then_some(wft_stream) - .flatten(); - let act_poller = config - .task_types - .enable_remote_activities - .then_some(act_poller) - .flatten(); - let nexus_poller = config - .task_types - .enable_nexus - .then_some(nexus_poller) - .flatten(); - - let ap = act_poller - .map(|ap| MockPermittedPollBuffer::new(Arc::new(act_slots.clone()), ap)); - let np = nexus_poller - .map(|np| MockPermittedPollBuffer::new(Arc::new(nexus_slots.clone()), np)); - let wfs = wft_stream.map(|stream| { - let wft_semaphore = wft_slots.clone(); - let wfs = stream.then(move |s| { - let wft_semaphore = wft_semaphore.clone(); - async move { - let permit = wft_semaphore.acquire_owned().await; - s.map(|s| (s, permit)) - } - }); - wfs.right_stream() - }); - ( - wfs, - ap.map(|ap| Box::new(ap) as BoxedActPoller), - np.map(|np| Box::new(np) as BoxedNexusPoller), - ) - } - }; - let la_permit_dealer = MeteredPermitDealer::new( tuner.local_activity_slot_supplier(), metrics.with_new_attrs([local_activity_worker_type()]), @@ -845,32 +781,251 @@ impl Worker { (None, None, None) }; - let at_task_mgr = act_poller.map(|ap| { - WorkerActivityTasks::new( - act_slots.clone(), - ap, - client.clone(), - metrics.clone(), - config.max_heartbeat_throttle_interval, - config.default_heartbeat_throttle_interval, - config.graceful_shutdown_period, - config.local_timeout_buffer_for_activities, - ) - }); - let poll_on_non_local_activities = at_task_mgr.is_some(); + // Determine, before `task_pollers` is moved into the deferred build closure, whether we + // will poll for non-local activities. + let poll_on_non_local_activities = config.task_types.enable_remote_activities + && match &task_pollers { + TaskPollers::Real => true, + #[cfg(any(feature = "test-utilities", test))] + TaskPollers::Mocked { act_poller, .. } => act_poller.is_some(), + }; if !poll_on_non_local_activities && !shared_namespace_worker { info!("Activity polling is disabled for this worker"); }; - let nexus_mgr = nexus_poller.map(|poller| { - NexusManager::new( - poller, - metrics.clone(), - config.graceful_shutdown_period, - shutdown_token.child_token(), - ) - }); + let worker_instance_key = client.worker_instance_key(); + let worker_status = Arc::new(RwLock::new(WorkerStatus::Running)); + let sdk_name_and_ver = client.sdk_name_and_version(); + + // Slot filled by the deferred build with the activity cancellation callback, so the + // client-worker registrator can resolve it lazily (the activity task manager doesn't + // exist until the pollers are built). + let cancel_activity_slot: Arc> = Arc::new(OnceLock::new()); + + // Build the poller-dependent subsystems lazily. This closure runs once, after namespace + // capabilities have been fetched (see `Worker::validate`), so the effective poller behavior + // can be resolved with knowledge of those capabilities. Everything it needs is captured by + // clone (or moved) *before* the synchronous heartbeat manager below consumes the originals. + let task_subsystems_builder: Box TaskSubsystems + Send> = { + let config = config.clone(); + let client = client.clone(); + let capabilities = capabilities.clone(); + let shutdown_token = shutdown_token.clone(); + let metrics = metrics.clone(); + let wft_slots = wft_slots.clone(); + let act_slots = act_slots.clone(); + let nexus_slots = nexus_slots.clone(); + let wf_last_suc_poll_time = wf_last_suc_poll_time.clone(); + let wf_sticky_last_suc_poll_time = wf_sticky_last_suc_poll_time.clone(); + let act_last_suc_poll_time = act_last_suc_poll_time.clone(); + let nexus_last_suc_poll_time = nexus_last_suc_poll_time.clone(); + let worker_telemetry = worker_telemetry.clone(); + let local_act_mgr = local_act_mgr.clone(); + let cancel_activity_slot = cancel_activity_slot.clone(); + Box::new(move || { + let (wft_stream, act_poller, nexus_poller) = match task_pollers { + TaskPollers::Real => { + let wft_stream = if config.task_types.enable_workflows { + let stream = make_wft_poller( + &config, + &sticky_queue_name, + &client, + &metrics, + &shutdown_token, + &wft_slots, + wf_last_suc_poll_time.clone(), + wf_sticky_last_suc_poll_time.clone(), + capabilities.clone(), + ) + .boxed(); + let stream = if !client.is_mock() { + // Some replay tests combine a mock client with real pollers, + // and they don't need to use the external stream + stream::select( + stream, + UnboundedReceiverStream::new(external_wft_rx), + ) + .left_stream() + } else { + stream.right_stream() + }; + Some(stream) + } else { + None + }; + + let act_poll_buffer = if config.task_types.enable_remote_activities { + let act_metrics = metrics.with_new_attrs([activity_poller()]); + let ap = LongPollBuffer::new_activity_task( + client.clone(), + config.task_queue.clone(), + resolve_effective_behavior( + config.activity_task_poller_behavior, + &capabilities, + ), + act_slots.clone(), + shutdown_token.child_token(), + Some(move |np| act_metrics.record_num_pollers(np)), + ActivityTaskOptions { + max_worker_acts_per_second: config + .max_worker_activities_per_second, + max_tps: config.max_task_queue_activities_per_second, + }, + act_last_suc_poll_time.clone(), + capabilities.clone(), + ); + Some(Box::from(ap) as BoxedActPoller) + } else { + None + }; + + let nexus_poll_buffer = if config.task_types.enable_nexus { + let np_metrics = metrics.with_new_attrs([nexus_poller()]); + Some(Box::new(LongPollBuffer::new_nexus_task( + client.clone(), + config.task_queue.clone(), + resolve_effective_behavior( + config.nexus_task_poller_behavior, + &capabilities, + ), + nexus_slots.clone(), + shutdown_token.child_token(), + Some(move |np| np_metrics.record_num_pollers(np)), + nexus_last_suc_poll_time.clone(), + capabilities.clone(), + shared_namespace_worker, + )) as BoxedNexusPoller) + } else { + None + }; + + #[cfg(any(feature = "test-utilities", test))] + let wft_stream = wft_stream.map(|s| s.left_stream()); + (wft_stream, act_poll_buffer, nexus_poll_buffer) + } + #[cfg(any(feature = "test-utilities", test))] + TaskPollers::Mocked { + wft_stream, + act_poller, + nexus_poller, + } => { + let wft_stream = config + .task_types + .enable_workflows + .then_some(wft_stream) + .flatten(); + let act_poller = config + .task_types + .enable_remote_activities + .then_some(act_poller) + .flatten(); + let nexus_poller = config + .task_types + .enable_nexus + .then_some(nexus_poller) + .flatten(); + + let ap = act_poller.map(|ap| { + MockPermittedPollBuffer::new(Arc::new(act_slots.clone()), ap) + }); + let np = nexus_poller.map(|np| { + MockPermittedPollBuffer::new(Arc::new(nexus_slots.clone()), np) + }); + let wfs = wft_stream.map(|stream| { + let wft_semaphore = wft_slots.clone(); + let wfs = stream.then(move |s| { + let wft_semaphore = wft_semaphore.clone(); + async move { + let permit = wft_semaphore.acquire_owned().await; + s.map(|s| (s, permit)) + } + }); + wfs.right_stream() + }); + ( + wfs, + ap.map(|ap| Box::new(ap) as BoxedActPoller), + np.map(|np| Box::new(np) as BoxedNexusPoller), + ) + } + }; + let at_task_mgr = act_poller.map(|ap| { + WorkerActivityTasks::new( + act_slots.clone(), + ap, + client.clone(), + metrics.clone(), + config.max_heartbeat_throttle_interval, + config.default_heartbeat_throttle_interval, + config.graceful_shutdown_period, + config.local_timeout_buffer_for_activities, + ) + }); + + if let Some(mgr) = &at_task_mgr { + let _ = cancel_activity_slot.set(mgr.cancel_activity_callback()); + } + + let nexus_mgr = nexus_poller.map(|poller| { + NexusManager::new( + poller, + metrics.clone(), + config.graceful_shutdown_period, + shutdown_token.child_token(), + ) + }); + + let workflows = wft_stream.map(|stream| { + Workflows::new( + WorkflowBasics { + worker_config: Arc::new(config.clone()), + shutdown_token: shutdown_token.child_token(), + metrics, + server_capabilities: client.capabilities().unwrap_or_default(), + sdk_name: sdk_name_and_ver.0, + sdk_version: sdk_name_and_ver.1, + default_versioning_behavior: config + .versioning_strategy + .default_versioning_behavior(), + }, + sticky_queue_name.map(|sq| StickyExecutionAttributes { + worker_task_queue: Some(TaskQueue { + name: sq, + kind: TaskQueueKind::Sticky as i32, + normal_name: config.task_queue.clone(), + }), + schedule_to_start_timeout: Some( + config + .sticky_queue_schedule_to_start_timeout + .try_into() + .expect("timeout fits into proto"), + ), + }), + client, + wft_slots, + stream, + la_sink, + local_act_mgr.clone(), + hb_rx, + at_task_mgr.as_ref().and_then(|mgr| { + match config.max_task_queue_activities_per_second { + Some(persec) if persec > 0.0 => None, + _ => Some(mgr.get_handle_for_workflows()), + } + }), + worker_telemetry + .as_ref() + .and_then(|telem| telem.trace_subscriber.clone()), + ) + }); + TaskSubsystems { + workflows, + at_task_mgr, + nexus_mgr, + } + }) + }; let deployment_options = match &config.versioning_strategy { WorkerVersioningStrategy::WorkerDeploymentBased(opts) => Some(opts.clone()), _ => None, @@ -882,10 +1037,7 @@ impl Worker { external_wft_tx, deployment_options, ); - let worker_instance_key = client.worker_instance_key(); - let worker_status = Arc::new(RwLock::new(WorkerStatus::Running)); - let sdk_name_and_ver = client.sdk_name_and_version(); let worker_heartbeat = worker_heartbeat_interval.map(|hb_interval| { let heartbeat_sys_info = sys_info.unwrap_or_else(|| Arc::new(RealSysInfo::new(hb_interval))); @@ -908,17 +1060,15 @@ impl Worker { hb_interval, worker_telemetry.clone(), hb_metrics, + capabilities.clone(), ) }); - let cancel_activity_callback = at_task_mgr - .as_ref() - .map(|mgr| mgr.cancel_activity_callback()); let client_worker_registrator = Arc::new(ClientWorkerRegistrator { worker_instance_key, slot_provider: provider, heartbeat_manager: worker_heartbeat, - cancel_activity_callback, + cancel_activity_slot: cancel_activity_slot.clone(), client: RwLock::new(client.clone()), shared_namespace_worker, task_types: config.task_types, @@ -934,50 +1084,7 @@ impl Worker { Ok(Self { worker_instance_key, client: client.clone(), - workflows: wft_stream.map(|stream| { - Workflows::new( - WorkflowBasics { - worker_config: Arc::new(config.clone()), - shutdown_token: shutdown_token.child_token(), - metrics, - server_capabilities: client.capabilities().unwrap_or_default(), - sdk_name: sdk_name_and_ver.0, - sdk_version: sdk_name_and_ver.1, - default_versioning_behavior: config - .versioning_strategy - .default_versioning_behavior(), - }, - sticky_queue_name.map(|sq| StickyExecutionAttributes { - worker_task_queue: Some(TaskQueue { - name: sq, - kind: TaskQueueKind::Sticky as i32, - normal_name: config.task_queue.clone(), - }), - schedule_to_start_timeout: Some( - config - .sticky_queue_schedule_to_start_timeout - .try_into() - .expect("timeout fits into proto"), - ), - }), - client, - wft_slots, - stream, - la_sink, - local_act_mgr.clone(), - hb_rx, - at_task_mgr.as_ref().and_then(|mgr| { - match config.max_task_queue_activities_per_second { - Some(persec) if persec > 0.0 => None, - _ => Some(mgr.get_handle_for_workflows()), - } - }), - worker_telemetry - .as_ref() - .and_then(|telem| telem.trace_subscriber.clone()), - ) - }), - at_task_mgr, + task_subsystems: LazyLock::new(task_subsystems_builder), local_act_mgr, config, shutdown_token, @@ -990,7 +1097,6 @@ impl Worker { act_permits, la_permits, }), - nexus_mgr, client_worker_registrator, status: worker_status, capabilities, @@ -1025,19 +1131,21 @@ impl Worker { if let Some(la_mgr) = &self.local_act_mgr { la_mgr.wait_all_outstanding_tasks_finished().await; } - // Wait for workflows to finish - if let Some(workflows) = &self.workflows { + // Wait for workflows to finish. If a caller reached shutdown without ever polling, this + // deref builds the subsystems; the build closure sees the cancelled shutdown token and + // signals them, so this still terminates promptly. + if let Some(workflows) = self.task_subsystems.workflows.as_ref() { workflows .shutdown() .await .expect("Workflow processing terminates cleanly"); } // Wait for activities to finish - if let Some(acts) = self.at_task_mgr.as_ref() { + if let Some(acts) = self.task_subsystems.at_task_mgr.as_ref() { acts.shutdown().await; } // Wait for nexus tasks to finish - if let Some(nexus) = &self.nexus_mgr { + if let Some(nexus) = self.task_subsystems.nexus_mgr.as_ref() { nexus.shutdown().await; } // Wait for all permits to be released, but don't totally hang real-world shutdown. @@ -1056,8 +1164,8 @@ impl Worker { /// functions have returned `ShutDown` errors. pub async fn finalize_shutdown(self) { self.shutdown().await; - if let Some(b) = self.at_task_mgr { - b.shutdown().await; + if let Some(atm) = self.task_subsystems.at_task_mgr.as_ref() { + atm.shutdown().await; } // Only after worker is fully shutdown do we remove the heartbeat callback // from SharedNamespaceWorker, allowing for accurate worker shutdown @@ -1074,7 +1182,7 @@ impl Worker { /// Returns number of currently cached workflows pub async fn cached_workflows(&self) -> usize { - match &self.workflows { + match self.task_subsystems.workflows.as_ref() { Some(workflows) => workflows .get_state_info() .await @@ -1087,7 +1195,7 @@ impl Worker { /// Returns number of currently outstanding workflow tasks #[cfg(test)] pub(crate) async fn outstanding_workflow_tasks(&self) -> usize { - match &self.workflows { + match self.task_subsystems.workflows.as_ref() { Some(workflows) => workflows .get_state_info() .await @@ -1099,13 +1207,17 @@ impl Worker { #[allow(unused)] pub(crate) fn available_wft_permits(&self) -> Option { - self.workflows + self.task_subsystems + .workflows .as_ref() .and_then(|w| w.available_wft_permits()) } #[cfg(test)] pub(crate) fn unused_wft_permits(&self) -> Option { - self.workflows.as_ref().and_then(|w| w.unused_wft_permits()) + self.task_subsystems + .workflows + .as_ref() + .and_then(|w| w.unused_wft_permits()) } /// Ask the worker for some work, returning an [ActivityTask]. It is then the language SDK's @@ -1141,7 +1253,7 @@ impl Worker { unreachable!() } if self.config.task_types.enable_remote_activities { - if let Some(ref act_mgr) = self.at_task_mgr { + if let Some(act_mgr) = self.task_subsystems.at_task_mgr.as_ref() { let res = act_mgr.poll().await; if let Err(err) = res.as_ref() && matches!(err, PollError::ShutDown) @@ -1229,7 +1341,7 @@ impl Worker { /// the user as we don't want to break activity execution due to badly configured heartbeat /// options. pub fn record_activity_heartbeat(&self, details: ActivityHeartbeat) { - if let Some(at_mgr) = self.at_task_mgr.as_ref() { + if let Some(at_mgr) = self.task_subsystems.at_task_mgr.as_ref() { let tt = TaskToken(details.task_token.clone()); if let Err(e) = at_mgr.record_heartbeat(details) { warn!(task_token = %tt, details = ?e, "Activity heartbeat failed."); @@ -1266,7 +1378,7 @@ impl Worker { return Ok(()); } - if let Some(atm) = &self.at_task_mgr { + if let Some(atm) = self.task_subsystems.at_task_mgr.as_ref() { atm.complete(task_token, status, &*self.client).await; Ok(()) } else { @@ -1288,7 +1400,7 @@ impl Worker { /// Do not call poll concurrently. It handles polling the server concurrently internally. #[instrument(skip(self), fields(run_id, workflow_id, task_queue=%self.config.task_queue))] pub async fn poll_workflow_activation(&self) -> Result { - match &self.workflows { + match self.task_subsystems.workflows.as_ref() { Some(workflows) => { let r = workflows.next_workflow_activation().await; // In the event workflows are shutdown or erroring, begin shutdown of everything else. Once @@ -1320,7 +1432,7 @@ impl Worker { &self, completion: WorkflowActivationCompletion, ) -> Result<(), CompleteWfError> { - match &self.workflows { + match self.task_subsystems.workflows.as_ref() { Some(workflows) => { workflows .activation_completed( @@ -1346,7 +1458,7 @@ impl Worker { /// Do not call poll concurrently. It handles polling the server concurrently internally. #[instrument(skip(self))] pub async fn poll_nexus_task(&self) -> Result { - match &self.nexus_mgr { + match self.task_subsystems.nexus_mgr.as_ref() { Some(mgr) => mgr.next_nexus_task().await, None => Err(PollError::ShutDown), } @@ -1373,7 +1485,7 @@ impl Worker { tracing::Span::current().record("task_token", tt.to_string()); tracing::Span::current().record("status", status.to_string()); - match &self.nexus_mgr { + match self.task_subsystems.nexus_mgr.as_ref() { Some(mgr) => mgr.complete_task(tt, status, &*self.client).await, None => Err(CompleteNexusError::NexusNotEnabled), } @@ -1399,7 +1511,7 @@ impl Worker { message: impl Into, reason: EvictionReason, ) { - if let Some(workflows) = &self.workflows { + if let Some(workflows) = self.task_subsystems.workflows.as_ref() { workflows.request_eviction(run_id, message, reason); } else { dbg_panic!("trying to request wf eviction when workflows not enabled for this worker"); @@ -1448,12 +1560,12 @@ impl Worker { // Push a BumpStream message to the workflow activation queue. This ensures that // any pending workflow activation polls will resolve, even if there are no other inputs. - if let Some(workflows) = &self.workflows { + if let Some(workflows) = self.task_subsystems.workflows.as_ref() { workflows.bump_stream(); } // Second, we want to stop polling of both activity and workflow tasks - if let Some(atm) = self.at_task_mgr.as_ref() { + if let Some(atm) = self.task_subsystems.at_task_mgr.as_ref() { atm.initiate_shutdown(); } // Let the manager know that shutdown has been initiated to try to unblock the local @@ -1464,7 +1576,12 @@ impl Worker { // If workflows have never been polled, immediately tell the local activity manager // that workflows have shut down, so it can proceed with shutdown without waiting. // This is particularly important for activity-only workers. - if self.workflows.as_ref().is_none_or(|w| !w.ever_polled()) { + if self + .task_subsystems + .workflows + .as_ref() + .is_none_or(|w| !w.ever_polled()) + { la_mgr.workflows_have_shutdown(); } } @@ -1478,6 +1595,7 @@ impl Worker { let client = self.client.clone(); let sticky_name = self + .task_subsystems .workflows .as_ref() .and_then(|wf| wf.get_sticky_queue_name()) @@ -1560,7 +1678,7 @@ impl Worker { } fn notify_local_result(&self, run_id: &str, res: LocalResolution) { - if let Some(workflows) = &self.workflows { + if let Some(workflows) = self.task_subsystems.workflows.as_ref() { workflows.notify_of_local_result(run_id, res); } else { dbg_panic!("trying to notify local result when workflows not enabled for this worker"); @@ -2012,7 +2130,9 @@ struct ClientWorkerRegistrator { worker_instance_key: Uuid, slot_provider: SlotProvider, heartbeat_manager: Option, - cancel_activity_callback: Option, + /// Slot filled by the deferred poller build with the activity cancellation callback. Shared + /// with the worker so it can be resolved lazily once the activity task manager is constructed. + cancel_activity_slot: Arc>, client: RwLock>, shared_namespace_worker: bool, task_types: WorkerTaskTypes, @@ -2051,7 +2171,12 @@ impl ClientWorker for ClientWorkerRegistrator { } fn cancel_activity_callback(&self) -> Option { - self.cancel_activity_callback.clone() + // When activities are disabled the slot is never filled, so the wrapper returns false (a + // harmless no-op); no need to gate on whether activities are enabled. + let slot = self.cancel_activity_slot.clone(); + Some(Arc::new(move |tt| { + slot.get().map(|cb| cb(tt)).unwrap_or(false) + })) } fn new_shared_namespace_worker( @@ -2104,6 +2229,7 @@ impl WorkerHeartbeatManager { heartbeat_interval: Duration, telemetry_instance: Option, heartbeat_manager_metrics: HeartbeatMetrics, + capabilities: Arc, ) -> Self { let start_time = Some(SystemTime::now().into()); let worker_heartbeat_callback: HeartbeatFn = Arc::new(move || { @@ -2165,7 +2291,11 @@ impl WorkerHeartbeatManager { .wf_last_suc_poll_time .load() .map(|time| time.into()), - is_autoscaling: config.workflow_task_poller_behavior.is_autoscaling(), + is_autoscaling: resolve_effective_behavior( + config.workflow_task_poller_behavior, + &capabilities, + ) + .is_autoscaling(), }); worker_heartbeat.workflow_sticky_poller_info = Some(WorkerPollerInfo { current_pollers: in_mem @@ -2176,7 +2306,11 @@ impl WorkerHeartbeatManager { .wf_sticky_last_suc_poll_time .load() .map(|time| time.into()), - is_autoscaling: config.workflow_task_poller_behavior.is_autoscaling(), + is_autoscaling: resolve_effective_behavior( + config.workflow_task_poller_behavior, + &capabilities, + ) + .is_autoscaling(), }); worker_heartbeat.activity_poller_info = Some(WorkerPollerInfo { current_pollers: in_mem @@ -2187,7 +2321,11 @@ impl WorkerHeartbeatManager { .act_last_suc_poll_time .load() .map(|time| time.into()), - is_autoscaling: config.activity_task_poller_behavior.is_autoscaling(), + is_autoscaling: resolve_effective_behavior( + config.activity_task_poller_behavior, + &capabilities, + ) + .is_autoscaling(), }); worker_heartbeat.nexus_poller_info = Some(WorkerPollerInfo { current_pollers: in_mem @@ -2198,7 +2336,11 @@ impl WorkerHeartbeatManager { .nexus_last_suc_poll_time .load() .map(|time| time.into()), - is_autoscaling: config.nexus_task_poller_behavior.is_autoscaling(), + is_autoscaling: resolve_effective_behavior( + config.nexus_task_poller_behavior, + &capabilities, + ) + .is_autoscaling(), }); worker_heartbeat.workflow_task_slots_info = make_slots_info( @@ -2259,12 +2401,19 @@ pub(crate) enum TaskPollers { }, } -fn wft_poller_behavior(config: &WorkerConfig, is_sticky: bool) -> PollerBehavior { +/// Given an already-resolved workflow task poller `behavior` (see [resolve_effective_behavior]), +/// applies the sticky/non-sticky split when it is a [PollerBehavior::SimpleMaximum]; other +/// behaviors are returned unchanged. +pub(crate) fn wft_poller_behavior( + behavior: PollerBehavior, + config: &WorkerConfig, + is_sticky: bool, +) -> PollerBehavior { fn calc_max_nonsticky(max_polls: usize, ratio: f32) -> usize { ((max_polls as f32 * ratio) as usize).max(1) } - if let PollerBehavior::SimpleMaximum(m) = config.workflow_task_poller_behavior { + if let PollerBehavior::SimpleMaximum(m) = behavior { if !is_sticky { PollerBehavior::SimpleMaximum(calc_max_nonsticky( m, @@ -2277,7 +2426,7 @@ fn wft_poller_behavior(config: &WorkerConfig, is_sticky: bool) -> PollerBehavior ) } } else { - config.workflow_task_poller_behavior + behavior } } @@ -2339,7 +2488,12 @@ mod tests { let fut = worker.poll_activity_task(); advance_fut!(fut); assert_eq!( - worker.at_task_mgr.as_ref().unwrap().unused_permits(), + worker + .task_subsystems + .at_task_mgr + .as_ref() + .unwrap() + .unused_permits(), Some(5) ); } @@ -2363,22 +2517,30 @@ mod tests { .unwrap(); let worker = Worker::new_test(cfg, mock_client); assert!(worker.activity_poll().await.is_err()); - assert_eq!(worker.at_task_mgr.unwrap().unused_permits(), Some(5)); + assert_eq!( + worker + .task_subsystems + .at_task_mgr + .as_ref() + .unwrap() + .unused_permits(), + Some(5) + ); } #[test] fn max_polls_calculated_properly() { let cfg = { let mut cfg = test_worker_cfg().build().unwrap(); - cfg.workflow_task_poller_behavior = PollerBehavior::SimpleMaximum(5_usize); + cfg.workflow_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(5_usize)); cfg }; assert_eq!( - wft_poller_behavior(&cfg, false), + wft_poller_behavior(PollerBehavior::SimpleMaximum(5), &cfg, false), PollerBehavior::SimpleMaximum(1) ); assert_eq!( - wft_poller_behavior(&cfg, true), + wft_poller_behavior(PollerBehavior::SimpleMaximum(5), &cfg, true), PollerBehavior::SimpleMaximum(4) ); } @@ -2551,4 +2713,52 @@ mod tests { "Error should mention default_versioning_behavior: {err}", ); } + + fn auto_enroll_caps() -> NamespaceCapabilities { + let caps = NamespaceCapabilities::default(); + caps.poller_autoscaling_auto_enroll + .store(true, Ordering::Relaxed); + caps + } + + #[test] + fn resolve_effective_behavior_enrolls_unset_with_capability() { + let caps = auto_enroll_caps(); + assert_eq!( + resolve_effective_behavior(None, &caps), + PollerBehavior::Autoscaling { + minimum: 1, + maximum: 100, + initial: 5, + }, + ); + } + + #[test] + fn resolve_effective_behavior_unset_without_capability_defaults() { + let caps = NamespaceCapabilities::default(); + assert_eq!( + resolve_effective_behavior(None, &caps), + PollerBehavior::SimpleMaximum(5), + ); + } + + #[test] + fn resolve_effective_behavior_configured_unchanged() { + // A configured behavior is always used as-is, even when the capability is present. + let caps = auto_enroll_caps(); + assert_eq!( + resolve_effective_behavior(Some(PollerBehavior::SimpleMaximum(5)), &caps), + PollerBehavior::SimpleMaximum(5), + ); + let autoscaling = PollerBehavior::Autoscaling { + minimum: 2, + maximum: 20, + initial: 4, + }; + assert_eq!( + resolve_effective_behavior(Some(autoscaling), &caps), + autoscaling, + ); + } } diff --git a/crates/sdk-core/src/worker/workflow/wft_poller.rs b/crates/sdk-core/src/worker/workflow/wft_poller.rs index d5069a67a..06f78a6a8 100644 --- a/crates/sdk-core/src/worker/workflow/wft_poller.rs +++ b/crates/sdk-core/src/worker/workflow/wft_poller.rs @@ -38,7 +38,10 @@ pub(crate) fn make_wft_poller( > + Sized + 'static { let wft_metrics = metrics.with_new_attrs([workflow_poller()]); - let poller_behavior = wft_poller_behavior(config, false); + let effective_behavior = crate::worker::resolve_effective_behavior( + config.workflow_task_poller_behavior, + &capabilities, + ); let wft_poller_shared = if sticky_queue_name.is_some() { Some(Arc::new(WFTPollerShared::new( wft_slots.available_permits(), @@ -50,7 +53,7 @@ pub(crate) fn make_wft_poller( client.clone(), config.task_queue.clone(), None, - poller_behavior, + wft_poller_behavior(effective_behavior, config, false), wft_slots.clone(), shutdown_token.child_token(), Some(move |np| { @@ -68,7 +71,7 @@ pub(crate) fn make_wft_poller( client.clone(), config.task_queue.clone(), Some(sqn.clone()), - wft_poller_behavior(config, true), + wft_poller_behavior(effective_behavior, config, true), wft_slots.clone().into_sticky(), shutdown_token.child_token(), Some(move |np| { diff --git a/crates/sdk-core/tests/heavy_tests.rs b/crates/sdk-core/tests/heavy_tests.rs index 50a80a86e..9b96db4c6 100644 --- a/crates/sdk-core/tests/heavy_tests.rs +++ b/crates/sdk-core/tests/heavy_tests.rs @@ -83,7 +83,7 @@ async fn activity_load() { let mut starter = CoreWfStarter::new("activity_load"); starter.sdk_config.max_cached_workflows = CONCURRENCY; - starter.sdk_config.activity_task_poller_behavior = PollerBehavior::SimpleMaximum(10); + starter.sdk_config.activity_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(10)); starter.sdk_config.tuner = Arc::new(TunerHolder::fixed_size(CONCURRENCY, CONCURRENCY, 100, 100)); starter.sdk_config.register_activities(StdActivities); @@ -162,8 +162,10 @@ async fn chunky_activities_resource_based() { const WORKFLOWS: usize = 100; let mut starter = CoreWfStarter::new("chunky_activities_resource_based"); - starter.sdk_config.workflow_task_poller_behavior = PollerBehavior::SimpleMaximum(10_usize); - starter.sdk_config.activity_task_poller_behavior = PollerBehavior::SimpleMaximum(10_usize); + starter.sdk_config.workflow_task_poller_behavior = + Some(PollerBehavior::SimpleMaximum(10_usize)); + starter.sdk_config.activity_task_poller_behavior = + Some(PollerBehavior::SimpleMaximum(10_usize)); let mut tuner = ResourceBasedTuner::new(0.7, 0.7); tuner .with_workflow_slots_options(ResourceSlotOptions::new( @@ -245,7 +247,7 @@ async fn workflow_load() { let rt = CoreRuntime::new_assume_tokio(get_integ_runtime_options(telemopts)).unwrap(); let mut starter = CoreWfStarter::new_with_runtime("workflow_load", rt); starter.sdk_config.max_cached_workflows = 200; - starter.sdk_config.activity_task_poller_behavior = PollerBehavior::SimpleMaximum(10); + starter.sdk_config.activity_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(10)); starter.sdk_config.tuner = Arc::new(TunerHolder::fixed_size(5, 100, 100, 100)); starter.sdk_config.register_activities(StdActivities); let task_queue = starter.get_task_queue().to_owned(); @@ -456,16 +458,16 @@ async fn poller_autoscaling_basic_loadtest() { let mut starter = CoreWfStarter::new("poller_load"); starter.sdk_config.max_cached_workflows = 5000; starter.sdk_config.tuner = Arc::new(TunerHolder::fixed_size(1000, 1000, 100, 1)); - starter.sdk_config.workflow_task_poller_behavior = PollerBehavior::Autoscaling { + starter.sdk_config.workflow_task_poller_behavior = Some(PollerBehavior::Autoscaling { minimum: 1, maximum: 200, initial: 5, - }; - starter.sdk_config.activity_task_poller_behavior = PollerBehavior::Autoscaling { + }); + starter.sdk_config.activity_task_poller_behavior = Some(PollerBehavior::Autoscaling { minimum: 1, maximum: 200, initial: 5, - }; + }); starter.sdk_config.register_activities(JitteryActivities); let mut worker = starter.worker().await; diff --git a/crates/sdk-core/tests/integ_tests/metrics_tests.rs b/crates/sdk-core/tests/integ_tests/metrics_tests.rs index 234c75acd..c8b9e9e23 100644 --- a/crates/sdk-core/tests/integ_tests/metrics_tests.rs +++ b/crates/sdk-core/tests/integ_tests/metrics_tests.rs @@ -470,11 +470,11 @@ async fn idle_activity_worker_reports_zero_slots_used() { let rt = CoreRuntime::new_assume_tokio(get_integ_runtime_options(telemopts)).unwrap(); let mut starter = CoreWfStarter::new_with_runtime("idle_activity_worker_reports_zero_slots_used", rt); - starter.sdk_config.activity_task_poller_behavior = PollerBehavior::Autoscaling { + starter.sdk_config.activity_task_poller_behavior = Some(PollerBehavior::Autoscaling { minimum: 1, maximum: 1, initial: 1, - }; + }); let activity_slots = Arc::new(ReservationTrackingActivitySlotSupplier::new(3)); let mut tuner = TunerBuilder::default(); tuner.activity_slot_supplier(activity_slots.clone()); diff --git a/crates/sdk-core/tests/integ_tests/polling_tests.rs b/crates/sdk-core/tests/integ_tests/polling_tests.rs index 2acad2a58..2a1634b48 100644 --- a/crates/sdk-core/tests/integ_tests/polling_tests.rs +++ b/crates/sdk-core/tests/integ_tests/polling_tests.rs @@ -277,15 +277,15 @@ async fn small_workflow_slots_and_pollers(#[values(false, true)] use_autoscaling let wf_name = "only_one_workflow_slot_and_two_pollers"; let mut starter = CoreWfStarter::new(wf_name); if use_autoscaling { - starter.sdk_config.workflow_task_poller_behavior = PollerBehavior::Autoscaling { + starter.sdk_config.workflow_task_poller_behavior = Some(PollerBehavior::Autoscaling { minimum: 1, maximum: 5, initial: 1, - }; + }); } else { - starter.sdk_config.workflow_task_poller_behavior = PollerBehavior::SimpleMaximum(2); + starter.sdk_config.workflow_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(2)); } - starter.sdk_config.activity_task_poller_behavior = PollerBehavior::SimpleMaximum(1); + starter.sdk_config.activity_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(1)); starter.sdk_config.tuner = Arc::new(TunerHolder::fixed_size(2, 1, 1, 1)); starter.sdk_config.register_activities(StdActivities); let mut worker = starter.worker().await; diff --git a/crates/sdk-core/tests/integ_tests/worker_heartbeat_tests.rs b/crates/sdk-core/tests/integ_tests/worker_heartbeat_tests.rs index ca8b5f8ba..80b692f74 100644 --- a/crates/sdk-core/tests/integ_tests/worker_heartbeat_tests.rs +++ b/crates/sdk-core/tests/integ_tests/worker_heartbeat_tests.rs @@ -351,16 +351,16 @@ async fn docker_worker_heartbeat_tuner() { tuner .with_workflow_slots_options(ResourceSlotOptions::new(2, 10, Duration::from_millis(0))) .with_activity_slots_options(ResourceSlotOptions::new(5, 10, Duration::from_millis(50))); - starter.sdk_config.workflow_task_poller_behavior = PollerBehavior::Autoscaling { + starter.sdk_config.workflow_task_poller_behavior = Some(PollerBehavior::Autoscaling { minimum: 1, maximum: 200, initial: 5, - }; - starter.sdk_config.nexus_task_poller_behavior = PollerBehavior::Autoscaling { + }); + starter.sdk_config.nexus_task_poller_behavior = Some(PollerBehavior::Autoscaling { minimum: 1, maximum: 200, initial: 5, - }; + }); starter.sdk_config.tuner = Arc::new(tuner); starter.sdk_config.register_activities(StdActivities); let mut worker = starter.worker().await; diff --git a/crates/sdk-core/tests/integ_tests/worker_tests.rs b/crates/sdk-core/tests/integ_tests/worker_tests.rs index 7dbd68d01..b26dd9a5f 100644 --- a/crates/sdk-core/tests/integ_tests/worker_tests.rs +++ b/crates/sdk-core/tests/integ_tests/worker_tests.rs @@ -199,7 +199,7 @@ async fn resource_based_few_pollers_guarantees_non_sticky_poll() { let mut starter = CoreWfStarter::new(wf_name); starter.sdk_config.task_types = WorkerTaskTypes::workflow_only(); // 3 pollers so the minimum slots of 2 can both be handed out to a sticky poller - starter.sdk_config.workflow_task_poller_behavior = PollerBehavior::SimpleMaximum(3_usize); + starter.sdk_config.workflow_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(3_usize)); // Set the limits to zero so it's essentially unwilling to hand out slots let mut tuner = ResourceBasedTuner::new(0.0, 0.0); tuner.with_workflow_slots_options(ResourceSlotOptions::new(2, 10, Duration::from_millis(0))); diff --git a/crates/sdk-core/tests/integ_tests/workflow_tests.rs b/crates/sdk-core/tests/integ_tests/workflow_tests.rs index e93cf14ff..b8223fb59 100644 --- a/crates/sdk-core/tests/integ_tests/workflow_tests.rs +++ b/crates/sdk-core/tests/integ_tests/workflow_tests.rs @@ -388,7 +388,8 @@ async fn wft_timeout_doesnt_create_unsolvable_autocomplete() { // Test needs eviction on and a short timeout wf_starter.sdk_config.max_cached_workflows = 0_usize; wf_starter.sdk_config.tuner = Arc::new(TunerHolder::fixed_size(1, 1, 1, 1)); - wf_starter.sdk_config.workflow_task_poller_behavior = PollerBehavior::SimpleMaximum(1_usize); + wf_starter.sdk_config.workflow_task_poller_behavior = + Some(PollerBehavior::SimpleMaximum(1_usize)); wf_starter.workflow_options.task_timeout = Some(Duration::from_secs(1)); let core = wf_starter.get_worker().await; let client = wf_starter.get_client().await; diff --git a/crates/sdk-core/tests/integ_tests/workflow_tests/activities.rs b/crates/sdk-core/tests/integ_tests/workflow_tests/activities.rs index d4a7dde8c..bb2e5d92a 100644 --- a/crates/sdk-core/tests/integ_tests/workflow_tests/activities.rs +++ b/crates/sdk-core/tests/integ_tests/workflow_tests/activities.rs @@ -1990,16 +1990,16 @@ async fn activity_can_be_cancelled_by_local_timeout() { async fn long_activity_timeout_repro() { let wf_name = "long_activity_timeout_repro"; let mut starter = CoreWfStarter::new(wf_name); - starter.sdk_config.workflow_task_poller_behavior = PollerBehavior::Autoscaling { + starter.sdk_config.workflow_task_poller_behavior = Some(PollerBehavior::Autoscaling { minimum: 1, maximum: 10, initial: 5, - }; - starter.sdk_config.activity_task_poller_behavior = PollerBehavior::Autoscaling { + }); + starter.sdk_config.activity_task_poller_behavior = Some(PollerBehavior::Autoscaling { minimum: 1, maximum: 10, initial: 5, - }; + }); starter .set_core_cfg_mutator(|m| m.local_timeout_buffer_for_activities = Duration::from_secs(0)); starter.sdk_config.register_activities(StdActivities); diff --git a/crates/sdk-core/tests/integ_tests/workflow_tests/stickyness.rs b/crates/sdk-core/tests/integ_tests/workflow_tests/stickyness.rs index b8238d943..53dce4e03 100644 --- a/crates/sdk-core/tests/integ_tests/workflow_tests/stickyness.rs +++ b/crates/sdk-core/tests/integ_tests/workflow_tests/stickyness.rs @@ -125,7 +125,7 @@ async fn cache_miss_ok() { starter.sdk_config.task_types = WorkerTaskTypes::workflow_only(); starter.sdk_config.tuner = Arc::new(TunerHolder::fixed_size(2, 1, 1, 1)); starter.sdk_config.max_cached_workflows = 0_usize; - starter.sdk_config.workflow_task_poller_behavior = PollerBehavior::SimpleMaximum(1_usize); + starter.sdk_config.workflow_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(1_usize)); let mut worker = starter.worker().await; let barr = Arc::new(Barrier::new(2)); diff --git a/crates/sdk-core/tests/manual_tests.rs b/crates/sdk-core/tests/manual_tests.rs index 3cb3d9f86..e2ec78c75 100644 --- a/crates/sdk-core/tests/manual_tests.rs +++ b/crates/sdk-core/tests/manual_tests.rs @@ -136,16 +136,16 @@ async fn poller_load_spiky() { let mut starter = CoreWfStarter::new_with_runtime("poller_load", rt); starter.sdk_config.max_cached_workflows = 5000; starter.sdk_config.tuner = Arc::new(TunerHolder::fixed_size(1000, 1000, 100, 100)); - starter.sdk_config.workflow_task_poller_behavior = PollerBehavior::Autoscaling { + starter.sdk_config.workflow_task_poller_behavior = Some(PollerBehavior::Autoscaling { minimum: 1, maximum: 200, initial: 5, - }; - starter.sdk_config.activity_task_poller_behavior = PollerBehavior::Autoscaling { + }); + starter.sdk_config.activity_task_poller_behavior = Some(PollerBehavior::Autoscaling { minimum: 1, maximum: 200, initial: 5, - }; + }); let mut worker = starter.worker().await; let submitter = worker.get_submitter_handle(); @@ -278,11 +278,11 @@ async fn poller_load_sustained() { let mut starter = CoreWfStarter::new_with_runtime("poller_load", rt); starter.sdk_config.max_cached_workflows = 5000; starter.sdk_config.tuner = Arc::new(TunerHolder::fixed_size(1000, 100, 100, 100)); - starter.sdk_config.workflow_task_poller_behavior = PollerBehavior::Autoscaling { + starter.sdk_config.workflow_task_poller_behavior = Some(PollerBehavior::Autoscaling { minimum: 1, maximum: 200, initial: 5, - }; + }); starter.sdk_config.task_types = WorkerTaskTypes::workflow_only(); let mut worker = starter.worker().await; worker.register_workflow::().unwrap(); @@ -353,16 +353,16 @@ async fn poller_load_spike_then_sustained() { let mut starter = CoreWfStarter::new_with_runtime("poller_load", rt); starter.sdk_config.max_cached_workflows = 5000; starter.sdk_config.tuner = Arc::new(TunerHolder::fixed_size(1000, 100, 100, 100)); - starter.sdk_config.workflow_task_poller_behavior = PollerBehavior::Autoscaling { + starter.sdk_config.workflow_task_poller_behavior = Some(PollerBehavior::Autoscaling { minimum: 1, maximum: 200, initial: 5, - }; - starter.sdk_config.activity_task_poller_behavior = PollerBehavior::Autoscaling { + }); + starter.sdk_config.activity_task_poller_behavior = Some(PollerBehavior::Autoscaling { minimum: 1, maximum: 200, initial: 5, - }; + }); let mut worker = starter.worker().await; let submitter = worker.get_submitter_handle(); diff --git a/crates/sdk/src/lib.rs b/crates/sdk/src/lib.rs index cd1ab6ef6..51f2b4264 100644 --- a/crates/sdk/src/lib.rs +++ b/crates/sdk/src/lib.rs @@ -196,8 +196,10 @@ pub struct WorkerOptions { /// Controls how polling for Workflow tasks will happen on this worker's task queue. See also /// [WorkerConfig::nonsticky_to_sticky_poll_ratio]. If using SimpleMaximum, Must be at least 2 /// when `max_cached_workflows` > 0, or is an error. - #[builder(default = PollerBehavior::SimpleMaximum(5))] - pub workflow_task_poller_behavior: PollerBehavior, + /// + /// If left unset, the worker uses `SimpleMaximum(5)` and becomes eligible for automatic + /// enrollment into poller autoscaling when the namespace advertises support for it. + pub workflow_task_poller_behavior: Option, /// Only applies when using [PollerBehavior::SimpleMaximum] /// /// (max workflow task polls * this number) = the number of max pollers that will be allowed for @@ -208,11 +210,15 @@ pub struct WorkerOptions { #[builder(default = 0.2)] pub nonsticky_to_sticky_poll_ratio: f32, /// Controls how polling for Activity tasks will happen on this worker's task queue. - #[builder(default = PollerBehavior::SimpleMaximum(5))] - pub activity_task_poller_behavior: PollerBehavior, + /// + /// If left unset, the worker uses `SimpleMaximum(5)` and becomes eligible for automatic + /// enrollment into poller autoscaling when the namespace advertises support for it. + pub activity_task_poller_behavior: Option, /// Controls how polling for Nexus tasks will happen on this worker's task queue. - #[builder(default = PollerBehavior::SimpleMaximum(5))] - pub nexus_task_poller_behavior: PollerBehavior, + /// + /// If left unset, the worker uses `SimpleMaximum(5)` and becomes eligible for automatic + /// enrollment into poller autoscaling when the namespace advertises support for it. + pub nexus_task_poller_behavior: Option, // TODO [rust-sdk-branch]: Will go away once workflow registration can only happen in here. // Then it can be auto-determined. /// Specifies which task types this worker will poll for. @@ -433,9 +439,9 @@ impl WorkerOptions { })) .max_cached_workflows(self.max_cached_workflows) .tuner(self.tuner.clone()) - .workflow_task_poller_behavior(self.workflow_task_poller_behavior) - .activity_task_poller_behavior(self.activity_task_poller_behavior) - .nexus_task_poller_behavior(self.nexus_task_poller_behavior) + .maybe_workflow_task_poller_behavior(self.workflow_task_poller_behavior) + .maybe_activity_task_poller_behavior(self.activity_task_poller_behavior) + .maybe_nexus_task_poller_behavior(self.nexus_task_poller_behavior) .task_types(self.task_types) .sticky_queue_schedule_to_start_timeout(self.sticky_queue_schedule_to_start_timeout) .max_heartbeat_throttle_interval(self.max_heartbeat_throttle_interval) @@ -688,6 +694,9 @@ impl Worker { /// Runs the worker. Eventually resolves after the worker has been explicitly shut down, /// or may return early with an error in the event of some unresolvable problem. pub async fn run(&mut self) -> Result<(), anyhow::Error> { + // Perform the namespace check-in so poller behavior (e.g. autoscaling auto-enroll) is + // resolved before any polling begins. + self.common.worker.validate().await?; let shutdown_token = CancellationToken::new(); let (common, wf_half, act_half) = self.split_apart(); let (wf_future_tx, wf_future_rx) = From d40f3dfbe6c92d8f6f7746d2ef262de0a04a1d12 Mon Sep 17 00:00:00 2001 From: Veeral Patel Date: Mon, 27 Jul 2026 14:45:14 -0700 Subject: [PATCH 2/5] test: default describe_namespace expectation for mock clients Worker::run() now calls validate(), which describes the namespace to discover capabilities. Mock-based tests never wired up that expectation, so mockall panicked; during unwind this became a non-unwinding abort that killed the whole test process and cascaded into unrelated tests. Add a fallback describe_namespace returning an empty response (no capabilities, so poller behavior stays at its default) for both the automock (build_mock_pollers) and manual (mock_manual_worker_client, used by replay) paths. Registered with an open call count and after any test-provided expectation, so mockall's FIFO matching still prefers a test's own expectation and workers that never validate don't trip an unsatisfied expectation. --- crates/sdk-core/src/test_help/integ_helpers.rs | 14 ++++++++++++-- crates/sdk-core/src/worker/client/mocks.rs | 6 ++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/crates/sdk-core/src/test_help/integ_helpers.rs b/crates/sdk-core/src/test_help/integ_helpers.rs index 87edc58b8..5847f83df 100644 --- a/crates/sdk-core/src/test_help/integ_helpers.rs +++ b/crates/sdk-core/src/test_help/integ_helpers.rs @@ -56,8 +56,9 @@ use temporalio_common::{ protocol::{self, v1::message}, update, workflowservice::v1::{ - PollActivityTaskQueueResponse, PollNexusTaskQueueResponse, - PollWorkflowTaskQueueResponse, RespondWorkflowTaskCompletedResponse, + DescribeNamespaceResponse, PollActivityTaskQueueResponse, + PollNexusTaskQueueResponse, PollWorkflowTaskQueueResponse, + RespondWorkflowTaskCompletedResponse, }, }, utilities::pack_any, @@ -836,6 +837,15 @@ pub fn build_mock_pollers(mut cfg: MockPollCfg) -> MocksHolder { Ok(Default::default()) }); + // Fallback so worker validation (which describes the namespace to discover capabilities) works + // for mock-based tests that don't care about capabilities. Added last and with an open call + // count, so any test-provided `describe_namespace` expectation is matched first (mockall checks + // expectations FIFO) and tests that never validate don't trip an unsatisfied expectation. + cfg.mock_client + .expect_describe_namespace() + .times(0..) + .returning(|| Ok(DescribeNamespaceResponse::default())); + let mut mh = MocksHolder { client: Arc::new(cfg.mock_client), inputs: mock_worker, diff --git a/crates/sdk-core/src/worker/client/mocks.rs b/crates/sdk-core/src/worker/client/mocks.rs index a67fa3b95..c601f2ee8 100644 --- a/crates/sdk-core/src/worker/client/mocks.rs +++ b/crates/sdk-core/src/worker/client/mocks.rs @@ -56,6 +56,12 @@ pub(crate) fn mock_manual_worker_client() -> MockManualWorkerClient { r.expect_workers() .returning(|| DEFAULT_WORKERS_REGISTRY.clone()); r.expect_is_mock().returning(|| true); + // Fallback so worker validation (which describes the namespace to discover capabilities) works + // for mock-based workers (e.g. replay) that don't care about capabilities. Open call count so + // workers that never validate don't trip an unsatisfied expectation. + r.expect_describe_namespace() + .times(0..) + .returning(|| async { Ok(DescribeNamespaceResponse::default()) }.boxed()); r.expect_shutdown_worker() .returning(|_, _, _, _| async { Ok(ShutdownWorkerResponse {}) }.boxed()); r.expect_sdk_name_and_version() From 5f3cd12472a298c43eeab4a835c9118c0c2ff8e5 Mon Sep 17 00:00:00 2001 From: Spencer Judge Date: Mon, 27 Jul 2026 15:22:17 -0700 Subject: [PATCH 3/5] Reduce clones somewhat --- crates/sdk-core/src/worker/mod.rs | 156 ++++++++++++++---------------- 1 file changed, 73 insertions(+), 83 deletions(-) diff --git a/crates/sdk-core/src/worker/mod.rs b/crates/sdk-core/src/worker/mod.rs index b21b246d6..1c76156dc 100644 --- a/crates/sdk-core/src/worker/mod.rs +++ b/crates/sdk-core/src/worker/mod.rs @@ -802,26 +802,71 @@ impl Worker { // exist until the pollers are built). let cancel_activity_slot: Arc> = Arc::new(OnceLock::new()); + let deployment_options = match &config.versioning_strategy { + WorkerVersioningStrategy::WorkerDeploymentBased(opts) => Some(opts.clone()), + _ => None, + }; + let provider = SlotProvider::new( + config.namespace.clone(), + config.task_queue.clone(), + wft_slots.clone(), + external_wft_tx, + deployment_options, + ); + + let worker_heartbeat = worker_heartbeat_interval.map(|hb_interval| { + let heartbeat_sys_info = + sys_info.unwrap_or_else(|| Arc::new(RealSysInfo::new(hb_interval))); + let hb_metrics = HeartbeatMetrics { + in_mem_metrics: metrics.in_memory_meter(), + wft_slots: wft_slots.clone(), + act_slots: act_slots.clone(), + nexus_slots: nexus_slots.clone(), + la_slots: la_permit_dealer, + wf_last_suc_poll_time: wf_last_suc_poll_time.clone(), + wf_sticky_last_suc_poll_time: wf_sticky_last_suc_poll_time.clone(), + act_last_suc_poll_time: act_last_suc_poll_time.clone(), + nexus_last_suc_poll_time: nexus_last_suc_poll_time.clone(), + status: worker_status.clone(), + sys_info: heartbeat_sys_info, + }; + WorkerHeartbeatManager::new( + config.clone(), + worker_instance_key, + hb_interval, + worker_telemetry.clone(), + hb_metrics, + capabilities.clone(), + ) + }); + + let client_worker_registrator = Arc::new(ClientWorkerRegistrator { + worker_instance_key, + slot_provider: provider, + heartbeat_manager: worker_heartbeat, + cancel_activity_slot: cancel_activity_slot.clone(), + client: RwLock::new(client.clone()), + shared_namespace_worker, + task_types: config.task_types, + }); + + if !shared_namespace_worker { + client.workers().register_worker( + client_worker_registrator.clone(), + config.skip_client_worker_set_check, + )?; + } + + let worker_config = config.clone(); + let worker_client = client.clone(); + let worker_shutdown_token = shutdown_token.clone(); + let worker_local_act_mgr = local_act_mgr.clone(); + let worker_capabilities = capabilities.clone(); + // Build the poller-dependent subsystems lazily. This closure runs once, after namespace // capabilities have been fetched (see `Worker::validate`), so the effective poller behavior - // can be resolved with knowledge of those capabilities. Everything it needs is captured by - // clone (or moved) *before* the synchronous heartbeat manager below consumes the originals. - let task_subsystems_builder: Box TaskSubsystems + Send> = { - let config = config.clone(); - let client = client.clone(); - let capabilities = capabilities.clone(); - let shutdown_token = shutdown_token.clone(); - let metrics = metrics.clone(); - let wft_slots = wft_slots.clone(); - let act_slots = act_slots.clone(); - let nexus_slots = nexus_slots.clone(); - let wf_last_suc_poll_time = wf_last_suc_poll_time.clone(); - let wf_sticky_last_suc_poll_time = wf_sticky_last_suc_poll_time.clone(); - let act_last_suc_poll_time = act_last_suc_poll_time.clone(); - let nexus_last_suc_poll_time = nexus_last_suc_poll_time.clone(); - let worker_telemetry = worker_telemetry.clone(); - let local_act_mgr = local_act_mgr.clone(); - let cancel_activity_slot = cancel_activity_slot.clone(); + // can be resolved with knowledge of those capabilities. + let task_subsystems_builder: Box TaskSubsystems + Send> = Box::new(move || { let (wft_stream, act_poller, nexus_poller) = match task_pollers { TaskPollers::Real => { @@ -891,8 +936,8 @@ impl Worker { nexus_slots.clone(), shutdown_token.child_token(), Some(move |np| np_metrics.record_num_pollers(np)), - nexus_last_suc_poll_time.clone(), - capabilities.clone(), + nexus_last_suc_poll_time, + capabilities, shared_namespace_worker, )) as BoxedNexusPoller) } else { @@ -993,7 +1038,7 @@ impl Worker { worker_task_queue: Some(TaskQueue { name: sq, kind: TaskQueueKind::Sticky as i32, - normal_name: config.task_queue.clone(), + normal_name: config.task_queue, }), schedule_to_start_timeout: Some( config @@ -1006,7 +1051,7 @@ impl Worker { wft_slots, stream, la_sink, - local_act_mgr.clone(), + local_act_mgr, hb_rx, at_task_mgr.as_ref().and_then(|mgr| { match config.max_task_queue_activities_per_second { @@ -1024,70 +1069,15 @@ impl Worker { at_task_mgr, nexus_mgr, } - }) - }; - let deployment_options = match &config.versioning_strategy { - WorkerVersioningStrategy::WorkerDeploymentBased(opts) => Some(opts.clone()), - _ => None, - }; - let provider = SlotProvider::new( - config.namespace.clone(), - config.task_queue.clone(), - wft_slots.clone(), - external_wft_tx, - deployment_options, - ); - - let worker_heartbeat = worker_heartbeat_interval.map(|hb_interval| { - let heartbeat_sys_info = - sys_info.unwrap_or_else(|| Arc::new(RealSysInfo::new(hb_interval))); - let hb_metrics = HeartbeatMetrics { - in_mem_metrics: metrics.in_memory_meter(), - wft_slots: wft_slots.clone(), - act_slots, - nexus_slots, - la_slots: la_permit_dealer, - wf_last_suc_poll_time, - wf_sticky_last_suc_poll_time, - act_last_suc_poll_time, - nexus_last_suc_poll_time, - status: worker_status.clone(), - sys_info: heartbeat_sys_info, - }; - WorkerHeartbeatManager::new( - config.clone(), - worker_instance_key, - hb_interval, - worker_telemetry.clone(), - hb_metrics, - capabilities.clone(), - ) - }); - - let client_worker_registrator = Arc::new(ClientWorkerRegistrator { - worker_instance_key, - slot_provider: provider, - heartbeat_manager: worker_heartbeat, - cancel_activity_slot: cancel_activity_slot.clone(), - client: RwLock::new(client.clone()), - shared_namespace_worker, - task_types: config.task_types, - }); - - if !shared_namespace_worker { - client.workers().register_worker( - client_worker_registrator.clone(), - config.skip_client_worker_set_check, - )?; - } + }); Ok(Self { worker_instance_key, - client: client.clone(), + client: worker_client, task_subsystems: LazyLock::new(task_subsystems_builder), - local_act_mgr, - config, - shutdown_token, + local_act_mgr: worker_local_act_mgr, + config: worker_config, + shutdown_token: worker_shutdown_token, post_activate_hook: None, // Non-local activities are already complete if configured not to poll for them. non_local_activities_complete: Arc::new(AtomicBool::new(!poll_on_non_local_activities)), @@ -1099,7 +1089,7 @@ impl Worker { }), client_worker_registrator, status: worker_status, - capabilities, + capabilities: worker_capabilities, shutdown_rpc_handle: Mutex::new(None), }) } From 721578e0fb51d9bd840d74241eb57cdeecb3ae78 Mon Sep 17 00:00:00 2001 From: Spencer Judge Date: Mon, 27 Jul 2026 15:56:57 -0700 Subject: [PATCH 4/5] Add regression for validation race --- CHANGELOG.md | 3 ++ crates/sdk-core/src/replay/mod.rs | 30 +++++++++++- crates/sdk-core/src/worker/client/mocks.rs | 6 --- crates/sdk-core/src/worker/heartbeat.rs | 55 ++++++++++++++++++++++ 4 files changed, 87 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3b21c70bb..60337bfe7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -54,6 +54,9 @@ to docs, or any other relevant information. ### Breaking Changes * The `ActivityContext` constructor now requires `ClientOptions`. +* `WorkerConfig::{workflow,activity,nexus}_task_poller_behavior` and the corresponding Rust SDK + `WorkerOptions` fields are now `Option`. `None` means the poller was not explicitly + configured and is eligible for automatic enrollment into poller autoscaling. ### Breaking Changes - Rust SDK `ApplicationFailure` and `WorkflowError` APIs now use boxed `std::error::Error` values instead of diff --git a/crates/sdk-core/src/replay/mod.rs b/crates/sdk-core/src/replay/mod.rs index c1cc6ef29..726e43d73 100644 --- a/crates/sdk-core/src/replay/mod.rs +++ b/crates/sdk-core/src/replay/mod.rs @@ -34,7 +34,8 @@ use temporalio_common::{ common::v1::WorkflowExecution, history::v1::History, workflowservice::v1::{ - RespondWorkflowTaskCompletedResponse, RespondWorkflowTaskFailedResponse, + DescribeNamespaceResponse, RespondWorkflowTaskCompletedResponse, + RespondWorkflowTaskFailedResponse, }, }, }, @@ -92,6 +93,12 @@ where } else { mock_manual_worker_client() }; + // Worker::run validates before polling. Installing this after an optional client override + // lets a test-provided describe expectation take precedence over the fallback. + client + .expect_describe_namespace() + .times(0..) + .returning(|| async { Ok(DescribeNamespaceResponse::default()) }.boxed()); let hist_allow_tx = historator.replay_done_tx.clone(); let historator = Arc::new(TokioMutex::new(historator)); @@ -138,6 +145,27 @@ where } } +#[cfg(test)] +mod tests { + use super::*; + use crate::test_help::test_worker_cfg; + use futures_util::{FutureExt, stream}; + + #[tokio::test] + async fn client_override_describe_namespace_precedes_fallback() { + let mut client = mock_manual_worker_client(); + client + .expect_describe_namespace() + .times(1) + .returning(|| async { Ok(DescribeNamespaceResponse::default()) }.boxed()); + + let mut input = ReplayWorkerInput::new(test_worker_cfg().build().unwrap(), stream::empty()); + input.client_override = Some(client); + + input.into_core_worker().unwrap().validate().await.unwrap(); + } +} + /// A history which will be used during replay verification. Since histories do not include the /// workflow id, it must be manually attached. #[derive(Debug, Clone)] diff --git a/crates/sdk-core/src/worker/client/mocks.rs b/crates/sdk-core/src/worker/client/mocks.rs index c601f2ee8..a67fa3b95 100644 --- a/crates/sdk-core/src/worker/client/mocks.rs +++ b/crates/sdk-core/src/worker/client/mocks.rs @@ -56,12 +56,6 @@ pub(crate) fn mock_manual_worker_client() -> MockManualWorkerClient { r.expect_workers() .returning(|| DEFAULT_WORKERS_REGISTRY.clone()); r.expect_is_mock().returning(|| true); - // Fallback so worker validation (which describes the namespace to discover capabilities) works - // for mock-based workers (e.g. replay) that don't care about capabilities. Open call count so - // workers that never validate don't trip an unsatisfied expectation. - r.expect_describe_namespace() - .times(0..) - .returning(|| async { Ok(DescribeNamespaceResponse::default()) }.boxed()); r.expect_shutdown_worker() .returning(|_, _, _, _| async { Ok(ShutdownWorkerResponse {}) }.boxed()); r.expect_sdk_name_and_version() diff --git a/crates/sdk-core/src/worker/heartbeat.rs b/crates/sdk-core/src/worker/heartbeat.rs index d07ad81d0..1bdec5728 100644 --- a/crates/sdk-core/src/worker/heartbeat.rs +++ b/crates/sdk-core/src/worker/heartbeat.rs @@ -443,6 +443,61 @@ mod tests { assert_eq!(3, heartbeat_count.load(Ordering::Relaxed)); } + #[tokio::test] + async fn heartbeat_reports_auto_enroll_before_worker_validation() { + let mut mock = mock_worker_client(); + let (reported_tx, reported_rx) = tokio::sync::oneshot::channel(); + let reported_tx = Mutex::new(Some(reported_tx)); + mock.expect_record_worker_heartbeat() + .returning(move |_, heartbeats| { + if let Some(tx) = reported_tx.lock().unwrap().take() { + let is_autoscaling = heartbeats[0] + .activity_poller_info + .as_ref() + .unwrap() + .is_autoscaling; + let _ = tx.send(is_autoscaling); + } + Ok(RecordWorkerHeartbeatResponse {}) + }); + mock.expect_describe_namespace().times(2).returning(|| { + Ok(DescribeNamespaceResponse { + namespace_info: Some(NamespaceInfo { + capabilities: Some(Capabilities { + worker_heartbeats: true, + poller_autoscaling_auto_enroll: true, + ..Capabilities::default() + }), + ..NamespaceInfo::default() + }), + ..DescribeNamespaceResponse::default() + }) + }); + + let mut config = test_worker_cfg().build().unwrap(); + config.task_types = WorkerTaskTypes::activity_only(); + let worker = worker::Worker::new( + config, + None, + Arc::new(mock), + None, + Some(Duration::from_secs(60)), + ) + .unwrap(); + + let reported_autoscaling = tokio::time::timeout(Duration::from_secs(5), reported_rx) + .await + .expect("worker heartbeat was not recorded in time") + .expect("heartbeat sender was dropped"); + worker.validate().await.unwrap(); + worker.drain_activity_poller_and_shutdown().await; + + assert!( + reported_autoscaling, + "heartbeat emitted before validation must reflect namespace auto-enrollment" + ); + } + #[tokio::test] async fn worker_commands_not_polled_when_capability_disabled() { let mut mock = mock_worker_client(); From 1a847940c9ac15dad327fbd7ca824d32d1640725 Mon Sep 17 00:00:00 2001 From: Spencer Judge Date: Mon, 27 Jul 2026 16:52:56 -0700 Subject: [PATCH 5/5] Ensure various calls to describe namespace are synchronized --- crates/client/src/worker.rs | 132 ++++++++++++++++++- crates/sdk-core/src/core_tests/workers.rs | 12 +- crates/sdk-core/src/pollers/poll_buffer.rs | 20 ++- crates/sdk-core/src/worker/client/mocks.rs | 13 +- crates/sdk-core/src/worker/heartbeat.rs | 18 ++- crates/sdk-core/src/worker/mod.rs | 142 ++++++++++++--------- 6 files changed, 247 insertions(+), 90 deletions(-) diff --git a/crates/client/src/worker.rs b/crates/client/src/worker.rs index 1886b34f4..c9cf39e62 100644 --- a/crates/client/src/worker.rs +++ b/crates/client/src/worker.rs @@ -8,17 +8,21 @@ use std::{ HashMap, hash_map::Entry::{Occupied, Vacant}, }, - sync::Arc, + future::Future, + sync::{Arc, Weak}, }; use temporalio_common::{ protos::{ TaskToken, temporal::api::{ - worker::v1::WorkerHeartbeat, workflowservice::v1::PollWorkflowTaskQueueResponse, + worker::v1::WorkerHeartbeat, + workflowservice::v1::{DescribeNamespaceResponse, PollWorkflowTaskQueueResponse}, }, }, worker::{WorkerDeploymentOptions, WorkerTaskTypes}, }; +use tokio::sync::OnceCell; +use tonic::Code; use uuid::Uuid; /// This trait represents a slot reserved for processing a WFT by a worker. @@ -83,6 +87,8 @@ struct ClientWorkerSetImpl { all_workers: HashMap>, /// Maps namespace to shared worker for worker heartbeating shared_worker: HashMap>, + // Avoid retaining namespace limits and capabilities after the last worker using them is gone. + namespace_descriptions: HashMap>, } impl ClientWorkerSetImpl { @@ -92,9 +98,25 @@ impl ClientWorkerSetImpl { slot_providers: Default::default(), all_workers: Default::default(), shared_worker: Default::default(), + namespace_descriptions: Default::default(), } } + fn namespace_description_source(&mut self, namespace: &str) -> Arc { + if let Some(description) = self + .namespace_descriptions + .get(namespace) + .and_then(Weak::upgrade) + { + return description; + } + + let description = Arc::new(NamespaceDescriptionSource::unresolved()); + self.namespace_descriptions + .insert(namespace.to_owned(), Arc::downgrade(&description)); + description + } + fn try_reserve_wft_slot( &self, namespace: String, @@ -292,6 +314,55 @@ impl ClientWorkerSetImpl { } } +/// A connection-scoped source for a namespace description shared by all workers in that namespace. +#[derive(Debug)] +#[doc(hidden)] +pub struct NamespaceDescriptionSource { + description: OnceCell, +} + +impl NamespaceDescriptionSource { + /// Construct a source whose description has not yet been resolved. + pub fn unresolved() -> Self { + Self { + description: OnceCell::new(), + } + } + + /// Construct a source whose description has already been resolved. + pub fn resolved(description: DescribeNamespaceResponse) -> Self { + Self { + description: OnceCell::new_with(Some(description)), + } + } + + /// Resolve the namespace description once, allowing concurrent callers to await the same RPC. + pub async fn resolve( + &self, + fetch: F, + ) -> Result<&DescribeNamespaceResponse, tonic::Status> + where + F: FnOnce() -> Fut, + Fut: Future>, + { + self.description + .get_or_try_init(|| async { + match fetch().await { + Err(status) if status.code() == Code::Unimplemented => { + Ok(DescribeNamespaceResponse::default()) + } + result => result, + } + }) + .await + } + + /// Return the resolved namespace description, if resolution has completed successfully. + pub fn get(&self) -> Option<&DescribeNamespaceResponse> { + self.description.get() + } +} + /// This trait represents a shared namespace worker that sends worker heartbeats and /// receives worker commands. pub trait SharedNamespaceWorkerTrait { @@ -340,6 +411,14 @@ impl ClientWorkerSet { } } + /// Return the shared namespace description source for this connection and namespace. + #[doc(hidden)] + pub fn namespace_description_source(&self, namespace: &str) -> Arc { + self.worker_manager + .write() + .namespace_description_source(namespace) + } + /// Try to reserve a compatible processing slot in any of the registered workers. /// Returns the slot and the worker's deployment options (if using deployment-based versioning). pub(crate) fn try_reserve_wft_slot( @@ -476,6 +555,55 @@ pub trait ClientWorker: Send + Sync { #[cfg(test)] mod tests { use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + + #[tokio::test] + async fn namespace_description_source_resolves_once() { + let source = NamespaceDescriptionSource::unresolved(); + let calls = AtomicUsize::new(0); + + let first = source.resolve(|| async { + calls.fetch_add(1, Ordering::Relaxed); + tokio::task::yield_now().await; + Ok(DescribeNamespaceResponse::default()) + }); + let second = source.resolve(|| async { + calls.fetch_add(1, Ordering::Relaxed); + Ok(DescribeNamespaceResponse::default()) + }); + + let (first, second) = tokio::join!(first, second); + first.unwrap(); + second.unwrap(); + assert_eq!(calls.load(Ordering::Relaxed), 1); + } + + #[tokio::test] + async fn namespace_description_source_resolves_unimplemented_as_default() { + let source = NamespaceDescriptionSource::unresolved(); + + let description = source + .resolve(|| async { Err(tonic::Status::unimplemented("unsupported")) }) + .await + .unwrap(); + + assert_eq!(description, &DescribeNamespaceResponse::default()); + } + + #[test] + fn namespace_description_sources_are_scoped_by_namespace() { + let workers = ClientWorkerSet::new(); + let first = workers.namespace_description_source("first"); + + assert!(Arc::ptr_eq( + &first, + &workers.namespace_description_source("first") + )); + assert!(!Arc::ptr_eq( + &first, + &workers.namespace_description_source("second") + )); + } fn new_mock_slot(with_error: bool) -> Box { let mut mock_slot = MockSlot::new(); diff --git a/crates/sdk-core/src/core_tests/workers.rs b/crates/sdk-core/src/core_tests/workers.rs index efb2b207c..01d858910 100644 --- a/crates/sdk-core/src/core_tests/workers.rs +++ b/crates/sdk-core/src/core_tests/workers.rs @@ -10,7 +10,7 @@ use crate::{ self, PollerBehavior, client::{ MockWorkerClient, - mocks::{DEFAULT_TEST_CAPABILITIES, DEFAULT_WORKERS_REGISTRY, mock_worker_client}, + mocks::{DEFAULT_TEST_CAPABILITIES, mock_worker_client}, }, }, }; @@ -24,6 +24,7 @@ use std::{ }, time::Duration, }; +use temporalio_client::worker::ClientWorkerSet; use temporalio_common::{ protos::{ coresdk::{ @@ -343,10 +344,10 @@ async fn worker_shutdown_api(#[case] use_cache: bool, #[case] api_success: bool) // This will no longer be needed if // https://github.com/asomers/mockall/issues/283 is implemented. let mut mock = MockWorkerClient::new(); + let workers = Arc::new(ClientWorkerSet::new()); mock.expect_capabilities() .returning(|| Some(*DEFAULT_TEST_CAPABILITIES)); - mock.expect_workers() - .returning(|| DEFAULT_WORKERS_REGISTRY.clone()); + mock.expect_workers().returning(move || workers.clone()); mock.expect_is_mock().returning(|| true); mock.expect_sdk_name_and_version() .returning(|| ("test-core".to_string(), "0.0.0".to_string())); @@ -1239,12 +1240,13 @@ async fn graceful_shutdown_sends_shutdown_worker_rpc_during_initiate() { let poll_releaser_for_rpc = poll_releaser.clone(); let mut mock_client = MockWorkerClient::new(); + let workers = Arc::new(ClientWorkerSet::new()); mock_client .expect_capabilities() .returning(|| Some(*DEFAULT_TEST_CAPABILITIES)); mock_client .expect_workers() - .returning(|| DEFAULT_WORKERS_REGISTRY.clone()); + .returning(move || workers.clone()); mock_client.expect_is_mock().returning(|| true); mock_client .expect_sdk_name_and_version() @@ -1266,7 +1268,7 @@ async fn graceful_shutdown_sends_shutdown_worker_rpc_during_initiate() { hb.worker_identity = "test-identity".to_string(); hb.heartbeat_time = Some(std::time::SystemTime::now().into()); }); - // Return the worker_poll_complete_on_shutdown capability so validate() enables graceful mode + // Return the worker_poll_complete_on_shutdown capability so graceful mode is enabled. mock_client.expect_describe_namespace().returning(move || { Ok(DescribeNamespaceResponse { namespace_info: Some(NamespaceInfo { diff --git a/crates/sdk-core/src/pollers/poll_buffer.rs b/crates/sdk-core/src/pollers/poll_buffer.rs index 65c94dde3..45fe81d03 100644 --- a/crates/sdk-core/src/pollers/poll_buffer.rs +++ b/crates/sdk-core/src/pollers/poll_buffer.rs @@ -864,6 +864,7 @@ mod tests { use futures_util::FutureExt; use rstest::rstest; use std::time::Duration; + use temporalio_common::protos::temporal::api::namespace::v1::namespace_info::Capabilities; use tokio::{select, sync::Notify}; #[tokio::test] @@ -1213,11 +1214,10 @@ mod tests { wft_poller_shared: None, }, Arc::new(AtomicCell::new(None)), - Arc::new({ - let ns = NamespaceCapabilities::default(); - ns.graceful_poll_shutdown.store(graceful, Ordering::Relaxed); - ns - }), + Arc::new(NamespaceCapabilities::resolved(Capabilities { + worker_poll_complete_on_shutdown: graceful, + ..Default::default() + })), ); let first = pb.poll().await.unwrap().unwrap(); @@ -1269,12 +1269,10 @@ mod tests { min: minimum, target: AtomicUsize::new(10), ever_saw_scaling_decision: AtomicBool::new(false), - capabilities: Arc::new({ - let ns = NamespaceCapabilities::default(); - ns.poller_autoscaling - .store(supports_autoscaling, Ordering::Relaxed); - ns - }), + capabilities: Arc::new(NamespaceCapabilities::resolved(Capabilities { + poller_autoscaling: supports_autoscaling, + ..Default::default() + })), behavior: PollerBehavior::Autoscaling { minimum, maximum: 10, diff --git a/crates/sdk-core/src/worker/client/mocks.rs b/crates/sdk-core/src/worker/client/mocks.rs index a67fa3b95..3e649a9cb 100644 --- a/crates/sdk-core/src/worker/client/mocks.rs +++ b/crates/sdk-core/src/worker/client/mocks.rs @@ -1,11 +1,8 @@ use super::*; use futures_util::{Future, FutureExt}; -use std::sync::{Arc, LazyLock}; +use std::sync::Arc; use temporalio_client::worker::ClientWorkerSet; -pub(crate) static DEFAULT_WORKERS_REGISTRY: LazyLock> = - LazyLock::new(|| Arc::new(ClientWorkerSet::new())); - pub(crate) static DEFAULT_TEST_CAPABILITIES: &Capabilities = &Capabilities { signal_and_query_header: true, internal_error_differentiation: true, @@ -26,10 +23,10 @@ pub(crate) static DEFAULT_TEST_CAPABILITIES: &Capabilities = &Capabilities { /// Create a mock client primed with basic necessary expectations pub fn mock_worker_client() -> MockWorkerClient { let mut r = MockWorkerClient::new(); + let workers = Arc::new(ClientWorkerSet::new()); r.expect_capabilities() .returning(|| Some(*DEFAULT_TEST_CAPABILITIES)); - r.expect_workers() - .returning(|| DEFAULT_WORKERS_REGISTRY.clone()); + r.expect_workers().returning(move || workers.clone()); r.expect_is_mock().returning(|| true); r.expect_shutdown_worker() .returning(|_, _, _, _| Ok(ShutdownWorkerResponse {})); @@ -51,10 +48,10 @@ pub fn mock_worker_client() -> MockWorkerClient { /// Create a mock manual client primed with basic necessary expectations pub(crate) fn mock_manual_worker_client() -> MockManualWorkerClient { let mut r = MockManualWorkerClient::new(); + let workers = Arc::new(ClientWorkerSet::new()); r.expect_capabilities() .returning(|| Some(*DEFAULT_TEST_CAPABILITIES)); - r.expect_workers() - .returning(|| DEFAULT_WORKERS_REGISTRY.clone()); + r.expect_workers().returning(move || workers.clone()); r.expect_is_mock().returning(|| true); r.expect_shutdown_worker() .returning(|_, _, _, _| async { Ok(ShutdownWorkerResponse {}) }.boxed()); diff --git a/crates/sdk-core/src/worker/heartbeat.rs b/crates/sdk-core/src/worker/heartbeat.rs index 1bdec5728..e83417389 100644 --- a/crates/sdk-core/src/worker/heartbeat.rs +++ b/crates/sdk-core/src/worker/heartbeat.rs @@ -15,7 +15,9 @@ use std::{ }, time::Duration, }; -use temporalio_client::worker::{SharedNamespaceWorkerTrait, WorkerCallbacks}; +use temporalio_client::worker::{ + NamespaceDescriptionSource, SharedNamespaceWorkerTrait, WorkerCallbacks, +}; use temporalio_common::{ protos::{ TaskToken, @@ -55,6 +57,7 @@ impl SharedNamespaceWorker { namespace: String, heartbeat_interval: Duration, telemetry: Option, + namespace_description: Arc, ) -> Result { let reset_notify = Arc::new(Notify::new()); let cancel = CancellationToken::new(); @@ -68,12 +71,16 @@ impl SharedNamespaceWorker { let cancel = cancel.clone(); let worker_control_task_queue_enabled = worker_control_task_queue_enabled.clone(); async move { - let worker_commands_supported = match client.describe_namespace().await { + let worker_commands_supported = match namespace_description + .resolve(|| client.describe_namespace()) + .await + { Ok(namespace_resp) => { let caps = namespace_resp .namespace_info - .and_then(|info| info.capabilities); - if caps.as_ref().map(|c| c.worker_heartbeats) != Some(true) { + .as_ref() + .and_then(|info| info.capabilities.as_ref()); + if caps.map(|c| c.worker_heartbeats) != Some(true) { debug!( "Worker heartbeating configured for runtime, but server version does not support it." ); @@ -460,7 +467,7 @@ mod tests { } Ok(RecordWorkerHeartbeatResponse {}) }); - mock.expect_describe_namespace().times(2).returning(|| { + mock.expect_describe_namespace().times(1).returning(|| { Ok(DescribeNamespaceResponse { namespace_info: Some(NamespaceInfo { capabilities: Some(Capabilities { @@ -538,6 +545,7 @@ mod tests { namespace, Duration::from_millis(100), None, + Arc::new(temporalio_client::worker::NamespaceDescriptionSource::unresolved()), ) .unwrap(); diff --git a/crates/sdk-core/src/worker/mod.rs b/crates/sdk-core/src/worker/mod.rs index 1c76156dc..739e01d98 100644 --- a/crates/sdk-core/src/worker/mod.rs +++ b/crates/sdk-core/src/worker/mod.rs @@ -10,7 +10,7 @@ mod workflow; /// the error limit, shared so every conversion site reports the same identifier. pub(crate) const PAYLOADS_TOO_LARGE_FAILURE_TYPE: &str = "PayloadsTooLarge"; -use temporalio_client::{Connection, PayloadErrorLimits}; +use temporalio_client::{Connection, PayloadErrorLimits, worker::NamespaceDescriptionSource}; use temporalio_common::{ protos::{ coresdk::{ @@ -85,6 +85,8 @@ use temporalio_client::worker::{ CancelActivityCallback, ClientWorker, HeartbeatCallback, SharedNamespaceWorkerTrait, Slot as SlotTrait, }; +#[cfg(test)] +use temporalio_common::protos::temporal::api::namespace::v1::NamespaceInfo as ApiNamespaceInfo; use temporalio_common::{ protos::{ TaskToken, @@ -98,8 +100,10 @@ use temporalio_common::{ temporal::api::{ deployment, enums::v1::{TaskQueueKind, TaskQueueType, WorkerStatus}, + namespace::v1::namespace_info::Capabilities as ApiNamespaceCapabilities, taskqueue::v1::{StickyExecutionAttributes, TaskQueue}, worker::v1::{WorkerHeartbeat, WorkerHostInfo, WorkerPollerInfo, WorkerSlotsInfo}, + workflowservice::v1::DescribeNamespaceResponse, }, }, telemetry::metrics::TemporalMeter, @@ -450,44 +454,77 @@ pub struct Worker { client_worker_registrator: Arc, /// Status of the worker status: Arc>, - /// Capabilities as returned by a describe namespace rpc. Not set until after validate() is - /// called. + /// Capabilities from the namespace description shared by workers on this connection. capabilities: Arc, /// Handle for the spawned ShutdownWorker RPC task, awaited during shutdown. shutdown_rpc_handle: Mutex>>, } -/// Namespace capabilities discovered via `describe_namespace` during worker validation. -#[derive(Default)] +/// Namespace capabilities discovered via `describe_namespace`. pub struct NamespaceCapabilities { - pub(crate) graceful_poll_shutdown: AtomicBool, - pub(crate) poller_autoscaling: AtomicBool, - pub(crate) poller_autoscaling_auto_enroll: AtomicBool, - pub(crate) worker_commands: AtomicBool, + description: Arc, +} + +impl Default for NamespaceCapabilities { + fn default() -> Self { + Self { + description: Arc::new(NamespaceDescriptionSource::resolved( + DescribeNamespaceResponse::default(), + )), + } + } } impl NamespaceCapabilities { + fn new(description: Arc) -> Self { + Self { description } + } + + #[cfg(test)] + pub(crate) fn resolved(capabilities: ApiNamespaceCapabilities) -> Self { + Self::new(Arc::new(NamespaceDescriptionSource::resolved( + DescribeNamespaceResponse { + namespace_info: Some(ApiNamespaceInfo { + capabilities: Some(capabilities), + ..Default::default() + }), + ..Default::default() + }, + ))) + } + + fn capabilities(&self) -> Option<&ApiNamespaceCapabilities> { + self.description + .get() + .and_then(|description| description.namespace_info.as_ref()) + .and_then(|namespace_info| namespace_info.capabilities.as_ref()) + } + /// Returns true if the server supports graceful poll cancellation on shutdown, so pollers /// can let in-flight polls complete rather than hard-killing them. pub fn graceful_poll_shutdown(&self) -> bool { - self.graceful_poll_shutdown.load(Ordering::Relaxed) + self.capabilities() + .is_some_and(|capabilities| capabilities.worker_poll_complete_on_shutdown) } /// Returns true if pollers may scale down on poll timeout even without an explicit scaling /// decision from the server. pub fn poller_autoscaling(&self) -> bool { - self.poller_autoscaling.load(Ordering::Relaxed) + self.capabilities() + .is_some_and(|capabilities| capabilities.poller_autoscaling) } /// Returns true if the namespace opts workers into poller autoscaling by default. Poller types /// left at their default are automatically enrolled into autoscaling when this is set. pub fn poller_autoscaling_auto_enroll(&self) -> bool { - self.poller_autoscaling_auto_enroll.load(Ordering::Relaxed) + self.capabilities() + .is_some_and(|capabilities| capabilities.poller_autoscaling_auto_enroll) } /// Returns true if worker commands are supported in this namespace. pub fn worker_commands(&self) -> bool { - self.worker_commands.load(Ordering::Relaxed) + self.capabilities() + .is_some_and(|capabilities| capabilities.worker_commands) } } @@ -580,14 +617,21 @@ impl Worker { /// needs to be done asynchronously. Lang SDKs should call this function once before calling /// any others. pub async fn validate(&self) -> Result { - match self.client.describe_namespace().await { + match self + .capabilities + .description + .resolve(|| self.client.describe_namespace()) + .await + { Ok(info) => { - let ns_info = info.namespace_info; - let limits = ns_info.as_ref().and_then(|ns_info| { - ns_info.limits.map(|api_limits| namespace_info::Limits { - blob_size_limit_error: api_limits.blob_size_limit_error, - memo_size_limit_error: api_limits.memo_size_limit_error, - }) + let limits = info.namespace_info.as_ref().and_then(|api_namespace_info| { + api_namespace_info + .limits + .as_ref() + .map(|api_limits| namespace_info::Limits { + blob_size_limit_error: api_limits.blob_size_limit_error, + memo_size_limit_error: api_limits.memo_size_limit_error, + }) }); // Install the namespace error limits on the client (enforced on completions) unless // opted out; warn-level enforcement is always on, configured on the connection. @@ -600,38 +644,11 @@ impl Worker { memo: limits.memo_size_limit_error.max(0) as usize, })); } - if let Some(caps) = ns_info.and_then(|ns| ns.capabilities) { - if caps.worker_poll_complete_on_shutdown { - self.capabilities - .graceful_poll_shutdown - .store(true, Ordering::Relaxed); - } - if caps.poller_autoscaling { - self.capabilities - .poller_autoscaling - .store(true, Ordering::Relaxed); - } - if caps.poller_autoscaling_auto_enroll { - self.capabilities - .poller_autoscaling_auto_enroll - .store(true, Ordering::Relaxed); - } - if caps.worker_commands { - self.capabilities - .worker_commands - .store(true, Ordering::Relaxed); - } - } // Now that capabilities are known, eagerly build the pollers so effective poller // behavior is resolved during the normal check-in path. LazyLock::force(&self.task_subsystems); Ok(NamespaceInfo { limits }) } - Err(e) if e.code() == tonic::Code::Unimplemented => { - // Ignore if unimplemented since we wouldn't want to fail against an old server, for - // example. - Ok(NamespaceInfo::default()) - } Err(e) => Err(WorkerValidationError::NamespaceDescribeError { source: e, namespace: self.config.namespace.clone(), @@ -744,12 +761,10 @@ impl Worker { let wf_sticky_last_suc_poll_time = Arc::new(AtomicCell::new(None)); let act_last_suc_poll_time = Arc::new(AtomicCell::new(None)); let nexus_last_suc_poll_time = Arc::new(AtomicCell::new(None)); - let capabilities = Arc::new(NamespaceCapabilities { - graceful_poll_shutdown: AtomicBool::new(false), - poller_autoscaling: AtomicBool::new(false), - poller_autoscaling_auto_enroll: AtomicBool::new(false), - worker_commands: AtomicBool::new(false), - }); + let namespace_description = client + .workers() + .namespace_description_source(config.namespace.as_str()); + let capabilities = Arc::new(NamespaceCapabilities::new(namespace_description.clone())); let nexus_slots = MeteredPermitDealer::new( tuner.nexus_task_slot_supplier(), @@ -846,6 +861,7 @@ impl Worker { heartbeat_manager: worker_heartbeat, cancel_activity_slot: cancel_activity_slot.clone(), client: RwLock::new(client.clone()), + namespace_description, shared_namespace_worker, task_types: config.task_types, }); @@ -1518,7 +1534,10 @@ impl Worker { self.client.connection() } - /// Returns the namespace capabilities discovered during [Worker::validate]. + /// Returns the namespace capabilities discovered during worker initialization. + /// + /// Lang SDKs should call [Worker::validate] before reading these capabilities when worker + /// heartbeating is disabled, since validation is then what resolves the namespace description. pub fn get_namespace_capabilities(&self) -> &NamespaceCapabilities { &self.capabilities } @@ -2124,6 +2143,7 @@ struct ClientWorkerRegistrator { /// with the worker so it can be resolved lazily once the activity task manager is constructed. cancel_activity_slot: Arc>, client: RwLock>, + namespace_description: Arc, shared_namespace_worker: bool, task_types: WorkerTaskTypes, } @@ -2178,6 +2198,7 @@ impl ClientWorker for ClientWorkerRegistrator { self.namespace().to_string(), hb_mgr.heartbeat_interval, hb_mgr.telemetry.clone(), + self.namespace_description.clone(), )?)) } else { bail!("Shared namespace worker creation never be called without a heartbeat manager"); @@ -2461,7 +2482,10 @@ mod tests { }, }; use futures_util::FutureExt; - use temporalio_common::protos::temporal::api::workflowservice::v1::PollActivityTaskQueueResponse; + use temporalio_common::protos::temporal::api::{ + namespace::v1::namespace_info::Capabilities, + workflowservice::v1::PollActivityTaskQueueResponse, + }; #[tokio::test] async fn activity_timeouts_maintain_permit() { @@ -2705,10 +2729,10 @@ mod tests { } fn auto_enroll_caps() -> NamespaceCapabilities { - let caps = NamespaceCapabilities::default(); - caps.poller_autoscaling_auto_enroll - .store(true, Ordering::Relaxed); - caps + NamespaceCapabilities::resolved(Capabilities { + poller_autoscaling_auto_enroll: true, + ..Default::default() + }) } #[test]