From 13629af5903bcf8b0bf8f11eae14fdea5288eec1 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 13 Oct 2026 19:01:29 -0700 Subject: [PATCH 1/3] 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 3b1d2ebc6def7f0c55edc5d132eed6e349be4d66 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 13 Oct 2026 19:19:27 -0700 Subject: [PATCH 2/3] 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 0bbeb6eab120000f231cef057d858000031684a2 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Wed, 14 Oct 2026 02:17:36 -0700 Subject: [PATCH 3/3] 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();