From 748842dfe7a2586afee70c24ef2aff6ca96f4f30 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Thu, 15 Oct 2026 05:27:00 -0700 Subject: [PATCH 1/3] Closed a Workflow Activity's streams when the Workflow resolves it. Only History knows every way an Activity ends, and the Workflow's Core gets that outcome when it resolves the Activity. A header on the schedule command marks the Activities that own streams, so replay rebuilds the mark and unmarked Activities cost nothing. --- crates/sdk-core/CHANGELOG.md | 10 + .../src/core_tests/external_streams.rs | 349 +++++++++++++++++- crates/sdk-core/src/streams.rs | 39 +- .../machines/activity_state_machine.rs | 11 +- .../workflow/machines/workflow_machines.rs | 26 ++ .../src/worker/workflow/managed_run.rs | 7 +- crates/sdk-core/src/worker/workflow/mod.rs | 20 +- .../src/worker/workflow/stream_output.rs | 146 +++++++- crates/streams/src/subscribe.rs | 52 +++ 9 files changed, 652 insertions(+), 8 deletions(-) diff --git a/crates/sdk-core/CHANGELOG.md b/crates/sdk-core/CHANGELOG.md index e7f6c06dda..7467883963 100644 --- a/crates/sdk-core/CHANGELOG.md +++ b/crates/sdk-core/CHANGELOG.md @@ -190,6 +190,16 @@ relevant information. set, and every store and wrapper must implement it. A Workflow's `DeleteOwner` also deletes its Activities' streams. `OwnerId` gains `WorkflowActivity` and `OwnerStatus` gains `ActivityGone`. The documented Redis ACL gains `+sadd +smembers`. +* Experimental: the scheduling Workflow's Core closes a Workflow Activity's streams when it + resolves the Activity. Workflow code marks the Activity with the `ScheduleActivity` header + `__temporal_stream_owner`. When a marked Activity resolves outside replay, Core closes each + topic in the Activity's topic set in the background, once the activation completed and its + stream output was appended. A completed Activity closes with a null result, and a failed, + timed-out or canceled one with the failure its resolution holds, cut to `MAX_CLOSE_BYTES`. The + close names `@` as `closed_by` and reaches the stream's notifier. A close + that fails is logged and does not fail the Workflow Task. Unmarked Activities cost no store + call. Notifier calls name a Workflow Activity stream by its Workflow, scheduling run and + Activity id. ### Breaking Changes * Replay detects changes to the service or operation of scheduled Nexus operations, reporting diff --git a/crates/sdk-core/src/core_tests/external_streams.rs b/crates/sdk-core/src/core_tests/external_streams.rs index f373a5bc91..2ed32321c9 100644 --- a/crates/sdk-core/src/core_tests/external_streams.rs +++ b/crates/sdk-core/src/core_tests/external_streams.rs @@ -3,7 +3,7 @@ use crate::{ TaskToken, - replay::{TestHistoryBuilder, canned_histories, default_wes_attribs}, + replay::{TestHistoryBuilder, canned_histories, default_act_sched, default_wes_attribs}, test_help::{ MockPollCfg, ResponseType, WorkerExt, WorkerTestHelpers, build_mock_pollers, hist_to_poll_resp, mock_worker, schedule_local_activity_cmd, start_timer_cmd, @@ -22,7 +22,8 @@ use temporalio_common::{ workflow_activation::{WorkflowActivation, workflow_activation_job}, workflow_commands::{ ActivityCancellationType, CompleteWorkflowExecution, OutputRecord, - WorkflowOutputStreamCommit, workflow_command, + RequestCancelActivity, ScheduleActivity, WorkflowOutputStreamCommit, + workflow_command, }, workflow_completion::WorkflowActivationCompletion, }, @@ -48,7 +49,7 @@ use temporalio_streams::{ proto::{ ChainId, DeleteOwnerRequest, DeleteOwnerResponse, ReadRequest, StoreAppendRequest, StoreAppendResponse, StoreLatestRequest, StoreLatestResponse, StoreReadRequest, - StoreReadResponse, read_record, + StoreReadResponse, StoredClose, StreamOwnerKind, read_record, }, read_page, stream_hash, workflow_batch_digest, }; @@ -68,6 +69,10 @@ struct RecordingStore { refuse: Option, /// How long each append takes. append_delay: Option, + /// Each stream close as `(chain, topic, close)`. + closes: Mutex>, + /// Refuses every topic close. + refuse_closes: bool, } impl RecordingStore { @@ -191,6 +196,22 @@ impl StreamStore for RecordingStore { self.inner.close_topic(chain, topic).await } + async fn close_topic_with( + &self, + chain: &ChainId, + topic: &str, + close: StoredClose, + ) -> StreamResult { + self.calls.lock().push("close_topic_with"); + if self.refuse_closes { + return Err(StreamError::storage("the store is down")); + } + self.closes + .lock() + .push((chain.clone(), topic.to_string(), close.clone())); + self.inner.close_topic_with(chain, topic, close).await + } + async fn retire_producer( &self, chain: &ChainId, @@ -1290,3 +1311,325 @@ async fn a_rerun_after_the_run_closed_its_own_topic_drops_its_records_and_goes_o "the rerun's records were dropped" ); } + +/// What a Workflow Activity's stream holds once the agent `fakeid` scheduled it in `run_id`. +fn activity_chain(run_id: &str, activity_id: &str) -> ChainId { + ChainId { + namespace: "default".to_string(), + workflow_id: "fakeid".to_string(), + first_run_id: run_id.to_string(), + owner_kind: StreamOwnerKind::WorkflowActivity as i32, + activity_id: activity_id.to_string(), + } +} + +/// A `ScheduleActivity` that marks Activity `activity_id` as the owner of its streams, or not. +fn schedule_activity( + activity_id: &str, + marked: bool, + cancellation_type: ActivityCancellationType, +) -> workflow_command::Variant { + let mut schedule = ScheduleActivity { + seq: 1, + activity_id: activity_id.to_string(), + cancellation_type: cancellation_type as i32, + ..default_act_sched() + }; + if marked { + schedule + .headers + .insert("__temporal_stream_owner".to_string(), Payload::default()); + } + schedule.into() +} + +/// Appends one record to each of `topics` of the Activity's stream, as its producer would. +async fn activity_wrote(store: &RecordingStore, run_id: &str, activity_id: &str, topics: &[&str]) { + for topic in topics { + store + .inner + .append(StoreAppendRequest { + chain: Some(activity_chain(run_id, activity_id)), + topic: topic.to_string(), + producer_id: format!("{activity_id}@{run_id}"), + attempt: 1, + sequence: 1, + digest: vec![1; 32], + records: vec![b"record".to_vec()], + }) + .await + .unwrap(); + } +} + +/// Waits a bounded time for `count` stream closes, which run in the background. +async fn closes_within_deadline( + store: &RecordingStore, + count: usize, +) -> Vec<(String, StoredClose)> { + for _ in 0..200 { + let closes = store.closes.lock().clone(); + if closes.len() >= count { + return closes + .into_iter() + .map(|(chain, topic, close)| { + assert_eq!(chain.owner_kind, StreamOwnerKind::WorkflowActivity as i32); + (topic, close) + }) + .collect(); + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + panic!( + "expected {count} stream closes, got {:?}", + store.closes.lock() + ); +} + +/// Runs a Workflow that schedules Activity `tool`, which writes `topics`, and completes once the +/// Activity resolves. Returns the run id. +async fn run_activity_owner( + history: TestHistoryBuilder, + store: Arc, + marked: bool, + topics: &[&str], +) -> String { + let worker = worker_appending_to( + history, + vec![1.into(), 2.into()], + store.clone(), + Default::default(), + ); + let start = worker.poll_workflow_activation().await.unwrap(); + let run_id = start.run_id.clone(); + activity_wrote(&store, &run_id, "tool", topics).await; + complete( + &worker, + &run_id, + vec![schedule_activity( + "tool", + marked, + ActivityCancellationType::TryCancel, + )], + ) + .await; + let resolved = worker.poll_workflow_activation().await.unwrap(); + assert!(!resolved.is_replaying); + assert!(resolved.jobs.iter().any(|job| matches!( + job.variant, + Some(workflow_activation_job::Variant::ResolveActivity(_)) + ))); + worker.complete_execution(&resolved.run_id).await; + worker.drain_pollers_and_shutdown().await; + run_id +} + +#[tokio::test] +async fn a_marked_activity_that_completes_closes_each_topic_it_wrote_with_null() { + let store = Arc::new(RecordingStore::default()); + let run_id = run_activity_owner( + canned_histories::single_activity("tool"), + store.clone(), + true, + &["status", "tokens"], + ) + .await; + let closes = closes_within_deadline(&store, 2).await; + assert_eq!( + closes + .iter() + .map(|(topic, _)| topic.as_str()) + .collect::>(), + ["status", "tokens"] + ); + for (_, close) in &closes { + assert_eq!(close.result, Some(temporalio_streams::empty_close_result())); + assert!(close.failure.is_none()); + assert_eq!(close.closed_by, format!("tool@{run_id}")); + } + let chain = store.closes.lock()[0].0.clone(); + assert_eq!(chain, activity_chain(&run_id, "tool")); + // A reader of the stream finds the close. + let read = store + .inner + .read(StoreReadRequest { + chain: Some(chain), + topic: "tokens".to_string(), + ..Default::default() + }) + .await + .unwrap(); + assert!(read.closed); +} + +#[tokio::test] +async fn a_marked_activity_that_fails_closes_with_the_failure_the_workflow_saw() { + let store = Arc::new(RecordingStore::default()); + run_activity_owner( + canned_histories::single_failed_activity("tool"), + store.clone(), + true, + &["tokens"], + ) + .await; + let closes = closes_within_deadline(&store, 1).await; + let failure = closes[0].1.failure.clone().unwrap(); + assert!(closes[0].1.result.is_none()); + assert!( + matches!( + failure.failure_info, + Some(temporalio_common::protos::temporal::api::failure::v1::failure::FailureInfo::ActivityFailureInfo(_)) + ), + "{failure:?}" + ); +} + +#[tokio::test] +async fn a_marked_activity_that_times_out_closes_with_its_timeout() { + let store = Arc::new(RecordingStore::default()); + run_activity_owner( + canned_histories::started_activity_timeout("tool"), + store.clone(), + true, + &["tokens"], + ) + .await; + let closes = closes_within_deadline(&store, 1).await; + let failure = closes[0].1.failure.clone().unwrap(); + assert_eq!(failure.message, "Activity task timed out", "{failure:?}"); +} + +#[tokio::test] +async fn an_unmarked_activity_costs_no_store_call() { + let store = Arc::new(RecordingStore::default()); + run_activity_owner( + canned_histories::single_activity("tool"), + store.clone(), + false, + &[], + ) + .await; + tokio::time::sleep(Duration::from_millis(50)).await; + assert!(store.calls().is_empty(), "{:?}", store.calls()); +} + +#[tokio::test] +async fn a_store_that_refuses_the_close_leaves_the_task_completed() { + let store = Arc::new(RecordingStore { + refuse_closes: true, + ..Default::default() + }); + let completions: Arc>> = Default::default(); + let worker = worker_appending_to( + canned_histories::single_activity("tool"), + vec![1.into(), 2.into()], + store.clone(), + completions.clone(), + ); + let start = worker.poll_workflow_activation().await.unwrap(); + activity_wrote(&store, &start.run_id, "tool", &["tokens"]).await; + complete( + &worker, + &start.run_id, + vec![schedule_activity( + "tool", + true, + ActivityCancellationType::TryCancel, + )], + ) + .await; + let resolved = worker.poll_workflow_activation().await.unwrap(); + worker.complete_execution(&resolved.run_id).await; + worker.drain_pollers_and_shutdown().await; + for _ in 0..200 { + if store.calls().contains(&"close_topic_with") { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + assert!( + store.calls().contains(&"close_topic_with"), + "{:?}", + store.calls() + ); + assert_eq!(completions.lock().len(), 2, "both tasks completed"); +} + +#[tokio::test] +async fn replay_of_a_resolution_makes_no_close() { + let mut history = canned_histories::single_activity("tool"); + // Another task after the resolution, so the resolution replays. + history.add_workflow_task_completed(); + let timer_started = history.add_by_type(EventType::TimerStarted); + history.add_timer_fired(timer_started, "1".to_string()); + history.add_workflow_task_scheduled_and_started(); + let store = Arc::new(RecordingStore::default()); + let worker = worker_appending_to(history, vec![3.into()], store.clone(), Default::default()); + let start = worker.poll_workflow_activation().await.unwrap(); + assert!(start.is_replaying); + complete( + &worker, + &start.run_id, + vec![schedule_activity( + "tool", + true, + ActivityCancellationType::TryCancel, + )], + ) + .await; + let resolved = worker.poll_workflow_activation().await.unwrap(); + assert!(resolved.is_replaying); + complete( + &worker, + &resolved.run_id, + vec![start_timer_cmd(1, Duration::from_secs(1))], + ) + .await; + let fired = worker.poll_workflow_activation().await.unwrap(); + worker.complete_execution(&fired.run_id).await; + worker.drain_pollers_and_shutdown().await; + tokio::time::sleep(Duration::from_millis(50)).await; + assert!(store.calls().is_empty(), "{:?}", store.calls()); +} + +#[tokio::test] +async fn a_try_cancel_closes_the_streams_at_once_with_the_canceled_failure() { + let store = Arc::new(RecordingStore::default()); + let worker = worker_appending_to( + canned_histories::cancel_scheduled_activity("tool", "signal"), + vec![1.into(), 2.into()], + store.clone(), + Default::default(), + ); + let start = worker.poll_workflow_activation().await.unwrap(); + activity_wrote(&store, &start.run_id, "tool", &["tokens"]).await; + complete( + &worker, + &start.run_id, + vec![schedule_activity( + "tool", + true, + ActivityCancellationType::TryCancel, + )], + ) + .await; + let signaled = worker.poll_workflow_activation().await.unwrap(); + complete( + &worker, + &signaled.run_id, + vec![RequestCancelActivity { seq: 1 }.into()], + ) + .await; + let closes = closes_within_deadline(&store, 1).await; + let failure = closes[0].1.failure.clone().unwrap(); + assert!( + matches!( + failure.cause.as_deref().and_then(|cause| cause.failure_info.as_ref()), + Some(temporalio_common::protos::temporal::api::failure::v1::failure::FailureInfo::CanceledFailureInfo(_)) + ), + "{failure:?}" + ); + let resolved = worker.poll_workflow_activation().await.unwrap(); + worker.complete_execution(&resolved.run_id).await; + worker.drain_pollers_and_shutdown().await; +} diff --git a/crates/sdk-core/src/streams.rs b/crates/sdk-core/src/streams.rs index e599ae1f07..48c532ae7f 100644 --- a/crates/sdk-core/src/streams.rs +++ b/crates/sdk-core/src/streams.rs @@ -140,8 +140,18 @@ fn poll_answer(response: PollStreamNotifierResponse) -> NotifierPoll { /// 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. An Activity run is its own chain, so its notifier is keyed by the run. +/// its own. An Activity run is its own chain, so its notifier is keyed by the run. A Workflow +/// Activity's notifier is keyed by its Workflow, its scheduling run and its Activity id. fn chain_reference(chain: &ChainId, topic: String) -> StreamReference { + if chain.owner_kind == proto::StreamOwnerKind::WorkflowActivity as i32 { + return StreamReference { + owner_kind: StreamOwnerKind::WorkflowActivity as i32, + workflow_id: chain.workflow_id.clone(), + run_id: chain.first_run_id.clone(), + topic, + activity_id: chain.activity_id.clone(), + }; + } if chain.owner_kind == proto::StreamOwnerKind::Activity as i32 { return StreamReference { owner_kind: StreamOwnerKind::Activity as i32, @@ -566,6 +576,33 @@ mod tests { assert!(request.close_result.is_none()); } + #[test] + fn a_notification_of_a_workflow_activity_stream_names_its_scheduling_run() { + let request = notify_request(StreamNotification { + chain: proto::ChainId { + namespace: "ns".to_string(), + workflow_id: "agent".to_string(), + first_run_id: "run-2".to_string(), + owner_kind: proto::StreamOwnerKind::WorkflowActivity as i32, + activity_id: "tool-42".to_string(), + }, + topic: "tokens".to_string(), + position: "cursor".to_string(), + counter: 3, + close: None, + }); + assert_eq!( + request.stream_ref, + Some(StreamReference { + owner_kind: StreamOwnerKind::WorkflowActivity as i32, + workflow_id: "agent".to_string(), + run_id: "run-2".to_string(), + topic: "tokens".to_string(), + activity_id: "tool-42".to_string(), + }) + ); + } + #[test] fn a_notification_of_an_activity_stream_names_the_activity_run() { let request = notify_request(StreamNotification { diff --git a/crates/sdk-core/src/worker/workflow/machines/activity_state_machine.rs b/crates/sdk-core/src/worker/workflow/machines/activity_state_machine.rs index 0d9a2dc5f4..3287386bc8 100644 --- a/crates/sdk-core/src/worker/workflow/machines/activity_state_machine.rs +++ b/crates/sdk-core/src/worker/workflow/machines/activity_state_machine.rs @@ -9,7 +9,7 @@ use crate::{ internal_flags::CoreInternalFlags, worker::workflow::{ CommandAnnotations, InternalFlagsRef, ProtoCommandExt, fatal, machines::HistEventData, - nondeterminism, + nondeterminism, stream_output::STREAM_OWNER_HEADER, }, }; use std::convert::{TryFrom, TryInto}; @@ -224,6 +224,15 @@ impl ActivityMachine { pub(super) fn was_cancelled_before_sent_to_server(&self) -> bool { self.shared_state().cancelled_before_sent } + + /// The Activity's id when Workflow code marked it as the owner of its own streams. + pub(super) fn stream_owner_activity_id(&self) -> Option<&str> { + let attrs = &self.shared_state.attrs; + attrs + .headers + .contains_key(STREAM_OWNER_HEADER) + .then_some(attrs.activity_id.as_str()) + } } impl TryFrom for ActivityMachineEvents { 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 d290ded89b..e2af9231b2 100644 --- a/crates/sdk-core/src/worker/workflow/machines/workflow_machines.rs +++ b/crates/sdk-core/src/worker/workflow/machines/workflow_machines.rs @@ -37,6 +37,7 @@ use crate::{ upsert_search_attributes_state_machine::upsert_search_attrs_internal, }, nondeterminism, + stream_output::ActivityStreamClose, }, }, }; @@ -137,6 +138,9 @@ pub(crate) struct WorkflowMachines { pub(crate) last_processed_event: i64, /// True if the workflow is replaying from history pub(crate) replaying: bool, + /// The streams of Workflow Activities this run resolved live, to close once the activation + /// completes. + activity_stream_closes: Vec, /// Workflow identifier pub(crate) workflow_id: String, /// Workflow type identifier. (Function name, class, etc) @@ -313,6 +317,7 @@ impl WorkflowMachines { last_history_from_server: basics.history, protocol_msgs: vec![], pending_stream_wakes: vec![], + activity_stream_closes: vec![], workflow_id: basics.workflow_id, workflow_type: basics.workflow_type, run_id: basics.run_id, @@ -352,6 +357,11 @@ impl WorkflowMachines { } } + /// Takes the stream closes of the Workflow Activities resolved since the last call. + pub(crate) fn take_activity_stream_closes(&mut self) -> Vec { + std::mem::take(&mut self.activity_stream_closes) + } + /// Returns true if workflow has seen a terminal command pub(crate) const fn workflow_is_finished(&self) -> bool { self.workflow_end_time.is_some() @@ -1269,6 +1279,22 @@ impl WorkflowMachines { for response in machine_responses { match response { MachineResponse::PushWFJob(a) => { + // Replay does no store calls, so only a live resolution closes the streams. + if !self.replaying + && let workflow_activation_job::Variant::ResolveActivity(resolved) = + &a.variant + && !resolved.is_local + && let Machines::ActivityMachine(machine) = self.machine(smk) + && let Some(activity_id) = machine.stream_owner_activity_id() + && let Some(close) = ActivityStreamClose::of( + &self.workflow_id, + &self.run_id, + activity_id, + resolved.result.as_ref(), + ) + { + self.activity_stream_closes.push(close); + } // We don't need to notify lang about jobs created by core-internal machines if !self.machine_is_core_created.contains_key(smk) { self.drive_me.send_job(a); diff --git a/crates/sdk-core/src/worker/workflow/managed_run.rs b/crates/sdk-core/src/worker/workflow/managed_run.rs index 2d28ff7a38..3a5dea6ab7 100644 --- a/crates/sdk-core/src/worker/workflow/managed_run.rs +++ b/crates/sdk-core/src/worker/workflow/managed_run.rs @@ -907,6 +907,7 @@ impl ManagedRun { replaying: self.wfm.machines.replaying, output: None, retired: None, + activity_stream_closes: vec![], }, resp_chan: completion.resp_chan, })) @@ -917,10 +918,12 @@ impl ManagedRun { complete_resp_chan: completion.resp_chan, }), }; - // Only a completion the run took has output to append. + // Only a completion the run took has output to append and streams to close. + let activity_stream_closes = self.wfm.machines.take_activity_stream_closes(); completed.map(|done| { done.map(|mut done| { done.result.output = output; + done.result.activity_stream_closes = activity_stream_closes; done }) }) @@ -1421,6 +1424,7 @@ impl ManagedRun { replaying, output: None, retired, + activity_stream_closes: vec![], }, resp_chan, } @@ -1457,6 +1461,7 @@ impl ManagedRun { replaying: self.wfm.machines.replaying, output: None, retired: None, + activity_stream_closes: vec![], }) .is_err() { diff --git a/crates/sdk-core/src/worker/workflow/mod.rs b/crates/sdk-core/src/worker/workflow/mod.rs index 51db47abe5..5f83ce8aa6 100644 --- a/crates/sdk-core/src/worker/workflow/mod.rs +++ b/crates/sdk-core/src/worker/workflow/mod.rs @@ -638,7 +638,7 @@ impl Workflows { ); } - let completion_outcome = if let Ok(c) = rx.await { + let mut completion_outcome = if let Ok(c) = rx.await { c } else if is_empty_completion { // The stream may shut down after accepting the message but before processing @@ -654,8 +654,12 @@ impl Workflows { ); return Ok(()); }; + let activity_stream_closes = std::mem::take(&mut completion_outcome.activity_stream_closes); let (completion_outcome, output_appended) = self.append_output(&run_id, completion_outcome).await; + if output_appended { + self.close_activity_streams(activity_stream_closes); + } let replaying = completion_outcome.replaying; let mut wft_from_complete = None; @@ -765,6 +769,18 @@ impl Workflows { )) } + /// Closes the streams of the Workflow Activities a completion resolved, in the background. + fn close_activity_streams(&self, closes: Vec) { + #[cfg(feature = "streams")] + if let Some(store) = &self.output_store { + store.close_activity_streams(closes); + return; + } + if !closes.is_empty() { + debug!("This Worker has no stream store to close Workflow Activity streams in"); + } + } + async fn retire_stream_sessions(&self, retired: &stream_output::RetiredSessions) { #[cfg(feature = "streams")] if let Some(store) = &self.output_store { @@ -1407,6 +1423,8 @@ struct ActivationCompleteResult { output: Option>, /// The store sessions to retire once the server accepts the reported task. retired: Option>, + /// The streams of Workflow Activities resolved live, to close once the output is appended. + activity_stream_closes: Vec, } /// One activation's stream output, and what failing its Workflow Task takes. diff --git a/crates/sdk-core/src/worker/workflow/stream_output.rs b/crates/sdk-core/src/worker/workflow/stream_output.rs index 0d995f0850..712af39d8c 100644 --- a/crates/sdk-core/src/worker/workflow/stream_output.rs +++ b/crates/sdk-core/src/worker/workflow/stream_output.rs @@ -5,13 +5,57 @@ use super::WFMachinesError; use std::collections::{BTreeMap, BTreeSet, HashMap}; use temporalio_common::{ protos::{ - coresdk::workflow_commands::OutputRecord, temporal::sdk::streams::v1::StreamRecordKind, + coresdk::{ + activity_result::{ActivityResolution, activity_resolution}, + workflow_commands::OutputRecord, + }, + temporal::{api::failure::v1::Failure, sdk::streams::v1::StreamRecordKind}, }, streams::check_topic, }; const CONTENT_HASH_BYTES: usize = 32; +/// The `ScheduleActivity` header that marks an Activity as the owner of its own streams. The +/// scheduling Workflow's Core closes them when it resolves the Activity. The command holds it, so +/// replay rebuilds the mark with no History read, and the Activity's Worker sees it too. +pub(crate) const STREAM_OWNER_HEADER: &str = "__temporal_stream_owner"; + +/// The streams of one Workflow Activity to close, with the outcome its scheduling Workflow saw. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct ActivityStreamClose { + pub(crate) workflow_id: String, + /// The run that scheduled the Activity, which keys its streams. + pub(crate) run_id: String, + pub(crate) activity_id: String, + /// `None` for a completed Activity, which closes with a null result. + pub(crate) failure: Option, +} + +impl ActivityStreamClose { + /// The close for `resolution`, or `None` for one that does not end the Activity. + pub(crate) fn of( + workflow_id: &str, + run_id: &str, + activity_id: &str, + resolution: Option<&ActivityResolution>, + ) -> Option { + let failure = match resolution?.status.as_ref()? { + activity_resolution::Status::Completed(_) => None, + activity_resolution::Status::Failed(failed) => Some(failed.failure.clone()?), + activity_resolution::Status::Cancelled(cancelled) => Some(cancelled.failure.clone()?), + // Only Local Activities back off, and they own no streams. + activity_resolution::Status::Backoff(_) => return None, + }; + Some(Self { + workflow_id: workflow_id.to_string(), + run_id: run_id.to_string(), + activity_id: activity_id.to_string(), + failure, + }) + } +} + /// The records one activation committed, appended when its completion reaches Core. #[derive(Debug, Clone, PartialEq)] pub(crate) struct OutputBatch { @@ -222,6 +266,22 @@ impl OutputStore { Ok(dropped) } + /// Closes every topic each Workflow Activity in `closes` wrote, in the background. A close + /// that fails is logged and dropped, since it must not fail the Workflow Task. The owner + /// checks then close the streams with a generic failure. + pub(crate) fn close_activity_streams(&self, closes: Vec) { + if closes.is_empty() { + return; + } + let store = self.store.clone(); + let namespace = self.namespace.clone(); + tokio::spawn(async move { + for close in closes { + close_one(store.as_ref(), &namespace, close).await; + } + }); + } + /// Drops the store's dedupe sessions of attempts that will not write again. A failure only /// leaves a session in the store until retention, so it is logged and not returned. pub(crate) async fn retire(&self, retired: &RetiredSessions) { @@ -241,6 +301,90 @@ impl OutputStore { } } +#[cfg(feature = "streams")] +async fn close_one( + store: &dyn temporalio_streams::StreamStore, + namespace: &str, + close: ActivityStreamClose, +) { + use temporalio_streams::proto::{ChainId, StoredClose, StreamOwnerKind}; + let chain = ChainId { + namespace: namespace.to_string(), + workflow_id: close.workflow_id, + first_run_id: close.run_id, + owner_kind: StreamOwnerKind::WorkflowActivity as i32, + activity_id: close.activity_id, + }; + let closed_by = format!("{}@{}", chain.activity_id, chain.first_run_id); + let (result, failure) = match close.failure { + None => (Some(temporalio_streams::empty_close_result()), None), + Some(failure) => (None, Some(fit_failure(failure, closed_by.len()))), + }; + let topics = match store.written_topics(&chain).await { + Ok(topics) => topics, + Err(error) => { + warn!( + workflow_id = %chain.workflow_id, + activity_id = %chain.activity_id, + %error, + "Could not list the topics of an Activity's streams to close them" + ); + return; + } + }; + for topic in topics { + let stored = StoredClose { + result: result.clone(), + failure: failure.clone(), + closed_by: closed_by.clone(), + }; + if let Err(error) = store.close_stream(&chain, &topic, stored).await { + warn!( + workflow_id = %chain.workflow_id, + activity_id = %chain.activity_id, + %topic, + %error, + "Could not close an Activity's stream when its Workflow resolved it" + ); + } + } +} + +/// `failure` cut to what a stored close holds. The causes and the stack trace go first, then the +/// details, and a failure still too large keeps only the start of its message. +#[cfg(feature = "streams")] +fn fit_failure(failure: Failure, closed_by_bytes: usize) -> Failure { + use prost::Message; + // Room for the close's other fields and their tags. + let limit = temporalio_streams::MAX_CLOSE_BYTES - closed_by_bytes - 64; + if failure.encoded_len() <= limit { + return failure; + } + let mut cut = Failure { + cause: None, + stack_trace: String::new(), + ..failure + }; + if cut.encoded_len() <= limit { + return cut; + } + cut.encoded_attributes = None; + cut.failure_info = None; + if cut.encoded_len() <= limit { + return cut; + } + let mut end = limit / 2; + while !cut.message.is_char_boundary(end) { + end -= 1; + } + cut.message.truncate(end); + Failure { + message: cut.message, + source: cut.source, + ..Default::default() + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/streams/src/subscribe.rs b/crates/streams/src/subscribe.rs index cd3f50d350..e39b1254a2 100644 --- a/crates/streams/src/subscribe.rs +++ b/crates/streams/src/subscribe.rs @@ -658,6 +658,58 @@ mod tests { assert_eq!(setup.server.detached.lock().unwrap()[0].0, activity_chain); } + #[tokio::test] + async fn a_workflow_reader_subscribes_to_a_stream_that_a_workflow_activity_owns() { + let setup = setup(); + let tool = |run_id: &str| StreamAddress { + owner_kind: StreamOwnerKind::WorkflowActivity as i32, + run_id: run_id.to_string(), + activity_id: "1".to_string(), + ..stream() + }; + let tool_chain = |run_id: &str| ChainId { + namespace: "ns".to_string(), + workflow_id: "owner".to_string(), + first_run_id: run_id.to_string(), + owner_kind: StreamOwnerKind::WorkflowActivity as i32, + activity_id: "1".to_string(), + }; + // Activity `1` of two runs, and the Workflow's own stream, are three streams. + for (address, id) in [(tool("run-a"), "1"), (tool("run-b"), "2")] { + let request = SubscribeRequest { + stream: Some(address), + ..subscribe_request(id) + }; + let answer = subscribe(&setup.service, request).await.unwrap(); + assert!(answer.attached); + } + assert!( + subscribe(&setup.service, subscribe_request("3")) + .await + .unwrap() + .attached + ); + let attached = setup.server.attached.lock().unwrap().clone(); + assert_eq!(attached[0].chain, tool_chain("run-a")); + assert_eq!(attached[1].chain, tool_chain("run-b")); + assert_eq!(attached[2].chain, chain("chain-1")); + let unsubscribe = UnsubscribeRequest { + stream: Some(tool("run-b")), + run_id: "reader-run".to_string(), + subscription_id: "2".to_string(), + ..Default::default() + }; + setup + .service + .call("Unsubscribe", &unsubscribe.encode_to_vec()) + .await + .unwrap(); + assert_eq!( + setup.server.detached.lock().unwrap()[0].0, + tool_chain("run-b") + ); + } + async fn wait( setup: &Setup, after_counter: i64, From 90cd6c0ed5d76d25845aa42dffb132bd58709522 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Fri, 16 Oct 2026 05:05:48 -0700 Subject: [PATCH 2/3] Closed a Workflow Activity's streams when its resolution has no failure. A Failed or Cancelled resolution without a failure closed nothing, so readers waited for the owner checks. --- .../src/worker/workflow/stream_output.rs | 43 ++++++++++++++++++- 1 file changed, 41 insertions(+), 2 deletions(-) diff --git a/crates/sdk-core/src/worker/workflow/stream_output.rs b/crates/sdk-core/src/worker/workflow/stream_output.rs index f2408384c8..ec84f1679f 100644 --- a/crates/sdk-core/src/worker/workflow/stream_output.rs +++ b/crates/sdk-core/src/worker/workflow/stream_output.rs @@ -42,8 +42,20 @@ impl ActivityStreamClose { ) -> Option { let failure = match resolution?.status.as_ref()? { activity_resolution::Status::Completed(_) => None, - activity_resolution::Status::Failed(failed) => Some(failed.failure.clone()?), - activity_resolution::Status::Cancelled(cancelled) => Some(cancelled.failure.clone()?), + // A resolution with no failure still ends the Activity, so its streams close with a + // failure of their own. + activity_resolution::Status::Failed(failed) => Some( + failed + .failure + .clone() + .unwrap_or_else(|| no_failure("failed")), + ), + activity_resolution::Status::Cancelled(cancelled) => Some( + cancelled + .failure + .clone() + .unwrap_or_else(|| no_failure("was cancelled")), + ), // Only Local Activities back off, and they own no streams. activity_resolution::Status::Backoff(_) => return None, }; @@ -56,6 +68,13 @@ impl ActivityStreamClose { } } +fn no_failure(what: &str) -> Failure { + Failure { + message: format!("The Activity {what} with no failure"), + ..Default::default() + } +} + /// The records one activation committed, appended when its completion reaches Core. #[derive(Debug, Clone, PartialEq)] pub(crate) struct OutputBatch { @@ -416,6 +435,26 @@ mod tests { use super::*; use temporalio_common::protos::temporal::api::common::v1::Payload; + #[test] + fn a_failed_or_cancelled_resolution_with_no_failure_still_closes() { + use temporalio_common::protos::coresdk::activity_result::{ + Cancellation, Failure as Failed, + }; + for status in [ + activity_resolution::Status::Failed(Failed { + failure: None, + ..Default::default() + }), + activity_resolution::Status::Cancelled(Cancellation { failure: None }), + ] { + let resolution = ActivityResolution { + status: Some(status), + }; + let close = ActivityStreamClose::of("wf", "run-1", "tool", Some(&resolution)).unwrap(); + assert!(close.failure.is_some()); + } + } + fn data(topic: &str, hash: u8) -> OutputRecord { OutputRecord { topic: topic.to_string(), From b5d5ca126ca75c8916d953bb3306b13b62f920db Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Fri, 16 Oct 2026 05:05:49 -0700 Subject: [PATCH 3/3] Tracked the Workflow Activity stream closes in the shutdown flush. They ran in an untracked task, so a Worker stop lost them, and replay never sends them again. --- .../src/core_tests/external_streams.rs | 25 +++++++++++++++++++ .../src/worker/workflow/stream_output.rs | 3 ++- 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/crates/sdk-core/src/core_tests/external_streams.rs b/crates/sdk-core/src/core_tests/external_streams.rs index 7d35e096fc..03ced1c1c2 100644 --- a/crates/sdk-core/src/core_tests/external_streams.rs +++ b/crates/sdk-core/src/core_tests/external_streams.rs @@ -73,6 +73,8 @@ struct RecordingStore { closes: Mutex>, /// Refuses every topic close. refuse_closes: bool, + /// How long each topic close takes. + close_delay: Option, /// A retire waits for a permit here, when set. retire_gate: Option>, @@ -209,6 +211,9 @@ impl StreamStore for RecordingStore { if self.refuse_closes { return Err(StreamError::storage("the store is down")); } + if let Some(delay) = self.close_delay { + tokio::time::sleep(delay).await; + } self.closes .lock() .push((chain.clone(), topic.to_string(), close.clone())); @@ -1609,6 +1614,26 @@ async fn a_marked_activity_that_completes_closes_each_topic_it_wrote_with_null() assert!(read.closed); } +#[tokio::test] +async fn a_worker_shutdown_waits_for_the_activity_stream_closes() { + let store = Arc::new(RecordingStore { + close_delay: Some(Duration::from_millis(200)), + ..Default::default() + }); + run_activity_owner( + canned_histories::single_activity("tool"), + store.clone(), + true, + &["tokens"], + ) + .await; + assert_eq!( + store.closes.lock().len(), + 1, + "closed before the shutdown returned" + ); +} + #[tokio::test] async fn a_marked_activity_that_fails_closes_with_the_failure_the_workflow_saw() { let store = Arc::new(RecordingStore::default()); diff --git a/crates/sdk-core/src/worker/workflow/stream_output.rs b/crates/sdk-core/src/worker/workflow/stream_output.rs index ec84f1679f..0e68399536 100644 --- a/crates/sdk-core/src/worker/workflow/stream_output.rs +++ b/crates/sdk-core/src/worker/workflow/stream_output.rs @@ -311,7 +311,8 @@ impl OutputStore { } let store = self.store.clone(); let namespace = self.namespace.clone(); - tokio::spawn(async move { + // Replay never sends the close again, so a Worker shutdown waits for it. + self.tasks.spawn(async move { for close in closes { close_one(store.as_ref(), &namespace, close).await; }