From 58d62d287207eddbf97864a176706e348fca0215 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 13 Oct 2026 18:54:25 -0700 Subject: [PATCH 1/8] Let lang complete a workflow task while local activities run. A reader that finds records on every read never blocks, so its task would complete only at the local activity heartbeat. With heartbeat_task, Core completes the task now as it does at that heartbeat, with the commands and markers so far and a request for a new task. Lang sets the flag from a count or a clock, so Core ignores it while replaying. --- .../workflow_completion.proto | 6 + crates/protos/src/protos/mod.rs | 1 + crates/sdk-core/CHANGELOG.md | 4 + crates/sdk-core/src/core_tests/mod.rs | 1 + .../sdk-core/src/core_tests/task_heartbeat.rs | 370 ++++++++++++++++++ .../src/worker/workflow/managed_run.rs | 8 + crates/sdk-core/src/worker/workflow/mod.rs | 2 + .../src/worker/workflow/workflow_stream.rs | 2 + crates/sdk/src/workflow_future.rs | 1 + 9 files changed, 395 insertions(+) create mode 100644 crates/sdk-core/src/core_tests/task_heartbeat.rs diff --git a/crates/protos/protos/local/temporal/sdk/core/workflow_completion/workflow_completion.proto b/crates/protos/protos/local/temporal/sdk/core/workflow_completion/workflow_completion.proto index afe7d64427..f70c206e5a 100644 --- a/crates/protos/protos/local/temporal/sdk/core/workflow_completion/workflow_completion.proto +++ b/crates/protos/protos/local/temporal/sdk/core/workflow_completion/workflow_completion.proto @@ -33,6 +33,12 @@ message Success { repeated uint32 used_internal_flags = 6; // The versioning behavior this workflow is currently using temporal.api.enums.v1.VersioningBehavior versioning_behavior = 7; + // Asks Core to complete the workflow task now if local activities are still running, as it + // does at a local activity heartbeat. The completion carries the commands and markers so far + // and asks the server for a new task, and the local activities run on. Lang sets it from a + // count or a clock, so Core ignores it while replaying. Without running local activities, + // the task completes as it would anyway. + bool heartbeat_task = 8; } // Failure to activate or execute a workflow diff --git a/crates/protos/src/protos/mod.rs b/crates/protos/src/protos/mod.rs index ffd7640263..427215e933 100644 --- a/crates/protos/src/protos/mod.rs +++ b/crates/protos/src/protos/mod.rs @@ -81,6 +81,7 @@ pub mod coresdk { commands: v, used_internal_flags: vec![], versioning_behavior: VersioningBehavior::Unspecified.into(), + heartbeat_task: false, } } } diff --git a/crates/sdk-core/CHANGELOG.md b/crates/sdk-core/CHANGELOG.md index 9aedcb07cd..a119bbd945 100644 --- a/crates/sdk-core/CHANGELOG.md +++ b/crates/sdk-core/CHANGELOG.md @@ -40,6 +40,10 @@ relevant information. requests. * Experimental: the C bridge has `temporal_core_stream_store_new`, `_call` and `_free` for the stream service, built with the Redis store. +* Experimental: `Success.heartbeat_task` on a workflow activation completion asks Core to + complete the workflow task now while local activities run, as at a local activity heartbeat. + The task carries the commands and markers so far and asks for a new task. Core ignores the + flag while replaying, since lang sets it from a count or a clock. ### Breaking Changes * Replay detects changes to the service or operation of scheduled Nexus operations, reporting diff --git a/crates/sdk-core/src/core_tests/mod.rs b/crates/sdk-core/src/core_tests/mod.rs index 6dd4479bb7..88f444fb9a 100644 --- a/crates/sdk-core/src/core_tests/mod.rs +++ b/crates/sdk-core/src/core_tests/mod.rs @@ -2,6 +2,7 @@ mod activity_tasks; mod event_groups; mod queries; mod replay_flag; +mod task_heartbeat; mod updates; mod workers; mod workflow_cancels; diff --git a/crates/sdk-core/src/core_tests/task_heartbeat.rs b/crates/sdk-core/src/core_tests/task_heartbeat.rs new file mode 100644 index 0000000000..47a6223530 --- /dev/null +++ b/crates/sdk-core/src/core_tests/task_heartbeat.rs @@ -0,0 +1,370 @@ +//! Lang's early workflow task heartbeat: `Success.heartbeat_task` completes the task while local +//! activities run, and replay must not depend on it. + +use crate::{ + Worker, + replay::TestHistoryBuilder, + test_help::{ + MockPollCfg, ResponseType, WorkerTestHelpers, build_mock_pollers, + drain_pollers_and_shutdown, mock_worker, schedule_local_activity_cmd, start_timer_cmd, + }, + worker::client::{WorkflowTaskCompletion, mocks::mock_worker_client}, +}; +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; +use temporalio_common::protos::{ + coresdk::{ + ActivityTaskCompletion, + activity_result::ActivityExecutionResult, + activity_task::activity_task, + workflow_activation::{WorkflowActivation, WorkflowActivationJob, workflow_activation_job}, + workflow_commands::{ + ActivityCancellationType, CompleteWorkflowExecution, workflow_command, + }, + workflow_completion::{WorkflowActivationCompletion, workflow_activation_completion}, + }, + temporal::api::{ + command::v1::command::Attributes, + enums::v1::{CommandType, EventType}, + history::v1::{MarkerRecordedEventAttributes, history_event}, + workflowservice::v1::RespondWorkflowTaskCompletedResponse, + }, +}; + +// Long enough that Core's own heartbeat, at 80% of it, never fires during a test. +const WFT_TIMEOUT: Duration = Duration::from_secs(60); + +fn first_commands() -> Vec { + vec![ + start_timer_cmd(1, Duration::from_secs(3600)), + schedule_local_activity_cmd( + 1, + "1", + ActivityCancellationType::TryCancel, + Duration::from_secs(60), + ), + ] +} + +fn completion( + run_id: &str, + commands: Vec, + heartbeat_task: bool, +) -> WorkflowActivationCompletion { + let mut completion = WorkflowActivationCompletion::from_cmds(run_id, commands); + let Some(workflow_activation_completion::Status::Successful(success)) = + completion.status.as_mut() + else { + unreachable!("from_cmds builds a successful completion"); + }; + success.heartbeat_task = heartbeat_task; + completion +} + +fn is_init(activation: &WorkflowActivation) -> bool { + matches!( + activation.jobs.as_slice(), + [WorkflowActivationJob { + variant: Some(workflow_activation_job::Variant::InitializeWorkflow(_)), + }] + ) +} + +fn resolves_la_1(activation: &WorkflowActivation) -> bool { + matches!( + activation.jobs.as_slice(), + [WorkflowActivationJob { + variant: Some(workflow_activation_job::Variant::ResolveActivity(resolve)), + }] if resolve.seq == 1 && resolve.is_local + ) +} + +fn command_types(completion: &WorkflowTaskCompletion) -> Vec { + completion + .commands + .iter() + .map(|command| command.command_type()) + .collect() +} + +/// A worker over `history` that records every workflow task completion it reports. +fn recording_worker( + history: TestHistoryBuilder, + responses: impl IntoIterator, + expected_completions: usize, +) -> (Worker, Arc>>) { + let reported = Arc::new(Mutex::new(vec![])); + let mut cfg = + MockPollCfg::from_resp_batches("fake_wf_id", history, responses, mock_worker_client()); + cfg.num_expected_fails = 0; + cfg.num_expected_completions = Some(expected_completions.into()); + let recorder = reported.clone(); + cfg.completion_mock_fn = Some(Box::new(move |completion| { + recorder.lock().unwrap().push(completion.clone()); + Ok(RespondWorkflowTaskCompletedResponse::default()) + })); + let mut mock = build_mock_pollers(cfg); + mock.worker_cfg(|wc| wc.max_cached_workflows = 1); + (mock_worker(mock), reported) +} + +/// What a cut produces on the server: the first task completes with the timer and asks for a new +/// task while the local activity runs, and the next task records the marker. +fn cut_history_up_to_second_task() -> TestHistoryBuilder { + let mut t = TestHistoryBuilder::default(); + t.add_wfe_started_with_wft_timeout(WFT_TIMEOUT); + t.add_full_wf_task(); + t.add_timer_started("1".to_string()); + t.add_workflow_task_scheduled_and_started(); + t +} + +/// The history the server writes for `completions`, one workflow task each. +fn history_of(completions: &[WorkflowTaskCompletion]) -> TestHistoryBuilder { + let mut t = TestHistoryBuilder::default(); + t.add_wfe_started_with_wft_timeout(WFT_TIMEOUT); + for completion in completions { + t.add_workflow_task_scheduled_and_started(); + t.add_workflow_task_completed(); + let completed_id = t.current_event_id(); + if let Some(history_event::Attributes::WorkflowTaskCompletedEventAttributes(attrs)) = + t.last_event().and_then(|event| event.attributes.as_mut()) + { + attrs.sdk_metadata = Some(completion.sdk_metadata.clone()); + } + for command in &completion.commands { + match command.attributes.clone() { + Some(Attributes::StartTimerCommandAttributes(timer)) => { + t.add_timer_started(timer.timer_id); + } + Some(Attributes::RecordMarkerCommandAttributes(marker)) => { + t.add(MarkerRecordedEventAttributes { + marker_name: marker.marker_name, + details: marker.details, + workflow_task_completed_event_id: completed_id, + header: marker.header, + failure: marker.failure, + }); + } + Some(Attributes::CompleteWorkflowExecutionCommandAttributes(_)) => { + t.add_workflow_execution_completed(); + } + other => panic!("no history event for {other:?}"), + } + } + } + t +} + +/// Runs the cut live and answers the two completions it reported. +async fn run_cut_live() -> Vec { + let (core, reported) = recording_worker( + cut_history_up_to_second_task(), + [ResponseType::ToTaskNum(1), ResponseType::ToTaskNum(2)], + 2, + ); + let activation = core.poll_workflow_activation().await.unwrap(); + assert!(is_init(&activation), "{activation:?}"); + core.complete_workflow_activation(completion(&activation.run_id, first_commands(), true)) + .await + .unwrap(); + let task = core.poll_activity_task().await.unwrap(); + let Some(activity_task::Variant::Start(start)) = &task.variant else { + panic!("expected the local activity, got {task:?}"); + }; + assert!(start.is_local); + { + let reported = reported.lock().unwrap(); + let [cut] = reported.as_slice() else { + panic!("the cut reports one completion while the local activity runs: {reported:?}"); + }; + assert!(cut.force_create_new_workflow_task); + assert_eq!(command_types(cut), [CommandType::StartTimer]); + } + core.complete_activity_task(ActivityTaskCompletion { + task_token: task.task_token, + result: Some(ActivityExecutionResult::ok(b"done".to_vec().into())), + }) + .await + .unwrap(); + let activation = core.poll_workflow_activation().await.unwrap(); + assert!(resolves_la_1(&activation), "{activation:?}"); + core.complete_execution(&activation.run_id).await; + drain_pollers_and_shutdown(&core).await; + let reported = reported.lock().unwrap().clone(); + let [_, last] = reported.as_slice() else { + panic!("expected two completions, got {reported:?}"); + }; + assert!(!last.force_create_new_workflow_task); + assert_eq!( + command_types(last), + [ + CommandType::RecordMarker, + CommandType::CompleteWorkflowExecution + ] + ); + reported +} + +#[tokio::test] +async fn a_heartbeat_task_completes_the_task_while_the_local_activity_runs() { + run_cut_live().await; +} + +#[tokio::test] +async fn a_heartbeat_task_without_local_activities_completes_the_task_as_usual() { + let mut t = TestHistoryBuilder::default(); + t.add_wfe_started_with_wft_timeout(WFT_TIMEOUT); + t.add_workflow_task_scheduled_and_started(); + let (core, reported) = recording_worker(t, [ResponseType::ToTaskNum(1)], 1); + let activation = core.poll_workflow_activation().await.unwrap(); + core.complete_workflow_activation(completion( + &activation.run_id, + vec![start_timer_cmd(1, Duration::from_secs(3600))], + true, + )) + .await + .unwrap(); + drain_pollers_and_shutdown(&core).await; + let reported = reported.lock().unwrap(); + let [only] = reported.as_slice() else { + panic!("expected one completion, got {reported:?}"); + }; + assert!( + !only.force_create_new_workflow_task, + "with nothing running, there is nothing to heartbeat for" + ); + assert_eq!(command_types(only), [CommandType::StartTimer]); +} + +/// Replays `history` to its end with lang deciding the cut by `heartbeat_task`, and checks that +/// lang sees the activations of the live run and that nothing goes to the server. +async fn replay(history: TestHistoryBuilder, heartbeat_task: bool) { + let (core, reported) = recording_worker(history, [ResponseType::AllHistory], 0); + let activation = core.poll_workflow_activation().await.unwrap(); + assert!(is_init(&activation), "{activation:?}"); + core.complete_workflow_activation(completion( + &activation.run_id, + first_commands(), + heartbeat_task, + )) + .await + .unwrap(); + let activation = tokio::time::timeout(Duration::from_secs(10), core.poll_workflow_activation()) + .await + .expect("replay resolves the local activity from its marker") + .unwrap(); + assert!(resolves_la_1(&activation), "{activation:?}"); + core.complete_execution(&activation.run_id).await; + drain_pollers_and_shutdown(&core).await; + assert!(reported.lock().unwrap().is_empty()); +} + +#[rstest::rstest] +#[case::with_the_cut(true)] +#[case::without_the_cut(false)] +#[tokio::test] +async fn a_cut_history_replays_whatever_lang_decides(#[case] heartbeat_task: bool) { + let live = run_cut_live().await; + let history = history_of(&live); + let tasks = history + .get_full_history_info() + .unwrap() + .events() + .iter() + .filter(|event| event.event_type() == EventType::WorkflowTaskCompleted) + .count(); + assert_eq!(tasks, 2, "the cut adds a workflow task"); + replay(history, heartbeat_task).await; +} + +#[tokio::test] +async fn an_uncut_history_replays_when_lang_asks_for_a_cut() { + // The local activity resolved within the first task, so its marker sits in that task. A + // replay whose clock says to cut must not ask for the task that the live run never made. + let mut t = TestHistoryBuilder::default(); + t.add_wfe_started_with_wft_timeout(WFT_TIMEOUT); + t.add_full_wf_task(); + t.add_timer_started("1".to_string()); + t.add_local_activity_result_marker(1, "1", b"done".to_vec().into()); + t.add_workflow_execution_completed(); + replay(t, true).await; +} + +#[tokio::test] +async fn a_cut_task_runs_its_local_activity_after_a_cache_miss() { + // A new Worker gets the second task with the whole history. It replays the first task, where + // the cut is ignored, and then runs the local activity that the cut left without a marker. + let (core, reported) = recording_worker( + cut_history_up_to_second_task(), + [ResponseType::AllHistory], + 1, + ); + let activation = core.poll_workflow_activation().await.unwrap(); + assert!(is_init(&activation), "{activation:?}"); + core.complete_workflow_activation(completion(&activation.run_id, first_commands(), true)) + .await + .unwrap(); + let task = core.poll_activity_task().await.unwrap(); + assert!(matches!( + &task.variant, + Some(activity_task::Variant::Start(start)) if start.is_local + )); + core.complete_activity_task(ActivityTaskCompletion { + task_token: task.task_token, + result: Some(ActivityExecutionResult::ok(b"done".to_vec().into())), + }) + .await + .unwrap(); + let activation = core.poll_workflow_activation().await.unwrap(); + assert!(resolves_la_1(&activation), "{activation:?}"); + core.complete_workflow_activation(completion( + &activation.run_id, + vec![CompleteWorkflowExecution { result: None }.into()], + true, + )) + .await + .unwrap(); + drain_pollers_and_shutdown(&core).await; + let reported = reported.lock().unwrap(); + let [only] = reported.as_slice() else { + panic!("expected one completion, got {reported:?}"); + }; + assert!(!only.force_create_new_workflow_task); + assert_eq!( + command_types(only), + [ + CommandType::RecordMarker, + CommandType::CompleteWorkflowExecution + ] + ); +} + +#[tokio::test] +async fn a_cut_history_has_the_events_the_server_writes() { + let live = run_cut_live().await; + let events: Vec = history_of(&live) + .get_full_history_info() + .unwrap() + .events() + .iter() + .map(|event| event.event_type()) + .collect(); + assert_eq!( + events, + [ + EventType::WorkflowExecutionStarted, + EventType::WorkflowTaskScheduled, + EventType::WorkflowTaskStarted, + EventType::WorkflowTaskCompleted, + EventType::TimerStarted, + EventType::WorkflowTaskScheduled, + EventType::WorkflowTaskStarted, + EventType::WorkflowTaskCompleted, + EventType::MarkerRecorded, + EventType::WorkflowExecutionCompleted, + ] + ); +} diff --git a/crates/sdk-core/src/worker/workflow/managed_run.rs b/crates/sdk-core/src/worker/workflow/managed_run.rs index bb8dd3399c..d1c790dd19 100644 --- a/crates/sdk-core/src/worker/workflow/managed_run.rs +++ b/crates/sdk-core/src/worker/workflow/managed_run.rs @@ -425,6 +425,7 @@ impl ManagedRun { versioning_behavior: VersioningBehavior, resp_chan: Option>, is_forced_failure: bool, + heartbeat_task: bool, ) -> Result> { let activation_was_only_eviction = self.activation_is_eviction(); if activation_was_only_eviction && !commands.is_empty() { @@ -504,6 +505,7 @@ impl ManagedRun { resp_chan, is_forced_failure, versioning_behavior, + heartbeat_task, }; // Verify we can actually apply the next workflow task, which will happen as part of @@ -673,6 +675,7 @@ impl ManagedRun { VersioningBehavior::Unspecified, // Doesn't matter since we're failing wf resp_chan, true, + false, ) .unwrap_or_else(|e| { dbg_panic!("Got next page request when auto-failing workflow: {e:?}"); @@ -748,7 +751,10 @@ impl ManagedRun { let completing_heartbeat_autocomplete = matches!(self.activation, Some(OutstandingActivation::Autocomplete)) && self.waiting_on_la.is_some(); + // Lang decides a task heartbeat from a count or a clock, so replay must not depend on it. + let lang_heartbeat = completion.heartbeat_task && !self.wfm.machines.replaying; let completing_la_heartbeat = completing_heartbeat_autocomplete + || lang_heartbeat || self .waiting_on_la .as_ref() @@ -1634,6 +1640,8 @@ struct RunActivationCompletion { /// unblock. Must always be `Some` when initialized. resp_chan: Option>, versioning_behavior: VersioningBehavior, + /// Lang asked to complete the workflow task now while local activities run. + heartbeat_task: bool, } #[derive(Debug, derive_more::From)] enum ActOrFulfill { diff --git a/crates/sdk-core/src/worker/workflow/mod.rs b/crates/sdk-core/src/worker/workflow/mod.rs index 4dc05d1556..27d43696f5 100644 --- a/crates/sdk-core/src/worker/workflow/mod.rs +++ b/crates/sdk-core/src/worker/workflow/mod.rs @@ -1444,6 +1444,7 @@ fn validate_completion( commands, used_flags: success.used_internal_flags, versioning_behavior: success.versioning_behavior.try_into().unwrap_or_default(), + heartbeat_task: success.heartbeat_task, }) } Some(workflow_activation_completion::Status::Failed(failure)) => { @@ -1468,6 +1469,7 @@ enum ValidatedCompletion { commands: Vec, used_flags: Vec, versioning_behavior: VersioningBehavior, + heartbeat_task: bool, }, Fail { run_id: String, diff --git a/crates/sdk-core/src/worker/workflow/workflow_stream.rs b/crates/sdk-core/src/worker/workflow/workflow_stream.rs index 04207f5d8a..15bf05d70d 100644 --- a/crates/sdk-core/src/worker/workflow/workflow_stream.rs +++ b/crates/sdk-core/src/worker/workflow/workflow_stream.rs @@ -306,6 +306,7 @@ impl WFStream { commands, used_flags, versioning_behavior, + heartbeat_task, .. } => match rh.successful_completion( commands, @@ -313,6 +314,7 @@ impl WFStream { versioning_behavior, complete.response_tx, false, + heartbeat_task, ) { Ok(acts) => acts, Err(npr) => { diff --git a/crates/sdk/src/workflow_future.rs b/crates/sdk/src/workflow_future.rs index f18604ce0b..cb96af1e58 100644 --- a/crates/sdk/src/workflow_future.rs +++ b/crates/sdk/src/workflow_future.rs @@ -238,6 +238,7 @@ impl WorkflowFuture { commands: activation_cmds, used_internal_flags: vec![], versioning_behavior: VersioningBehavior::Unspecified.into(), + heartbeat_task: false, }, )), ..Default::default() From 13629af5903bcf8b0bf8f11eae14fdea5288eec1 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 13 Oct 2026 19:01:29 -0700 Subject: [PATCH 2/8] Attached a Workflow's stream reader to the stream's notifier. Subscribe attaches the reader as a workflow target before its first read, with a request id from its run and subscription id, so a rerun attaches as the same reader. The process remembers what it attached, so later reads send no attach. The attach names the chain that holds the Workflow id now and arms an empty stream, so the first append wakes the Workflow. Unsubscribe detaches. --- .../temporal/sdk/core/streams/streams.proto | 34 ++ crates/sdk-core/CHANGELOG.md | 6 + crates/sdk-core/src/streams.rs | 112 +++- crates/streams/src/handle.rs | 17 +- crates/streams/src/lib.rs | 4 +- crates/streams/src/notify.rs | 41 +- crates/streams/src/service.rs | 45 +- crates/streams/src/subscribe.rs | 526 ++++++++++++++++++ 8 files changed, 761 insertions(+), 24 deletions(-) create mode 100644 crates/streams/src/subscribe.rs diff --git a/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto b/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto index 2d1cb59e6e..89f45bde98 100644 --- a/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto +++ b/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto @@ -28,6 +28,13 @@ service StreamService { // Waits until the stream notifications this process sent are out, for at most 10 seconds. A // notification goes in the background, so a process calls it before it stops. rpc FlushNotifications(FlushNotificationsRequest) returns (FlushNotificationsResponse); + // Attaches a Workflow's reader of a stream to the stream's notifier, so the notifier wakes the + // Workflow with a `StreamWake` on its next Workflow Task. A reader calls it before its first + // read. The process attaches each subscription once and answers a repeat from memory, and the + // server takes a repeat from another process as the same attach. + rpc Subscribe(SubscribeRequest) returns (SubscribeResponse); + // Detaches a Workflow's reader that stops before the stream ends. + rpc Unsubscribe(UnsubscribeRequest) returns (UnsubscribeResponse); } // Which store holds a process's streams, and how to reach it. Lang builds one store from it and @@ -239,6 +246,33 @@ message FlushNotificationsRequest {} message FlushNotificationsResponse {} +message SubscribeRequest { + // The stream to read. The reader's Workflow is in the stream's namespace. + StreamAddress stream = 1; + // The reader's Workflow and run. + string workflow_id = 2; + string run_id = 3; + // Names the reader in its run. Lang derives it from the reader's call number, so replay gives + // the same one. At most 256 bytes. + string subscription_id = 4; +} + +message SubscribeResponse { + // The first run of the stream's run chain, which keys the notifier the reader attached to. + string chain_first_run_id = 1; + // Whether this call sent the attach. False when this process attached the reader before. + bool attached = 2; +} + +message UnsubscribeRequest { + // The stream, the reader's run and the subscription, as the Subscribe named them. + StreamAddress stream = 1; + string run_id = 2; + string subscription_id = 3; +} + +message UnsubscribeResponse {} + message DeleteOwnerRequest { string namespace = 1; StreamOwnerKind owner_kind = 2; diff --git a/crates/sdk-core/CHANGELOG.md b/crates/sdk-core/CHANGELOG.md index 83f83c220d..4c29b9d231 100644 --- a/crates/sdk-core/CHANGELOG.md +++ b/crates/sdk-core/CHANGELOG.md @@ -66,6 +66,12 @@ relevant information. finds the owner ended closes only the store. The notifier's owner check closes the notifier. * Experimental: `StreamStore` has `close_topic_with` and `close_stream`, which carry the close outcome to a wrapper that tells the notifier. +* Experimental: the new `Subscribe` and `Unsubscribe` stream calls attach a Workflow's reader to + a stream's notifier as a `workflow_target`, and detach it. The attach's request id is + `jetty//`, so a rerun attaches as the same reader. A process attaches + each subscription once and answers a repeat from memory. The attach names the chain that holds + the Workflow id now, and on an empty stream it arms the stream, so the first append wakes the + Workflow. ### Breaking Changes * Replay detects changes to the service or operation of scheduled Nexus operations, reporting diff --git a/crates/sdk-core/src/streams.rs b/crates/sdk-core/src/streams.rs index 2ccc324276..1c133d90e5 100644 --- a/crates/sdk-core/src/streams.rs +++ b/crates/sdk-core/src/streams.rs @@ -6,20 +6,22 @@ use temporalio_client::Connection; use temporalio_common::protos::temporal::api::{ common::v1::WorkflowExecution, enums::v1::StreamOwnerKind, - stream::v1::StreamReference, + stream::v1::{StreamReference, WorkflowWakeTarget}, workflowservice::v1::{ - DescribeWorkflowExecutionRequest, DescribeWorkflowExecutionResponse, NotifyStreamRequest, + AttachStreamCallbackRequest, DescribeWorkflowExecutionRequest, + DescribeWorkflowExecutionResponse, DetachStreamCallbackRequest, NotifyStreamRequest, }, }; use temporalio_streams::{ CloseOutcome, NotifierClient, NotifyError, OwnerClient, OwnerDescription, OwnerError, - StreamNotification, + StreamNotification, WorkflowAttachment, proto::ChainId, }; pub use temporalio_streams::{StreamError, StreamResult, StreamService, StreamStore, proto}; use tonic::IntoRequest; /// Connects to the store `config` names. The service asks `connection`'s server about streams' -/// owners and sends the stream notifications the config turns on to it. +/// owners, sends the stream notifications the config turns on to it, and attaches Workflow +/// readers there. pub async fn connect_stream_service( config: proto::StreamStoreConfig, connection: Connection, @@ -49,19 +51,69 @@ impl NotifierClient for ConnectionNotifierClient { .map(|_| ()) .map_err(notify_error) } + + async fn attach_workflow(&self, attachment: WorkflowAttachment) -> Result<(), NotifyError> { + self.connection + .workflow_service() + .attach_stream_callback(attach_request(attachment).into_request()) + .await + .map(|_| ()) + .map_err(notify_error) + } + + async fn detach( + &self, + chain: ChainId, + topic: String, + request_id: String, + ) -> Result<(), NotifyError> { + self.connection + .workflow_service() + .detach_stream_callback( + DetachStreamCallbackRequest { + namespace: chain.namespace.clone(), + stream_ref: Some(chain_reference(&chain, topic)), + request_id, + ..Default::default() + } + .into_request(), + ) + .await + .map(|_| ()) + .map_err(notify_error) + } +} + +/// The notifier of `chain`'s stream on `topic`. The run chain's first run keys the notifier, so one +/// notifier serves the stream across Continue-as-New and a new chain on the same Workflow id gets +/// its own. +fn chain_reference(chain: &ChainId, topic: String) -> StreamReference { + StreamReference { + owner_kind: StreamOwnerKind::Workflow as i32, + workflow_id: chain.workflow_id.clone(), + run_id: chain.first_run_id.clone(), + topic, + } +} + +fn attach_request(attachment: WorkflowAttachment) -> AttachStreamCallbackRequest { + AttachStreamCallbackRequest { + namespace: attachment.chain.namespace.clone(), + stream_ref: Some(chain_reference(&attachment.chain, attachment.topic)), + request_id: attachment.request_id, + workflow_target: Some(WorkflowWakeTarget { + workflow_id: attachment.workflow_id, + run_id: attachment.run_id, + subscription_id: attachment.subscription_id, + }), + ..Default::default() + } } fn notify_request(notification: StreamNotification) -> NotifyStreamRequest { NotifyStreamRequest { + stream_ref: Some(chain_reference(¬ification.chain, notification.topic)), namespace: notification.chain.namespace, - // The run chain's first run keys the notifier, so one notifier serves the stream across - // Continue-as-New and a new chain on the same Workflow id gets its own. - stream_ref: Some(StreamReference { - owner_kind: StreamOwnerKind::Workflow as i32, - workflow_id: notification.chain.workflow_id, - run_id: notification.chain.first_run_id, - topic: notification.topic, - }), position: notification.position, counter: notification.counter, close: notification.close.is_some(), @@ -261,6 +313,42 @@ mod tests { assert!(request.close_result.is_none()); } + #[test] + fn an_attach_names_the_chains_notifier_and_wakes_the_reader() { + let request = attach_request(WorkflowAttachment { + chain: proto::ChainId { + namespace: "ns".to_string(), + workflow_id: "owner".to_string(), + first_run_id: "chain-1".to_string(), + }, + topic: "tokens".to_string(), + request_id: "jetty/reader-run/1".to_string(), + workflow_id: "reader".to_string(), + run_id: "reader-run".to_string(), + subscription_id: "1".to_string(), + }); + assert_eq!(request.namespace, "ns"); + assert_eq!( + request.stream_ref, + Some(StreamReference { + owner_kind: StreamOwnerKind::Workflow as i32, + workflow_id: "owner".to_string(), + run_id: "chain-1".to_string(), + topic: "tokens".to_string(), + }) + ); + assert_eq!(request.request_id, "jetty/reader-run/1"); + assert_eq!( + request.workflow_target, + Some(WorkflowWakeTarget { + workflow_id: "reader".to_string(), + run_id: "reader-run".to_string(), + subscription_id: "1".to_string(), + }) + ); + assert!(request.callback.is_none(), "exactly one target"); + } + #[test] fn only_answers_a_retry_cant_change_refuse_a_notification() { assert!(matches!( diff --git a/crates/streams/src/handle.rs b/crates/streams/src/handle.rs index 8c2df3f34b..31a09c393b 100644 --- a/crates/streams/src/handle.rs +++ b/crates/streams/src/handle.rs @@ -484,6 +484,21 @@ impl Streams { )) } + /// The run chain a checked stream address names now. A pinned run's chain comes from the + /// owner cache, since it never changes. The latest run's chain comes from a fresh describe, + /// for a caller that must not land on a chain that already ended. + pub(crate) async fn chain_now(&self, stream: &StreamAddress) -> StreamResult { + let key = owner_key(stream); + if !key.run_id.is_empty() { + return self.chain(stream).await; + } + let described = self.describe(&key, "").await.map_err(not_found)?; + let first_run_id = first_run_id(&key, &described)?; + self.chains + .put(key, &first_run_id, self.options.owner_recheck); + Ok(chain_of(stream, &first_run_id)) + } + /// The chain a stream address names now. async fn chain(&self, stream: &StreamAddress) -> StreamResult { let key = owner_key(stream); @@ -526,7 +541,7 @@ impl Streams { } /// The address of a stream this release keeps, checked. -fn address(stream: Option<&StreamAddress>) -> StreamResult<&StreamAddress> { +pub(crate) fn address(stream: Option<&StreamAddress>) -> StreamResult<&StreamAddress> { let stream = stream.ok_or_else(|| StreamError::invalid_argument("a stream call needs its stream"))?; if stream.owner_kind != StreamOwnerKind::Workflow as i32 { diff --git a/crates/streams/src/lib.rs b/crates/streams/src/lib.rs index 52bfa93e18..93cbc42f82 100644 --- a/crates/streams/src/lib.rs +++ b/crates/streams/src/lib.rs @@ -20,6 +20,7 @@ mod record; mod redis; mod service; mod store; +mod subscribe; #[cfg(feature = "redis")] pub use self::redis::{RedisStore, RedisStoreOptions}; @@ -29,7 +30,8 @@ pub use handle::{OWNER_RECHECK, Streams, StreamsOptions}; pub use memory::MemoryStore; pub use notify::{ Arm, CLOSE_RETRY_CAP, CLOSE_RETRY_FIRST, DEFAULT_MAX_NOTIFIERS, FLUSH_LIMIT, NOTIFIER_REFRESH, - Notifier, NotifierClient, NotifyError, NotifyingStore, StreamNotification, progress_counter, + Notifier, NotifierClient, NotifyError, NotifyingStore, StreamNotification, WorkflowAttachment, + progress_counter, }; pub use owner::{OwnerClient, OwnerDescription, OwnerError}; pub use reader::{DEFAULT_MAX_RECORDS, MAX_TRACKED_PRODUCERS, ReadTarget, read_page}; diff --git a/crates/streams/src/notify.rs b/crates/streams/src/notify.rs index 4b86c05543..89d7a11bf8 100644 --- a/crates/streams/src/notify.rs +++ b/crates/streams/src/notify.rs @@ -83,11 +83,50 @@ pub enum NotifyError { Unsupported(String), } -/// Sends notifications to the stream notifier on the server. +/// A Workflow's reader, to attach to a stream's notifier as a wake target. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct WorkflowAttachment { + /// The run chain whose stream it reads. Its first run keys the notifier. + pub chain: ChainId, + /// The stream's topic. + pub topic: String, + /// The attach is idempotent by it, and a detach names it. + pub request_id: String, + /// The reader's Workflow, in the chain's namespace. + pub workflow_id: String, + /// The reader's run. + pub run_id: String, + /// Names the reader in its run. + pub subscription_id: String, +} + +/// Sends notifications to the stream notifier on the server, and attaches and detaches the +/// Workflows it wakes. #[async_trait::async_trait] pub trait NotifierClient: Send + Sync { /// Sends one notification. async fn notify(&self, notification: StreamNotification) -> Result<(), NotifyError>; + + /// Attaches a Workflow's reader to a stream's notifier. Idempotent by its request id. + async fn attach_workflow(&self, attachment: WorkflowAttachment) -> Result<(), NotifyError> { + let _ = attachment; + Err(NotifyError::Unsupported( + "this client can't attach readers".to_string(), + )) + } + + /// Detaches what the attach with `request_id` attached to `chain`'s stream on `topic`. + async fn detach( + &self, + chain: ChainId, + topic: String, + request_id: String, + ) -> Result<(), NotifyError> { + let _ = (chain, topic, request_id); + Err(NotifyError::Unsupported( + "this client can't detach readers".to_string(), + )) + } } /// The notifier counter for a store position. diff --git a/crates/streams/src/service.rs b/crates/streams/src/service.rs index 99134a7309..81df8e92d5 100644 --- a/crates/streams/src/service.rs +++ b/crates/streams/src/service.rs @@ -6,8 +6,9 @@ use crate::{ proto::{ AppendRequest, CloseRequest, DeleteOwnerRequest, FlushNotificationsRequest, FlushNotificationsResponse, LatestRequest, ReadRequest, StreamFailure, StreamStoreConfig, - stream_store_config, + SubscribeRequest, UnsubscribeRequest, stream_store_config, }, + subscribe::Subscriptions, }; use prost::Message; use std::{sync::Arc, time::Duration}; @@ -27,6 +28,7 @@ const MAX_READ_WAIT: Duration = Duration::from_secs(60); pub struct StreamService { store: Arc, streams: Streams, + subscriptions: Subscriptions, } impl std::fmt::Debug for StreamService { @@ -48,7 +50,8 @@ impl StreamService { } /// Like [StreamService::connect], closing each stream's notifier through `notifier` after - /// the store closes the stream, and sending progress too when the config turns it on. + /// the store closes the stream, sending progress too when the config turns it on, and + /// attaching and detaching Workflow readers through it. pub async fn connect_with_notifier( config: StreamStoreConfig, owner: Arc, @@ -70,7 +73,7 @@ impl StreamService { "stream notifications need a client to send them with", )); } - let store: Arc = match notifier { + let store: Arc = match notifier.clone() { Some(client) => { let max_notifiers = match config.max_notifiers { 0 => DEFAULT_MAX_NOTIFIERS, @@ -84,11 +87,11 @@ impl StreamService { } None => store, }; - Ok(Self::new( - store, - owner, - StreamsOptions::for_retention(retention), - )) + let service = Self::new(store, owner, StreamsOptions::for_retention(retention)); + Ok(match notifier { + Some(client) => service.with_subscriptions(client), + None => service, + }) } /// Serves calls on `store`. @@ -98,7 +101,17 @@ impl StreamService { options: StreamsOptions, ) -> Self { let streams = Streams::new(store.clone(), owner, options); - Self { store, streams } + Self { + store, + streams, + subscriptions: Subscriptions::default(), + } + } + + /// Attaches and detaches Workflow readers through `client`. + pub fn with_subscriptions(mut self, client: Arc) -> Self { + self.subscriptions = Subscriptions::new(Some(client)); + self } /// Serves one call: `rpc` is a `StreamService` method name and `request` its serialized @@ -138,6 +151,20 @@ impl StreamService { .await; Ok(FlushNotificationsResponse {}.encode_to_vec()) } + "Subscribe" => self + .subscriptions + .subscribe( + &self.streams, + self.store.as_ref(), + decode::(rpc, request)?, + ) + .await + .map(|r| r.encode_to_vec()), + "Unsubscribe" => self + .subscriptions + .unsubscribe(&self.streams, decode::(rpc, request)?) + .await + .map(|r| r.encode_to_vec()), other => Err(StreamError::unsupported(format!( "the stream service has no call {other:?}" ))), diff --git a/crates/streams/src/subscribe.rs b/crates/streams/src/subscribe.rs new file mode 100644 index 0000000000..761120713e --- /dev/null +++ b/crates/streams/src/subscribe.rs @@ -0,0 +1,526 @@ +//! Workflow readers that sleep on a stream's notifier. +//! +//! A Workflow's reader attaches to the stream's notifier before its first read, so the notifier +//! wakes the Workflow on its next Workflow Task when the stream moves. The attach names the reader +//! by its run and its subscription id, which replay gives back, so a rerun attaches as the same +//! reader and the server keeps one callback. + +use crate::{ + NotifierClient, NotifyError, StreamError, StreamResult, StreamStore, Streams, + WorkflowAttachment, + handle::address, + proto::{ + StreamAddress, SubscribeRequest, SubscribeResponse, UnsubscribeRequest, UnsubscribeResponse, + }, +}; +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, +}; + +/// The longest subscription id the server takes. +const MAX_SUBSCRIPTION_ID_BYTES: usize = 256; + +/// The subscriptions this process attached, so a reader's later reads send no attach. +#[derive(Default)] +pub(crate) struct Subscriptions { + client: Option>, + attached: Mutex>, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct SubscriptionKey { + namespace: String, + owner_workflow_id: String, + owner_run_id: String, + topic: String, + request_id: String, +} + +impl SubscriptionKey { + fn new(stream: &StreamAddress, request_id: &str) -> Self { + Self { + namespace: stream.namespace.clone(), + owner_workflow_id: stream.workflow_id.clone(), + owner_run_id: stream.run_id.clone(), + topic: stream.topic.clone(), + request_id: request_id.to_string(), + } + } +} + +/// The request id of a reader's attach. The server dedupes attaches by it, so it depends only on +/// what replay gives back. +fn request_id(run_id: &str, subscription_id: &str) -> String { + format!("jetty/{run_id}/{subscription_id}") +} + +fn check_subscription(run_id: &str, subscription_id: &str) -> StreamResult<()> { + if run_id.is_empty() || subscription_id.is_empty() { + return Err(StreamError::invalid_argument( + "a subscription needs the reader's run id and its subscription id", + )); + } + if subscription_id.len() > MAX_SUBSCRIPTION_ID_BYTES { + return Err(StreamError::invalid_argument(format!( + "a subscription id is at most {MAX_SUBSCRIPTION_ID_BYTES} bytes, not {}", + subscription_id.len() + ))); + } + Ok(()) +} + +fn subscribe_error(error: NotifyError) -> StreamError { + match error { + NotifyError::Refused(message) => StreamError::refused(message), + NotifyError::Unsupported(message) => StreamError::unsupported(message), + NotifyError::Failed(message) => StreamError::storage(message), + } +} + +impl Subscriptions { + pub(crate) fn new(client: Option>) -> Self { + Self { + client, + attached: Mutex::default(), + } + } + + fn client(&self) -> StreamResult<&Arc> { + self.client.as_ref().ok_or_else(|| { + StreamError::unsupported("a Workflow reader needs a client connected to the server") + }) + } + + /// Attaches the reader once per process. On an empty stream the attach also arms the stream, + /// in the same store step that finds it empty, so the first append wakes the Workflow. + pub(crate) async fn subscribe( + &self, + streams: &Streams, + store: &dyn StreamStore, + request: SubscribeRequest, + ) -> StreamResult { + let stream = address(request.stream.as_ref())?; + if request.workflow_id.is_empty() { + return Err(StreamError::invalid_argument( + "a subscription needs the reader's Workflow id", + )); + } + check_subscription(&request.run_id, &request.subscription_id)?; + let request_id = request_id(&request.run_id, &request.subscription_id); + let key = SubscriptionKey::new(stream, &request_id); + if let Some(first_run_id) = self.attached.lock().unwrap().get(&key) { + return Ok(SubscribeResponse { + chain_first_run_id: first_run_id.clone(), + attached: false, + }); + } + let client = self.client()?; + // Fresh for the latest run, since the Workflow id may have moved on to a new chain since a + // read warmed the cache, and an attach to the ended chain would never wake the reader. + let chain = streams.chain_now(stream).await?; + client + .attach_workflow(WorkflowAttachment { + chain: chain.clone(), + topic: stream.topic.clone(), + request_id, + workflow_id: request.workflow_id, + run_id: request.run_id, + subscription_id: request.subscription_id, + }) + .await + .map_err(subscribe_error)?; + // The reader's first read arms an empty stream too, so a failed arm costs nothing. + if let Err(error) = store.newest_or_arm(&chain, &stream.topic).await { + tracing::warn!( + workflow_id = %chain.workflow_id, + topic = %stream.topic, + %error, + "Could not arm the stream a Workflow reader subscribed to" + ); + } + self.attached + .lock() + .unwrap() + .insert(key, chain.first_run_id.clone()); + Ok(SubscribeResponse { + chain_first_run_id: chain.first_run_id, + attached: true, + }) + } + + /// Detaches the reader from the chain it attached to, or from the chain the stream names now + /// when this process never attached it. + pub(crate) async fn unsubscribe( + &self, + streams: &Streams, + request: UnsubscribeRequest, + ) -> StreamResult { + let stream = address(request.stream.as_ref())?; + check_subscription(&request.run_id, &request.subscription_id)?; + let request_id = request_id(&request.run_id, &request.subscription_id); + let client = self.client()?; + let key = SubscriptionKey::new(stream, &request_id); + let known = self.attached.lock().unwrap().get(&key).cloned(); + let chain = match known { + Some(first_run_id) => crate::proto::ChainId { + namespace: stream.namespace.clone(), + workflow_id: stream.workflow_id.clone(), + first_run_id, + }, + None => streams.chain_now(stream).await?, + }; + client + .detach(chain, stream.topic.clone(), request_id) + .await + .map_err(subscribe_error)?; + self.attached.lock().unwrap().remove(&key); + Ok(UnsubscribeResponse {}) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + MemoryStore, OwnerClient, OwnerDescription, OwnerError, StreamNotification, StreamService, + StreamsOptions, + proto::{ChainId, StoreAppendRequest, StreamFailureKind, StreamOwnerKind}, + }; + use prost::Message; + use temporalio_common::protos::temporal::api::enums::v1::WorkflowExecutionStatus; + + /// The owner's chain is `first_run_id`, which a test moves to a new chain. + struct Owner(Mutex); + + #[async_trait::async_trait] + impl OwnerClient for Owner { + async fn describe( + &self, + _: &str, + _: &str, + _: &str, + ) -> Result { + let first_run_id = self.0.lock().unwrap().clone(); + Ok(OwnerDescription { + run_id: first_run_id.clone(), + first_run_id, + status: WorkflowExecutionStatus::Running, + }) + } + } + + /// Records attaches and detaches, and answers each with `answer`. + #[derive(Default)] + struct Server { + attached: Mutex>, + detached: Mutex>, + answer: Mutex>, + } + + #[async_trait::async_trait] + impl NotifierClient for Server { + async fn notify(&self, _: StreamNotification) -> Result<(), NotifyError> { + Ok(()) + } + + async fn attach_workflow(&self, attachment: WorkflowAttachment) -> Result<(), NotifyError> { + if let Some(error) = self.answer.lock().unwrap().clone() { + return Err(error); + } + self.attached.lock().unwrap().push(attachment); + Ok(()) + } + + async fn detach( + &self, + chain: ChainId, + topic: String, + request_id: String, + ) -> Result<(), NotifyError> { + self.detached + .lock() + .unwrap() + .push((chain, topic, request_id)); + Ok(()) + } + } + + struct Setup { + service: StreamService, + store: Arc, + owner: Arc, + server: Arc, + } + + fn setup() -> Setup { + let store = Arc::new(MemoryStore::new()); + let owner = Arc::new(Owner(Mutex::new("chain-1".to_string()))); + let server = Arc::new(Server::default()); + let service = StreamService::new(store.clone(), owner.clone(), StreamsOptions::default()) + .with_subscriptions(server.clone()); + Setup { + service, + store, + owner, + server, + } + } + + fn stream() -> StreamAddress { + StreamAddress { + namespace: "ns".to_string(), + owner_kind: StreamOwnerKind::Workflow as i32, + workflow_id: "owner".to_string(), + run_id: String::new(), + topic: "tokens".to_string(), + } + } + + fn subscribe_request(subscription_id: &str) -> SubscribeRequest { + SubscribeRequest { + stream: Some(stream()), + workflow_id: "reader".to_string(), + run_id: "reader-run".to_string(), + subscription_id: subscription_id.to_string(), + } + } + + async fn subscribe( + service: &StreamService, + request: SubscribeRequest, + ) -> Result { + service + .call("Subscribe", &request.encode_to_vec()) + .await + .map(|answer| SubscribeResponse::decode(answer.as_slice()).unwrap()) + } + + fn chain(first_run_id: &str) -> ChainId { + ChainId { + namespace: "ns".to_string(), + workflow_id: "owner".to_string(), + first_run_id: first_run_id.to_string(), + } + } + + /// The arm the next append on the stream reports. + async fn arm_of_next_append(store: &MemoryStore, first_run_id: &str) -> Option { + store + .append(StoreAppendRequest { + chain: Some(chain(first_run_id)), + topic: "tokens".to_string(), + producer_id: "p".to_string(), + attempt: 1, + sequence: 1, + digest: vec![1], + records: vec![b"record".to_vec()], + }) + .await + .unwrap() + .arm + } + + #[tokio::test] + async fn a_subscription_attaches_the_workflow_target_once_per_process() { + let setup = setup(); + let first = subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + assert!(first.attached); + assert_eq!(first.chain_first_run_id, "chain-1"); + let again = subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + assert!(!again.attached, "a later read sends no attach"); + assert_eq!(again.chain_first_run_id, "chain-1"); + let attached = setup.server.attached.lock().unwrap().clone(); + assert_eq!( + attached, + [WorkflowAttachment { + chain: chain("chain-1"), + topic: "tokens".to_string(), + request_id: "jetty/reader-run/1".to_string(), + workflow_id: "reader".to_string(), + run_id: "reader-run".to_string(), + subscription_id: "1".to_string(), + }] + ); + } + + #[tokio::test] + async fn readers_of_one_run_attach_apart() { + let setup = setup(); + subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + subscribe(&setup.service, subscribe_request("2")) + .await + .unwrap(); + let ids: Vec = setup + .server + .attached + .lock() + .unwrap() + .iter() + .map(|a| a.request_id.clone()) + .collect(); + assert_eq!(ids, ["jetty/reader-run/1", "jetty/reader-run/2"]); + } + + #[tokio::test] + async fn a_subscription_to_an_empty_stream_arms_it_for_the_first_append() { + let setup = setup(); + subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + assert!( + arm_of_next_append(&setup.store, "chain-1") + .await + .is_some_and(|arm| arm > 0), + "the first append wakes the Workflow" + ); + } + + #[tokio::test] + async fn a_subscription_to_a_stream_with_records_does_not_arm() { + let setup = setup(); + arm_of_next_append(&setup.store, "chain-1").await; + subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + let landed = setup + .store + .append(StoreAppendRequest { + chain: Some(chain("chain-1")), + topic: "tokens".to_string(), + producer_id: "other".to_string(), + attempt: 1, + sequence: 1, + digest: vec![1], + records: vec![b"record".to_vec()], + }) + .await + .unwrap(); + assert_eq!(landed.arm, Some(0), "the reader reads the records first"); + } + + #[tokio::test] + async fn a_subscription_attaches_to_the_chain_that_holds_the_workflow_id_now() { + let setup = setup(); + // A read warms the owner cache with the first chain. + setup + .service + .call( + "Latest", + &crate::proto::LatestRequest { + stream: Some(stream()), + } + .encode_to_vec(), + ) + .await + .unwrap(); + *setup.owner.0.lock().unwrap() = "chain-2".to_string(); + let answer = subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + assert_eq!(answer.chain_first_run_id, "chain-2"); + } + + #[tokio::test] + async fn a_refused_attach_fails_the_subscription_and_is_not_remembered() { + let setup = setup(); + *setup.server.answer.lock().unwrap() = Some(NotifyError::Refused( + "the execution has 32 stream wake sources".to_string(), + )); + let error = subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap_err(); + assert_eq!(error.kind, StreamFailureKind::Refused as i32); + *setup.server.answer.lock().unwrap() = None; + let answer = subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + assert!(answer.attached, "the next read tries the attach again"); + } + + #[tokio::test] + async fn an_attach_that_may_land_later_fails_as_storage() { + let setup = setup(); + *setup.server.answer.lock().unwrap() = Some(NotifyError::Failed("down".to_string())); + let error = subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap_err(); + assert_eq!(error.kind, StreamFailureKind::Storage as i32); + } + + #[tokio::test] + async fn a_subscription_needs_its_reader() { + let setup = setup(); + for request in [ + SubscribeRequest { + workflow_id: String::new(), + ..subscribe_request("1") + }, + SubscribeRequest { + run_id: String::new(), + ..subscribe_request("1") + }, + subscribe_request(""), + subscribe_request(&"x".repeat(MAX_SUBSCRIPTION_ID_BYTES + 1)), + SubscribeRequest { + stream: None, + ..subscribe_request("1") + }, + ] { + let error = subscribe(&setup.service, request).await.unwrap_err(); + assert_eq!(error.kind, StreamFailureKind::InvalidArgument as i32); + } + assert!(setup.server.attached.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn a_subscription_without_a_client_is_unsupported() { + let service = StreamService::new( + Arc::new(MemoryStore::new()), + Arc::new(Owner(Mutex::new("chain-1".to_string()))), + StreamsOptions::default(), + ); + let error = subscribe(&service, subscribe_request("1")) + .await + .unwrap_err(); + assert_eq!(error.kind, StreamFailureKind::Unsupported as i32); + } + + #[tokio::test] + async fn an_unsubscribe_detaches_from_the_chain_it_attached_to_and_forgets_it() { + let setup = setup(); + subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + *setup.owner.0.lock().unwrap() = "chain-2".to_string(); + let unsubscribe = UnsubscribeRequest { + stream: Some(stream()), + run_id: "reader-run".to_string(), + subscription_id: "1".to_string(), + }; + setup + .service + .call("Unsubscribe", &unsubscribe.encode_to_vec()) + .await + .unwrap(); + assert_eq!( + setup.server.detached.lock().unwrap().clone(), + [( + chain("chain-1"), + "tokens".to_string(), + "jetty/reader-run/1".to_string() + )] + ); + let answer = subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + assert!(answer.attached, "a new subscription attaches again"); + assert_eq!(answer.chain_first_run_id, "chain-2"); + } +} From f2b808bec0d5435ce414fe177df3c1ce282965e9 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 13 Oct 2026 19:04:48 -0700 Subject: [PATCH 3/8] Let a client stream reader wait on the stream's notifier. WaitForNotification sends the server's PollStreamNotifier long poll with the wake_after of the read that armed the stream. It returns when the counter passes it, when the stream closed, or with timed_out at the wait limit. The reader then reads again with arm, which also finds records whose notification a producer never sent. --- .../temporal/sdk/core/streams/streams.proto | 27 +++ crates/sdk-core/CHANGELOG.md | 4 + crates/sdk-core/src/streams.rs | 43 ++++- crates/streams/src/handle.rs | 2 +- crates/streams/src/lib.rs | 4 +- crates/streams/src/notify.rs | 34 +++- crates/streams/src/service.rs | 16 +- crates/streams/src/subscribe.rs | 161 +++++++++++++++++- 8 files changed, 275 insertions(+), 16 deletions(-) diff --git a/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto b/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto index 89f45bde98..d50ef6728b 100644 --- a/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto +++ b/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto @@ -35,6 +35,11 @@ service StreamService { rpc Subscribe(SubscribeRequest) returns (SubscribeResponse); // Detaches a Workflow's reader that stops before the stream ends. rpc Unsubscribe(UnsubscribeRequest) returns (UnsubscribeResponse); + // Waits on the stream's notifier, for a client reader whose read armed the stream. Returns when + // the notifier's counter is above `after_counter`, when the stream closed, or at the server's + // wait limit (`streamnotifier.pollMaxWait`). The reader then reads again with `arm`. After the + // wait limit, that read finds records whose notification was lost, or arms again. + rpc WaitForNotification(WaitForNotificationRequest) returns (WaitForNotificationResponse); } // Which store holds a process's streams, and how to reach it. Lang builds one store from it and @@ -273,6 +278,28 @@ message UnsubscribeRequest { message UnsubscribeResponse {} +message WaitForNotificationRequest { + // The stream the reader reads. + StreamAddress stream = 1; + // The `wake_after` of the read that armed the stream. + int64 after_counter = 2; +} + +message WaitForNotificationResponse { + // The highest counter the notifier has seen. + int64 counter = 1; + // The position that came with `counter`. + string position = 2; + // The wait reached its limit with no counter above `after_counter` and no close. + bool timed_out = 3; + // The stream closed. The reader reads to its end. + bool closed = 4; + // The result the stream closed with, when it gave one. + temporal.api.common.v1.Payload close_result = 5; + // The failure the stream closed with, when it closed with one. + temporal.api.failure.v1.Failure close_failure = 6; +} + message DeleteOwnerRequest { string namespace = 1; StreamOwnerKind owner_kind = 2; diff --git a/crates/sdk-core/CHANGELOG.md b/crates/sdk-core/CHANGELOG.md index 4c29b9d231..b1b0eca8e9 100644 --- a/crates/sdk-core/CHANGELOG.md +++ b/crates/sdk-core/CHANGELOG.md @@ -72,6 +72,10 @@ relevant information. each subscription once and answers a repeat from memory. The attach names the chain that holds the Workflow id now, and on an empty stream it arms the stream, so the first append wakes the Workflow. +* Experimental: the new `WaitForNotification` stream call waits on a stream's notifier with the + server's `PollStreamNotifier` long poll, for a client reader whose read armed the stream. It + returns when the notifier's counter passes the read's `wake_after`, when the stream closed, or + with `timed_out` at the server's wait limit. The reader then reads again with `arm`. ### Breaking Changes * Replay detects changes to the service or operation of scheduled Nexus operations, reporting diff --git a/crates/sdk-core/src/streams.rs b/crates/sdk-core/src/streams.rs index 1c133d90e5..11382f23f6 100644 --- a/crates/sdk-core/src/streams.rs +++ b/crates/sdk-core/src/streams.rs @@ -10,11 +10,12 @@ use temporalio_common::protos::temporal::api::{ workflowservice::v1::{ AttachStreamCallbackRequest, DescribeWorkflowExecutionRequest, DescribeWorkflowExecutionResponse, DetachStreamCallbackRequest, NotifyStreamRequest, + PollStreamNotifierRequest, PollStreamNotifierResponse, }, }; use temporalio_streams::{ - CloseOutcome, NotifierClient, NotifyError, OwnerClient, OwnerDescription, OwnerError, - StreamNotification, WorkflowAttachment, proto::ChainId, + CloseOutcome, NotifierClient, NotifierPoll, NotifyError, OwnerClient, OwnerDescription, + OwnerError, StreamNotification, WorkflowAttachment, proto::ChainId, }; pub use temporalio_streams::{StreamError, StreamResult, StreamService, StreamStore, proto}; use tonic::IntoRequest; @@ -82,6 +83,44 @@ impl NotifierClient for ConnectionNotifierClient { .map(|_| ()) .map_err(notify_error) } + + async fn poll( + &self, + chain: ChainId, + topic: String, + after_counter: i64, + ) -> Result, NotifyError> { + let answer = self + .connection + .workflow_service() + .poll_stream_notifier( + PollStreamNotifierRequest { + namespace: chain.namespace.clone(), + stream_ref: Some(chain_reference(&chain, topic)), + after_counter, + identity: self.connection.identity().to_string(), + } + .into_request(), + ) + .await; + match answer { + Ok(response) => Ok(Some(poll_answer(response.into_inner()))), + // The server answers at its own wait limit, before the client's deadline. A deadline + // that passes anyway asks for the same re-read. + Err(status) if status.code() == tonic::Code::DeadlineExceeded => Ok(None), + Err(status) => Err(notify_error(status)), + } + } +} + +fn poll_answer(response: PollStreamNotifierResponse) -> NotifierPoll { + NotifierPoll { + counter: response.counter, + position: response.position, + closed: response.closed, + close_result: response.close_result, + close_failure: response.close_failure, + } } /// The notifier of `chain`'s stream on `topic`. The run chain's first run keys the notifier, so one diff --git a/crates/streams/src/handle.rs b/crates/streams/src/handle.rs index 31a09c393b..2078901637 100644 --- a/crates/streams/src/handle.rs +++ b/crates/streams/src/handle.rs @@ -500,7 +500,7 @@ impl Streams { } /// The chain a stream address names now. - async fn chain(&self, stream: &StreamAddress) -> StreamResult { + pub(crate) async fn chain(&self, stream: &StreamAddress) -> StreamResult { let key = owner_key(stream); // A pinned run's chain never changes. The latest run's chain changes when a new chain // takes the Workflow id. diff --git a/crates/streams/src/lib.rs b/crates/streams/src/lib.rs index 93cbc42f82..327ce96c73 100644 --- a/crates/streams/src/lib.rs +++ b/crates/streams/src/lib.rs @@ -30,8 +30,8 @@ pub use handle::{OWNER_RECHECK, Streams, StreamsOptions}; pub use memory::MemoryStore; pub use notify::{ Arm, CLOSE_RETRY_CAP, CLOSE_RETRY_FIRST, DEFAULT_MAX_NOTIFIERS, FLUSH_LIMIT, NOTIFIER_REFRESH, - Notifier, NotifierClient, NotifyError, NotifyingStore, StreamNotification, WorkflowAttachment, - progress_counter, + Notifier, NotifierClient, NotifierPoll, NotifyError, NotifyingStore, StreamNotification, + WorkflowAttachment, progress_counter, }; pub use owner::{OwnerClient, OwnerDescription, OwnerError}; pub use reader::{DEFAULT_MAX_RECORDS, MAX_TRACKED_PRODUCERS, ReadTarget, read_page}; diff --git a/crates/streams/src/notify.rs b/crates/streams/src/notify.rs index 89d7a11bf8..7a33d54af5 100644 --- a/crates/streams/src/notify.rs +++ b/crates/streams/src/notify.rs @@ -100,8 +100,23 @@ pub struct WorkflowAttachment { pub subscription_id: String, } -/// Sends notifications to the stream notifier on the server, and attaches and detaches the -/// Workflows it wakes. +/// What a stream's notifier answered a client's wait with. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct NotifierPoll { + /// The highest counter the notifier has seen. + pub counter: i64, + /// The position that came with `counter`. + pub position: String, + /// Whether the stream closed. + pub closed: bool, + /// The result the stream closed with, when it gave one. + pub close_result: Option, + /// The failure the stream closed with, when it closed with one. + pub close_failure: Option, +} + +/// Sends notifications to the stream notifier on the server, attaches and detaches the +/// Workflows it wakes, and waits on it for clients. #[async_trait::async_trait] pub trait NotifierClient: Send + Sync { /// Sends one notification. @@ -127,6 +142,21 @@ pub trait NotifierClient: Send + Sync { "this client can't detach readers".to_string(), )) } + + /// Waits until the notifier of `chain`'s stream on `topic` has a counter above + /// `after_counter`, or the stream closed. Answers `None` when the wait reached its limit with + /// no answer from the server. + async fn poll( + &self, + chain: ChainId, + topic: String, + after_counter: i64, + ) -> Result, NotifyError> { + let _ = (chain, topic, after_counter); + Err(NotifyError::Unsupported( + "this client can't wait on a notifier".to_string(), + )) + } } /// The notifier counter for a store position. diff --git a/crates/streams/src/service.rs b/crates/streams/src/service.rs index 81df8e92d5..aca1a43e00 100644 --- a/crates/streams/src/service.rs +++ b/crates/streams/src/service.rs @@ -6,7 +6,7 @@ use crate::{ proto::{ AppendRequest, CloseRequest, DeleteOwnerRequest, FlushNotificationsRequest, FlushNotificationsResponse, LatestRequest, ReadRequest, StreamFailure, StreamStoreConfig, - SubscribeRequest, UnsubscribeRequest, stream_store_config, + SubscribeRequest, UnsubscribeRequest, WaitForNotificationRequest, stream_store_config, }, subscribe::Subscriptions, }; @@ -50,8 +50,8 @@ impl StreamService { } /// Like [StreamService::connect], closing each stream's notifier through `notifier` after - /// the store closes the stream, sending progress too when the config turns it on, and - /// attaching and detaching Workflow readers through it. + /// the store closes the stream, sending progress too when the config turns it on, attaching + /// and detaching Workflow readers through it, and waiting on it for client readers. pub async fn connect_with_notifier( config: StreamStoreConfig, owner: Arc, @@ -108,7 +108,7 @@ impl StreamService { } } - /// Attaches and detaches Workflow readers through `client`. + /// Attaches and detaches Workflow readers, and waits for client readers, through `client`. pub fn with_subscriptions(mut self, client: Arc) -> Self { self.subscriptions = Subscriptions::new(Some(client)); self @@ -165,6 +165,14 @@ impl StreamService { .unsubscribe(&self.streams, decode::(rpc, request)?) .await .map(|r| r.encode_to_vec()), + "WaitForNotification" => self + .subscriptions + .wait( + &self.streams, + decode::(rpc, request)?, + ) + .await + .map(|r| r.encode_to_vec()), other => Err(StreamError::unsupported(format!( "the stream service has no call {other:?}" ))), diff --git a/crates/streams/src/subscribe.rs b/crates/streams/src/subscribe.rs index 761120713e..2c8d643eb4 100644 --- a/crates/streams/src/subscribe.rs +++ b/crates/streams/src/subscribe.rs @@ -1,16 +1,18 @@ -//! Workflow readers that sleep on a stream's notifier. +//! Readers that sleep on a stream's notifier. //! //! A Workflow's reader attaches to the stream's notifier before its first read, so the notifier //! wakes the Workflow on its next Workflow Task when the stream moves. The attach names the reader //! by its run and its subscription id, which replay gives back, so a rerun attaches as the same -//! reader and the server keeps one callback. +//! reader and the server keeps one callback. A client reader has no task to wake, so it waits on +//! the notifier with a long poll. use crate::{ NotifierClient, NotifyError, StreamError, StreamResult, StreamStore, Streams, WorkflowAttachment, handle::address, proto::{ - StreamAddress, SubscribeRequest, SubscribeResponse, UnsubscribeRequest, UnsubscribeResponse, + StreamAddress, SubscribeRequest, SubscribeResponse, UnsubscribeRequest, + UnsubscribeResponse, WaitForNotificationRequest, WaitForNotificationResponse, }, }; use std::{ @@ -177,14 +179,45 @@ impl Subscriptions { self.attached.lock().unwrap().remove(&key); Ok(UnsubscribeResponse {}) } + + /// Waits on the notifier of the chain the stream names. A wait that reached its limit answers + /// `timed_out`, and the reader's next read checks the store itself, which covers a + /// notification that a producer never sent. + pub(crate) async fn wait( + &self, + streams: &Streams, + request: WaitForNotificationRequest, + ) -> StreamResult { + let stream = address(request.stream.as_ref())?; + let client = self.client()?; + let chain = streams.chain(stream).await?; + let polled = client + .poll(chain, stream.topic.clone(), request.after_counter) + .await + .map_err(subscribe_error)?; + let Some(polled) = polled else { + return Ok(WaitForNotificationResponse { + timed_out: true, + ..Default::default() + }); + }; + Ok(WaitForNotificationResponse { + timed_out: !polled.closed && polled.counter <= request.after_counter, + counter: polled.counter, + position: polled.position, + closed: polled.closed, + close_result: polled.close_result, + close_failure: polled.close_failure, + }) + } } #[cfg(test)] mod tests { use super::*; use crate::{ - MemoryStore, OwnerClient, OwnerDescription, OwnerError, StreamNotification, StreamService, - StreamsOptions, + MemoryStore, NotifierPoll, OwnerClient, OwnerDescription, OwnerError, StreamNotification, + StreamService, StreamsOptions, proto::{ChainId, StoreAppendRequest, StreamFailureKind, StreamOwnerKind}, }; use prost::Message; @@ -216,6 +249,8 @@ mod tests { attached: Mutex>, detached: Mutex>, answer: Mutex>, + polled: Mutex>, + poll_answer: Mutex, NotifyError>>>, } #[async_trait::async_trait] @@ -244,6 +279,23 @@ mod tests { .push((chain, topic, request_id)); Ok(()) } + + async fn poll( + &self, + chain: ChainId, + topic: String, + after_counter: i64, + ) -> Result, NotifyError> { + self.polled + .lock() + .unwrap() + .push((chain, topic, after_counter)); + self.poll_answer + .lock() + .unwrap() + .clone() + .expect("the test sets the poll's answer") + } } struct Setup { @@ -523,4 +575,103 @@ mod tests { assert!(answer.attached, "a new subscription attaches again"); assert_eq!(answer.chain_first_run_id, "chain-2"); } + + async fn wait( + setup: &Setup, + after_counter: i64, + answer: Result, NotifyError>, + ) -> Result { + *setup.server.poll_answer.lock().unwrap() = Some(answer); + setup + .service + .call( + "WaitForNotification", + &WaitForNotificationRequest { + stream: Some(stream()), + after_counter, + } + .encode_to_vec(), + ) + .await + .map(|answer| WaitForNotificationResponse::decode(answer.as_slice()).unwrap()) + } + + #[tokio::test] + async fn a_wait_asks_the_notifier_of_the_streams_chain_for_a_counter_above_the_arm() { + let setup = setup(); + let answer = wait( + &setup, + 41, + Ok(Some(NotifierPoll { + counter: 42, + position: "cursor-42".to_string(), + ..Default::default() + })), + ) + .await + .unwrap(); + assert_eq!( + setup.server.polled.lock().unwrap().clone(), + [(chain("chain-1"), "tokens".to_string(), 41)] + ); + assert!(!answer.timed_out && !answer.closed); + assert_eq!( + (answer.counter, answer.position.as_str()), + (42, "cursor-42") + ); + } + + #[tokio::test] + async fn a_wait_with_no_new_counter_times_out_for_a_re_read() { + let setup = setup(); + let unchanged = wait( + &setup, + 41, + Ok(Some(NotifierPoll { + counter: 41, + ..Default::default() + })), + ) + .await + .unwrap(); + assert!(unchanged.timed_out); + let deadline = wait(&setup, 41, Ok(None)).await.unwrap(); + assert!(deadline.timed_out && !deadline.closed); + } + + #[tokio::test] + async fn a_wait_on_a_closed_stream_carries_the_outcome() { + let setup = setup(); + let failure = temporalio_common::protos::temporal::api::failure::v1::Failure { + message: "the agent gave up".to_string(), + ..Default::default() + }; + let answer = wait( + &setup, + 41, + Ok(Some(NotifierPoll { + counter: 41, + closed: true, + close_failure: Some(failure.clone()), + ..Default::default() + })), + ) + .await + .unwrap(); + assert!(answer.closed && !answer.timed_out); + assert_eq!(answer.close_failure, Some(failure)); + } + + #[tokio::test] + async fn a_wait_on_a_server_without_the_poll_is_unsupported() { + let setup = setup(); + let error = wait( + &setup, + 0, + Err(NotifyError::Unsupported("unimplemented".to_string())), + ) + .await + .unwrap_err(); + assert_eq!(error.kind, StreamFailureKind::Unsupported as i32); + } } From aaf79bd2a18447673c665d090bb392f2674cdb02 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 13 Oct 2026 19:11:02 -0700 Subject: [PATCH 4/8] Handed stream wakes to lang as a ResolveStreamWake job. The server folds a reader's wakes onto the next Workflow Task scheduled event. Core folds them per subscription across the task's scheduled events and issues one job each, from History, so replay issues the same jobs. A task scheduled with wakes ends a heartbeat chain, as N10 does for progress. --- .../workflow_activation.proto | 31 +++ crates/protos/src/protos/mod.rs | 3 + crates/sdk-core/CHANGELOG.md | 5 + crates/sdk-core/src/core_tests/mod.rs | 1 + .../sdk-core/src/core_tests/stream_wakes.rs | 214 ++++++++++++++++++ crates/sdk-core/src/replay/history_builder.rs | 11 + .../src/worker/workflow/history_update.rs | 50 +++- .../workflow/machines/workflow_machines.rs | 129 ++++++++++- crates/sdk/src/workflow_future.rs | 5 + crates/workflow/src/runtime/instance.rs | 3 + 10 files changed, 450 insertions(+), 2 deletions(-) create mode 100644 crates/sdk-core/src/core_tests/stream_wakes.rs diff --git a/crates/protos/protos/local/temporal/sdk/core/workflow_activation/workflow_activation.proto b/crates/protos/protos/local/temporal/sdk/core/workflow_activation/workflow_activation.proto index 7287ab1c1a..b0114c6ff7 100644 --- a/crates/protos/protos/local/temporal/sdk/core/workflow_activation/workflow_activation.proto +++ b/crates/protos/protos/local/temporal/sdk/core/workflow_activation/workflow_activation.proto @@ -140,6 +140,9 @@ message WorkflowActivationJob { ResolveNexusOperationStart resolve_nexus_operation_start = 15; // A nexus operation resolved. ResolveNexusOperation resolve_nexus_operation = 16; + // A stream that one of the Workflow's readers subscribed to moved or closed. Numbers 17 + // and 18 are left to the external stream replay job and to Nexus operation progress. + ResolveStreamWake resolve_stream_wake = 19; // Remove the workflow identified by the [WorkflowActivation] containing this job from the // cache after performing the activation. It is guaranteed that this will be the only job // in the activation if present. @@ -384,6 +387,34 @@ message ResolveNexusOperation { nexus.NexusOperationResult result = 2; } +// A wake for one of the Workflow's stream readers, as the server folded it onto this Workflow +// Task's scheduled event. +// +// Core issues at most one per subscription per Workflow Task, in the task's first activation. It +// comes from History, so replay issues the same job at the same point. A wake says that the +// stream moved, not what it holds, so the reader reads the stream itself. +message ResolveStreamWake { + // The subscription id the reader subscribed with. + string subscription_id = 1; + // The position that came with `counter`, in the producer's terms. + string position = 2; + // The highest counter the stream's notifier has seen. The reader reads when it is above the + // `wake_after` of the read that armed the stream. + int64 counter = 3; + // The metadata that came with `counter`. + map metadata = 4; + // The notifier asks the reader to read again whatever the counter, after a quiet period. + bool recheck = 5; + // The stream closed. The reader reads to its end, then sees the outcome. + bool closed = 6; + // The result the stream closed with, when it gave one. + temporal.api.common.v1.Payload close_result = 7; + // The failure the stream closed with, when it closed with one. + temporal.api.failure.v1.Failure close_failure = 8; + // The notifier dropped the subscription. The reader subscribes again. + bool resubscribe = 9; +} + message RemoveFromCache { string message = 1; diff --git a/crates/protos/src/protos/mod.rs b/crates/protos/src/protos/mod.rs index a292e7462b..1391aa508b 100644 --- a/crates/protos/src/protos/mod.rs +++ b/crates/protos/src/protos/mod.rs @@ -1254,6 +1254,9 @@ pub mod coresdk { workflow_activation_job::Variant::ResolveNexusOperation(_) => { write!(f, "ResolveNexusOperation") } + workflow_activation_job::Variant::ResolveStreamWake(w) => { + write!(f, "ResolveStreamWake({}, {})", w.subscription_id, w.counter) + } } } } diff --git a/crates/sdk-core/CHANGELOG.md b/crates/sdk-core/CHANGELOG.md index b1b0eca8e9..cc9b04f209 100644 --- a/crates/sdk-core/CHANGELOG.md +++ b/crates/sdk-core/CHANGELOG.md @@ -76,6 +76,11 @@ relevant information. server's `PollStreamNotifier` long poll, for a client reader whose read armed the stream. It returns when the notifier's counter passes the read's `wake_after`, when the stream closed, or with `timed_out` at the server's wait limit. The reader then reads again with `arm`. +* Experimental: a stream wake for one of a Workflow's readers reaches lang as a + `ResolveStreamWake` job, at most one per subscription per workflow task, folded from the + `stream_wakes` on the task's scheduled events. Replay delivers the same jobs in the same + activations. A task scheduled with wakes ends a workflow task heartbeat chain, so it keeps an + activation of its own on replay. ### Breaking Changes * Replay detects changes to the service or operation of scheduled Nexus operations, reporting diff --git a/crates/sdk-core/src/core_tests/mod.rs b/crates/sdk-core/src/core_tests/mod.rs index 6dd4479bb7..f3ea9c8630 100644 --- a/crates/sdk-core/src/core_tests/mod.rs +++ b/crates/sdk-core/src/core_tests/mod.rs @@ -2,6 +2,7 @@ mod activity_tasks; mod event_groups; mod queries; mod replay_flag; +mod stream_wakes; mod updates; mod workers; mod workflow_cancels; diff --git a/crates/sdk-core/src/core_tests/stream_wakes.rs b/crates/sdk-core/src/core_tests/stream_wakes.rs new file mode 100644 index 0000000000..6d0ec452aa --- /dev/null +++ b/crates/sdk-core/src/core_tests/stream_wakes.rs @@ -0,0 +1,214 @@ +//! Stream wakes: the server folds them onto a Workflow Task's scheduled event, and Core hands them +//! to lang as one job per subscription per task, live and on replay alike. + +use crate::{ + replay::TestHistoryBuilder, + test_help::{MockPollCfg, ResponseType, build_mock_pollers, mock_worker}, + worker::client::mocks::mock_worker_client, +}; +use temporalio_common::{ + protos::{ + coresdk::{ + workflow_activation::{ResolveStreamWake, WorkflowActivation, workflow_activation_job}, + workflow_completion::WorkflowActivationCompletion, + }, + temporal::api::{ + common::v1::Payload, + enums::v1::{EventType, WorkflowTaskFailedCause}, + failure::v1::Failure, + stream::v1::StreamWake, + }, + }, + worker::WorkerTaskTypes, +}; + +fn wake(subscription_id: &str, counter: i64) -> StreamWake { + StreamWake { + subscription_id: subscription_id.to_string(), + position: format!("position-{counter}"), + counter, + metadata: [("records".to_string(), counter.to_string())].into(), + ..Default::default() + } +} + +/// The job lang should get for `wake`. +fn job_for(wake: &StreamWake) -> workflow_activation_job::Variant { + workflow_activation_job::Variant::ResolveStreamWake(ResolveStreamWake { + subscription_id: wake.subscription_id.clone(), + position: wake.position.clone(), + counter: wake.counter, + metadata: wake.metadata.clone(), + recheck: wake.recheck, + closed: wake.closed, + close_result: wake.close_result.clone(), + close_failure: wake.close_failure.clone(), + resubscribe: wake.resubscribe, + }) +} + +fn wake_jobs(activation: &WorkflowActivation) -> Vec { + activation + .jobs + .iter() + .filter_map(|job| match &job.variant { + v @ Some(workflow_activation_job::Variant::ResolveStreamWake(_)) => v.clone(), + _ => None, + }) + .collect() +} + +fn job_names(activation: &WorkflowActivation) -> Vec { + activation + .jobs + .iter() + .map(|job| match &job.variant { + Some(workflow_activation_job::Variant::ResolveStreamWake(w)) => { + format!("Wake({}, {})", w.subscription_id, w.counter) + } + Some(workflow_activation_job::Variant::SignalWorkflow(s)) => { + format!("Signal({})", s.signal_name) + } + Some(other) => other.to_string(), + None => "none".to_string(), + }) + .collect() +} + +fn worker(history: TestHistoryBuilder, batches: Vec) -> crate::Worker { + let mut mock = build_mock_pollers(MockPollCfg::from_resp_batches( + "fake_wf_id", + history, + batches, + mock_worker_client(), + )); + mock.worker_cfg(|w| { + w.task_types = WorkerTaskTypes::workflow_only(); + w.max_cached_workflows = 1; + }); + mock_worker(mock) +} + +async fn complete_empty(worker: &crate::Worker, run_id: &str) { + worker + .complete_workflow_activation(WorkflowActivationCompletion::empty(run_id.to_string())) + .await + .unwrap(); +} + +/// A first task with nothing for lang to do, then two tasks scheduled only because wakes arrived. +/// The first and the second complete with no commands, which is the shape of a heartbeat chain, so +/// replay must still give each wake task an activation of its own. +fn wakes_after_a_quiet_task() -> (TestHistoryBuilder, [StreamWake; 2]) { + let mut t = TestHistoryBuilder::default(); + t.add_by_type(EventType::WorkflowExecutionStarted); + t.add_full_wf_task(); + let first = wake("1", 41); + t.add_workflow_task_scheduled_with_stream_wakes(vec![first.clone()]); + t.add_workflow_task_started(); + t.add_workflow_task_completed(); + let second = wake("1", 42); + t.add_workflow_task_scheduled_with_stream_wakes(vec![second.clone()]); + t.add_workflow_task_started(); + (t, [first, second]) +} + +async fn run_wakes_after_a_quiet_task(batches: Vec, replaying: bool) { + let (history, [first, second]) = wakes_after_a_quiet_task(); + let worker = worker(history, batches); + + let init = worker.poll_workflow_activation().await.unwrap(); + let run_id = init.run_id.clone(); + assert!(wake_jobs(&init).is_empty(), "{:?}", job_names(&init)); + complete_empty(&worker, &run_id).await; + + let first_wake = worker.poll_workflow_activation().await.unwrap(); + assert_eq!( + wake_jobs(&first_wake), + vec![job_for(&first)], + "the wake-only task gets an activation of its own: {:?}", + job_names(&first_wake) + ); + assert_eq!(first_wake.is_replaying, replaying); + complete_empty(&worker, &run_id).await; + + let second_wake = worker.poll_workflow_activation().await.unwrap(); + assert_eq!(wake_jobs(&second_wake), vec![job_for(&second)]); + assert!(!second_wake.is_replaying); +} + +#[tokio::test] +async fn each_wake_task_gets_its_own_activation_live() { + run_wakes_after_a_quiet_task( + vec![ + ResponseType::ToTaskNum(1), + ResponseType::ToTaskNum(2), + ResponseType::ToTaskNum(3), + ], + false, + ) + .await; +} + +#[tokio::test] +async fn each_wake_task_gets_its_own_activation_on_replay() { + run_wakes_after_a_quiet_task(vec![ResponseType::AllHistory], true).await; +} + +#[tokio::test] +async fn a_failed_task_and_its_retry_fold_their_wakes_per_subscription() { + let mut t = TestHistoryBuilder::default(); + t.add_by_type(EventType::WorkflowExecutionStarted); + t.add_full_wf_task(); + t.add_workflow_task_scheduled_with_stream_wakes(vec![wake("a", 5), wake("b", 1)]); + t.add_workflow_task_started(); + t.add_workflow_task_failed_with_failure( + WorkflowTaskFailedCause::WorkflowWorkerUnhandledFailure, + Failure::default(), + ); + let closed = StreamWake { + closed: true, + close_result: Some(Payload { + data: b"summary".to_vec(), + ..Default::default() + }), + ..wake("a", 7) + }; + t.add_workflow_task_scheduled_with_stream_wakes(vec![closed.clone()]); + t.add_workflow_task_started(); + let worker = worker(t, vec![ResponseType::AllHistory]); + + let init = worker.poll_workflow_activation().await.unwrap(); + complete_empty(&worker, &init.run_id).await; + let woken = worker.poll_workflow_activation().await.unwrap(); + assert_eq!( + wake_jobs(&woken), + vec![job_for(&closed), job_for(&wake("b", 1))], + "{:?}", + job_names(&woken) + ); +} + +#[tokio::test] +async fn a_wake_and_a_signal_of_one_task_share_its_activation() { + let mut t = TestHistoryBuilder::default(); + t.add_by_type(EventType::WorkflowExecutionStarted); + t.add_full_wf_task(); + t.add_we_signaled("go", vec![]); + t.add_workflow_task_scheduled_with_stream_wakes(vec![StreamWake { + recheck: true, + ..wake("1", 0) + }]); + t.add_workflow_task_started(); + let worker = worker( + t, + vec![ResponseType::ToTaskNum(1), ResponseType::ToTaskNum(2)], + ); + + let init = worker.poll_workflow_activation().await.unwrap(); + complete_empty(&worker, &init.run_id).await; + let woken = worker.poll_workflow_activation().await.unwrap(); + let mut names = job_names(&woken); + names.sort(); + assert_eq!(names, ["Signal(go)", "Wake(1, 0)"]); +} diff --git a/crates/sdk-core/src/replay/history_builder.rs b/crates/sdk-core/src/replay/history_builder.rs index 51472482c9..3b259aab21 100644 --- a/crates/sdk-core/src/replay/history_builder.rs +++ b/crates/sdk-core/src/replay/history_builder.rs @@ -111,6 +111,17 @@ impl TestHistoryBuilder { self.workflow_task_scheduled_event_id = self.add_by_type(EventType::WorkflowTaskScheduled); } + /// Add a workflow task scheduled event carrying folded stream wakes. + pub fn add_workflow_task_scheduled_with_stream_wakes( + &mut self, + stream_wakes: Vec, + ) { + self.workflow_task_scheduled_event_id = self.add(WorkflowTaskScheduledEventAttributes { + stream_wakes, + ..Default::default() + }); + } + /// Add a workflow task started event. pub fn add_workflow_task_started(&mut self) { self.final_workflow_task_started_event_id = self.add(WorkflowTaskStartedEventAttributes { diff --git a/crates/sdk-core/src/worker/workflow/history_update.rs b/crates/sdk-core/src/worker/workflow/history_update.rs index 8565899465..c60fdf39c1 100644 --- a/crates/sdk-core/src/worker/workflow/history_update.rs +++ b/crates/sdk-core/src/worker/workflow/history_update.rs @@ -741,10 +741,14 @@ fn find_end_index_of_next_wft_seq( if let Some(next_next_event) = events.get(ix + 2) { if !saw_command && next_next_event.event_type() == EventType::WorkflowTaskScheduled + && !scheduled_with_stream_wakes(next_next_event) { // If we've never seen an interesting event and the next two events are // a completion followed immediately again by scheduled, then this is a - // WFT heartbeat and also doesn't conclude the sequence. + // WFT heartbeat and also doesn't conclude the sequence. A task + // scheduled with stream wakes had an activation of its own live, so + // folding it into this one on replay would hand lang the wakes too + // early. continue; } else { // If we see an update accepted command after WFT completed, we want to @@ -801,6 +805,13 @@ fn find_end_index_of_next_wft_seq( NextWFTSeqEndIndex::Incomplete(last_index) } +fn scheduled_with_stream_wakes(e: &HistoryEvent) -> bool { + matches!( + &e.attributes, + Some(Attributes::WorkflowTaskScheduledEventAttributes(a)) if !a.stream_wakes.is_empty() + ) +} + #[cfg(test)] mod tests { use super::*; @@ -941,6 +952,43 @@ mod tests { assert_eq!(seq.len(), 2); } + #[test] + fn a_task_scheduled_with_stream_wakes_ends_a_heartbeat_chain() { + let history = |with_wakes: bool| { + let mut t = TestHistoryBuilder::default(); + t.add_by_type(EventType::WorkflowExecutionStarted); + t.add_full_wf_task(); + t.add_full_wf_task(); // wft started 6 + if with_wakes { + t.add_workflow_task_scheduled_with_stream_wakes(vec![Default::default()]); + t.add_workflow_task_started(); + } else { + t.add_workflow_task_scheduled_and_started(); + } + t.add_workflow_task_completed(); // wft started 9 + t.add_we_signaled("whee", vec![]); + t.add_full_wf_task(); + t.add_workflow_execution_completed(); + t.as_history_update() + }; + + let mut heartbeat = history(false); + next_check_peek(&mut heartbeat, 0); + let seq = next_check_peek(&mut heartbeat, 6); + assert_eq!(seq.last().unwrap().event_id, 13); + + let mut woken = history(true); + next_check_peek(&mut woken, 0); + let seq = next_check_peek(&mut woken, 6); + assert_eq!( + seq.last().unwrap().event_id, + 9, + "the wake task must end the sequence at its own started event" + ); + let seq = next_check_peek(&mut woken, 9); + assert_eq!(seq.last().unwrap().event_id, 13); + } + #[test] fn heartbeat_marker_end() { let mut t = TestHistoryBuilder::default(); 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 6c1b1a59f7..122d403372 100644 --- a/crates/sdk-core/src/worker/workflow/machines/workflow_machines.rs +++ b/crates/sdk-core/src/worker/workflow/machines/workflow_machines.rs @@ -59,7 +59,8 @@ use temporalio_common::{ coresdk::{ common::{NamespacedWorkflowExecution, VersioningIntent}, workflow_activation::{ - self, NotifyHasPatch, UpdateRandomSeed, WorkflowActivation, workflow_activation_job, + self, NotifyHasPatch, ResolveStreamWake, UpdateRandomSeed, WorkflowActivation, + workflow_activation_job, }, workflow_commands::ContinueAsNewWorkflowExecution, }, @@ -70,6 +71,7 @@ use temporalio_common::{ history::v1::{HistoryEvent, history_event}, protocol::v1::{Message as ProtocolMessage, message::SequencingId}, sdk::v1::WorkflowTaskCompletedMetadata, + stream::v1::StreamWake, }, }, worker::WorkerDeploymentVersion, @@ -88,6 +90,10 @@ pub(crate) struct WorkflowMachines { last_history_from_server: HistoryUpdate, /// Protocol messages that have yet to be processed for the current WFT. protocol_msgs: Vec, + /// Stream wakes read from the scheduled events of the task being applied, one per + /// subscription. A failed task's scheduled event and its retry's are applied together, and + /// the server clears what it put on the first, so both count. + pending_stream_wakes: Vec, /// EventId of the last handled WorkflowTaskStarted event current_started_event_id: i64, /// The event id of the next workflow task started event that the machines need to process. @@ -272,6 +278,7 @@ impl WorkflowMachines { Self { last_history_from_server: basics.history, protocol_msgs: vec![], + pending_stream_wakes: vec![], workflow_id: basics.workflow_id, workflow_type: basics.workflow_type, run_id: basics.run_id, @@ -486,6 +493,27 @@ impl WorkflowMachines { } } + /// Hands lang one job per subscription for the wakes folded from this task's scheduled + /// events. The wakes come from History, so replay hands lang the same jobs at the same point. + fn deliver_stream_wakes(&mut self) { + for wake in std::mem::take(&mut self.pending_stream_wakes) { + self.drive_me.send_job( + workflow_activation_job::Variant::ResolveStreamWake(ResolveStreamWake { + subscription_id: wake.subscription_id, + position: wake.position, + counter: wake.counter, + metadata: wake.metadata, + recheck: wake.recheck, + closed: wake.closed, + close_result: wake.close_result, + close_failure: wake.close_failure, + resubscribe: wake.resubscribe, + }) + .into(), + ); + } + } + pub(crate) fn has_pending_jobs(&self) -> bool { !self.drive_me.peek_pending_jobs().is_empty() } @@ -759,6 +787,7 @@ impl WorkflowMachines { } self.last_processed_event = eid; } + self.deliver_stream_wakes(); // Needed to delay mutation of self until after we've iterated over peeked events. #[allow(clippy::large_enum_variant)] @@ -1035,6 +1064,14 @@ impl WorkflowMachines { } } Ok(EventType::WorkflowTaskScheduled) => { + if let Some(history_event::Attributes::WorkflowTaskScheduledEventAttributes( + ref attrs, + )) = event_dat.event.attributes + { + for wake in &attrs.stream_wakes { + fold_stream_wake(&mut self.pending_stream_wakes, wake); + } + } let wf_task_sm = WorkflowTaskMachine::new(self.next_started_event_id); let key = self.all_machines.insert(wf_task_sm.into()); self.submachine_handle_event(key, event_dat)?; @@ -1861,3 +1898,93 @@ enum CommandIdKind { /// A command which is fire-and-forget (ex: Upsert search attribs) NeverResolves, } + +/// Keeps one wake per subscription, in order of first appearance. The highest counter keeps its +/// position and metadata. A recheck, a close or a resubscribe on any of the folded wakes stays, +/// since each one asks the reader to act whatever the counter. +fn fold_stream_wake(folded: &mut Vec, wake: &StreamWake) { + let Some(held) = folded + .iter_mut() + .find(|held| held.subscription_id == wake.subscription_id) + else { + folded.push(wake.clone()); + return; + }; + if wake.counter > held.counter { + held.counter = wake.counter; + held.position = wake.position.clone(); + held.metadata = wake.metadata.clone(); + } + held.recheck |= wake.recheck; + held.resubscribe |= wake.resubscribe; + if wake.closed && !held.closed { + held.closed = true; + held.close_result = wake.close_result.clone(); + held.close_failure = wake.close_failure.clone(); + } +} + +#[cfg(test)] +mod stream_wake_tests { + use super::fold_stream_wake; + use temporalio_common::protos::temporal::api::{common::v1::Payload, stream::v1::StreamWake}; + + fn wake(subscription_id: &str, counter: i64, position: &str) -> StreamWake { + StreamWake { + subscription_id: subscription_id.to_string(), + position: position.to_string(), + counter, + ..Default::default() + } + } + + #[test] + fn the_highest_counter_per_subscription_is_kept_in_first_seen_order() { + let mut folded = vec![]; + for w in [ + wake("b", 2, "a"), + wake("a", 1, "b"), + wake("b", 4, "c"), + wake("b", 3, "d"), + wake("a", 1, "e"), + ] { + fold_stream_wake(&mut folded, &w); + } + assert_eq!(folded, vec![wake("b", 4, "c"), wake("a", 1, "b")]); + } + + #[test] + fn a_recheck_a_close_or_a_resubscribe_survives_a_higher_counter() { + let result = Payload { + data: b"done".to_vec(), + ..Default::default() + }; + let mut folded = vec![]; + fold_stream_wake( + &mut folded, + &StreamWake { + recheck: true, + closed: true, + close_result: Some(result.clone()), + ..wake("a", 5, "closed") + }, + ); + fold_stream_wake( + &mut folded, + &StreamWake { + resubscribe: true, + ..wake("a", 7, "later") + }, + ); + assert_eq!( + folded, + vec![StreamWake { + recheck: true, + closed: true, + close_result: Some(result), + resubscribe: true, + ..wake("a", 7, "later") + }] + ); + } +} diff --git a/crates/sdk/src/workflow_future.rs b/crates/sdk/src/workflow_future.rs index f18604ce0b..6c7a7f4cd1 100644 --- a/crates/sdk/src/workflow_future.rs +++ b/crates/sdk/src/workflow_future.rs @@ -336,6 +336,11 @@ impl WorkflowFuture { .context("Nexus operation must have result")?; push_polled_context!(ActivationJobContext::Passive); } + // This SDK has no stream reader yet. A wake carries no records, only a hint to + // read the stream, so dropping it loses nothing. + Variant::ResolveStreamWake(_) => { + push_polled_context!(ActivationJobContext::Passive); + } Variant::RemoveFromCache(_) => { unreachable!("Cache removal should happen higher up"); } diff --git a/crates/workflow/src/runtime/instance.rs b/crates/workflow/src/runtime/instance.rs index 7c6b3a0332..3baaf46741 100644 --- a/crates/workflow/src/runtime/instance.rs +++ b/crates/workflow/src/runtime/instance.rs @@ -1119,6 +1119,9 @@ where ActivationJobResult::None } Some(ActivationVariant::RemoveFromCache(_)) => ActivationJobResult::None, + // This runtime has no stream reader yet. A wake carries no records, only a hint + // to read the stream, so dropping it loses nothing. + Some(ActivationVariant::ResolveStreamWake(_)) => ActivationJobResult::None, None => { return Err(Box::new(Failure { message: "Activation job missing variant".to_string(), From 410e62915a474211e1ce7dfb58b6049fb6630584 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 13 Oct 2026 19:13:44 -0700 Subject: [PATCH 5/8] Tested a stream wake on the task that a cut asked for. The wake reaches lang while the local activity runs, and replay gives the wake and then the resolution in the live order, whatever lang decides about the cut. --- .../sdk-core/src/core_tests/task_heartbeat.rs | 120 +++++++++++++++++- 1 file changed, 118 insertions(+), 2 deletions(-) diff --git a/crates/sdk-core/src/core_tests/task_heartbeat.rs b/crates/sdk-core/src/core_tests/task_heartbeat.rs index 47a6223530..33c1950ebf 100644 --- a/crates/sdk-core/src/core_tests/task_heartbeat.rs +++ b/crates/sdk-core/src/core_tests/task_heartbeat.rs @@ -29,6 +29,7 @@ use temporalio_common::protos::{ command::v1::command::Attributes, enums::v1::{CommandType, EventType}, history::v1::{MarkerRecordedEventAttributes, history_event}, + stream::v1::StreamWake, workflowservice::v1::RespondWorkflowTaskCompletedResponse, }, }; @@ -123,10 +124,24 @@ fn cut_history_up_to_second_task() -> TestHistoryBuilder { /// The history the server writes for `completions`, one workflow task each. fn history_of(completions: &[WorkflowTaskCompletion]) -> TestHistoryBuilder { + history_with_wakes(completions, &[]) +} + +/// Like [history_of], with `wakes[i]` on the scheduled event of task `i`. +fn history_with_wakes( + completions: &[WorkflowTaskCompletion], + wakes: &[Vec], +) -> TestHistoryBuilder { let mut t = TestHistoryBuilder::default(); t.add_wfe_started_with_wft_timeout(WFT_TIMEOUT); - for completion in completions { - t.add_workflow_task_scheduled_and_started(); + for (task, completion) in completions.iter().enumerate() { + match wakes.get(task) { + Some(wakes) if !wakes.is_empty() => { + t.add_workflow_task_scheduled_with_stream_wakes(wakes.clone()); + t.add_workflow_task_started(); + } + _ => t.add_workflow_task_scheduled_and_started(), + } t.add_workflow_task_completed(); let completed_id = t.current_event_id(); if let Some(history_event::Attributes::WorkflowTaskCompletedEventAttributes(attrs)) = @@ -368,3 +383,104 @@ async fn a_cut_history_has_the_events_the_server_writes() { ] ); } + +fn wake() -> StreamWake { + StreamWake { + subscription_id: "1".to_string(), + position: "position-42".to_string(), + counter: 42, + ..Default::default() + } +} + +fn wakes_reader_1(activation: &WorkflowActivation) -> bool { + matches!( + activation.jobs.as_slice(), + [WorkflowActivationJob { + variant: Some(workflow_activation_job::Variant::ResolveStreamWake(wake)), + }] if wake.subscription_id == "1" && wake.counter == 42 + ) +} + +/// A wake that lands while the cut task is open comes on the task the cut asked for, with the +/// local activity still running, and the activity resolves in a later activation of that task. +async fn run_cut_then_wake_live() -> Vec { + let mut t = TestHistoryBuilder::default(); + t.add_wfe_started_with_wft_timeout(WFT_TIMEOUT); + t.add_full_wf_task(); + t.add_timer_started("1".to_string()); + t.add_workflow_task_scheduled_with_stream_wakes(vec![wake()]); + t.add_workflow_task_started(); + let (core, reported) = recording_worker( + t, + [ResponseType::ToTaskNum(1), ResponseType::ToTaskNum(2)], + 2, + ); + let activation = core.poll_workflow_activation().await.unwrap(); + core.complete_workflow_activation(completion(&activation.run_id, first_commands(), true)) + .await + .unwrap(); + let task = core.poll_activity_task().await.unwrap(); + let woken = core.poll_workflow_activation().await.unwrap(); + assert!(wakes_reader_1(&woken), "{woken:?}"); + core.complete_workflow_activation(completion(&woken.run_id, vec![], false)) + .await + .unwrap(); + core.complete_activity_task(ActivityTaskCompletion { + task_token: task.task_token, + result: Some(ActivityExecutionResult::ok(b"done".to_vec().into())), + }) + .await + .unwrap(); + let activation = core.poll_workflow_activation().await.unwrap(); + assert!(resolves_la_1(&activation), "{activation:?}"); + core.complete_execution(&activation.run_id).await; + drain_pollers_and_shutdown(&core).await; + let reported = reported.lock().unwrap().clone(); + let [cut, last] = reported.as_slice() else { + panic!("expected two completions, got {reported:?}"); + }; + assert!(cut.force_create_new_workflow_task); + assert!(!last.force_create_new_workflow_task); + reported +} + +#[tokio::test] +async fn a_wake_on_the_task_after_a_cut_reaches_lang_while_the_local_activity_runs() { + run_cut_then_wake_live().await; +} + +#[rstest::rstest] +#[case::with_the_cut(true)] +#[case::without_the_cut(false)] +#[tokio::test] +async fn a_cut_then_a_wake_replays_in_the_live_order(#[case] heartbeat_task: bool) { + let live = run_cut_then_wake_live().await; + let history = history_with_wakes(&live, &[vec![], vec![wake()]]); + let (core, reported) = recording_worker(history, [ResponseType::AllHistory], 0); + let activation = core.poll_workflow_activation().await.unwrap(); + assert!(is_init(&activation), "{activation:?}"); + core.complete_workflow_activation(completion( + &activation.run_id, + first_commands(), + heartbeat_task, + )) + .await + .unwrap(); + let woken = core.poll_workflow_activation().await.unwrap(); + assert!( + wakes_reader_1(&woken), + "the wake comes first, as live: {woken:?}" + ); + core.complete_workflow_activation(completion(&woken.run_id, vec![], false)) + .await + .unwrap(); + let activation = tokio::time::timeout(Duration::from_secs(10), core.poll_workflow_activation()) + .await + .expect("replay resolves the local activity from its marker") + .unwrap(); + assert!(resolves_la_1(&activation), "{activation:?}"); + core.complete_execution(&activation.run_id).await; + drain_pollers_and_shutdown(&core).await; + assert!(reported.lock().unwrap().is_empty()); +} From 3b1d2ebc6def7f0c55edc5d132eed6e349be4d66 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 13 Oct 2026 19:19:27 -0700 Subject: [PATCH 6/8] Let a dropped Workflow reader subscribe again with a new attach. The server answers a repeat of a dropped attach as done, so a reader that gets a resubscribe wake raises its generation, which gives a new request id. Lang counts the generation from its wakes, so replay gives the same one. --- .../temporal/sdk/core/streams/streams.proto | 8 ++- crates/sdk-core/CHANGELOG.md | 9 ++-- crates/streams/src/subscribe.rs | 49 +++++++++++++++++-- 3 files changed, 56 insertions(+), 10 deletions(-) diff --git a/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto b/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto index 89f45bde98..9544693459 100644 --- a/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto +++ b/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto @@ -255,6 +255,10 @@ message SubscribeRequest { // Names the reader in its run. Lang derives it from the reader's call number, so replay gives // the same one. At most 256 bytes. string subscription_id = 4; + // Counts the times the reader subscribed again after a wake with `resubscribe`. Each value is a + // new attach, since the server answers a repeat of a dropped attach as done. Lang counts it from + // the wakes, so replay gives the same one. + uint32 generation = 5; } message SubscribeResponse { @@ -265,10 +269,12 @@ message SubscribeResponse { } message UnsubscribeRequest { - // The stream, the reader's run and the subscription, as the Subscribe named them. + // The stream, the reader's run, the subscription and its generation, as the last Subscribe + // named them. StreamAddress stream = 1; string run_id = 2; string subscription_id = 3; + uint32 generation = 4; } message UnsubscribeResponse {} diff --git a/crates/sdk-core/CHANGELOG.md b/crates/sdk-core/CHANGELOG.md index 4c29b9d231..ec96129efb 100644 --- a/crates/sdk-core/CHANGELOG.md +++ b/crates/sdk-core/CHANGELOG.md @@ -68,10 +68,11 @@ relevant information. outcome to a wrapper that tells the notifier. * Experimental: the new `Subscribe` and `Unsubscribe` stream calls attach a Workflow's reader to a stream's notifier as a `workflow_target`, and detach it. The attach's request id is - `jetty//`, so a rerun attaches as the same reader. A process attaches - each subscription once and answers a repeat from memory. The attach names the chain that holds - the Workflow id now, and on an empty stream it arms the stream, so the first append wakes the - Workflow. + `jetty//`, so a rerun attaches as the same reader. After a wake with + `resubscribe`, lang raises `SubscribeRequest.generation`, which makes a new attach. A process + attaches each subscription once and answers a repeat from memory. The attach names the chain + that holds the Workflow id now, and on an empty stream it arms the stream, so the first append + wakes the Workflow. ### Breaking Changes * Replay detects changes to the service or operation of scheduled Nexus operations, reporting diff --git a/crates/streams/src/subscribe.rs b/crates/streams/src/subscribe.rs index 761120713e..6a3e66ad0a 100644 --- a/crates/streams/src/subscribe.rs +++ b/crates/streams/src/subscribe.rs @@ -50,9 +50,13 @@ impl SubscriptionKey { } /// The request id of a reader's attach. The server dedupes attaches by it, so it depends only on -/// what replay gives back. -fn request_id(run_id: &str, subscription_id: &str) -> String { - format!("jetty/{run_id}/{subscription_id}") +/// what replay gives back. A later generation is a new attach, after the notifier dropped the +/// reader and the server would answer the old id as already done. +fn request_id(run_id: &str, subscription_id: &str, generation: u32) -> String { + match generation { + 0 => format!("jetty/{run_id}/{subscription_id}"), + generation => format!("jetty/{run_id}/{subscription_id}/{generation}"), + } } fn check_subscription(run_id: &str, subscription_id: &str) -> StreamResult<()> { @@ -107,7 +111,11 @@ impl Subscriptions { )); } check_subscription(&request.run_id, &request.subscription_id)?; - let request_id = request_id(&request.run_id, &request.subscription_id); + let request_id = request_id( + &request.run_id, + &request.subscription_id, + request.generation, + ); let key = SubscriptionKey::new(stream, &request_id); if let Some(first_run_id) = self.attached.lock().unwrap().get(&key) { return Ok(SubscribeResponse { @@ -158,7 +166,11 @@ impl Subscriptions { ) -> StreamResult { let stream = address(request.stream.as_ref())?; check_subscription(&request.run_id, &request.subscription_id)?; - let request_id = request_id(&request.run_id, &request.subscription_id); + let request_id = request_id( + &request.run_id, + &request.subscription_id, + request.generation, + ); let client = self.client()?; let key = SubscriptionKey::new(stream, &request_id); let known = self.attached.lock().unwrap().get(&key).cloned(); @@ -283,6 +295,7 @@ mod tests { workflow_id: "reader".to_string(), run_id: "reader-run".to_string(), subscription_id: subscription_id.to_string(), + generation: 0, } } @@ -348,6 +361,31 @@ mod tests { ); } + #[tokio::test] + async fn a_new_generation_attaches_again_after_the_notifier_dropped_the_reader() { + let setup = setup(); + subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + let again = subscribe( + &setup.service, + SubscribeRequest { + generation: 1, + ..subscribe_request("1") + }, + ) + .await + .unwrap(); + assert!(again.attached, "a resubscribe sends a new attach"); + let attached = setup.server.attached.lock().unwrap().clone(); + let ids: Vec<&str> = attached.iter().map(|a| a.request_id.as_str()).collect(); + assert_eq!(ids, ["jetty/reader-run/1", "jetty/reader-run/1/1"]); + assert!( + attached.iter().all(|a| a.subscription_id == "1"), + "the wakes still name the same reader" + ); + } + #[tokio::test] async fn readers_of_one_run_attach_apart() { let setup = setup(); @@ -503,6 +541,7 @@ mod tests { stream: Some(stream()), run_id: "reader-run".to_string(), subscription_id: "1".to_string(), + ..Default::default() }; setup .service From c653bf42b1e8821f06e700cf00675fd8ca11839c Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 13 Oct 2026 19:20:19 -0700 Subject: [PATCH 7/8] Kept one changelog entry for Subscribe after the merge. --- crates/sdk-core/CHANGELOG.md | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/crates/sdk-core/CHANGELOG.md b/crates/sdk-core/CHANGELOG.md index db8346cd73..99272874c6 100644 --- a/crates/sdk-core/CHANGELOG.md +++ b/crates/sdk-core/CHANGELOG.md @@ -68,19 +68,15 @@ relevant information. outcome to a wrapper that tells the notifier. * Experimental: the new `Subscribe` and `Unsubscribe` stream calls attach a Workflow's reader to a stream's notifier as a `workflow_target`, and detach it. The attach's request id is - `jetty//`, so a rerun attaches as the same reader. A process attaches - each subscription once and answers a repeat from memory. The attach names the chain that holds - the Workflow id now, and on an empty stream it arms the stream, so the first append wakes the - Workflow. -* Experimental: the new `WaitForNotification` stream call waits on a stream's notifier with the - server's `PollStreamNotifier` long poll, for a client reader whose read armed the stream. It - returns when the notifier's counter passes the read's `wake_after`, when the stream closed, or - with `timed_out` at the server's wait limit. The reader then reads again with `arm`. `jetty//`, so a rerun attaches as the same reader. After a wake with `resubscribe`, lang raises `SubscribeRequest.generation`, which makes a new attach. A process attaches each subscription once and answers a repeat from memory. The attach names the chain that holds the Workflow id now, and on an empty stream it arms the stream, so the first append wakes the Workflow. +* Experimental: the new `WaitForNotification` stream call waits on a stream's notifier with the + server's `PollStreamNotifier` long poll, for a client reader whose read armed the stream. It + returns when the notifier's counter passes the read's `wake_after`, when the stream closed, or + with `timed_out` at the server's wait limit. The reader then reads again with `arm`. ### Breaking Changes * Replay detects changes to the service or operation of scheduled Nexus operations, reporting From 0bbeb6eab120000f231cef057d858000031684a2 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Wed, 14 Oct 2026 02:17:36 -0700 Subject: [PATCH 8/8] Bounded the subscriptions a process remembers. A reader that reads to its end, or a run that closes, never detaches, so the map grew for the life of the Worker. It now forgets the oldest entries like the owner cache. A forgotten reader attaches again on its next read, and the server dedupes that attach. --- crates/streams/src/handle.rs | 16 +++++--- crates/streams/src/subscribe.rs | 71 +++++++++++++++++++++++---------- 2 files changed, 59 insertions(+), 28 deletions(-) diff --git a/crates/streams/src/handle.rs b/crates/streams/src/handle.rs index 066175c978..b649156a7d 100644 --- a/crates/streams/src/handle.rs +++ b/crates/streams/src/handle.rs @@ -80,31 +80,35 @@ struct ProducerKey { } #[derive(Debug, Clone)] -struct Checked { - first_run_id: String, +pub(crate) struct Checked { + pub(crate) first_run_id: String, at: Instant, } /// Remembers when each key was last checked, forgetting the oldest when full. #[derive(Debug)] -struct Remembered { +pub(crate) struct Remembered { entries: Mutex>, limit: usize, } impl Remembered { - fn new(limit: usize) -> Self { + pub(crate) fn new(limit: usize) -> Self { Self { entries: Mutex::default(), limit: limit.max(1), } } - fn get(&self, key: &K) -> Option { + pub(crate) fn get(&self, key: &K) -> Option { self.entries.lock().unwrap().get(key).cloned() } - fn put(&self, key: K, first_run_id: &str, fresh_for: Duration) { + pub(crate) fn forget(&self, key: &K) { + self.entries.lock().unwrap().remove(key); + } + + pub(crate) fn put(&self, key: K, first_run_id: &str, fresh_for: Duration) { let mut entries = self.entries.lock().unwrap(); if entries.len() >= self.limit && !entries.contains_key(&key) { entries.retain(|_, checked| checked.at.elapsed() < fresh_for); diff --git a/crates/streams/src/subscribe.rs b/crates/streams/src/subscribe.rs index 6a3e66ad0a..af75629a4b 100644 --- a/crates/streams/src/subscribe.rs +++ b/crates/streams/src/subscribe.rs @@ -6,26 +6,30 @@ //! reader and the server keeps one callback. use crate::{ - NotifierClient, NotifyError, StreamError, StreamResult, StreamStore, Streams, - WorkflowAttachment, - handle::address, + NotifierClient, NotifyError, OWNER_RECHECK, StreamError, StreamResult, StreamStore, Streams, + StreamsOptions, WorkflowAttachment, + handle::{Remembered, address}, proto::{ StreamAddress, SubscribeRequest, SubscribeResponse, UnsubscribeRequest, UnsubscribeResponse, }, }; -use std::{ - collections::HashMap, - sync::{Arc, Mutex}, -}; +use std::sync::Arc; /// The longest subscription id the server takes. const MAX_SUBSCRIPTION_ID_BYTES: usize = 256; -/// The subscriptions this process attached, so a reader's later reads send no attach. -#[derive(Default)] +/// The subscriptions this process attached, so a reader's later reads send no attach. A reader +/// that ends or a run that closes never detaches, so the oldest are forgotten. A reader this +/// process forgot attaches again, and the server dedupes that attach. pub(crate) struct Subscriptions { client: Option>, - attached: Mutex>, + attached: Remembered, +} + +impl Default for Subscriptions { + fn default() -> Self { + Self::new(None) + } } #[derive(Debug, Clone, PartialEq, Eq, Hash)] @@ -86,7 +90,7 @@ impl Subscriptions { pub(crate) fn new(client: Option>) -> Self { Self { client, - attached: Mutex::default(), + attached: Remembered::new(StreamsOptions::default().remembered), } } @@ -117,9 +121,9 @@ impl Subscriptions { request.generation, ); let key = SubscriptionKey::new(stream, &request_id); - if let Some(first_run_id) = self.attached.lock().unwrap().get(&key) { + if let Some(attached) = self.attached.get(&key) { return Ok(SubscribeResponse { - chain_first_run_id: first_run_id.clone(), + chain_first_run_id: attached.first_run_id, attached: false, }); } @@ -147,10 +151,7 @@ impl Subscriptions { "Could not arm the stream a Workflow reader subscribed to" ); } - self.attached - .lock() - .unwrap() - .insert(key, chain.first_run_id.clone()); + self.attached.put(key, &chain.first_run_id, OWNER_RECHECK); Ok(SubscribeResponse { chain_first_run_id: chain.first_run_id, attached: true, @@ -173,12 +174,11 @@ impl Subscriptions { ); let client = self.client()?; let key = SubscriptionKey::new(stream, &request_id); - let known = self.attached.lock().unwrap().get(&key).cloned(); - let chain = match known { - Some(first_run_id) => crate::proto::ChainId { + let chain = match self.attached.get(&key) { + Some(attached) => crate::proto::ChainId { namespace: stream.namespace.clone(), workflow_id: stream.workflow_id.clone(), - first_run_id, + first_run_id: attached.first_run_id, }, None => streams.chain_now(stream).await?, }; @@ -186,7 +186,7 @@ impl Subscriptions { .detach(chain, stream.topic.clone(), request_id) .await .map_err(subscribe_error)?; - self.attached.lock().unwrap().remove(&key); + self.attached.forget(&key); Ok(UnsubscribeResponse {}) } } @@ -200,6 +200,7 @@ mod tests { proto::{ChainId, StoreAppendRequest, StreamFailureKind, StreamOwnerKind}, }; use prost::Message; + use std::sync::Mutex; use temporalio_common::protos::temporal::api::enums::v1::WorkflowExecutionStatus; /// The owner's chain is `first_run_id`, which a test moves to a new chain. @@ -334,6 +335,32 @@ mod tests { .arm } + #[tokio::test] + async fn a_full_process_forgets_its_subscriptions_and_attaches_again() { + let setup = setup(); + let subscriptions = Subscriptions { + client: Some(setup.server.clone()), + attached: Remembered::new(2), + }; + let streams = Streams::new( + setup.store.clone(), + setup.owner.clone(), + StreamsOptions::default(), + ); + let subscribe = |id: &str| { + subscriptions.subscribe(&streams, setup.store.as_ref(), subscribe_request(id)) + }; + for id in ["1", "2", "3"] { + assert!(subscribe(id).await.unwrap().attached); + } + assert!(!subscribe("3").await.unwrap().attached); + assert!( + subscribe("1").await.unwrap().attached, + "the oldest was forgotten" + ); + assert_eq!(setup.server.attached.lock().unwrap().len(), 4); + } + #[tokio::test] async fn a_subscription_attaches_the_workflow_target_once_per_process() { let setup = setup();