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 32b8468c69..d547b4a780 100644 --- a/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto +++ b/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto @@ -31,6 +31,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 @@ -247,6 +254,39 @@ 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; + // 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 { + // 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, 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 {} + message DeleteOwnerRequest { string namespace = 1; StreamOwnerKind owner_kind = 2; diff --git a/crates/sdk-core/CHANGELOG.md b/crates/sdk-core/CHANGELOG.md index 95a142a712..4963e74c27 100644 --- a/crates/sdk-core/CHANGELOG.md +++ b/crates/sdk-core/CHANGELOG.md @@ -66,6 +66,13 @@ 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. 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 stream service's `RetireProducer` call drops one producer attempt's dedupe state on a stream, which the Redis store keeps as one field of the topic's meta hash until retention. A producer that writes once per short-lived attempt calls it after the attempt's diff --git a/crates/sdk-core/src/streams.rs b/crates/sdk-core/src/streams.rs index 97e8694aef..0039972c17 100644 --- a/crates/sdk-core/src/streams.rs +++ b/crates/sdk-core/src/streams.rs @@ -8,20 +8,22 @@ use temporalio_common::protos::google::rpc::Status as RpcStatus; 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, @@ -51,19 +53,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(), @@ -295,6 +347,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 e49faead7a..a0752b4889 100644 --- a/crates/streams/src/handle.rs +++ b/crates/streams/src/handle.rs @@ -87,31 +87,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); @@ -132,10 +136,6 @@ impl Remembered { }, ); } - - fn forget(&self, key: &K) { - self.entries.lock().unwrap().remove(key); - } } /// The stream layer over one store. @@ -625,6 +625,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); @@ -676,7 +691,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 6b5eead95b..39a895795c 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}; @@ -30,7 +31,7 @@ pub use memory::MemoryStore; pub use notify::{ Arm, CLOSE_RETRY_CAP, CLOSE_RETRY_FIRST, CLOSE_RETRY_LIMIT, DEFAULT_MAX_NOTIFIERS, FLUSH_LIMIT, NOTIFIER_REFRESH, Notifier, NotifierClient, NotifyError, NotifyingStore, StreamNotification, - progress_counter, + WorkflowAttachment, progress_counter, }; pub use owner::{OwnerClient, OwnerDescription, OwnerError}; pub use reader::{ diff --git a/crates/streams/src/notify.rs b/crates/streams/src/notify.rs index d6793d4a36..2df95bf6e0 100644 --- a/crates/streams/src/notify.rs +++ b/crates/streams/src/notify.rs @@ -86,11 +86,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 35dca41fd6..9b30dc1cb4 100644 --- a/crates/streams/src/service.rs +++ b/crates/streams/src/service.rs @@ -6,8 +6,10 @@ use crate::{ proto::{ AppendRequest, CloseRequest, DeleteOwnerRequest, FlushNotificationsRequest, FlushNotificationsResponse, LatestRequest, ReadRequest, RetireProducerRequest, - StreamFailure, StreamStoreConfig, stream_store_config, + StreamFailure, StreamStoreConfig, SubscribeRequest, UnsubscribeRequest, + stream_store_config, }, + subscribe::Subscriptions, }; use prost::Message; use std::{sync::Arc, time::Duration}; @@ -27,6 +29,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 +51,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 +74,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, @@ -90,7 +94,11 @@ impl StreamService { StreamError::refused("the stream store's owner start grace must not be negative") })?; } - Ok(Self::new(store, owner, options)) + let service = Self::new(store, owner, options); + Ok(match notifier { + Some(client) => service.with_subscriptions(client), + None => service, + }) } /// Serves calls on `store`. @@ -100,7 +108,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 } /// The store, for the Workers that append Workflow output to it. @@ -150,6 +168,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..af75629a4b --- /dev/null +++ b/crates/streams/src/subscribe.rs @@ -0,0 +1,592 @@ +//! 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, OWNER_RECHECK, StreamError, StreamResult, StreamStore, Streams, + StreamsOptions, WorkflowAttachment, + handle::{Remembered, address}, + proto::{ + StreamAddress, SubscribeRequest, SubscribeResponse, UnsubscribeRequest, UnsubscribeResponse, + }, +}; +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. 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: Remembered, +} + +impl Default for Subscriptions { + fn default() -> Self { + Self::new(None) + } +} + +#[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. 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<()> { + 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: Remembered::new(StreamsOptions::default().remembered), + } + } + + 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, + request.generation, + ); + let key = SubscriptionKey::new(stream, &request_id); + if let Some(attached) = self.attached.get(&key) { + return Ok(SubscribeResponse { + chain_first_run_id: attached.first_run_id, + 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.put(key, &chain.first_run_id, OWNER_RECHECK); + 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, + request.generation, + ); + let client = self.client()?; + let key = SubscriptionKey::new(stream, &request_id); + 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: attached.first_run_id, + }, + None => streams.chain_now(stream).await?, + }; + client + .detach(chain, stream.topic.clone(), request_id) + .await + .map_err(subscribe_error)?; + self.attached.forget(&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 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. + 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(), + generation: 0, + } + } + + 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_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(); + 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 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(); + 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(), + ..Default::default() + }; + 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"); + } +}