From 2db6b4302232f0c3c1e1049c36cac729a4eee240 Mon Sep 17 00:00:00 2001 From: Thomas Hardy Date: Fri, 16 Oct 2026 14:17:20 -0400 Subject: [PATCH 1/3] Update API subtree from sticky-activity Import temporalio/api sticky-activity at c28c289a9dbcaba8297e5a61108c451b7d20777c. Preserve current Core request handling, serde generation, and payload-limit policy. git-subtree-dir: crates/protos/protos/api_upstream git-subtree-split: c28c289a9dbcaba8297e5a61108c451b7d20777c --- crates/protos/build.rs | 1 + .../protos/protos/api_upstream/buf.gen.yaml | 20 ------------- .../api_upstream/openapi/openapiv2.json | 12 ++++++++ .../api_upstream/openapi/openapiv3.yaml | 19 +++++++++++++ .../temporal/api/command/v1/message.proto | 3 ++ .../temporal/api/schedule/v1/message.proto | 8 ++++++ .../temporal/api/session/v1/message.proto | 18 ++++++++++++ .../workflowservice/v1/request_response.proto | 28 +++++++++++++++++++ crates/protos/src/protos/mod.rs | 6 ++++ crates/sdk-core/src/worker/client.rs | 2 ++ 10 files changed, 97 insertions(+), 20 deletions(-) delete mode 100644 crates/protos/protos/api_upstream/buf.gen.yaml create mode 100644 crates/protos/protos/api_upstream/temporal/api/session/v1/message.proto diff --git a/crates/protos/build.rs b/crates/protos/build.rs index ddcdebf38..68f59ffcb 100644 --- a/crates/protos/build.rs +++ b/crates/protos/build.rs @@ -39,6 +39,7 @@ const SERDE_DERIVE_PREFIXES: &[&str] = &[ ".temporal.api.replication", ".temporal.api.rules", ".temporal.api.schedule", + ".temporal.api.session", ".temporal.api.taskqueue", ".temporal.api.testservice", ".temporal.api.update", diff --git a/crates/protos/protos/api_upstream/buf.gen.yaml b/crates/protos/protos/api_upstream/buf.gen.yaml deleted file mode 100644 index 555cd008d..000000000 --- a/crates/protos/protos/api_upstream/buf.gen.yaml +++ /dev/null @@ -1,20 +0,0 @@ -version: v1 -plugins: - - plugin: buf.build/protocolbuffers/go:v1.31.0 - out: ./ - opt: - - paths=source_relative - - plugin: buf.build/grpc/go:v1.3.0 - out: ./ - opt: - - paths=source_relative - - plugin: buf.build/grpc-ecosystem/gateway:v2.18.0 - out: ./ - opt: - - paths=source_relative - - allow_patch_feature=false - - name: go-helpers - out: ./ - path: ["go", "run", "./protoc-gen-go-helpers"] - opt: - - paths=source_relative diff --git a/crates/protos/protos/api_upstream/openapi/openapiv2.json b/crates/protos/protos/api_upstream/openapi/openapiv2.json index cbb73e756..a1da400b1 100644 --- a/crates/protos/protos/api_upstream/openapi/openapiv2.json +++ b/crates/protos/protos/api_upstream/openapi/openapiv2.json @@ -12917,6 +12917,10 @@ "propagatedNexusSerializationContext": { "$ref": "#/definitions/v1PropagatedSerializationContext", "description": "Serialization context propagated from the Nexus caller that started this workflow." + }, + "timeSkippingStatePropagation": { + "$ref": "#/definitions/v1TimeSkippingStatePropagation", + "description": "Time-skipping state propagated by the execution that started this workflow.\nThis field is intended for use by the Temporal server, for example when a schedule\nstarts a workflow, and is not intended to be set by external clients." } } }, @@ -19591,6 +19595,10 @@ "type": "string", "format": "int64", "description": "Size of the schedule's internal state (including payloads) in bytes." + }, + "timeSkippingInfo": { + "$ref": "#/definitions/v1TimeSkippingInfo", + "description": "Information about time skipping for this schedule.\nIf the schedule has never enabled time skipping, this will be unset." } } }, @@ -19688,6 +19696,10 @@ }, "unpause": { "type": "string" + }, + "timeSkippingConfig": { + "$ref": "#/definitions/v1TimeSkippingConfig", + "description": "If set, applies the specified time-skipping configuration to this schedule,\nreplacing any existing time-skipping configuration." } } }, diff --git a/crates/protos/protos/api_upstream/openapi/openapiv3.yaml b/crates/protos/protos/api_upstream/openapi/openapiv3.yaml index a6465fa4b..ce576907c 100644 --- a/crates/protos/protos/api_upstream/openapi/openapiv3.yaml +++ b/crates/protos/protos/api_upstream/openapi/openapiv3.yaml @@ -16446,6 +16446,12 @@ components: stateSizeBytes: type: string description: Size of the schedule's internal state (including payloads) in bytes. + timeSkippingInfo: + allOf: + - $ref: '#/components/schemas/TimeSkippingInfo' + description: |- + Information about time skipping for this schedule. + If the schedule has never enabled time skipping, this will be unset. ScheduleListEntry: type: object properties: @@ -16517,6 +16523,12 @@ components: notes field to the value of the string. unpause: type: string + timeSkippingConfig: + allOf: + - $ref: '#/components/schemas/TimeSkippingConfig' + description: |- + If set, applies the specified time-skipping configuration to this schedule, + replacing any existing time-skipping configuration. SchedulePolicies: type: object properties: @@ -17846,6 +17858,13 @@ components: allOf: - $ref: '#/components/schemas/PropagatedSerializationContext' description: Serialization context propagated from the Nexus caller that started this workflow. + timeSkippingStatePropagation: + allOf: + - $ref: '#/components/schemas/TimeSkippingStatePropagation' + description: |- + Time-skipping state propagated by the execution that started this workflow. + This field is intended for use by the Temporal server, for example when a schedule + starts a workflow, and is not intended to be set by external clients. StartWorkflowExecutionResponse: type: object properties: diff --git a/crates/protos/protos/api_upstream/temporal/api/command/v1/message.proto b/crates/protos/protos/api_upstream/temporal/api/command/v1/message.proto index ee839115b..b8be64315 100644 --- a/crates/protos/protos/api_upstream/temporal/api/command/v1/message.proto +++ b/crates/protos/protos/api_upstream/temporal/api/command/v1/message.proto @@ -16,6 +16,7 @@ import "temporal/api/enums/v1/command_type.proto"; import "temporal/api/common/v1/message.proto"; import "temporal/api/failure/v1/message.proto"; import "temporal/api/taskqueue/v1/message.proto"; +import "temporal/api/session/v1/message.proto"; import "temporal/api/workflow/v1/message.proto"; import "temporal/api/sdk/v1/user_metadata.proto"; import "temporal/api/sdk/v1/event_group_marker.proto"; @@ -68,6 +69,8 @@ message ScheduleActivityTaskCommandAttributes { // Priority metadata. If this message is not present, or any fields are not // present, they inherit the values from the workflow. temporal.api.common.v1.Priority priority = 14; + + temporal.api.session.v1.ActivitySessionOptions session_options = 15; } message RequestCancelActivityTaskCommandAttributes { diff --git a/crates/protos/protos/api_upstream/temporal/api/schedule/v1/message.proto b/crates/protos/protos/api_upstream/temporal/api/schedule/v1/message.proto index 3a0a9344c..45d5cb1c8 100644 --- a/crates/protos/protos/api_upstream/temporal/api/schedule/v1/message.proto +++ b/crates/protos/protos/api_upstream/temporal/api/schedule/v1/message.proto @@ -310,6 +310,10 @@ message SchedulePatch { // notes field to the value of the string. string pause = 3; string unpause = 4; + + // If set, applies the specified time-skipping configuration to this schedule, + // replacing any existing time-skipping configuration. + temporal.api.common.v1.TimeSkippingConfig time_skipping_config = 5; } message ScheduleInfo { @@ -352,6 +356,10 @@ message ScheduleInfo { // Size of the schedule's internal state (including payloads) in bytes. int64 state_size_bytes = 12; + + // Information about time skipping for this schedule. + // If the schedule has never enabled time skipping, this will be unset. + temporal.api.common.v1.TimeSkippingInfo time_skipping_info = 13; } message Schedule { diff --git a/crates/protos/protos/api_upstream/temporal/api/session/v1/message.proto b/crates/protos/protos/api_upstream/temporal/api/session/v1/message.proto new file mode 100644 index 000000000..e1e6d567e --- /dev/null +++ b/crates/protos/protos/api_upstream/temporal/api/session/v1/message.proto @@ -0,0 +1,18 @@ +syntax = "proto3"; + +package temporal.api.session.v1; + +option go_package = "go.temporal.io/api/session/v1;session"; +option java_package = "io.temporal.api.session.v1"; +option java_multiple_files = true; +option java_outer_classname = "MessageProto"; +option ruby_package = "Temporalio::Api::Session::V1"; +option csharp_namespace = "Temporalio.Api.Session.V1"; + +message ActivitySessionOptions { + oneof session { + bool workflow_sticky_session = 1; + + // later session_id, etc + } +} diff --git a/crates/protos/protos/api_upstream/temporal/api/workflowservice/v1/request_response.proto b/crates/protos/protos/api_upstream/temporal/api/workflowservice/v1/request_response.proto index bfdac483a..892c1f3cf 100644 --- a/crates/protos/protos/api_upstream/temporal/api/workflowservice/v1/request_response.proto +++ b/crates/protos/protos/api_upstream/temporal/api/workflowservice/v1/request_response.proto @@ -219,8 +219,14 @@ message StartWorkflowExecutionRequest { // Time-skipping configuration. If not set, time skipping is disabled. temporal.api.common.v1.TimeSkippingConfig time_skipping_config = 29; + // Serialization context propagated from the Nexus caller that started this workflow. temporal.api.nexus.v1.PropagatedSerializationContext propagated_nexus_serialization_context = 30; + + // Time-skipping state propagated by the execution that started this workflow. + // This field is intended for use by the Temporal server, for example when a schedule + // starts a workflow, and is not intended to be set by external clients. + temporal.api.common.v1.TimeSkippingStatePropagation time_skipping_state_propagation = 31; } message StartWorkflowExecutionResponse { @@ -317,6 +323,17 @@ message PollWorkflowTaskQueueRequest { // Removed in 1.55.0; was temporal.api.worker.v1.WorkerHeartbeat worker_heartbeat reserved 7; reserved "worker_heartbeat"; + + // EXPERIMENTAL. Subject to removal in the future. + // This flag is only applicable to sticky polls. In this mode, server will forward the workflow task sticky polls + // to the normal queue only once. The forwarding to normal queue will stop once a single task is dispatched from the + // normal queue to pollers of this particular sticky queue. + // Server persists the workflow ID of the first matched task in the sticky task queue's metadata. This will allow + // system to enforce single-workflow mode even when the server crashes, or a worker with stable identity restarts + // (when worker's sticky queue name remains the same after restart.) + // For activity sticky polls, this mode will disable forwarding to normal queue completely. + // (-- api-linter: core::0140::prepositions=disabled + bool single_workflow_mode = 11; } message PollWorkflowTaskQueueResponse { @@ -547,6 +564,17 @@ message PollActivityTaskQueueRequest { // Removed in 1.55.0; was temporal.api.worker.v1.WorkerHeartbeat worker_heartbeat reserved 7; reserved "worker_heartbeat"; + + // EXPERIMENTAL. Subject to removal in the future. + // This flag is only applicable to sticky polls. In this mode, server will forward the workflow task sticky polls + // to the normal queue only once. The forwarding to normal queue will stop once a single task is dispatched from the + // normal queue to pollers of this particular sticky queue. + // Server persists the workflow ID of the first matched task in the sticky task queue's metadata. This will allow + // system to enforce single-workflow mode even when the server crashes, or a worker with stable identity restarts + // (when worker's sticky queue name remains the same after restart.) + // For activity sticky polls, this mode will disable forwarding to normal queue completely. + // (-- api-linter: core::0140::prepositions=disabled + bool single_workflow_mode = 11; } message PollActivityTaskQueueResponse { diff --git a/crates/protos/src/protos/mod.rs b/crates/protos/src/protos/mod.rs index 3fca443f2..0ae9d430d 100644 --- a/crates/protos/src/protos/mod.rs +++ b/crates/protos/src/protos/mod.rs @@ -1910,6 +1910,7 @@ pub mod temporal { request_eager_execution: !s.do_not_eagerly_execute, use_workflow_build_id, priority: s.priority, + session_options: None, }, ) } @@ -2614,6 +2615,11 @@ pub mod temporal { include_proto_with_serde!("temporal.api.sdk.v1"); } } + pub mod session { + pub mod v1 { + tonic::include_proto!("temporal.api.session.v1"); + } + } pub mod taskqueue { pub mod v1 { tonic::include_proto!("temporal.api.taskqueue.v1"); diff --git a/crates/sdk-core/src/worker/client.rs b/crates/sdk-core/src/worker/client.rs index 856effc18..6f32b4422 100644 --- a/crates/sdk-core/src/worker/client.rs +++ b/crates/sdk-core/src/worker/client.rs @@ -496,6 +496,7 @@ impl WorkerClient for WorkerClientBag { worker_instance_key: self.worker_instance_key.to_string(), worker_control_task_queue: self.worker_control_task_queue(), poller_group_id: Default::default(), + single_workflow_mode: false, } .into_request(); request.extensions_mut().insert(IsWorkerTaskLongPoll); @@ -536,6 +537,7 @@ impl WorkerClient for WorkerClientBag { worker_instance_key: self.worker_instance_key.to_string(), worker_control_task_queue: self.worker_control_task_queue(), poller_group_id: Default::default(), + single_workflow_mode: false, } .into_request(); request.extensions_mut().insert(IsWorkerTaskLongPoll); From 6135ff34918ef1091f2ff32462574ebfedd7d698 Mon Sep 17 00:00:00 2001 From: Thomas Hardy Date: Fri, 16 Oct 2026 14:03:02 -0400 Subject: [PATCH 2/3] Add workflow-sticky activity session support Share workflow-style polling admission, activity capacity, and rate limiting. Protect built-in paired reservations, preserve workflow shutdown semantics, and scope command normalization and validation to session activities. --- .../workflow_commands/workflow_commands.proto | 2 + crates/protos/src/protos/mod.rs | 2 +- crates/sdk-core/CHANGELOG.md | 2 + crates/sdk-core/src/abstractions.rs | 17 +- .../sdk-core/src/core_tests/activity_tasks.rs | 63 ++ crates/sdk-core/src/lib.rs | 13 +- crates/sdk-core/src/pollers/mod.rs | 7 +- crates/sdk-core/src/pollers/poll_buffer.rs | 538 ++++++++++++++++-- crates/sdk-core/src/pollers/sticky.rs | 141 +++++ crates/sdk-core/src/replay/mod.rs | 66 ++- crates/sdk-core/src/telemetry/metrics.rs | 41 ++ crates/sdk-core/src/worker/activities.rs | 7 +- crates/sdk-core/src/worker/client.rs | 139 ++++- crates/sdk-core/src/worker/mod.rs | 145 +++-- .../sdk-core/src/worker/tuner/fixed_size.rs | 6 + .../src/worker/tuner/resource_based.rs | 76 ++- .../workflow/machines/workflow_machines.rs | 105 +++- .../src/worker/workflow/wft_poller.rs | 154 +---- 18 files changed, 1271 insertions(+), 253 deletions(-) create mode 100644 crates/sdk-core/src/pollers/sticky.rs diff --git a/crates/protos/protos/local/temporal/sdk/core/workflow_commands/workflow_commands.proto b/crates/protos/protos/local/temporal/sdk/core/workflow_commands/workflow_commands.proto index 2db4ccfb6..9fe20a71c 100644 --- a/crates/protos/protos/local/temporal/sdk/core/workflow_commands/workflow_commands.proto +++ b/crates/protos/protos/local/temporal/sdk/core/workflow_commands/workflow_commands.proto @@ -17,6 +17,7 @@ import "temporal/api/failure/v1/message.proto"; import "temporal/api/sdk/v1/user_metadata.proto"; import "temporal/api/sdk/v1/event_group_marker.proto"; import "temporal/api/workflow/v1/message.proto"; +import "temporal/api/session/v1/message.proto"; import "temporal/sdk/core/child_workflow/child_workflow.proto"; import "temporal/sdk/core/nexus/nexus.proto"; import "temporal/sdk/core/common/common.proto"; @@ -106,6 +107,7 @@ message ScheduleActivity { coresdk.common.VersioningIntent versioning_intent = 15; // The Priority to use for this activity temporal.api.common.v1.Priority priority = 16; + temporal.api.session.v1.ActivitySessionOptions session_options = 17; } message ScheduleLocalActivity { diff --git a/crates/protos/src/protos/mod.rs b/crates/protos/src/protos/mod.rs index 0ae9d430d..a8e3f09e2 100644 --- a/crates/protos/src/protos/mod.rs +++ b/crates/protos/src/protos/mod.rs @@ -1910,7 +1910,7 @@ pub mod temporal { request_eager_execution: !s.do_not_eagerly_execute, use_workflow_build_id, priority: s.priority, - session_options: None, + session_options: s.session_options, }, ) } diff --git a/crates/sdk-core/CHANGELOG.md b/crates/sdk-core/CHANGELOG.md index 7c34c69a2..11db13415 100644 --- a/crates/sdk-core/CHANGELOG.md +++ b/crates/sdk-core/CHANGELOG.md @@ -60,6 +60,8 @@ relevant information. make progress after unpause. ### Added +* Experimental workflow-affine activity support allows opted-in activities to use the workflow + worker's sticky destination while sharing activity concurrency and rate limits. * The OTLP metrics exporter can export histograms as OpenTelemetry exponential histograms, which Prometheus-compatible backends store as native histograms, when an SDK enables the option, including through the C bridge. It applies to every histogram, including custom ones, and names diff --git a/crates/sdk-core/src/abstractions.rs b/crates/sdk-core/src/abstractions.rs index d08186ec5..81148ccc5 100644 --- a/crates/sdk-core/src/abstractions.rs +++ b/crates/sdk-core/src/abstractions.rs @@ -35,9 +35,9 @@ pub(crate) struct MeteredPermitDealer { max_permits: Option, metrics_ctx: MetricsContext, meter: Option, - /// Only applies to permit dealers for workflow tasks. True if this permit dealer is associated - /// with a sticky queue poller. + /// Reservations must identify their source before a task has been received. is_sticky_poller: bool, + is_paired_activity_poller: bool, context_data: Arc, } @@ -100,6 +100,7 @@ where meter, max_permits, is_sticky_poller: false, + is_paired_activity_poller: false, context_data, } } @@ -205,6 +206,14 @@ impl MeteredPermitDealer { } } +impl MeteredPermitDealer { + pub(crate) fn for_activity_poller(mut self, sticky: bool, paired: bool) -> Self { + self.is_sticky_poller = sticky; + self.is_paired_activity_poller = paired; + self + } +} + impl SlotReservationContext for MeteredPermitDealer { fn task_queue(&self) -> &str { &self.context_data.task_queue @@ -226,6 +235,10 @@ impl SlotReservationContext for MeteredPermitDealer { self.is_sticky_poller } + fn is_paired_activity_poller(&self) -> bool { + self.is_paired_activity_poller + } + fn get_metrics_meter(&self) -> Option { self.meter.clone() } diff --git a/crates/sdk-core/src/core_tests/activity_tasks.rs b/crates/sdk-core/src/core_tests/activity_tasks.rs index 85156e1da..512ddf236 100644 --- a/crates/sdk-core/src/core_tests/activity_tasks.rs +++ b/crates/sdk-core/src/core_tests/activity_tasks.rs @@ -58,6 +58,7 @@ use temporalio_common::{ command::v1::{ScheduleActivityTaskCommandAttributes, command::Attributes}, enums::v1::{ApplicationErrorCategory, EventType}, failure::v1::{ApplicationFailureInfo, Failure, failure::FailureInfo}, + session::v1::{ActivitySessionOptions, activity_session_options}, workflowservice::v1::{ GetSystemInfoResponse, PollActivityTaskQueueResponse, RecordActivityTaskHeartbeatRequest, RecordActivityTaskHeartbeatResponse, @@ -1345,6 +1346,68 @@ async fn oversized_heartbeat_fails_activity() { core.drain_activity_poller_and_shutdown().await; } +#[rstest::rstest] +#[tokio::test] +async fn workflow_sticky_activity_session_commands( + #[values("enabled", "default_queue", "false_session")] scenario: &'static str, +) { + let mut history = TestHistoryBuilder::default(); + history.add_by_type(EventType::WorkflowExecutionStarted); + history.add_full_wf_task(); + let mut client = mock_worker_client(); + client + .expect_complete_workflow_task() + .times(1) + .returning(move |request, _| { + assert!(matches!( + &request.commands[0].attributes, + Some(Attributes::ScheduleActivityTaskCommandAttributes(attrs)) + if attrs.session_options.as_ref().unwrap().session + == Some(activity_session_options::Session::WorkflowStickySession( + scenario != "false_session" + )) + && !attrs.request_eager_execution + && !attrs.task_queue.as_ref().unwrap().name.is_empty() + )); + Ok(RespondWorkflowTaskCompletedResponse::default()) + }); + let mut mocks = single_hist_mock_sg("sticky_activity_session", history, [1], client, true); + mocks.worker_cfg(|config| { + config.max_cached_workflows = 2; + config.session_worker_options = (scenario != "false_session").then(Default::default); + config.max_eager_activity_reservations_per_workflow_task = 0; + }); + let core = mock_worker(mocks); + let activation = core.poll_workflow_activation().await.unwrap(); + let result = core + .complete_workflow_activation(WorkflowActivationCompletion::from_cmds( + activation.run_id.clone(), + vec![ + ScheduleActivity { + seq: 1, + activity_id: "affine_activity".to_owned(), + activity_type: "affine_activity".to_owned(), + task_queue: if scenario == "default_queue" { + String::new() + } else { + core.get_config().task_queue.clone() + }, + start_to_close_timeout: Some(prost_dur!(from_secs(10))), + session_options: Some(ActivitySessionOptions { + session: Some(activity_session_options::Session::WorkflowStickySession( + scenario != "false_session", + )), + }), + ..Default::default() + } + .into(), + ], + )) + .await; + result.unwrap(); + core.drain_pollers_and_shutdown().await; +} + #[tokio::test] async fn max_tq_acts_set_passed_to_poll_properly() { let rate = 9.28; diff --git a/crates/sdk-core/src/lib.rs b/crates/sdk-core/src/lib.rs index e7608a603..751469455 100644 --- a/crates/sdk-core/src/lib.rs +++ b/crates/sdk-core/src/lib.rs @@ -36,7 +36,8 @@ mod core_tests; pub mod test_help; pub use crate::worker::client::{ - PollActivityOptions, PollOptions, PollWorkflowOptions, WorkerClient, WorkflowTaskCompletion, + PollActivityOptions, PollActivityOptionsBuilder, PollOptions, PollWorkflowOptions, + PollWorkflowOptionsBuilder, WorkerClient, WorkflowTaskCompletion, }; pub use pollers::{Client, ClientOptions, ClientTlsOptions, RetryOptions, TlsOptions}; pub use temporalio_common::protos::TaskToken; @@ -46,11 +47,11 @@ pub use worker::{ FixedSizeSlotSupplier, LocalActivitySlotKind, NamespaceCapabilities, NexusSlotKind, PollError, PollerBehavior, ResourceBasedSlotsOptions, ResourceBasedSlotsOptionsBuilder, ResourceBasedTuner, ResourceBasedTunerConfig, ResourceController, ResourceSlotOptions, - SlotInfo, SlotInfoTrait, SlotKind, SlotKindType, SlotMarkUsedContext, SlotReleaseContext, - SlotReservationContext, SlotSupplier, SlotSupplierOptions, SlotSupplierPermit, TunerBuilder, - TunerHolder, TunerHolderOptions, TunerHolderOptionsBuilder, Worker, WorkerConfig, - WorkerConfigBuilder, WorkerTuner, WorkerValidationError, WorkerVersioningStrategy, - WorkflowErrorType, WorkflowSlotKind, + SessionWorkerOptions, SessionWorkerOptionsBuilder, SlotInfo, SlotInfoTrait, SlotKind, + SlotKindType, SlotMarkUsedContext, SlotReleaseContext, SlotReservationContext, SlotSupplier, + SlotSupplierOptions, SlotSupplierPermit, TunerBuilder, TunerHolder, TunerHolderOptions, + TunerHolderOptionsBuilder, Worker, WorkerConfig, WorkerConfigBuilder, WorkerTuner, + WorkerValidationError, WorkerVersioningStrategy, WorkflowErrorType, WorkflowSlotKind, }; use crate::{ diff --git a/crates/sdk-core/src/pollers/mod.rs b/crates/sdk-core/src/pollers/mod.rs index 0a1e2f616..0ed96fb78 100644 --- a/crates/sdk-core/src/pollers/mod.rs +++ b/crates/sdk-core/src/pollers/mod.rs @@ -1,7 +1,12 @@ mod poll_buffer; +mod sticky; +pub(crate) use sticky::StickyPollerShared; + +#[cfg(test)] +pub(crate) use poll_buffer::{ActivityTaskOptions, PollRateLimiter}; pub(crate) use poll_buffer::{ - ActivityTaskOptions, LongPollBuffer, WorkflowTaskOptions, WorkflowTaskPoller, + LongPollBuffer, StickyTaskPoller, WorkflowTaskOptions, make_activity_task_buffer, }; pub use temporalio_client::{Client, ClientOptions, ClientTlsOptions, RetryOptions, TlsOptions}; diff --git a/crates/sdk-core/src/pollers/poll_buffer.rs b/crates/sdk-core/src/pollers/poll_buffer.rs index b35249bc1..ff933e2fd 100644 --- a/crates/sdk-core/src/pollers/poll_buffer.rs +++ b/crates/sdk-core/src/pollers/poll_buffer.rs @@ -1,9 +1,10 @@ use crate::{ + WorkerConfig, abstractions::{ActiveCounter, MeteredPermitDealer, OwnedMeteredSemPermit, dbg_panic}, - pollers::{self, Poller}, + pollers::{self, Poller, StickyPollerShared}, worker::{ ActivitySlotKind, NamespaceCapabilities, NexusSlotKind, PollerBehavior, SlotKind, - WFTPollerShared, WorkflowSlotKind, + WorkflowSlotKind, client::{ PollActivityOptions, PollNexusOptions, PollOptions, PollWorkflowOptions, WorkerClient, }, @@ -68,13 +69,13 @@ const PERSISTENT_POLL_ERROR_WARN_BACKOFF: ExponentialBuilder = ExponentialBuilde type PollReceiver = Mutex)>>>; -struct PollRateLimiter { +pub(crate) struct PollRateLimiter { interval: Duration, next_allowed_at: Mutex, } impl PollRateLimiter { - fn new(polls_per_second: f64) -> Self { + pub(crate) fn new(polls_per_second: f64) -> Self { Self { interval: Duration::from_secs_f64(polls_per_second.recip()), next_allowed_at: Mutex::new(Instant::now()), @@ -102,15 +103,78 @@ pub(crate) struct LongPollBuffer { did_start: AtomicBool, } +#[derive(Default)] pub(crate) struct WorkflowTaskOptions { - pub(crate) wft_poller_shared: Option>, + pub(crate) wft_poller_shared: Option>, } +#[derive(Default)] pub(crate) struct ActivityTaskOptions { /// Local ratelimit - pub(crate) max_worker_acts_per_second: Option, + pub(crate) rate_limiter: Option>, /// Server ratelimit pub(crate) max_tps: Option, + pub(crate) sticky_queue_name: Option, + pub(crate) poller_shared: Option>, +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn make_activity_task_buffer( + config: &WorkerConfig, + sticky_queue_name: &Option, + client: &Arc, + metrics: &crate::MetricsContext, + shutdown: &CancellationToken, + slots: &MeteredPermitDealer, + last_successful_poll_time: Arc>>, + capabilities: Arc, +) -> StickyTaskPoller { + let sticky_queue_name = sticky_queue_name + .as_ref() + .filter(|_| config.session_worker_options.is_some()); + let shared = + sticky_queue_name.map(|_| Arc::new(StickyPollerShared::new(slots.available_permits()))); + let rate_limiter = config + .max_worker_activities_per_second + .map(|rate| Arc::new(PollRateLimiter::new(rate))); + let behavior = crate::worker::resolve_effective_behavior( + config.activity_task_poller_behavior, + &capabilities, + ); + let make = |sticky: Option| { + let behavior = if sticky_queue_name.is_some() { + crate::worker::sticky_poller_behavior(behavior, config, sticky.is_some()) + } else { + behavior + }; + let metrics = metrics.with_new_attrs([if sticky.is_some() { + crate::telemetry::metrics::activity_sticky_poller() + } else { + crate::telemetry::metrics::activity_poller() + }]); + LongPollBuffer::new_activity_task( + client.clone(), + config.task_queue.clone(), + behavior, + slots + .clone() + .for_activity_poller(sticky.is_some(), sticky_queue_name.is_some()), + shutdown.child_token(), + Some(move |np| metrics.record_num_pollers(np)), + ActivityTaskOptions { + rate_limiter: rate_limiter.clone(), + max_tps: config.max_task_queue_activities_per_second, + sticky_queue_name: sticky, + poller_shared: shared.clone(), + }, + last_successful_poll_time.clone(), + capabilities.clone(), + ) + }; + StickyTaskPoller::new( + Some(make(None)), + sticky_queue_name.map(|name| make(Some(name.clone()))), + ) } impl LongPollBuffer { @@ -180,7 +244,9 @@ impl LongPollBuffer { no_retry, timeout_override, }, - PollWorkflowOptions { sticky_queue_name }, + PollWorkflowOptions::builder() + .maybe_sticky_queue_name(sticky_queue_name) + .build(), ) .await } @@ -210,15 +276,36 @@ impl LongPollBuffer { last_successful_poll_time: Arc>>, capabilities: Arc, ) -> Self { - let pre_permit_delay = options - .max_worker_acts_per_second - .map(|ps| Arc::new(PollRateLimiter::new(ps))) - .map(|rl| { - move || { - let rl = rl.clone(); - async move { rl.wait().await }.boxed() + let is_sticky = options.sticky_queue_name.is_some(); + let poll_scaler = PollScaler::new( + poller_behavior, + num_pollers_handler, + shutdown.clone(), + last_successful_poll_time, + capabilities.clone(), + ); + if let Some(shared) = &options.poller_shared { + if is_sticky { + shared.set_sticky_active( + poll_scaler.active_rx.clone(), + poll_scaler.report_handle.target.subscribe(), + ); + } else { + shared.set_non_sticky_active(poll_scaler.active_rx.clone()); + } + } + let pre_permit_delay = Some(move || { + let shared = options.poller_shared.clone(); + let rate_limiter = options.rate_limiter.clone(); + async move { + if let Some(shared) = shared { + shared.wait_if_needed(is_sticky).await; } - }); + if let Some(rate_limiter) = rate_limiter { + rate_limiter.wait().await; + } + } + }); let no_retry = if matches!(poller_behavior, PollerBehavior::Autoscaling { .. }) { Some(NoRetryOnMatching { predicate: poll_scaling_error_matcher, @@ -229,6 +316,7 @@ impl LongPollBuffer { let poll_fn = move |timeout_override: Option| { let client = client.clone(); let task_queue = task_queue.clone(); + let sticky_queue_name = options.sticky_queue_name.clone(); async move { client .poll_activity_task( @@ -237,21 +325,15 @@ impl LongPollBuffer { no_retry, timeout_override, }, - PollActivityOptions { - max_tasks_per_sec: options.max_tps, - }, + PollActivityOptions::builder() + .maybe_max_tasks_per_sec(options.max_tps) + .maybe_sticky_queue_name(sticky_queue_name) + .build(), ) .await } }; - let poll_scaler = PollScaler::new( - poller_behavior, - num_pollers_handler, - shutdown.clone(), - last_successful_poll_time, - capabilities.clone(), - ); Self::new( poll_fn, permit_dealer, @@ -749,48 +831,43 @@ impl PollScalerReportHandle { } } -/// A poller capable of polling on a sticky and a nonsticky queue simultaneously for workflow tasks. +/// Preserves workflow polling's first-result termination semantics for both task types. #[derive(derive_more::Constructor)] -pub(crate) struct WorkflowTaskPoller { - normal_poller: PollWorkflowTaskBuffer, - sticky_poller: Option, +pub(crate) struct StickyTaskPoller { + normal_poller: Option>, + sticky_poller: Option>, } -type PollWorkflowTaskBuffer = LongPollBuffer; #[async_trait::async_trait] -impl - Poller<( - PollWorkflowTaskQueueResponse, - OwnedMeteredSemPermit, - )> for WorkflowTaskPoller +impl Poller<(T, OwnedMeteredSemPermit)> for StickyTaskPoller +where + T: TaskPollerResult + Send + Sync + Debug + 'static, + SK: SlotKind + 'static, { - async fn poll( - &self, - ) -> Option< - pollers::Result<( - PollWorkflowTaskQueueResponse, - OwnedMeteredSemPermit, - )>, - > { - if let Some(sq) = self.sticky_poller.as_ref() { - tokio::select! { - r = self.normal_poller.poll() => r, + async fn poll(&self) -> Option)>> { + match (self.normal_poller.as_ref(), self.sticky_poller.as_ref()) { + (Some(normal), Some(sq)) => tokio::select! { + r = normal.poll() => r, r = sq.poll() => r, - } - } else { - self.normal_poller.poll().await + }, + (Some(source), None) | (None, Some(source)) => source.poll().await, + (None, None) => None, } } fn notify_shutdown(&self) { - self.normal_poller.notify_shutdown(); + if let Some(normal) = self.normal_poller.as_ref() { + normal.notify_shutdown(); + } if let Some(sq) = self.sticky_poller.as_ref() { sq.notify_shutdown(); } } - async fn shutdown(mut self) { - self.normal_poller.shutdown().await; + async fn shutdown(self) { + if let Some(normal) = self.normal_poller { + normal.shutdown().await; + } if let Some(sq) = self.sticky_poller { sq.shutdown().await; } @@ -915,8 +992,13 @@ impl PersistentPollErrorWarningState { mod tests { use super::*; use crate::{ + WorkerVersioningStrategy, abstractions::tests::fixed_size_permit_dealer, - worker::client::mocks::mock_manual_worker_client, + test_help::test_worker_cfg, + worker::{ + ResourceBasedSlotsOptions, ResourceController, ResourceSlotOptions, SystemResourceInfo, + client::mocks::mock_manual_worker_client, + }, }; use futures_util::FutureExt; use rstest::rstest; @@ -924,6 +1006,348 @@ mod tests { use temporalio_common::protos::temporal::api::namespace::v1::namespace_info::Capabilities; use tokio::{select, sync::Notify}; + struct ExhaustedResources; + + impl SystemResourceInfo for ExhaustedResources { + fn total_mem(&self) -> u64 { + 100 + } + + fn used_mem(&self) -> u64 { + 100 + } + + fn used_cpu_percent(&self) -> f64 { + 1.0 + } + } + + #[tokio::test(start_paused = true)] + async fn paired_activity_resource_capacity_allows_both_poll_sources() { + let (tx, mut rx) = unbounded_channel(); + let mut client = mock_manual_worker_client(); + client + .expect_poll_activity_task() + .times(2) + .returning(move |_, activity| { + tx.send(activity.sticky_queue_name.is_some()).unwrap(); + std::future::pending().boxed() + }); + let client: Arc = Arc::new(client); + let mut config = test_worker_cfg().build().unwrap(); + config.session_worker_options = Some(crate::SessionWorkerOptions::default()); + config.activity_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(2)); + let supplier = Arc::new(ResourceController::new_with_sysinfo( + ResourceBasedSlotsOptions::builder() + .target_mem_usage(0.5) + .target_cpu_usage(0.5) + .build(), + Arc::new(ExhaustedResources), + )) + .as_kind(ResourceSlotOptions::new(0, 1, Duration::from_secs(60))); + let slots = MeteredPermitDealer::new( + supplier, + crate::MetricsContext::no_op(), + None, + Arc::new(Default::default()), + None, + ); + assert!(slots.try_acquire_owned().is_err()); + let poller = Arc::new(make_activity_task_buffer( + &config, + &Some("sticky".to_owned()), + &client, + &crate::MetricsContext::no_op(), + &CancellationToken::new(), + &slots, + Arc::new(AtomicCell::new(None)), + Arc::new(NamespaceCapabilities::default()), + )); + let polling = tokio::spawn({ + let poller = poller.clone(); + async move { poller.poll().await } + }); + let first = tokio::time::timeout(Duration::from_secs(5), rx.recv()) + .await + .unwrap() + .unwrap(); + let second = tokio::time::timeout(Duration::from_secs(5), rx.recv()) + .await + .unwrap() + .unwrap(); + assert_ne!(first, second); + poller.notify_shutdown(); + assert!(polling.await.unwrap().is_none()); + match Arc::try_unwrap(poller) { + Ok(poller) => poller.shutdown().await, + Err(_) => panic!("poller still referenced"), + } + let _normal = slots + .clone() + .for_activity_poller(false, true) + .try_acquire_owned() + .unwrap(); + let _sticky = slots + .clone() + .for_activity_poller(true, true) + .try_acquire_owned() + .unwrap(); + } + + #[rstest] + #[case(false)] + #[case(true)] + #[tokio::test] + async fn activity_sticky_polling_shares_capacity_and_options(#[case] enabled: bool) { + let (tx, mut rx) = unbounded_channel(); + let mut client = mock_manual_worker_client(); + client + .expect_poll_activity_task() + .times(2) + .returning(move |poll, activity| { + tx.send((poll, activity)).unwrap(); + std::future::pending().boxed() + }); + let client: Arc = Arc::new(client); + let mut config = test_worker_cfg().build().unwrap(); + config.session_worker_options = enabled.then(crate::SessionWorkerOptions::default); + config.activity_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(10)); + config.max_task_queue_activities_per_second = Some(7.0); + let slots = fixed_size_permit_dealer(2); + let poller = Arc::new(make_activity_task_buffer( + &config, + &Some("sticky".to_owned()), + &client, + &crate::MetricsContext::no_op(), + &CancellationToken::new(), + &slots, + Arc::new(AtomicCell::new(None)), + Arc::new(NamespaceCapabilities::default()), + )); + let polling = tokio::spawn({ + let poller = poller.clone(); + async move { poller.poll().await } + }); + let mut sticky_polls = 0; + for _ in 0..2 { + let (poll, activity) = tokio::time::timeout(Duration::from_secs(5), rx.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(poll.task_queue, config.task_queue); + assert_eq!(activity.max_tasks_per_sec, Some(7.0)); + if let Some(sticky) = activity.sticky_queue_name { + assert_eq!(sticky, "sticky"); + sticky_polls += 1; + } + } + assert_eq!(sticky_polls, usize::from(enabled)); + assert_eq!(slots.available_permits(), Some(0)); + assert!(rx.try_recv().is_err()); + poller.notify_shutdown(); + assert!(polling.await.unwrap().is_none()); + match Arc::try_unwrap(poller) { + Ok(poller) => poller.shutdown().await, + Err(_) => panic!("poller still referenced"), + } + assert_eq!(slots.available_permits(), Some(2)); + } + + #[tokio::test(start_paused = true)] + async fn activity_sticky_polling_uses_aggregate_worker_rate_limit() { + let (tx, mut rx) = unbounded_channel(); + let mut client = mock_manual_worker_client(); + client + .expect_poll_activity_task() + .times(2) + .returning(move |_, options| { + tx.send(( + options.sticky_queue_name.is_some(), + tokio::time::Instant::now(), + )) + .unwrap(); + std::future::pending().boxed() + }); + let client: Arc = Arc::new(client); + let mut config = test_worker_cfg().build().unwrap(); + config.session_worker_options = Some(crate::SessionWorkerOptions::default()); + config.activity_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(2)); + config.max_worker_activities_per_second = Some(2.0); + let poller = Arc::new(make_activity_task_buffer( + &config, + &Some("sticky".to_owned()), + &client, + &crate::MetricsContext::no_op(), + &CancellationToken::new(), + &fixed_size_permit_dealer(2), + Arc::new(AtomicCell::new(None)), + Arc::new(NamespaceCapabilities::default()), + )); + let polling = tokio::spawn({ + let poller = poller.clone(); + async move { poller.poll().await } + }); + let first = rx.recv().await.unwrap(); + let second = rx.recv().await.unwrap(); + assert_ne!(first.0, second.0); + assert!(second.1.duration_since(first.1) >= Duration::from_millis(500)); + poller.notify_shutdown(); + assert!(polling.await.unwrap().is_none()); + match Arc::try_unwrap(poller) { + Ok(poller) => poller.shutdown().await, + Err(_) => panic!("poller still referenced"), + } + } + + #[rstest] + #[case(false)] + #[case(true)] + #[tokio::test] + async fn activity_sticky_polling_forwards_shutdown_empty(#[case] sticky_first: bool) { + let (tx, mut rx) = unbounded_channel(); + let mut client = mock_manual_worker_client(); + client + .expect_poll_activity_task() + .times(2) + .returning(move |_, options| { + let (response_tx, response_rx) = tokio::sync::oneshot::channel(); + tx.send((options.sticky_queue_name.is_some(), response_tx)) + .unwrap(); + async move { Ok(response_rx.await.unwrap()) }.boxed() + }); + let client: Arc = Arc::new(client); + let mut config = test_worker_cfg().build().unwrap(); + config.session_worker_options = Some(crate::SessionWorkerOptions::default()); + config.activity_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(2)); + let slots = fixed_size_permit_dealer(2); + let poller = Arc::new(make_activity_task_buffer( + &config, + &Some("sticky".to_owned()), + &client, + &crate::MetricsContext::no_op(), + &CancellationToken::new(), + &slots, + Arc::new(AtomicCell::new(None)), + Arc::new(NamespaceCapabilities::resolved(Capabilities { + worker_poll_complete_on_shutdown: true, + ..Default::default() + })), + )); + let polling = tokio::spawn({ + let poller = poller.clone(); + async move { poller.poll().await } + }); + let first = tokio::time::timeout(Duration::from_secs(5), rx.recv()) + .await + .unwrap() + .unwrap(); + let second = tokio::time::timeout(Duration::from_secs(5), rx.recv()) + .await + .unwrap() + .unwrap(); + let (empty, task) = if first.0 == sticky_first { + (first.1, second.1) + } else { + (second.1, first.1) + }; + poller.notify_shutdown(); + empty + .send(PollActivityTaskQueueResponse::default()) + .unwrap(); + let (response, permit) = polling.await.unwrap().unwrap().unwrap(); + assert!(response.task_token.is_empty()); + drop(permit); + let other_source = if sticky_first { + poller.normal_poller.as_ref().unwrap() + } else { + poller.sticky_poller.as_ref().unwrap() + }; + task.send(PollActivityTaskQueueResponse { + task_token: vec![1], + ..Default::default() + }) + .unwrap(); + let (response, permit) = other_source.poll().await.unwrap().unwrap(); + assert_eq!(response.task_token, vec![1]); + drop(permit); + assert!(poller.poll().await.is_none()); + match Arc::try_unwrap(poller) { + Ok(poller) => poller.shutdown().await, + Err(_) => panic!("poller still referenced"), + } + assert_eq!(slots.available_permits(), Some(2)); + } + + #[rstest] + #[case(false)] + #[case(true)] + #[tokio::test] + async fn activity_sticky_polling_ends_with_either_source(#[case] sticky_first: bool) { + let (tx, mut rx) = unbounded_channel(); + let mut client = mock_manual_worker_client(); + client + .expect_poll_activity_task() + .times(2) + .returning(move |_, _| { + tx.send(()).unwrap(); + std::future::pending().boxed() + }); + let client: Arc = Arc::new(client); + let config = WorkerConfig::builder() + .namespace("default") + .task_queue("normal") + .versioning_strategy(WorkerVersioningStrategy::None { + build_id: "test".to_owned(), + }) + .task_types(temporalio_common::worker::WorkerTaskTypes::all()) + .max_cached_workflows(2_usize) + .workflow_task_poller_behavior(PollerBehavior::SimpleMaximum(2)) + .session_worker_options(crate::SessionWorkerOptions::default()) + .activity_task_poller_behavior(PollerBehavior::SimpleMaximum(2)) + .build() + .unwrap(); + let slots = fixed_size_permit_dealer(2); + let poller = Arc::new(make_activity_task_buffer( + &config, + &Some("sticky".to_owned()), + &client, + &crate::MetricsContext::no_op(), + &CancellationToken::new(), + &slots, + Arc::new(AtomicCell::new(None)), + Arc::new(NamespaceCapabilities::default()), + )); + let polling = tokio::spawn({ + let poller = poller.clone(); + async move { poller.poll().await } + }); + for _ in 0..2 { + tokio::time::timeout(Duration::from_secs(5), rx.recv()) + .await + .unwrap() + .unwrap(); + } + let first = if sticky_first { + poller.sticky_poller.as_ref().unwrap() + } else { + poller.normal_poller.as_ref().unwrap() + }; + first.notify_shutdown(); + assert!( + tokio::time::timeout(Duration::from_secs(5), polling) + .await + .unwrap() + .unwrap() + .is_none() + ); + poller.notify_shutdown(); + match Arc::try_unwrap(poller) { + Ok(poller) => poller.shutdown().await, + Err(_) => panic!("poller still referenced"), + } + assert_eq!(slots.available_permits(), Some(2)); + } + #[tokio::test] async fn only_polls_once_with_1_poller() { let mut mock_client = mock_manual_worker_client(); @@ -947,7 +1371,7 @@ mod tests { CancellationToken::new(), None::, WorkflowTaskOptions { - wft_poller_shared: Some(Arc::new(WFTPollerShared::new(Some(10)))), + wft_poller_shared: Some(Arc::new(StickyPollerShared::new(Some(10)))), }, Arc::new(AtomicCell::new(None)), Arc::new(NamespaceCapabilities::default()), @@ -1004,7 +1428,7 @@ mod tests { CancellationToken::new(), None::, WorkflowTaskOptions { - wft_poller_shared: Some(Arc::new(WFTPollerShared::new(Some(1)))), + wft_poller_shared: Some(Arc::new(StickyPollerShared::new(Some(1)))), }, Arc::new(AtomicCell::new(None)), Arc::new(NamespaceCapabilities::default()), @@ -1081,7 +1505,7 @@ mod tests { shutdown_token.clone(), None::, WorkflowTaskOptions { - wft_poller_shared: Some(Arc::new(WFTPollerShared::new(Some(10)))), + wft_poller_shared: Some(Arc::new(StickyPollerShared::new(Some(10)))), }, Arc::new(AtomicCell::new(None)), Arc::new(NamespaceCapabilities::default()), @@ -1190,7 +1614,7 @@ mod tests { CancellationToken::new(), None::, WorkflowTaskOptions { - wft_poller_shared: Some(Arc::new(WFTPollerShared::new(Some(INITIAL_POLLERS)))), + wft_poller_shared: Some(Arc::new(StickyPollerShared::new(Some(INITIAL_POLLERS)))), }, Arc::new(AtomicCell::new(None)), Arc::new(NamespaceCapabilities::default()), diff --git a/crates/sdk-core/src/pollers/sticky.rs b/crates/sdk-core/src/pollers/sticky.rs new file mode 100644 index 000000000..1bc458245 --- /dev/null +++ b/crates/sdk-core/src/pollers/sticky.rs @@ -0,0 +1,141 @@ +use std::sync::OnceLock; +use tokio::sync::watch; + +/// Coordinates sticky and normal polling before either source can monopolize slot reservations. +/// +/// Admission is checked before reserving a slot, in this order: +/// 1. Let a source with no active polls attempt its first poll, so neither queue is excluded. +/// 2. Hold back extra sticky polls at the sticky scaler's target, and extra polls that could prevent +/// the other source's first poll. +/// 3. Favor sticky polls when the reported sticky backlog exceeds both one and the active sticky +/// poll count and the sticky scaler can grow, so known sticky work gets usable capacity. +/// 4. Otherwise allow growth while the combined active poll count is below the captured bound. +/// +/// Activity poll responses have no sticky backlog hint, so rule 3 currently only affects workflows. +/// These are polling heuristics, not slot reservations or a guarantee of fairness: the slot supplier +/// still admits each reservation and the poll scaler still bounds each source's concurrency. +pub(crate) struct StickyPollerShared { + last_seen_sticky_backlog: (watch::Receiver, watch::Sender), + sticky_active: OnceLock>, + sticky_target: OnceLock>, + non_sticky_active: OnceLock>, + // This construction-time bound is not live capacity; suppliers without a known bound enforce + // admission alone. + max_slots: Option, +} + +impl StickyPollerShared { + pub(crate) fn new(max_slots: Option) -> Self { + let (tx, rx) = watch::channel(0); + Self { + last_seen_sticky_backlog: (rx, tx), + sticky_active: OnceLock::new(), + sticky_target: OnceLock::new(), + non_sticky_active: OnceLock::new(), + max_slots, + } + } + + pub(crate) fn set_sticky_active( + &self, + active_rx: watch::Receiver, + target_rx: watch::Receiver, + ) { + let _ = self.sticky_active.set(active_rx); + let _ = self.sticky_target.set(target_rx); + } + + pub(crate) fn set_non_sticky_active(&self, rx: watch::Receiver) { + let _ = self.non_sticky_active.set(rx); + } + + /// Makes either the sticky or non-sticky poller wait pre-permit-acquisition so that we can + /// balance which kind of queue we poll appropriately. + pub(crate) async fn wait_if_needed(&self, is_sticky: bool) { + // We need to make sure there's at least one poller of both kinds available. So, we check + // that we won't end up using every available permit with one kind of poller. In practice + // this is only ever likely to be an issue with very small numbers of slots. + if let Some(max_slots) = self.max_slots + && let Some(sticky_active) = self.sticky_active.get() + && let Some(sticky_target) = self.sticky_target.get() + && let Some(non_sticky_active) = self.non_sticky_active.get() + { + let mut sticky_active = sticky_active.clone(); + let mut sticky_target = sticky_target.clone(); + let mut non_sticky_active = non_sticky_active.clone(); + let mut sticky_backlog = self.last_seen_sticky_backlog.0.clone(); + + loop { + let num_sticky_active = *sticky_active.borrow_and_update(); + let num_sticky_target = *sticky_target.borrow_and_update(); + let num_non_sticky_active = *non_sticky_active.borrow_and_update(); + let num_sticky_backlog = *sticky_backlog.borrow_and_update(); + let sticky_should_be_chosen = num_sticky_backlog > 1 + && num_sticky_backlog > num_sticky_active + && num_sticky_active < num_sticky_target; + + let allow = || { + if !is_sticky { + // There should always be at least one non-sticky poller. + if num_non_sticky_active == 0 { + return true; + } + + // Do not allow an additional non-sticky poller to prevent starting a first sticky poller. + if num_sticky_active == 0 && num_non_sticky_active + 1 >= max_slots { + return false; + } + + // If there's a meaningful sticky backlog, prioritize sticky. + if sticky_should_be_chosen { + return false; + } + } else { + // There should always be at least one sticky poller. + if num_sticky_active == 0 { + return true; + } + + // Do not reserve a slot that the sticky scaler cannot use. + if num_sticky_active >= num_sticky_target { + return false; + } + + // Do not allow an additional sticky poller to prevent starting a first non-sticky poller. + if num_non_sticky_active == 0 && num_sticky_active + 1 >= max_slots { + return false; + } + + // If there's a meaningful sticky backlog, prioritize sticky. + if sticky_should_be_chosen { + return true; + } + } + + // Just balance the two poller types. + // FIXME: Do we need anything more here, to ensure proper balancing? + if num_sticky_active + num_non_sticky_active < max_slots { + return true; + } + + false + }; + + if allow() { + return; + } + + tokio::select! { + _ = sticky_active.changed() => (), + _ = sticky_target.changed() => (), + _ = non_sticky_active.changed() => (), + _ = sticky_backlog.changed() => (), + } + } + } + } + + pub(crate) fn record_sticky_backlog(&self, v: usize) { + let _ = self.last_seen_sticky_backlog.1.send(v); + } +} diff --git a/crates/sdk-core/src/replay/mod.rs b/crates/sdk-core/src/replay/mod.rs index 6e4cadc50..388304723 100644 --- a/crates/sdk-core/src/replay/mod.rs +++ b/crates/sdk-core/src/replay/mod.rs @@ -75,6 +75,8 @@ where } pub(crate) fn into_core_worker(mut self) -> Result { + // Replay consumes supplied history rather than this worker's server placement. + self.config.session_worker_options = None; self.config.max_cached_workflows = 1; self.config.workflow_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(1)); self.config.task_types = WorkerTaskTypes::workflow_only(); @@ -150,8 +152,70 @@ where #[cfg(test)] mod tests { use super::*; - use crate::test_help::test_worker_cfg; + use crate::test_help::{WorkerExt, test_worker_cfg}; use futures_util::{FutureExt, stream}; + use temporalio_common::protos::{ + coresdk::{ + workflow_activation::workflow_activation_job, + workflow_commands::{CompleteWorkflowExecution, ScheduleActivity}, + workflow_completion::WorkflowActivationCompletion, + }, + temporal::api::session::v1::{ActivitySessionOptions, activity_session_options}, + }; + + #[tokio::test] + async fn replay_workflow_sticky_activity_without_session_worker() { + let mut history = canned_histories::single_activity("activity"); + history.add_workflow_task_completed(); + history.add_workflow_execution_completed(); + let mut client = mock_manual_worker_client(); + client.expect_fail_workflow_task().times(0); + let config = test_worker_cfg().build().unwrap(); + assert!(config.session_worker_options.is_none()); + let mut input = ReplayWorkerInput::new(config, stream::iter([history.into()])); + input.client_override = Some(client); + let worker = input.into_core_worker().unwrap(); + + let start = worker.poll_workflow_activation().await.unwrap(); + assert!(start.is_replaying); + assert!(matches!( + start.jobs[0].variant, + Some(workflow_activation_job::Variant::InitializeWorkflow(_)) + )); + worker + .complete_workflow_activation(WorkflowActivationCompletion::from_cmds( + start.run_id, + vec![ + ScheduleActivity { + seq: 1, + activity_id: "activity".to_owned(), + activity_type: DEFAULT_ACTIVITY_TYPE.to_owned(), + session_options: Some(ActivitySessionOptions { + session: Some( + activity_session_options::Session::WorkflowStickySession(true), + ), + }), + ..Default::default() + } + .into(), + ], + )) + .await + .unwrap(); + let resolved = worker.poll_workflow_activation().await.unwrap(); + assert!(matches!( + resolved.jobs[0].variant, + Some(workflow_activation_job::Variant::ResolveActivity(_)) + )); + worker + .complete_workflow_activation(WorkflowActivationCompletion::from_cmds( + resolved.run_id, + vec![CompleteWorkflowExecution::default().into()], + )) + .await + .unwrap(); + worker.drain_pollers_and_shutdown().await; + } #[tokio::test] async fn client_override_describe_namespace_precedes_fallback() { diff --git a/crates/sdk-core/src/telemetry/metrics.rs b/crates/sdk-core/src/telemetry/metrics.rs index 5b3445c77..48048fc1e 100644 --- a/crates/sdk-core/src/telemetry/metrics.rs +++ b/crates/sdk-core/src/telemetry/metrics.rs @@ -647,6 +647,7 @@ pub(crate) struct NumPollersMetric { pub wft_current_pollers: HeartbeatMetric, pub sticky_wft_current_pollers: HeartbeatMetric, pub activity_current_pollers: HeartbeatMetric, + pub sticky_activity_current_pollers: HeartbeatMetric, pub nexus_current_pollers: HeartbeatMetric, } @@ -659,6 +660,10 @@ impl NumPollersMetric { self.sticky_wft_current_pollers.clone(), ), ("activity_task", self.activity_current_pollers.clone()), + ( + "sticky_activity_task", + self.sticky_activity_current_pollers.clone(), + ), ("nexus_task", self.nexus_current_pollers.clone()), ]) } @@ -760,6 +765,9 @@ pub(crate) fn workflow_sticky_poller() -> MetricKeyValue { pub(crate) fn activity_poller() -> MetricKeyValue { MetricKeyValue::new(KEY_POLLER_TYPE, "activity_task") } +pub(crate) fn activity_sticky_poller() -> MetricKeyValue { + MetricKeyValue::new(KEY_POLLER_TYPE, "sticky_activity_task") +} pub(crate) fn nexus_poller() -> MetricKeyValue { MetricKeyValue::new(KEY_POLLER_TYPE, "nexus_task") } @@ -1206,6 +1214,39 @@ mod tests { struct DummyInstrumentRef(usize); impl BufferInstrumentRef for DummyInstrumentRef {} + #[test] + fn heartbeat_tracks_normal_and_sticky_activity_pollers_independently() { + let metrics = MetricsContext::no_op(); + let normal = metrics.with_new_attrs([activity_poller()]); + let sticky = metrics.with_new_attrs([activity_sticky_poller()]); + let heartbeat = metrics.in_memory_meter().unwrap(); + normal.record_num_pollers(3); + sticky.record_num_pollers(5); + normal.record_num_pollers(1); + assert_eq!( + heartbeat + .num_pollers + .activity_current_pollers + .load(Ordering::Relaxed), + 1 + ); + assert_eq!( + heartbeat + .num_pollers + .sticky_activity_current_pollers + .load(Ordering::Relaxed), + 5 + ); + sticky.record_num_pollers(0); + assert_eq!( + heartbeat + .num_pollers + .activity_current_pollers + .load(Ordering::Relaxed), + 1 + ); + } + #[test] fn with_new_attrs_adds_attrs_to_returned_context() { let mc = MetricsContext::no_op(); diff --git a/crates/sdk-core/src/worker/activities.rs b/crates/sdk-core/src/worker/activities.rs index c41991103..5b6076e09 100644 --- a/crates/sdk-core/src/worker/activities.rs +++ b/crates/sdk-core/src/worker/activities.rs @@ -1031,7 +1031,7 @@ mod tests { use super::*; use crate::{ abstractions::tests::fixed_size_permit_dealer, - pollers::{ActivityTaskOptions, LongPollBuffer}, + pollers::{ActivityTaskOptions, LongPollBuffer, PollRateLimiter}, prost_dur, worker::{ NamespaceCapabilities, PollerBehavior, @@ -1055,8 +1055,8 @@ mod tests { shutdown_token, None::, ActivityTaskOptions { - max_worker_acts_per_second: None, max_tps: None, + ..Default::default() }, Arc::new(AtomicCell::new(None)), Arc::new(NamespaceCapabilities::default()), @@ -1142,8 +1142,9 @@ mod tests { shutdown_token.clone(), None::, ActivityTaskOptions { - max_worker_acts_per_second: Some(2.0), + rate_limiter: Some(Arc::new(PollRateLimiter::new(2.0))), max_tps: None, + ..Default::default() }, Arc::new(AtomicCell::new(None)), Arc::new(NamespaceCapabilities::default()), diff --git a/crates/sdk-core/src/worker/client.rs b/crates/sdk-core/src/worker/client.rs index 6f32b4422..1a1e989fa 100644 --- a/crates/sdk-core/src/worker/client.rs +++ b/crates/sdk-core/src/worker/client.rs @@ -446,16 +446,20 @@ pub struct PollOptions { pub timeout_override: Option, } /// Additional options specific to workflow task polling -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Default, bon::Builder)] +#[non_exhaustive] pub struct PollWorkflowOptions { /// Optional sticky queue name for session‐based workflow polling pub sticky_queue_name: Option, } /// Additional options specific to activity task polling -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Default, bon::Builder)] +#[non_exhaustive] pub struct PollActivityOptions { /// Optional rate limit (tasks per second) for activity polling pub max_tasks_per_sec: Option, + /// Poll the worker's sticky queue instead of the normal activity queue. + pub sticky_queue_name: Option, } /// Additional options specific to Nexus task polling #[derive(Debug, Clone, Default)] @@ -524,9 +528,20 @@ impl WorkerClient for WorkerClientBag { let mut request = PollActivityTaskQueueRequest { namespace: self.namespace.clone(), task_queue: Some(TaskQueue { - name: poll_options.task_queue, - kind: TaskQueueKind::Normal as i32, - normal_name: "".to_string(), + name: act_options + .sticky_queue_name + .clone() + .unwrap_or_else(|| poll_options.task_queue.clone()), + kind: if act_options.sticky_queue_name.is_some() { + TaskQueueKind::Sticky + } else { + TaskQueueKind::Normal + } as i32, + normal_name: if act_options.sticky_queue_name.is_some() { + poll_options.task_queue + } else { + String::new() + }, }), identity: self.identity(), task_queue_metadata: act_options.max_tasks_per_sec.map(|tps| TaskQueueMetadata { @@ -1240,6 +1255,120 @@ mod tests { }; use temporalio_common::worker::{WorkerDeploymentOptions, WorkerDeploymentVersion}; + #[tokio::test] + async fn poll_requests_encode_affinity_and_queue_names() { + let requests = Arc::new(Mutex::new(Vec::new())); + let captured = requests.clone(); + let service_override = CallbackBasedGrpcService { + callback: Arc::new(move |request| { + let captured = captured.clone(); + Box::pin(async move { + let proto = match request.rpc.as_str() { + "GetSystemInfo" => GetSystemInfoResponse::default().encode_to_vec(), + "PollActivityTaskQueue" => { + let request = + PollActivityTaskQueueRequest::decode(request.proto).unwrap(); + assert_eq!(request.namespace, "namespace"); + assert_eq!( + request.task_queue_metadata.unwrap().max_tasks_per_second, + Some(7.0) + ); + captured.lock().unwrap().push(( + "activity", + request.task_queue.unwrap(), + request.single_workflow_mode, + )); + PollActivityTaskQueueResponse::default().encode_to_vec() + } + "PollWorkflowTaskQueue" => { + let request = + PollWorkflowTaskQueueRequest::decode(request.proto).unwrap(); + assert_eq!(request.namespace, "namespace"); + captured.lock().unwrap().push(( + "workflow", + request.task_queue.unwrap(), + request.single_workflow_mode, + )); + PollWorkflowTaskQueueResponse::default().encode_to_vec() + } + rpc => panic!("unexpected RPC: {rpc}"), + }; + Ok(GrpcSuccessResponse { + headers: Default::default(), + proto, + }) + }) + }), + }; + let connection = Connection::connect( + ConnectionOptions::new(url::Url::parse("http://localhost:7233").unwrap()) + .service_override(service_override) + .dns_load_balancing(None) + .build(), + ) + .await + .unwrap(); + let client = WorkerClientBag::new( + SharedReplaceableClient::new(connection), + "namespace".to_owned(), + WorkerVersioningStrategy::None { + build_id: String::new(), + }, + Uuid::new_v4(), + ); + for sticky_queue_name in [None, Some("sticky".to_owned())] { + client + .poll_workflow_task( + PollOptions { + task_queue: "normal".to_owned(), + no_retry: None, + timeout_override: None, + }, + PollWorkflowOptions::builder() + .maybe_sticky_queue_name(sticky_queue_name.clone()) + .build(), + ) + .await + .unwrap(); + client + .poll_activity_task( + PollOptions { + task_queue: "normal".to_owned(), + no_retry: None, + timeout_override: None, + }, + PollActivityOptions::builder() + .maybe_sticky_queue_name(sticky_queue_name) + .max_tasks_per_sec(7.0) + .build(), + ) + .await + .unwrap(); + } + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), 4); + for (i, (rpc, queue, single_workflow_mode)) in requests.iter().enumerate() { + assert_eq!(*rpc, if i % 2 == 0 { "workflow" } else { "activity" }); + assert!(!*single_workflow_mode); + assert_eq!( + *queue, + if i < 2 { + TaskQueue { + name: "normal".to_owned(), + kind: TaskQueueKind::Normal as i32, + normal_name: String::new(), + } + } else { + TaskQueue { + name: "sticky".to_owned(), + kind: TaskQueueKind::Sticky as i32, + normal_name: "normal".to_owned(), + } + } + ); + } + } + #[tokio::test] async fn activity_failure_request_includes_last_heartbeat_details() { let requests = Arc::new(Mutex::new(Vec::new())); diff --git a/crates/sdk-core/src/worker/mod.rs b/crates/sdk-core/src/worker/mod.rs index d31addaf4..57da2d166 100644 --- a/crates/sdk-core/src/worker/mod.rs +++ b/crates/sdk-core/src/worker/mod.rs @@ -38,7 +38,6 @@ pub(crate) use activities::{ ExecutingLAId, LocalActRequest, LocalActivityExecutionResult, LocalActivityResolution, NewLocalAct, }; -pub(crate) use wft_poller::WFTPollerShared; #[allow(unreachable_pub)] // re-exported in test_help::integ_helpers pub use workflow::LEGACY_QUERY_ID; @@ -46,11 +45,11 @@ pub use workflow::LEGACY_QUERY_ID; use crate::{ ActivityHeartbeat, abstractions::{MeteredPermitDealer, PermitDealerContextData, dbg_panic}, - pollers::{ActivityTaskOptions, BoxedActPoller, BoxedNexusPoller, LongPollBuffer}, + pollers::{BoxedActPoller, BoxedNexusPoller, LongPollBuffer, make_activity_task_buffer}, protosext::validate_activity_completion, telemetry::metrics::{ - MetricsContext, WorkerHeartbeatMetrics, activity_poller, activity_worker_type, - local_activity_worker_type, nexus_poller, nexus_worker_type, workflow_worker_type, + MetricsContext, WorkerHeartbeatMetrics, activity_worker_type, local_activity_worker_type, + nexus_poller, nexus_worker_type, workflow_worker_type, }, worker::{ activities::{LACompleteAction, LocalActivityManager, NextPendingLAAction}, @@ -58,8 +57,7 @@ use crate::{ heartbeat::SharedNamespaceWorker, nexus::NexusManager, workflow::{ - LAReqSink, LocalResolution, WorkflowBasics, Workflows, wft_poller, - wft_poller::make_wft_poller, + LAReqSink, LocalResolution, WorkflowBasics, Workflows, wft_poller::make_wft_poller, }, }, }; @@ -128,6 +126,11 @@ use { }, }; +/// Experimental options enabling sticky activity support for session workers. +#[derive(Clone, Debug, Default, bon::Builder)] +#[non_exhaustive] +pub struct SessionWorkerOptions {} + /// Defines per-worker configuration options #[derive(Clone, bon::Builder)] #[builder(on(String, into), state_mod(vis = "pub"), finish_fn(vis = "", name = build_internal))] @@ -149,12 +152,16 @@ pub struct WorkerConfig { /// or failures. #[builder(default = 0)] pub max_cached_workflows: usize, + /// Experimental: enables session activity polling on this worker's workflow sticky queue. + /// Requires workflows, remote activities, and a nonzero workflow cache. + /// `None` leaves activity polling on the normal task queue only. + pub session_worker_options: Option, /// Set a [crate::WorkerTuner] for this worker. Either this or at least one of the /// `max_outstanding_*` fields must be set. pub tuner: Option>, /// 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. + /// If using SimpleMaximum, must be at least 2 when `max_cached_workflows` > 0. /// /// When left unset (`None`), the effective behavior is resolved at worker start: pollers are /// automatically enrolled into poller autoscaling when the namespace advertises the @@ -168,6 +175,9 @@ pub struct WorkerConfig { /// queue will allow 4 max pollers while the nonsticky queue will allow one. The minimum for /// either poller is 1, so if the maximum allowed is 1 and sticky queues are enabled, there will /// be 2 concurrent polls. + /// + /// Session workers also use this ratio to divide activity pollers between normal and sticky + /// queues. #[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 @@ -261,6 +271,7 @@ pub struct WorkerConfig { #[builder(into)] pub max_outstanding_workflow_tasks: Option, /// The maximum number of activity tasks that will ever be given to this worker concurrently. + /// Must be at least 2 for session workers polling both normal and sticky activity queues. /// /// Mutually exclusive with `tuner` #[builder(into)] @@ -348,6 +359,16 @@ impl WorkerConfigBuilder { "`task_types` cannot enable local activities without workflows".to_string(), ); } + if config.session_worker_options.is_some() + && (!task_types.enable_workflows + || !task_types.enable_remote_activities + || config.max_cached_workflows == 0) + { + return Err( + "`session_worker_options` requires workflows, remote activities, and `max_cached_workflows` > 0" + .to_string(), + ); + } if let Some(b) = config.workflow_task_poller_behavior { b.validate()?; @@ -373,6 +394,14 @@ impl WorkerConfigBuilder { if matches!(config.max_outstanding_activities, Some(v) if v == 0) { return Err("`max_outstanding_activities` must be > 0".to_string()); } + if config.session_worker_options.is_some() + && matches!(config.max_outstanding_activities, Some(v) if v < 2) + { + return Err( + "Paired session activity polling requires `max_outstanding_activities` >= 2" + .to_string(), + ); + } if matches!(config.max_outstanding_local_activities, Some(v) if v == 0) { return Err("`max_outstanding_local_activities` must be > 0".to_string()); } @@ -767,6 +796,17 @@ impl Worker { Arc::new(tuner_builder.build()) }); + let activity_supplier = tuner.activity_task_slot_supplier(); + if config.session_worker_options.is_some() + && activity_supplier + .fixed_capacity() + .is_some_and(|size| size < 2) + { + anyhow::bail!( + "Paired session activity polling requires at least 2 fixed activity slots" + ); + } + metrics.worker_registered(); let shutdown_token = CancellationToken::new(); let slot_context_data = Arc::new(PermitDealerContextData { @@ -789,7 +829,7 @@ impl Worker { ); let wft_permits = wft_slots.get_extant_count_rcv(); let act_slots = MeteredPermitDealer::new( - tuner.activity_task_slot_supplier(), + activity_supplier, metrics.with_new_attrs([activity_worker_type()]), None, slot_context_data.clone(), @@ -959,22 +999,13 @@ impl Worker { }; 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, - }, + let ap = make_activity_task_buffer( + &config, + &sticky_queue_name, + &client, + &metrics, + &shutdown_token, + &act_slots, act_last_suc_poll_time.clone(), capabilities.clone(), ); @@ -1855,7 +1886,10 @@ pub trait WorkerTuner { &self, ) -> Arc + Send + Sync>; - /// Return a [SlotSupplier] for activity tasks + /// Return a [SlotSupplier] for activity tasks. Session workers polling both normal and sticky + /// queues need at least one reservation per source, even under resource pressure, to avoid + /// becoming stuck polling only one queue. Custom suppliers must provide this protection; + /// [`SlotReservationContext::is_paired_activity_poller`] identifies these reservations. fn activity_task_slot_supplier( &self, ) -> Arc + Send + Sync>; @@ -1907,6 +1941,12 @@ pub trait SlotSupplier { None } + /// Returns the configured capacity for fixed-size suppliers, independent of reservations. + /// Custom and dynamically sized suppliers need not implement this. + fn fixed_capacity(&self) -> Option { + None + } + /// Returns a human-friendly identifier describing this supplier implementation for /// diagnostics and telemetry. fn slot_supplier_kind(&self) -> String { @@ -1928,9 +1968,16 @@ pub trait SlotReservationContext: Send + Sync { /// Returns the number of currently outstanding slot permits, whether used or un-used. fn num_issued_slots(&self) -> usize; - /// Returns true iff this is a sticky poll for a workflow task + /// Returns true iff this is a sticky poll for a workflow or activity task. fn is_sticky(&self) -> bool; + /// Whether this activity reservation belongs to a worker polling both normal and sticky + /// queues. Suppliers must allow at least one reservation for each source in this case. + /// Sticky-only single-workflow workers do not require two reservations. + fn is_paired_activity_poller(&self) -> bool { + false + } + /// Returns the metrics meter if metrics are enabled fn get_metrics_meter(&self) -> Option { None @@ -2096,7 +2143,7 @@ impl SlotKind for NexusSlotKind { #[derive(Clone, Copy, Debug, PartialEq)] pub enum PollerBehavior { /// Will attempt to poll as long as a slot is available, up to the provided maximum. Cannot - /// be less than two for workflow tasks, or one for other tasks. + /// be less than two for cached workflow tasks, or one otherwise. SimpleMaximum(usize), /// Will automatically scale the number of pollers based on feedback from the server. Still /// requires a slot to be available before beginning polling. @@ -2404,10 +2451,14 @@ impl WorkerHeartbeatManager { .is_autoscaling(), }); worker_heartbeat.activity_poller_info = Some(WorkerPollerInfo { - current_pollers: in_mem + current_pollers: (in_mem .num_pollers .activity_current_pollers - .load(Ordering::Relaxed) as i32, + .load(Ordering::Relaxed) + + in_mem + .num_pollers + .sticky_activity_current_pollers + .load(Ordering::Relaxed)) as i32, last_successful_poll_time: heartbeat_manager_metrics .act_last_suc_poll_time .load() @@ -2502,10 +2553,10 @@ pub(crate) enum TaskPollers { }, } -/// Given an already-resolved workflow task poller `behavior` (see [resolve_effective_behavior]), +/// Given an already-resolved 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( +pub(crate) fn sticky_poller_behavior( behavior: PollerBehavior, config: &WorkerConfig, is_sticky: bool, @@ -2640,11 +2691,11 @@ mod tests { cfg }; assert_eq!( - wft_poller_behavior(PollerBehavior::SimpleMaximum(5), &cfg, false), + sticky_poller_behavior(PollerBehavior::SimpleMaximum(5), &cfg, false), PollerBehavior::SimpleMaximum(1) ); assert_eq!( - wft_poller_behavior(PollerBehavior::SimpleMaximum(5), &cfg, true), + sticky_poller_behavior(PollerBehavior::SimpleMaximum(5), &cfg, true), PollerBehavior::SimpleMaximum(4) ); } @@ -2671,6 +2722,34 @@ mod tests { } } + #[test] + fn session_worker_requires_workflows_remote_activities_and_cache() { + for (workflows, activities, cache, valid) in [ + (true, true, 10, true), + (false, true, 10, false), + (true, false, 10, false), + (true, true, 0, false), + ] { + let result = WorkerConfig::builder() + .namespace("default") + .task_queue("test-queue") + .versioning_strategy(default_versioning_strategy()) + .task_types(WorkerTaskTypes { + enable_workflows: workflows, + enable_remote_activities: activities, + enable_local_activities: false, + enable_nexus: false, + }) + .max_cached_workflows(cache) + .session_worker_options(SessionWorkerOptions::default()) + .build(); + assert_eq!(result.is_ok(), valid); + if let Err(error) = result { + assert!(error.contains("session_worker_options")); + } + } + } + #[test] fn test_default_configuration_polls_all_types() { let config = WorkerConfig::builder() diff --git a/crates/sdk-core/src/worker/tuner/fixed_size.rs b/crates/sdk-core/src/worker/tuner/fixed_size.rs index b6b4fd557..79dc54956 100644 --- a/crates/sdk-core/src/worker/tuner/fixed_size.rs +++ b/crates/sdk-core/src/worker/tuner/fixed_size.rs @@ -8,6 +8,7 @@ use tokio::sync::Semaphore; /// Implements [SlotSupplier] with a fixed number of slots pub struct FixedSizeSlotSupplier { sem: Arc, + capacity: usize, _pd: PhantomData, } @@ -16,6 +17,7 @@ impl FixedSizeSlotSupplier { pub fn new(size: usize) -> Self { Self { sem: Arc::new(Semaphore::new(size)), + capacity: size, _pd: Default::default(), } } @@ -51,6 +53,10 @@ where Some(self.sem.available_permits()) } + fn fixed_capacity(&self) -> Option { + Some(self.capacity) + } + fn slot_supplier_kind(&self) -> String { "Fixed".to_string() } diff --git a/crates/sdk-core/src/worker/tuner/resource_based.rs b/crates/sdk-core/src/worker/tuner/resource_based.rs index e5fc9b6f6..7a1029c0f 100644 --- a/crates/sdk-core/src/worker/tuner/resource_based.rs +++ b/crates/sdk-core/src/worker/tuner/resource_based.rs @@ -325,6 +325,13 @@ where fn mark_slot_used(&self, _ctx: &dyn SlotMarkUsedContext) {} fn release_slot(&self, ctx: &dyn SlotReleaseContext) { + if let Some(source) = ctx.permit().user_data::() { + if source.is_sticky { + self.issued_sticky.fetch_sub(1, Ordering::Relaxed); + } else { + self.issued_nonsticky.fetch_sub(1, Ordering::Relaxed); + } + } // Really could use specialization here if let Some(SlotInfo::Workflow(info)) = ctx.info().map(|i| i.downcast()) { if info.is_sticky { @@ -340,6 +347,10 @@ where } } +struct PairedActivitySource { + is_sticky: bool, +} + impl ResourceBasedSlotsForType where MI: SystemResourceInfo + Send + Sync, @@ -367,7 +378,8 @@ where if ctx.num_issued_slots() < self.opts.min_slots { return Some(self.issue_slot(ctx)); } - if SK::kind() == SlotKindType::Workflow + if (SK::kind() == SlotKindType::Workflow + || (SK::kind() == SlotKindType::Activity && ctx.is_paired_activity_poller())) && (ctx.is_sticky() && self.issued_sticky.load(Ordering::Relaxed) == 0 || !ctx.is_sticky() && self.issued_nonsticky.load(Ordering::Relaxed) == 0) { @@ -380,7 +392,9 @@ where fn issue_slot(&self, ctx: &dyn SlotReservationContext) -> SlotSupplierPermit { // Always be willing to hand out at least 1 slot for sticky and 1 for non-sticky to avoid // getting stuck. - if SK::kind() == SlotKindType::Workflow { + let paired_activity = + SK::kind() == SlotKindType::Activity && ctx.is_paired_activity_poller(); + if SK::kind() == SlotKindType::Workflow || paired_activity { if ctx.is_sticky() { self.issued_sticky.fetch_add(1, Ordering::Relaxed); } else { @@ -388,7 +402,15 @@ where } } let _ = self.last_slot_issued_tx.send(Instant::now()); - SlotSupplierPermit::default() + if paired_activity { + // A poll may end without receiving a task, so slot info cannot identify its source + // when the reservation is released. + SlotSupplierPermit::with_user_data(PairedActivitySource { + is_sticky: ctx.is_sticky(), + }) + } else { + SlotSupplierPermit::default() + } } fn time_since_last_issued(&self) -> Duration { @@ -461,7 +483,7 @@ impl ResourceController { Arc::new(ResourceBasedSlotsForType::new(self.clone(), opts)) } - fn new_with_sysinfo(options: ResourceBasedSlotsOptions, sys_info: Arc) -> Self { + pub(crate) fn new_with_sysinfo(options: ResourceBasedSlotsOptions, sys_info: Arc) -> Self { Self { pids: Mutex::new(PidControllers::new(&options)), options, @@ -766,8 +788,9 @@ impl CGroupCpuFileSystem for CgroupV2CpuFileSystem { mod tests { use super::*; use crate::{ - abstractions::MeteredPermitDealer, telemetry::metrics::MetricsContext, - worker::WorkflowSlotKind, + abstractions::MeteredPermitDealer, + telemetry::metrics::MetricsContext, + worker::{ActivitySlotInfo, WorkflowSlotKind}, }; use std::{ cell::RefCell, @@ -889,6 +912,47 @@ mod tests { assert!(rbs.try_reserve_slot(&pd).is_some()); } + #[tokio::test(start_paused = true)] + async fn paired_activity_minimums_release_unused_and_used_reservations() { + let (fmis, used) = FakeMIS::new(); + used.store(90_000, Ordering::Release); + let rbs = Arc::new(ResourceController::new_with_sysinfo(test_options(), fmis)) + .as_kind::(ResourceSlotOptions { + min_slots: 0, + max_slots: 1, + ramp_throttle: Duration::from_secs(60), + }); + let ordinary = MeteredPermitDealer::new( + rbs.clone(), + MetricsContext::no_op(), + None, + Arc::new(Default::default()), + None, + ); + assert!(ordinary.try_acquire_owned().is_err()); + let normal = ordinary.clone().for_activity_poller(false, true); + let sticky = ordinary.clone().for_activity_poller(true, true); + let normal_permit = normal.acquire_owned().await; + let sticky_permit = sticky.acquire_owned().await; + assert!(normal.try_acquire_owned().is_err()); + assert!(sticky.try_acquire_owned().is_err()); + + // Canceled and empty polls return unused reservations without slot info. + drop(normal_permit); + let normal_permit = normal.try_acquire_owned().unwrap(); + assert!(sticky.try_acquire_owned().is_err()); + drop(sticky_permit); + let sticky_permit = sticky.try_acquire_owned().unwrap(); + assert!(normal.try_acquire_owned().is_err()); + + let normal_permit = normal_permit.into_used(ActivitySlotInfo::default()); + let sticky_permit = sticky_permit.into_used(ActivitySlotInfo::default()); + drop((normal_permit, sticky_permit)); + let _normal = normal.try_acquire_owned().unwrap(); + let _sticky = sticky.try_acquire_owned().unwrap(); + assert!(ordinary.try_acquire_owned().is_err()); + } + #[test] fn shared_controller_has_same_behavior_across_tuners() { let (fmis, used) = FakeMIS::new(); diff --git a/crates/sdk-core/src/worker/workflow/machines/workflow_machines.rs b/crates/sdk-core/src/worker/workflow/machines/workflow_machines.rs index 6c1b1a59f..8396d9d13 100644 --- a/crates/sdk-core/src/worker/workflow/machines/workflow_machines.rs +++ b/crates/sdk-core/src/worker/workflow/machines/workflow_machines.rs @@ -70,6 +70,7 @@ use temporalio_common::{ history::v1::{HistoryEvent, history_event}, protocol::v1::{Message as ProtocolMessage, message::SequencingId}, sdk::v1::WorkflowTaskCompletedMetadata, + session::v1::activity_session_options, }, }, worker::WorkerDeploymentVersion, @@ -1368,7 +1369,39 @@ impl WorkflowMachines { annotations ); } - WFCommandVariant::AddActivity(attrs) => { + WFCommandVariant::AddActivity(mut attrs) => { + if attrs.session_options.as_ref().is_some_and(|options| { + matches!( + options.session, + Some(activity_session_options::Session::WorkflowStickySession( + true + )) + ) + }) { + if attrs.task_queue.is_empty() { + attrs.task_queue.clone_from(&self.worker_config.task_queue); + } + // Historical commands must remain replayable without activity execution; + // only new dispatches depend on this worker's capabilities. + if !self.replaying && self.worker_config.session_worker_options.is_none() { + return Err(fatal!( + "Activity '{}' (seq {}) requests workflow affinity, but \ + session_worker_options is not enabled", + attrs.activity_id, + attrs.seq + )); + } + if !self.replaying && attrs.task_queue != self.worker_config.task_queue { + return Err(fatal!( + "Activity '{}' (seq {}) requests workflow affinity on queue '{}', \ + but must use the worker's task queue '{}'", + attrs.activity_id, + attrs.seq, + attrs.task_queue, + self.worker_config.task_queue + )); + } + } let seq = attrs.seq; let use_compat = self.determine_use_compatible_flag( attrs.versioning_intent(), @@ -1852,6 +1885,76 @@ fn patch_marker_handling( } } +#[cfg(test)] +mod session_tests { + use super::*; + use crate::test_help::test_worker_cfg; + use temporalio_common::protos::{ + coresdk::workflow_commands::ScheduleActivity, + temporal::api::session::v1::ActivitySessionOptions, + }; + + #[rstest::rstest] + fn activity_session_validation( + #[values(false, true)] replaying: bool, + #[values(false, true)] session_worker: bool, + #[values(false, true)] session_activity: bool, + #[values("", "q", "other")] task_queue: &str, + ) { + let mut config = test_worker_cfg().build().unwrap(); + config.task_queue = "q".to_owned(); + config.session_worker_options = session_worker.then(Default::default); + let (driven, _commands) = DrivenWorkflow::new(); + let mut machines = WorkflowMachines::new( + RunBasics { + worker_config: Arc::new(config), + workflow_id: "workflow".to_owned(), + workflow_type: "workflow_type".to_owned(), + run_id: "run".to_owned(), + history: HistoryUpdate::dummy(), + metrics: MetricsContext::no_op(), + capabilities: &Default::default(), + sdk_name: "test", + sdk_version: "test", + }, + driven, + ); + machines.replaying = replaying; + let result = machines.handle_driven_results(vec![WFCommand { + variant: WFCommandVariant::AddActivity(ScheduleActivity { + seq: 1, + activity_id: "activity".to_owned(), + activity_type: "activity_type".to_owned(), + task_queue: task_queue.to_owned(), + session_options: Some(ActivitySessionOptions { + session: Some(activity_session_options::Session::WorkflowStickySession( + session_activity, + )), + }), + ..Default::default() + }), + annotations: Default::default(), + }]); + let valid = !session_activity || replaying || (session_worker && task_queue != "other"); + assert_eq!(result.is_ok(), valid, "{result:?}"); + if valid { + assert!(matches!( + &machines.current_wf_task_commands[0].command.attributes, + Some(ProtoCmdAttrs::ScheduleActivityTaskCommandAttributes(attrs)) + if attrs.task_queue.as_ref().unwrap().name + == if session_activity && task_queue.is_empty() { "q" } else { task_queue } + )); + } else { + assert!( + result + .unwrap_err() + .to_string() + .contains("Activity 'activity' (seq 1)") + ); + } + } +} + #[derive(derive_more::From)] enum CommandIdKind { /// A normal command, requested by lang diff --git a/crates/sdk-core/src/worker/workflow/wft_poller.rs b/crates/sdk-core/src/worker/workflow/wft_poller.rs index ee6dee2db..c91277976 100644 --- a/crates/sdk-core/src/worker/workflow/wft_poller.rs +++ b/crates/sdk-core/src/worker/workflow/wft_poller.rs @@ -1,19 +1,20 @@ use crate::{ MetricsContext, WorkerConfig, abstractions::{MeteredPermitDealer, OwnedMeteredSemPermit}, - pollers::{BoxedWFPoller, LongPollBuffer, Poller, WorkflowTaskOptions, WorkflowTaskPoller}, + pollers::{ + BoxedWFPoller, LongPollBuffer, Poller, StickyPollerShared, StickyTaskPoller, + WorkflowTaskOptions, + }, protosext::ValidPollWFTQResponse, telemetry::metrics::{workflow_poller, workflow_sticky_poller}, - worker::{NamespaceCapabilities, WorkflowSlotKind, client::WorkerClient, wft_poller_behavior}, + worker::{ + NamespaceCapabilities, WorkflowSlotKind, client::WorkerClient, sticky_poller_behavior, + }, }; use crossbeam_utils::atomic::AtomicCell; use futures_util::{Stream, stream}; -use std::{ - sync::{Arc, OnceLock}, - time::SystemTime, -}; +use std::{sync::Arc, time::SystemTime}; use temporalio_common::protos::temporal::api::workflowservice::v1::PollWorkflowTaskQueueResponse; -use tokio::sync::watch; use tokio_util::sync::CancellationToken; #[allow(clippy::too_many_arguments)] @@ -49,15 +50,15 @@ pub(crate) fn make_wft_poller( .into_iter() .flatten() .min(); - Some(Arc::new(WFTPollerShared::new(balance_limit))) + Some(Arc::new(StickyPollerShared::new(balance_limit))) } else { None }; - let wf_task_poll_buffer = LongPollBuffer::new_workflow_task( + let wf_task_poll_buffer = Some(LongPollBuffer::new_workflow_task( client.clone(), config.task_queue.clone(), None, - wft_poller_behavior(effective_behavior, config, false), + sticky_poller_behavior(effective_behavior, config, false), wft_slots.clone(), shutdown_token.child_token(), Some(move |np| { @@ -68,14 +69,14 @@ pub(crate) fn make_wft_poller( }, last_successful_poll_time, capabilities.clone(), - ); + )); let sticky_queue_poller = sticky_queue_name.as_ref().map(|sqn| { let sticky_metrics = metrics.with_new_attrs([workflow_sticky_poller()]); LongPollBuffer::new_workflow_task( client.clone(), config.task_queue.clone(), Some(sqn.clone()), - wft_poller_behavior(effective_behavior, config, true), + sticky_poller_behavior(effective_behavior, config, true), wft_slots.clone().into_sticky(), shutdown_token.child_token(), Some(move |np| { @@ -86,135 +87,13 @@ pub(crate) fn make_wft_poller( capabilities, ) }); - let wf_task_poll_buffer = Box::new(WorkflowTaskPoller::new( + let wf_task_poll_buffer = Box::new(StickyTaskPoller::new( wf_task_poll_buffer, sticky_queue_poller, )); new_wft_poller(wf_task_poll_buffer, metrics.clone()) } -/// Info that needs to be shared across the sticky and non-sticky wft pollers to prioritize sticky -/// when appropriate -pub(crate) struct WFTPollerShared { - last_seen_sticky_backlog: (watch::Receiver, watch::Sender), - sticky_active: OnceLock>, - sticky_target: OnceLock>, - non_sticky_active: OnceLock>, - max_slots: Option, -} -impl WFTPollerShared { - pub(crate) fn new(max_slots: Option) -> Self { - let (tx, rx) = watch::channel(0); - Self { - last_seen_sticky_backlog: (rx, tx), - sticky_active: OnceLock::new(), - sticky_target: OnceLock::new(), - non_sticky_active: OnceLock::new(), - max_slots, - } - } - pub(crate) fn set_sticky_active( - &self, - active_rx: watch::Receiver, - target_rx: watch::Receiver, - ) { - let _ = self.sticky_active.set(active_rx); - let _ = self.sticky_target.set(target_rx); - } - pub(crate) fn set_non_sticky_active(&self, active_rx: watch::Receiver) { - let _ = self.non_sticky_active.set(active_rx); - } - /// Makes either the sticky or non-sticky poller wait pre-permit-acquisition so that we can - /// balance which kind of queue we poll appropriately. - pub(crate) async fn wait_if_needed(&self, is_sticky: bool) { - // We need to make sure there's at least one poller of both kinds available. So, we check - // that we won't end up using every available permit with one kind of poller. In practice - // this is only ever likely to be an issue with very small numbers of slots. - if let Some(max_slots) = self.max_slots - && let Some(sticky_active) = self.sticky_active.get() - && let Some(sticky_target) = self.sticky_target.get() - && let Some(non_sticky_active) = self.non_sticky_active.get() - { - let mut sticky_active = sticky_active.clone(); - let mut sticky_target = sticky_target.clone(); - let mut non_sticky_active = non_sticky_active.clone(); - let mut sticky_backlog = self.last_seen_sticky_backlog.0.clone(); - - loop { - let num_sticky_active = *sticky_active.borrow_and_update(); - let num_sticky_target = *sticky_target.borrow_and_update(); - let num_non_sticky_active = *non_sticky_active.borrow_and_update(); - let num_sticky_backlog = *sticky_backlog.borrow_and_update(); - let sticky_should_be_chosen = num_sticky_backlog > 1 - && num_sticky_backlog > num_sticky_active - && num_sticky_active < num_sticky_target; - - let allow = || { - if !is_sticky { - // There should always be at least one non-sticky poller. - if num_non_sticky_active == 0 { - return true; - } - - // Do not allow an additional non-sticky poller to prevent starting a first sticky poller. - if num_sticky_active == 0 && num_non_sticky_active + 1 >= max_slots { - return false; - } - - // If there's a meaningful sticky backlog, prioritize sticky. - if sticky_should_be_chosen { - return false; - } - } else { - // There should always be at least one sticky poller. - if num_sticky_active == 0 { - return true; - } - - // Do not reserve a slot that the sticky scaler cannot use. - if num_sticky_active >= num_sticky_target { - return false; - } - - // Do not allow an additional sticky poller to prevent starting a first non-sticky poller. - if num_non_sticky_active == 0 && num_sticky_active + 1 >= max_slots { - return false; - } - - // If there's a meaningful sticky backlog, prioritize sticky. - if sticky_should_be_chosen { - return true; - } - } - - // Just balance the two poller types. - // FIXME: Do we need anything more here, to ensure proper balancing? - if num_sticky_active + num_non_sticky_active < max_slots { - return true; - } - - false - }; - - if allow() { - return; - } - - tokio::select! { - _ = sticky_active.changed() => (), - _ = sticky_target.changed() => (), - _ = non_sticky_active.changed() => (), - _ = sticky_backlog.changed() => (), - } - } - } - } - - pub(crate) fn record_sticky_backlog(&self, v: usize) { - let _ = self.last_seen_sticky_backlog.1.send(v); - } -} - fn new_wft_poller( poller: BoxedWFPoller, metrics: MetricsContext, @@ -302,6 +181,7 @@ mod tests { task::{Context, Poll, Wake}, }; use temporalio_common::protos::temporal::api::enums::v1::EventType; + use tokio::sync::watch; struct WakeCounter(AtomicUsize); @@ -313,7 +193,7 @@ mod tests { #[tokio::test] async fn sticky_priority_respects_poller_target() { - let shared = WFTPollerShared::new(Some(4)); + let shared = StickyPollerShared::new(Some(4)); let (_sticky_tx, sticky_rx) = watch::channel(2); let (_non_sticky_tx, non_sticky_rx) = watch::channel(1); let (sticky_target_tx, sticky_target_rx) = watch::channel(3); @@ -330,7 +210,7 @@ mod tests { #[tokio::test] async fn target_change_wakes_balancer() { - let shared = WFTPollerShared::new(Some(4)); + let shared = StickyPollerShared::new(Some(4)); let (_sticky_tx, sticky_rx) = watch::channel(2); let (_non_sticky_tx, non_sticky_rx) = watch::channel(1); let (sticky_target_tx, sticky_target_rx) = watch::channel(3); From 242e4f82f4455c4cce13538bc427d948f2a80363 Mon Sep 17 00:00:00 2001 From: Thomas Hardy Date: Fri, 16 Oct 2026 14:07:49 -0400 Subject: [PATCH 3/3] Add experimental single-workflow affinity controls Keep stable sticky names, sticky-only worker construction, low-level single-workflow RPC options, and replay normalization of runtime affinity settings together in the short-lived experimental layer. --- crates/sdk-core/CHANGELOG.md | 2 + crates/sdk-core/src/lib.rs | 7 +- crates/sdk-core/src/pollers/poll_buffer.rs | 141 +++++++++- crates/sdk-core/src/replay/mod.rs | 44 ++- crates/sdk-core/src/worker/client.rs | 22 +- crates/sdk-core/src/worker/mod.rs | 261 +++++++++++++++++- .../src/worker/workflow/wft_poller.rs | 111 ++++++-- 7 files changed, 547 insertions(+), 41 deletions(-) diff --git a/crates/sdk-core/CHANGELOG.md b/crates/sdk-core/CHANGELOG.md index 11db13415..7e7c3b913 100644 --- a/crates/sdk-core/CHANGELOG.md +++ b/crates/sdk-core/CHANGELOG.md @@ -62,6 +62,8 @@ relevant information. ### Added * Experimental workflow-affine activity support allows opted-in activities to use the workflow worker's sticky destination while sharing activity concurrency and rate limits. +* Short-lived experimental affinity controls support stable sticky destinations and sticky-only + single-workflow polling. These integration controls may be redesigned or removed as sessions evolve. * The OTLP metrics exporter can export histograms as OpenTelemetry exponential histograms, which Prometheus-compatible backends store as native histograms, when an SDK enables the option, including through the C bridge. It applies to every histogram, including custom ones, and names diff --git a/crates/sdk-core/src/lib.rs b/crates/sdk-core/src/lib.rs index 751469455..97c991885 100644 --- a/crates/sdk-core/src/lib.rs +++ b/crates/sdk-core/src/lib.rs @@ -44,6 +44,7 @@ pub use temporalio_common::protos::TaskToken; pub use url::Url; pub use worker::{ ActivitySlotKind, CompleteActivityError, CompleteNexusError, CompleteWfError, + ExperimentalWorkerAffinityOptions, ExperimentalWorkerAffinityOptionsBuilder, FixedSizeSlotSupplier, LocalActivitySlotKind, NamespaceCapabilities, NexusSlotKind, PollError, PollerBehavior, ResourceBasedSlotsOptions, ResourceBasedSlotsOptionsBuilder, ResourceBasedTuner, ResourceBasedTunerConfig, ResourceController, ResourceSlotOptions, @@ -103,7 +104,11 @@ pub fn init_worker( if client_ident.is_empty() { bail!("Client identity cannot be empty. Either lang or user should be setting this value"); } - let sticky_q = sticky_q_name_for_worker(&client_ident, worker_config.max_cached_workflows); + let sticky_q = worker_config + .experimental_worker_affinity_options + .as_ref() + .and_then(|options| options.sticky_queue_name.clone()) + .or_else(|| sticky_q_name_for_worker(&client_ident, worker_config.max_cached_workflows)); let worker_instance_key = uuid::Uuid::new_v4(); let client_bag = Arc::new(WorkerClientBag::new( diff --git a/crates/sdk-core/src/pollers/poll_buffer.rs b/crates/sdk-core/src/pollers/poll_buffer.rs index ff933e2fd..c727951d9 100644 --- a/crates/sdk-core/src/pollers/poll_buffer.rs +++ b/crates/sdk-core/src/pollers/poll_buffer.rs @@ -106,6 +106,7 @@ pub(crate) struct LongPollBuffer { #[derive(Default)] pub(crate) struct WorkflowTaskOptions { pub(crate) wft_poller_shared: Option>, + pub(crate) single_workflow_mode: bool, } #[derive(Default)] @@ -116,6 +117,7 @@ pub(crate) struct ActivityTaskOptions { pub(crate) max_tps: Option, pub(crate) sticky_queue_name: Option, pub(crate) poller_shared: Option>, + pub(crate) single_workflow_mode: bool, } #[allow(clippy::too_many_arguments)] @@ -132,8 +134,10 @@ pub(crate) fn make_activity_task_buffer( let sticky_queue_name = sticky_queue_name .as_ref() .filter(|_| config.session_worker_options.is_some()); - let shared = - sticky_queue_name.map(|_| Arc::new(StickyPollerShared::new(slots.available_permits()))); + let single_workflow_mode = config.single_workflow_mode(); + let shared = sticky_queue_name + .filter(|_| !single_workflow_mode) + .map(|_| Arc::new(StickyPollerShared::new(slots.available_permits()))); let rate_limiter = config .max_worker_activities_per_second .map(|rate| Arc::new(PollRateLimiter::new(rate))); @@ -142,7 +146,7 @@ pub(crate) fn make_activity_task_buffer( &capabilities, ); let make = |sticky: Option| { - let behavior = if sticky_queue_name.is_some() { + let behavior = if sticky_queue_name.is_some() && !single_workflow_mode { crate::worker::sticky_poller_behavior(behavior, config, sticky.is_some()) } else { behavior @@ -156,9 +160,10 @@ pub(crate) fn make_activity_task_buffer( client.clone(), config.task_queue.clone(), behavior, - slots - .clone() - .for_activity_poller(sticky.is_some(), sticky_queue_name.is_some()), + slots.clone().for_activity_poller( + sticky.is_some(), + sticky_queue_name.is_some() && !single_workflow_mode, + ), shutdown.child_token(), Some(move |np| metrics.record_num_pollers(np)), ActivityTaskOptions { @@ -166,13 +171,14 @@ pub(crate) fn make_activity_task_buffer( max_tps: config.max_task_queue_activities_per_second, sticky_queue_name: sticky, poller_shared: shared.clone(), + single_workflow_mode, }, last_successful_poll_time.clone(), capabilities.clone(), ) }; StickyTaskPoller::new( - Some(make(None)), + (!single_workflow_mode).then(|| make(None)), sticky_queue_name.map(|name| make(Some(name.clone()))), ) } @@ -246,6 +252,7 @@ impl LongPollBuffer { }, PollWorkflowOptions::builder() .maybe_sticky_queue_name(sticky_queue_name) + .single_workflow_mode(options.single_workflow_mode) .build(), ) .await @@ -328,6 +335,7 @@ impl LongPollBuffer { PollActivityOptions::builder() .maybe_max_tasks_per_sec(options.max_tps) .maybe_sticky_queue_name(sticky_queue_name) + .single_workflow_mode(options.single_workflow_mode) .build(), ) .await @@ -1094,6 +1102,118 @@ mod tests { .unwrap(); } + #[tokio::test] + async fn single_workflow_activity_capacity_one_polls_sticky() { + let (tx, mut rx) = unbounded_channel(); + let mut client = mock_manual_worker_client(); + client + .expect_poll_activity_task() + .times(1) + .returning(move |_, activity| { + assert_eq!(activity.sticky_queue_name.as_deref(), Some("sticky")); + tx.send(()).unwrap(); + std::future::pending().boxed() + }); + let client: Arc = Arc::new(client); + let mut config = test_worker_cfg().build().unwrap(); + config.session_worker_options = Some(crate::SessionWorkerOptions::default()); + config.experimental_worker_affinity_options = Some( + crate::ExperimentalWorkerAffinityOptions::builder() + .single_workflow_mode(true) + .build(), + ); + let slots = fixed_size_permit_dealer(1); + let poller = Arc::new(make_activity_task_buffer( + &config, + &Some("sticky".to_owned()), + &client, + &crate::MetricsContext::no_op(), + &CancellationToken::new(), + &slots, + Arc::new(AtomicCell::new(None)), + Arc::new(NamespaceCapabilities::default()), + )); + let polling = tokio::spawn({ + let poller = poller.clone(); + async move { poller.poll().await } + }); + tokio::time::timeout(Duration::from_secs(5), rx.recv()) + .await + .unwrap() + .unwrap(); + poller.notify_shutdown(); + assert!(polling.await.unwrap().is_none()); + match Arc::try_unwrap(poller) { + Ok(poller) => poller.shutdown().await, + Err(_) => panic!("poller still referenced"), + } + assert_eq!(slots.available_permits(), Some(1)); + } + + #[rstest] + #[case(PollerBehavior::SimpleMaximum(1), 1)] + #[case(PollerBehavior::SimpleMaximum(5), 5)] + #[case(PollerBehavior::Autoscaling { minimum: 1, maximum: 8, initial: 3 }, 3)] + #[tokio::test] + async fn single_workflow_activity_polls_only_sticky_with_full_behavior( + #[case] behavior: PollerBehavior, + #[case] expected_polls: usize, + ) { + let (tx, mut rx) = unbounded_channel(); + let mut client = mock_manual_worker_client(); + client + .expect_poll_activity_task() + .times(expected_polls) + .returning(move |poll, activity| { + assert_eq!(poll.task_queue, "normal"); + assert_eq!(activity.sticky_queue_name.as_deref(), Some("stable-sticky")); + assert!(activity.single_workflow_mode); + tx.send(()).unwrap(); + std::future::pending().boxed() + }); + let client: Arc = Arc::new(client); + let mut config = test_worker_cfg().build().unwrap(); + config.task_queue = "normal".to_owned(); + config.session_worker_options = Some(crate::SessionWorkerOptions::default()); + config.experimental_worker_affinity_options = Some( + crate::ExperimentalWorkerAffinityOptions::builder() + .single_workflow_mode(true) + .build(), + ); + config.activity_task_poller_behavior = Some(behavior); + config.nonsticky_to_sticky_poll_ratio = 1.0; + let slots = fixed_size_permit_dealer(10); + let poller = Arc::new(make_activity_task_buffer( + &config, + &Some("stable-sticky".to_owned()), + &client, + &crate::MetricsContext::no_op(), + &CancellationToken::new(), + &slots, + Arc::new(AtomicCell::new(None)), + Arc::new(NamespaceCapabilities::default()), + )); + assert!(poller.normal_poller.is_none()); + assert!(poller.sticky_poller.is_some()); + let polling = tokio::spawn({ + let poller = poller.clone(); + async move { poller.poll().await } + }); + for _ in 0..expected_polls { + tokio::time::timeout(Duration::from_secs(5), rx.recv()) + .await + .unwrap() + .unwrap(); + } + poller.notify_shutdown(); + assert!(polling.await.unwrap().is_none()); + match Arc::try_unwrap(poller) { + Ok(poller) => poller.shutdown().await, + Err(_) => panic!("poller still referenced"), + } + assert_eq!(slots.available_permits(), Some(10)); + } + #[rstest] #[case(false)] #[case(true)] @@ -1136,6 +1256,7 @@ mod tests { .unwrap(); assert_eq!(poll.task_queue, config.task_queue); assert_eq!(activity.max_tasks_per_sec, Some(7.0)); + assert!(!activity.single_workflow_mode); if let Some(sticky) = activity.sticky_queue_name { assert_eq!(sticky, "sticky"); sticky_polls += 1; @@ -1372,6 +1493,7 @@ mod tests { None::, WorkflowTaskOptions { wft_poller_shared: Some(Arc::new(StickyPollerShared::new(Some(10)))), + ..Default::default() }, Arc::new(AtomicCell::new(None)), Arc::new(NamespaceCapabilities::default()), @@ -1429,6 +1551,7 @@ mod tests { None::, WorkflowTaskOptions { wft_poller_shared: Some(Arc::new(StickyPollerShared::new(Some(1)))), + ..Default::default() }, Arc::new(AtomicCell::new(None)), Arc::new(NamespaceCapabilities::default()), @@ -1506,6 +1629,7 @@ mod tests { None::, WorkflowTaskOptions { wft_poller_shared: Some(Arc::new(StickyPollerShared::new(Some(10)))), + ..Default::default() }, Arc::new(AtomicCell::new(None)), Arc::new(NamespaceCapabilities::default()), @@ -1615,6 +1739,7 @@ mod tests { None::, WorkflowTaskOptions { wft_poller_shared: Some(Arc::new(StickyPollerShared::new(Some(INITIAL_POLLERS)))), + ..Default::default() }, Arc::new(AtomicCell::new(None)), Arc::new(NamespaceCapabilities::default()), @@ -1696,6 +1821,7 @@ mod tests { }), WorkflowTaskOptions { wft_poller_shared: None, + ..Default::default() }, Arc::new(AtomicCell::new(None)), Arc::new(NamespaceCapabilities::resolved(Capabilities { @@ -1774,6 +1900,7 @@ mod tests { None::, WorkflowTaskOptions { wft_poller_shared: None, + ..Default::default() }, Arc::new(AtomicCell::new(None)), Arc::new(NamespaceCapabilities::resolved(Capabilities { diff --git a/crates/sdk-core/src/replay/mod.rs b/crates/sdk-core/src/replay/mod.rs index 388304723..2a02aa784 100644 --- a/crates/sdk-core/src/replay/mod.rs +++ b/crates/sdk-core/src/replay/mod.rs @@ -75,8 +75,10 @@ where } pub(crate) fn into_core_worker(mut self) -> Result { - // Replay consumes supplied history rather than this worker's server placement. + // Replay consumes supplied history, not server placement, even when callers reuse a + // production worker configuration with sticky-only polling enabled. self.config.session_worker_options = None; + self.config.experimental_worker_affinity_options = None; self.config.max_cached_workflows = 1; self.config.workflow_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(1)); self.config.task_types = WorkerTaskTypes::workflow_only(); @@ -152,7 +154,10 @@ where #[cfg(test)] mod tests { use super::*; - use crate::test_help::{WorkerExt, test_worker_cfg}; + use crate::{ + WorkerVersioningStrategy, + test_help::{WorkerExt, test_worker_cfg}, + }; use futures_util::{FutureExt, stream}; use temporalio_common::protos::{ coresdk::{ @@ -163,18 +168,47 @@ mod tests { temporal::api::session::v1::{ActivitySessionOptions, activity_session_options}, }; + #[rstest::rstest] #[tokio::test] - async fn replay_workflow_sticky_activity_without_session_worker() { + async fn replay_workflow_sticky_activity_without_session_worker( + #[values("ordinary", "single_generated_name", "single_explicit_name")] mode: &str, + ) { let mut history = canned_histories::single_activity("activity"); history.add_workflow_task_completed(); history.add_workflow_execution_completed(); let mut client = mock_manual_worker_client(); client.expect_fail_workflow_task().times(0); - let config = test_worker_cfg().build().unwrap(); - assert!(config.session_worker_options.is_none()); + let single_workflow = mode != "ordinary"; + let config = WorkerConfig::builder() + .namespace("default") + .task_queue("replay") + .versioning_strategy(WorkerVersioningStrategy::None { + build_id: "test".to_owned(), + }) + .task_types(WorkerTaskTypes::all()) + .max_cached_workflows(2_usize) + .workflow_task_poller_behavior(PollerBehavior::SimpleMaximum(2)) + .maybe_session_worker_options(single_workflow.then(Default::default)) + .maybe_experimental_worker_affinity_options(single_workflow.then(|| { + crate::ExperimentalWorkerAffinityOptions::builder() + .single_workflow_mode(true) + .maybe_sticky_queue_name( + (mode == "single_explicit_name").then(|| "stable-sticky".to_owned()), + ) + .build() + })) + .build() + .unwrap(); let mut input = ReplayWorkerInput::new(config, stream::iter([history.into()])); input.client_override = Some(client); let worker = input.into_core_worker().unwrap(); + assert!(worker.get_config().session_worker_options.is_none()); + assert!( + worker + .get_config() + .experimental_worker_affinity_options + .is_none() + ); let start = worker.poll_workflow_activation().await.unwrap(); assert!(start.is_replaying); diff --git a/crates/sdk-core/src/worker/client.rs b/crates/sdk-core/src/worker/client.rs index 1a1e989fa..0af042195 100644 --- a/crates/sdk-core/src/worker/client.rs +++ b/crates/sdk-core/src/worker/client.rs @@ -451,6 +451,9 @@ pub struct PollOptions { pub struct PollWorkflowOptions { /// Optional sticky queue name for session‐based workflow polling pub sticky_queue_name: Option, + /// Experimental: dedicate this sticky poller to a single workflow. + #[builder(default)] + pub single_workflow_mode: bool, } /// Additional options specific to activity task polling #[derive(Debug, Clone, Default, bon::Builder)] @@ -460,6 +463,9 @@ pub struct PollActivityOptions { pub max_tasks_per_sec: Option, /// Poll the worker's sticky queue instead of the normal activity queue. pub sticky_queue_name: Option, + /// Experimental: dedicate this sticky poller to a single workflow. + #[builder(default)] + pub single_workflow_mode: bool, } /// Additional options specific to Nexus task polling #[derive(Debug, Clone, Default)] @@ -500,7 +506,7 @@ impl WorkerClient for WorkerClientBag { worker_instance_key: self.worker_instance_key.to_string(), worker_control_task_queue: self.worker_control_task_queue(), poller_group_id: Default::default(), - single_workflow_mode: false, + single_workflow_mode: wf_options.single_workflow_mode, } .into_request(); request.extensions_mut().insert(IsWorkerTaskLongPoll); @@ -552,7 +558,7 @@ impl WorkerClient for WorkerClientBag { worker_instance_key: self.worker_instance_key.to_string(), worker_control_task_queue: self.worker_control_task_queue(), poller_group_id: Default::default(), - single_workflow_mode: false, + single_workflow_mode: act_options.single_workflow_mode, } .into_request(); request.extensions_mut().insert(IsWorkerTaskLongPoll); @@ -1316,7 +1322,11 @@ mod tests { }, Uuid::new_v4(), ); - for sticky_queue_name in [None, Some("sticky".to_owned())] { + for (sticky_queue_name, single_workflow_mode) in [ + (None, false), + (Some("sticky".to_owned()), false), + (Some("sticky".to_owned()), true), + ] { client .poll_workflow_task( PollOptions { @@ -1326,6 +1336,7 @@ mod tests { }, PollWorkflowOptions::builder() .maybe_sticky_queue_name(sticky_queue_name.clone()) + .single_workflow_mode(single_workflow_mode) .build(), ) .await @@ -1340,16 +1351,17 @@ mod tests { PollActivityOptions::builder() .maybe_sticky_queue_name(sticky_queue_name) .max_tasks_per_sec(7.0) + .single_workflow_mode(single_workflow_mode) .build(), ) .await .unwrap(); } let requests = requests.lock().unwrap(); - assert_eq!(requests.len(), 4); + assert_eq!(requests.len(), 6); for (i, (rpc, queue, single_workflow_mode)) in requests.iter().enumerate() { assert_eq!(*rpc, if i % 2 == 0 { "workflow" } else { "activity" }); - assert!(!*single_workflow_mode); + assert_eq!(*single_workflow_mode, i >= 4); assert_eq!( *queue, if i < 2 { diff --git a/crates/sdk-core/src/worker/mod.rs b/crates/sdk-core/src/worker/mod.rs index 57da2d166..92f949424 100644 --- a/crates/sdk-core/src/worker/mod.rs +++ b/crates/sdk-core/src/worker/mod.rs @@ -131,6 +131,28 @@ use { #[non_exhaustive] pub struct SessionWorkerOptions {} +/// Short-lived experimental worker affinity controls. +/// +/// This API is expected to be substantially redesigned or removed as sessions evolve. +/// Applications using it should expect migration work in future releases. +#[derive(Clone, Debug, Default, bon::Builder)] +#[non_exhaustive] +pub struct ExperimentalWorkerAffinityOptions { + /// Override the workflow sticky queue name, also used for sticky activity polling when + /// [`WorkerConfig::session_worker_options`] is enabled. + /// + /// The name must be unique to one concurrently live worker, but should remain stable across + /// that worker's restarts to retain the server's routing association. Sharing it between live + /// workers defeats worker affinity. Requires a nonzero workflow cache and a nonempty name. + pub sticky_queue_name: Option, + /// Poll only the sticky workflow and activity queues, dedicating this worker to one workflow. + /// Requires session worker support, workflows, remote activities, and a nonzero workflow cache. + /// Workflow completion does not rebind the sticky destination; serving another workflow requires + /// deliberately choosing a new destination and worker lifecycle. Nexus polling is unchanged. + #[builder(default)] + pub single_workflow_mode: bool, +} + /// Defines per-worker configuration options #[derive(Clone, bon::Builder)] #[builder(on(String, into), state_mod(vis = "pub"), finish_fn(vis = "", name = build_internal))] @@ -156,12 +178,15 @@ pub struct WorkerConfig { /// Requires workflows, remote activities, and a nonzero workflow cache. /// `None` leaves activity polling on the normal task queue only. pub session_worker_options: Option, + /// Short-lived experimental affinity controls, expected to require migration as sessions evolve. + pub experimental_worker_affinity_options: Option, /// Set a [crate::WorkerTuner] for this worker. Either this or at least one of the /// `max_outstanding_*` fields must be set. pub tuner: Option>, /// 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. + /// If using SimpleMaximum, must be at least 2 when `max_cached_workflows` > 0, except in + /// single-workflow mode. /// /// When left unset (`None`), the effective behavior is resolved at worker start: pollers are /// automatically enrolled into poller autoscaling when the namespace advertises the @@ -177,7 +202,7 @@ pub struct WorkerConfig { /// be 2 concurrent polls. /// /// Session workers also use this ratio to divide activity pollers between normal and sticky - /// queues. + /// queues. Single-workflow mode ignores the ratio and uses all configured pollers on sticky. #[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 @@ -272,6 +297,7 @@ pub struct WorkerConfig { pub max_outstanding_workflow_tasks: Option, /// The maximum number of activity tasks that will ever be given to this worker concurrently. /// Must be at least 2 for session workers polling both normal and sticky activity queues. + /// Sticky-only single-workflow workers may use 1. /// /// Mutually exclusive with `tuner` #[builder(into)] @@ -313,6 +339,12 @@ pub struct WorkerConfig { } impl WorkerConfig { + pub(crate) fn single_workflow_mode(&self) -> bool { + self.experimental_worker_affinity_options + .as_ref() + .is_some_and(|options| options.single_workflow_mode) + } + /// Returns true if the configuration specifies we should fail a workflow on a certain error /// type rather than failing the workflow task. pub fn should_fail_workflow( @@ -369,6 +401,21 @@ impl WorkerConfigBuilder { .to_string(), ); } + if let Some(options) = &config.experimental_worker_affinity_options { + if options.single_workflow_mode && config.session_worker_options.is_none() { + return Err("`single_workflow_mode` requires `session_worker_options`".to_string()); + } + if let Some(name) = &options.sticky_queue_name { + if name.is_empty() { + return Err("`sticky_queue_name` must not be empty".to_string()); + } + if config.max_cached_workflows == 0 { + return Err( + "`sticky_queue_name` requires `max_cached_workflows` > 0".to_string() + ); + } + } + } if let Some(b) = config.workflow_task_poller_behavior { b.validate()?; @@ -395,6 +442,7 @@ impl WorkerConfigBuilder { return Err("`max_outstanding_activities` must be > 0".to_string()); } if config.session_worker_options.is_some() + && !config.single_workflow_mode() && matches!(config.max_outstanding_activities, Some(v) if v < 2) { return Err( @@ -418,7 +466,8 @@ impl WorkerConfigBuilder { .to_string(), ); } - if matches!(config.workflow_task_poller_behavior, Some(PollerBehavior::SimpleMaximum(u)) if u < 2) + if !config.single_workflow_mode() + && 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()); } @@ -776,6 +825,11 @@ impl Worker { worker_heartbeat_options: Option, shared_namespace_worker: bool, ) -> Result { + let sticky_queue_name = config + .experimental_worker_affinity_options + .as_ref() + .and_then(|options| options.sticky_queue_name.clone()) + .or(sticky_queue_name); let (metrics, meter) = if let Some(wt) = worker_telemetry.as_ref() { ( MetricsContext::top_level_with_meter( @@ -798,6 +852,7 @@ impl Worker { let activity_supplier = tuner.activity_task_slot_supplier(); if config.session_worker_options.is_some() + && !config.single_workflow_mode() && activity_supplier .fixed_capacity() .is_some_and(|size| size < 2) @@ -2143,7 +2198,7 @@ impl SlotKind for NexusSlotKind { #[derive(Clone, Copy, Debug, PartialEq)] pub enum PollerBehavior { /// Will attempt to poll as long as a slot is available, up to the provided maximum. Cannot - /// be less than two for cached workflow tasks, or one otherwise. + /// be less than two for cached workflow tasks outside single-workflow mode, or one otherwise. SimpleMaximum(usize), /// Will automatically scale the number of pollers based on feedback from the server. Still /// requires a slot to be available before beginning polling. @@ -2750,6 +2805,204 @@ mod tests { } } + #[test] + fn affinity_options_validate_only_required_coherence() { + for (session, single, name, cache, workflows, activities, expected_error) in [ + (false, false, None, 0, true, true, None), + ( + false, + true, + None, + 10, + true, + true, + Some("session_worker_options"), + ), + (true, true, None, 10, true, true, None), + ( + true, + true, + None, + 0, + true, + true, + Some("session_worker_options"), + ), + ( + true, + true, + None, + 10, + false, + true, + Some("session_worker_options"), + ), + ( + true, + true, + None, + 10, + true, + false, + Some("session_worker_options"), + ), + ( + false, + false, + Some(""), + 10, + true, + true, + Some("must not be empty"), + ), + ( + false, + false, + Some("stable"), + 0, + true, + true, + Some("max_cached_workflows"), + ), + (false, false, Some("stable"), 10, true, false, None), + (true, true, Some("normal"), 10, true, true, None), + (true, true, Some("different"), 10, true, true, None), + ] { + let result = WorkerConfig::builder() + .namespace("default") + .task_queue("normal") + .versioning_strategy(default_versioning_strategy()) + .task_types(WorkerTaskTypes { + enable_workflows: workflows, + enable_remote_activities: activities, + enable_local_activities: false, + enable_nexus: false, + }) + .max_cached_workflows(cache) + .workflow_task_poller_behavior(PollerBehavior::SimpleMaximum(if single { + 1 + } else { + 2 + })) + .maybe_session_worker_options(session.then(SessionWorkerOptions::default)) + .experimental_worker_affinity_options(ExperimentalWorkerAffinityOptions { + sticky_queue_name: name.map(str::to_owned), + single_workflow_mode: single, + }) + .build(); + match expected_error { + Some(expected) => assert!(result.err().unwrap().contains(expected)), + None => assert!(result.is_ok(), "{:?}", result.err()), + } + } + } + + #[test] + fn session_activity_fixed_capacity_validation() { + for (session, single, capacity, valid) in [ + (false, false, 1_usize, true), + (true, false, 1, false), + (true, false, 2, true), + (true, true, 1, true), + ] { + let result = WorkerConfig::builder() + .namespace("default") + .task_queue("queue") + .versioning_strategy(default_versioning_strategy()) + .task_types(WorkerTaskTypes::all()) + .max_cached_workflows(10_usize) + .workflow_task_poller_behavior(PollerBehavior::SimpleMaximum(2)) + .maybe_session_worker_options(session.then(SessionWorkerOptions::default)) + .experimental_worker_affinity_options(ExperimentalWorkerAffinityOptions { + sticky_queue_name: None, + single_workflow_mode: single, + }) + .max_outstanding_activities(capacity) + .build(); + assert_eq!(result.is_ok(), valid); + if let Err(error) = result { + assert!(error.contains("max_outstanding_activities")); + } + } + } + + #[tokio::test] + async fn session_activity_tuner_fixed_capacity_validation() { + for (session, single, capacity, valid) in [ + (false, false, 1, true), + (true, false, 0, false), + (true, false, 1, false), + (true, false, 2, true), + (true, true, 1, true), + ] { + let config = WorkerConfig::builder() + .namespace("default") + .task_queue("queue") + .versioning_strategy(default_versioning_strategy()) + .task_types(WorkerTaskTypes::all()) + .max_cached_workflows(10_usize) + .workflow_task_poller_behavior(PollerBehavior::SimpleMaximum(2)) + .maybe_session_worker_options(session.then(SessionWorkerOptions::default)) + .experimental_worker_affinity_options(ExperimentalWorkerAffinityOptions { + sticky_queue_name: None, + single_workflow_mode: single, + }) + .tuner(Arc::new(TunerHolder::fixed_size(10, capacity, 10, 10))) + .build() + .unwrap(); + let result = Worker::new( + config, + Some("sticky".to_owned()), + Arc::new(mock_worker_client()), + None, + None, + None, + ); + assert_eq!(result.is_ok(), valid); + if let Err(error) = result { + assert!(error.to_string().contains("fixed activity slots")); + } + } + } + + #[tokio::test] + async fn explicit_sticky_name_survives_worker_recreation_and_test_construction() { + let config = test_worker_cfg() + .max_cached_workflows(10_usize) + .session_worker_options(SessionWorkerOptions::builder().build()) + .experimental_worker_affinity_options( + ExperimentalWorkerAffinityOptions::builder() + .sticky_queue_name("stable-sticky".to_owned()) + .single_workflow_mode(true) + .build(), + ) + .build() + .unwrap(); + for worker in [ + Worker::new_test(config.clone(), mock_worker_client()), + Worker::new( + config, + Some("randomly-generated".to_owned()), + Arc::new(mock_worker_client()), + None, + None, + None, + ) + .unwrap(), + ] { + assert_eq!( + worker + .task_subsystems + .workflows + .as_ref() + .unwrap() + .get_sticky_queue_name() + .as_deref(), + Some("stable-sticky") + ); + } + } + #[test] fn test_default_configuration_polls_all_types() { let config = WorkerConfig::builder() diff --git a/crates/sdk-core/src/worker/workflow/wft_poller.rs b/crates/sdk-core/src/worker/workflow/wft_poller.rs index c91277976..86af883e5 100644 --- a/crates/sdk-core/src/worker/workflow/wft_poller.rs +++ b/crates/sdk-core/src/worker/workflow/wft_poller.rs @@ -43,7 +43,8 @@ pub(crate) fn make_wft_poller( config.workflow_task_poller_behavior, &capabilities, ); - let wft_poller_shared = if sticky_queue_name.is_some() { + let single_workflow_mode = config.single_workflow_mode(); + let wft_poller_shared = if sticky_queue_name.is_some() && !single_workflow_mode { // Balance on the limit `acquire_owned` actually enforces (min of slot supplier and cache // size). Using only the slot supplier lets a small cache starve the non-sticky poller. let balance_limit = [wft_slots.available_permits(), wft_slots.max_permits()] @@ -54,35 +55,45 @@ pub(crate) fn make_wft_poller( } else { None }; - let wf_task_poll_buffer = Some(LongPollBuffer::new_workflow_task( - client.clone(), - config.task_queue.clone(), - None, - sticky_poller_behavior(effective_behavior, config, false), - wft_slots.clone(), - shutdown_token.child_token(), - Some(move |np| { - wft_metrics.record_num_pollers(np); - }), - WorkflowTaskOptions { - wft_poller_shared: wft_poller_shared.clone(), - }, - last_successful_poll_time, - capabilities.clone(), - )); + let wf_task_poll_buffer = (!single_workflow_mode).then(|| { + LongPollBuffer::new_workflow_task( + client.clone(), + config.task_queue.clone(), + None, + sticky_poller_behavior(effective_behavior, config, false), + wft_slots.clone(), + shutdown_token.child_token(), + Some(move |np| { + wft_metrics.record_num_pollers(np); + }), + WorkflowTaskOptions { + wft_poller_shared: wft_poller_shared.clone(), + single_workflow_mode, + }, + last_successful_poll_time, + capabilities.clone(), + ) + }); let sticky_queue_poller = sticky_queue_name.as_ref().map(|sqn| { let sticky_metrics = metrics.with_new_attrs([workflow_sticky_poller()]); LongPollBuffer::new_workflow_task( client.clone(), config.task_queue.clone(), Some(sqn.clone()), - sticky_poller_behavior(effective_behavior, config, true), + if single_workflow_mode { + effective_behavior + } else { + sticky_poller_behavior(effective_behavior, config, true) + }, wft_slots.clone().into_sticky(), shutdown_token.child_token(), Some(move |np| { sticky_metrics.record_num_pollers(np); }), - WorkflowTaskOptions { wft_poller_shared }, + WorkflowTaskOptions { + wft_poller_shared, + single_workflow_mode, + }, sticky_last_successful_poll_time, capabilities, ) @@ -172,6 +183,7 @@ mod tests { }, }; use futures_util::{FutureExt, StreamExt, pin_mut}; + use rstest::rstest; use std::{ future::Future, sync::{ @@ -179,6 +191,7 @@ mod tests { atomic::{AtomicUsize, Ordering}, }, task::{Context, Poll, Wake}, + time::Duration, }; use temporalio_common::protos::temporal::api::enums::v1::EventType; use tokio::sync::watch; @@ -230,6 +243,66 @@ mod tests { assert_eq!(waiter.as_mut().poll(&mut context), Poll::Ready(())); } + #[rstest] + #[case(PollerBehavior::SimpleMaximum(1), 1)] + #[case(PollerBehavior::SimpleMaximum(5), 5)] + #[case(PollerBehavior::Autoscaling { minimum: 1, maximum: 8, initial: 3 }, 3)] + #[tokio::test] + async fn single_workflow_polls_only_sticky_with_full_behavior( + #[case] behavior: PollerBehavior, + #[case] expected_polls: usize, + ) { + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + let mut client = mock_manual_worker_client(); + client + .expect_poll_workflow_task() + .times(expected_polls) + .returning(move |poll, workflow| { + assert_eq!(poll.task_queue, "normal"); + assert_eq!(workflow.sticky_queue_name.as_deref(), Some("stable-sticky")); + assert!(workflow.single_workflow_mode); + tx.send(()).unwrap(); + std::future::pending().boxed() + }); + let client: Arc = Arc::new(client); + let mut config = test_worker_cfg().build().unwrap(); + config.task_queue = "normal".to_owned(); + config.session_worker_options = Some(crate::SessionWorkerOptions::default()); + config.experimental_worker_affinity_options = Some( + crate::ExperimentalWorkerAffinityOptions::builder() + .single_workflow_mode(true) + .build(), + ); + config.workflow_task_poller_behavior = Some(behavior); + config.nonsticky_to_sticky_poll_ratio = 1.0; + let slots = fixed_size_permit_dealer(10); + let shutdown = CancellationToken::new(); + let stream = make_wft_poller( + &config, + &Some("stable-sticky".to_owned()), + &client, + &MetricsContext::no_op(), + &shutdown, + &slots, + Arc::new(AtomicCell::new(None)), + Arc::new(AtomicCell::new(None)), + Arc::new(NamespaceCapabilities::default()), + ); + let polling = tokio::spawn(async move { + pin_mut!(stream); + stream.next().await + }); + for _ in 0..expected_polls { + tokio::time::timeout(Duration::from_secs(5), rx.recv()) + .await + .unwrap() + .unwrap(); + } + shutdown.cancel(); + assert!(polling.await.unwrap().is_none()); + assert_eq!(slots.available_permits(), Some(10)); + } + #[tokio::test] async fn poll_timeouts_do_not_produce_responses() { let mut mock_poller = mock_poller();