From 4ad26e91e68465db274c31fd403b7f739985d21e Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Mon, 12 Oct 2026 13:50:36 -0700 Subject: [PATCH] Added the arm a stream reader sets to the store contract. An empty read that asks to arm sets a generation that the next append or promotion reports, so a producer can notify only while a reader sleeps. The memory store arms, other stores report that they can't tell. --- .../temporal/sdk/core/streams/streams.proto | 12 ++ .../src/core_tests/external_streams.rs | 1 + .../src/worker/workflow/stream_output.rs | 1 + crates/streams/src/lib.rs | 2 +- crates/streams/src/memory.rs | 101 ++++++++++-- crates/streams/src/notify.rs | 11 +- crates/streams/src/reader.rs | 3 +- crates/streams/src/redis/mod.rs | 3 + crates/streams/src/store.rs | 44 +++++- crates/streams/tests/conformance.rs | 149 +++++++++++++++++- 10 files changed, 310 insertions(+), 17 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 bca8cc975..cff30c83c 100644 --- a/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto +++ b/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto @@ -294,6 +294,9 @@ message StoreAppendRequest { message StoreAppendResponse { string first_position = 1; string last_position = 2; + // The topic's arm generation as the append found it, zero when no reader armed the topic. + // Unset when the store can't tell, and Core then notifies after the append. + optional uint64 arm = 3; } message StoreReadRequest { @@ -303,6 +306,9 @@ message StoreReadRequest { string after_position = 3; google.protobuf.Duration wait = 4; uint32 max_records = 5; + // When the read ends with no record and the topic is open, arm the topic in the same step that + // finds it empty, so the next append reports the arm. A store that can't arm ignores it. + bool arm = 6; } message StoredRecord { @@ -315,6 +321,9 @@ message StoreReadResponse { repeated StoredRecord records = 1; // The chain or this topic is marked closed. bool closed = 2; + // The read asked to arm, found no record and armed the topic. Never set with records or with + // `closed`. + bool armed = 3; } message StoreLatestRequest { @@ -368,6 +377,9 @@ enum PromoteOutcome { message PromoteResult { PromoteOutcome outcome = 1; uint32 records = 2; + // Each promoted topic's arm generation as the promotion found it, zero when no reader armed + // the topic. A topic left out means the store can't tell, and Core then notifies it. + map arms = 3; } // A stage no one has promoted or aborted yet. diff --git a/crates/sdk-core/src/core_tests/external_streams.rs b/crates/sdk-core/src/core_tests/external_streams.rs index 5adf6447d..b44435491 100644 --- a/crates/sdk-core/src/core_tests/external_streams.rs +++ b/crates/sdk-core/src/core_tests/external_streams.rs @@ -168,6 +168,7 @@ impl StreamStore for RecordingStore { Ok(PromoteResult { outcome: PromoteOutcome::Promoted as i32, records: 1, + ..Default::default() }) } diff --git a/crates/sdk-core/src/worker/workflow/stream_output.rs b/crates/sdk-core/src/worker/workflow/stream_output.rs index 808b40d90..3c93e5d02 100644 --- a/crates/sdk-core/src/worker/workflow/stream_output.rs +++ b/crates/sdk-core/src/worker/workflow/stream_output.rs @@ -1450,6 +1450,7 @@ mod tests { Ok(PromoteResult { outcome: PromoteOutcome::Promoted as i32, records: 1, + ..Default::default() }) } async fn abort(&self, stage: &StageRef) -> StreamResult<()> { diff --git a/crates/streams/src/lib.rs b/crates/streams/src/lib.rs index 3c7ae0088..b802b2c6f 100644 --- a/crates/streams/src/lib.rs +++ b/crates/streams/src/lib.rs @@ -43,7 +43,7 @@ pub use record::{ }; pub use repair::{Decision, decide_token}; pub use service::StreamService; -pub use store::{CloseOutcome, StreamStore, empty_close_result}; +pub use store::{CloseOutcome, NewestOrArm, StreamStore, empty_close_result}; /// The protos stores and lang exchange. pub mod proto { diff --git a/crates/streams/src/memory.rs b/crates/streams/src/memory.rs index 453106099..05f67306c 100644 --- a/crates/streams/src/memory.rs +++ b/crates/streams/src/memory.rs @@ -6,7 +6,7 @@ //! position keeps naming the same record after [MemoryStore::truncate] drops older ones. use crate::{ - StreamError, StreamResult, StreamStore, + NewestOrArm, StreamError, StreamResult, StreamStore, proto::{ ChainId, DeleteOwnerRequest, DeleteOwnerResponse, PendingStage, PromoteOutcome, PromoteResult, StageRef, StagedBatch, StoreAppendRequest, StoreAppendResponse, @@ -78,6 +78,29 @@ struct Chain { closed: bool, /// In staging order, which is the order a Worker commits them in. stages: Vec, + /// Kept apart from the topics, so an arm on a topic with no records leaves it unknown to + /// reads and to the topic list. + arms: HashMap, +} + +/// A topic's arm: the generation a reader set, zero once cleared, and the last one handed out. +#[derive(Debug, Default)] +struct Arm { + generation: u64, + last: u64, +} + +impl Chain { + fn arm(&self, topic: &str) -> u64 { + self.arms.get(topic).map_or(0, |arm| arm.generation) + } + + fn set_arm(&mut self, topic: &str) -> u64 { + let arm = self.arms.entry(topic.to_string()).or_default(); + arm.last += 1; + arm.generation = arm.last; + arm.generation + } } /// A [StreamStore] in process memory. For tests only. @@ -124,9 +147,33 @@ impl MemoryStore { self.written.send_modify(|version| *version += 1); } - /// One look at what a read gets now: `None` when it should wait for more. - fn read_now(&self, request: &StoreReadRequest) -> StreamResult> { - let chains = self.chains.lock().unwrap(); + /// One look at what a read gets now: `None` when it should wait for more. With `arm`, a look + /// that finds nothing arms the topic instead. + fn read_now( + &self, + request: &StoreReadRequest, + arm: bool, + ) -> StreamResult> { + let mut chains = self.chains.lock().unwrap(); + let found = Self::look(&chains, request)?; + if found.is_some() || !arm { + return Ok(found); + } + // Under the lock the look held, so no append lands between finding nothing and the arm. + chains + .entry(chain_key(chain_of(request.chain.as_ref())?)) + .or_default() + .set_arm(&request.topic); + Ok(Some(StoreReadResponse { + armed: true, + ..Default::default() + })) + } + + fn look( + chains: &HashMap, + request: &StoreReadRequest, + ) -> StreamResult> { let chain = chains.get(&chain_key(chain_of(request.chain.as_ref())?)); let topic = chain.and_then(|chain| chain.topics.get(&request.topic)); let closed = chain.is_some_and(|chain| chain.closed) || topic.is_some_and(|t| t.closed); @@ -140,8 +187,8 @@ impl MemoryStore { ))); } return Ok(closed.then(|| StoreReadResponse { - records: Vec::new(), closed, + ..Default::default() })); }; let start = after.map_or(topic.base, |after| after + 1); @@ -165,7 +212,13 @@ impl MemoryStore { record: topic.at(offset).cloned().unwrap_or_default(), }) .collect(); - Ok((!records.is_empty() || closed).then_some(StoreReadResponse { records, closed })) + Ok( + (!records.is_empty() || closed).then_some(StoreReadResponse { + records, + closed, + armed: false, + }), + ) } } @@ -228,6 +281,7 @@ impl StreamStore for MemoryStore { return Ok(StoreAppendResponse { first_position: held.first.to_string(), last_position: last.to_string(), + arm: Some(chain.arm(&request.topic)), }); } let next = held.sequence + held.count; @@ -247,6 +301,7 @@ impl StreamStore for MemoryStore { "the Workflow that owns this stream has closed", )); } + let arm = chain.arm(&request.topic); let topic = chain.topics.entry(request.topic.clone()).or_default(); let count = request.records.len() as i64; let mut positions = request.records.into_iter().map(|record| topic.push(record)); @@ -266,6 +321,7 @@ impl StreamStore for MemoryStore { Ok(StoreAppendResponse { first_position: first.to_string(), last_position: last.to_string(), + arm: Some(arm), }) } @@ -273,14 +329,14 @@ impl StreamStore for MemoryStore { let deadline = Instant::now() + wait_of(&request); let mut written = self.written.subscribe(); loop { - if let Some(response) = self.read_now(&request)? { + if let Some(response) = self.read_now(&request, false)? { return Ok(response); } if tokio::time::timeout_at(deadline, written.changed()) .await .is_err() { - return Ok(StoreReadResponse::default()); + return Ok(self.read_now(&request, request.arm)?.unwrap_or_default()); } } } @@ -388,11 +444,17 @@ impl StreamStore for MemoryStore { .or_default() .push(record.record); } + let arms = stage + .topics + .iter() + .map(|topic| (topic.clone(), chain.arm(topic))) + .collect(); drop(chains); self.wrote(); Ok(PromoteResult { outcome: PromoteOutcome::Promoted as i32, records, + arms, }) } @@ -461,6 +523,27 @@ impl StreamStore for MemoryStore { .unwrap_or_default()) } + async fn disarm(&self, chain: &ChainId, topic: &str, generation: u64) -> StreamResult<()> { + let mut chains = self.chains.lock().unwrap(); + if let Some(arm) = chains + .get_mut(&chain_key(chain)) + .and_then(|chain| chain.arms.get_mut(topic)) + && arm.generation == generation + { + arm.generation = 0; + } + Ok(()) + } + + async fn newest_or_arm(&self, chain: &ChainId, topic: &str) -> StreamResult { + let mut chains = self.chains.lock().unwrap(); + let chain = chains.entry(chain_key(chain)).or_default(); + if let Some(held) = chain.topics.get(topic).filter(|t| !t.records.is_empty()) { + return Ok(NewestOrArm::Newest((held.head() - 1).to_string())); + } + Ok(NewestOrArm::Armed(chain.set_arm(topic))) + } + async fn delete_owner(&self, request: DeleteOwnerRequest) -> StreamResult { if request.owner_kind != StreamOwnerKind::Workflow as i32 { return Err(StreamError::unsupported(format!( @@ -484,7 +567,7 @@ impl StreamStore for MemoryStore { fn settled() -> PromoteResult { PromoteResult { outcome: PromoteOutcome::Settled as i32, - records: 0, + ..Default::default() } } diff --git a/crates/streams/src/notify.rs b/crates/streams/src/notify.rs index e4c639f09..d12a100ac 100644 --- a/crates/streams/src/notify.rs +++ b/crates/streams/src/notify.rs @@ -4,7 +4,7 @@ //! notification is a hint that there is something new to read. The records stay on the read path. use crate::{ - CloseOutcome, StreamResult, StreamStore, WORKFLOW_OWNER_KIND, mint_cursor, + CloseOutcome, NewestOrArm, StreamResult, StreamStore, WORKFLOW_OWNER_KIND, mint_cursor, proto::{ ChainId, DeleteOwnerRequest, DeleteOwnerResponse, PendingStage, PromoteOutcome, PromoteResult, StageRef, StagedBatch, StoreAppendRequest, StoreAppendResponse, @@ -889,6 +889,14 @@ impl StreamStore for NotifyingStore { async fn flush_notifications(&self) { self.notifier.flush().await; } + + async fn disarm(&self, chain: &ChainId, topic: &str, generation: u64) -> StreamResult<()> { + self.inner.disarm(chain, topic, generation).await + } + + async fn newest_or_arm(&self, chain: &ChainId, topic: &str) -> StreamResult { + self.inner.newest_or_arm(chain, topic).await + } } #[cfg(test)] @@ -1249,6 +1257,7 @@ mod tests { Ok(StoreAppendResponse { first_position: "odd".to_string(), last_position: "odd".to_string(), + arm: None, }) } async fn read(&self, r: StoreReadRequest) -> StreamResult { diff --git a/crates/streams/src/reader.rs b/crates/streams/src/reader.rs index 85a5aa1fe..14cc19ab0 100644 --- a/crates/streams/src/reader.rs +++ b/crates/streams/src/reader.rs @@ -334,6 +334,7 @@ pub(crate) async fn read_with( 0 => DEFAULT_MAX_RECORDS, max => max, }, + arm: false, }) .await?; let closed = page.closed; @@ -494,7 +495,7 @@ mod tests { .collect(); Ok(StoreReadResponse { records, - closed: false, + ..Default::default() }) } diff --git a/crates/streams/src/redis/mod.rs b/crates/streams/src/redis/mod.rs index 11dfe1612..2fa148e0c 100644 --- a/crates/streams/src/redis/mod.rs +++ b/crates/streams/src/redis/mod.rs @@ -467,6 +467,7 @@ impl StreamStore for RedisStore { Ok(StoreAppendResponse { first_position, last_position, + arm: None, }) } @@ -511,6 +512,7 @@ impl StreamStore for RedisStore { Ok(StoreReadResponse { records: found, closed: state.chain_closed || state.topic_closed, + armed: false, }) } @@ -683,6 +685,7 @@ impl StreamStore for RedisStore { Ok(PromoteResult { outcome: outcome as i32, records, + arms: HashMap::new(), }) } diff --git a/crates/streams/src/store.rs b/crates/streams/src/store.rs index d962bd691..2f6f7e2dc 100644 --- a/crates/streams/src/store.rs +++ b/crates/streams/src/store.rs @@ -59,6 +59,7 @@ pub trait StreamStore: Send + Sync { /// positions and writes nothing. The same first sequence with another digest fails as /// producer divergent, and a first sequence below the end of the newest batch as producer /// stale. A closed chain or topic refuses a new batch as closed, but still answers a repeat. + /// The answer, a repeat's too, reports the topic's arm, so a lost notification goes again. async fn append(&self, request: StoreAppendRequest) -> StreamResult; /// Reads the records of a topic after a position, waiting up to the request's `wait` for one @@ -67,6 +68,10 @@ pub trait StreamStore: Send + Sync { /// Fails as expired when the position names a record retention dropped, or when retention /// dropped records the read had not delivered yet. Fails as not found when nothing is known /// about the topic and the position is not empty. + /// + /// A read that asks to arm and ends with no record on an open topic arms the topic in the + /// same step that finds it empty. The next append or promotion then reports the arm, and + /// the producer notifies the topic's callers. async fn read(&self, request: StoreReadRequest) -> StreamResult; /// The position of the newest record of a topic, or empty when it holds none. @@ -107,7 +112,7 @@ pub trait StreamStore: Send + Sync { async fn stage(&self, batch: StagedBatch) -> StreamResult<()>; /// Moves a stage's records into their topics in one step, so readers see all or none of - /// them. Idempotent, and the records become visible once. + /// them. Idempotent, and the records become visible once. Reports each promoted topic's arm. async fn promote(&self, stage: &StageRef) -> StreamResult; /// Drops a stage whose Workflow Task never committed. Idempotent. @@ -185,6 +190,43 @@ pub trait StreamStore: Send + Sync { /// Waits until the notifications this store sent are out. A store that sends none returns at /// once. async fn flush_notifications(&self) {} + + /// Clears the arm of a topic when it still holds `generation`, once the notification that + /// generation asked for landed. A reader that armed again since holds a higher generation, + /// and its arm stays for the next append. A store that can't arm keeps the default. + async fn disarm(&self, chain: &ChainId, topic: &str, generation: u64) -> StreamResult<()> { + let _ = (chain, topic, generation); + Ok(()) + } + + /// The position of the newest record of a topic, or, when it holds none, an arm set in the + /// same step that finds it empty, so the first append reports the arm. A store that can't + /// arm keeps the default, which answers [NewestOrArm::Unarmed] for an empty topic. + async fn newest_or_arm(&self, chain: &ChainId, topic: &str) -> StreamResult { + let position = self + .latest(StoreLatestRequest { + chain: Some(chain.clone()), + topic: topic.to_string(), + }) + .await? + .position; + Ok(if position.is_empty() { + NewestOrArm::Unarmed + } else { + NewestOrArm::Newest(position) + }) + } +} + +/// What [StreamStore::newest_or_arm] found. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum NewestOrArm { + /// The position of the topic's newest record. + Newest(String), + /// The topic held no record, and the store armed it with this generation. + Armed(u64), + /// The topic held no record, and the store can't arm. + Unarmed, } #[cfg(test)] diff --git a/crates/streams/tests/conformance.rs b/crates/streams/tests/conformance.rs index 374cff855..44b452d45 100644 --- a/crates/streams/tests/conformance.rs +++ b/crates/streams/tests/conformance.rs @@ -15,10 +15,10 @@ use std::{ }; use temporalio_common::protos::temporal::api::common::v1::Payload; use temporalio_streams::{ - BEGINNING, END, MemoryStore, Notifier, NotifierClient, NotifyError, NotifyingStore, - OwnerClient, OwnerDescription, OwnerError, ReadTarget, StreamNotification, StreamResult, - StreamStore, Streams, StreamsOptions, mint_cursor, proto::*, read_page, stored_append_record, - stored_output_record, stream_hash, + BEGINNING, END, MemoryStore, NewestOrArm, Notifier, NotifierClient, NotifyError, + NotifyingStore, OwnerClient, OwnerDescription, OwnerError, ReadTarget, StreamNotification, + StreamResult, StreamStore, Streams, StreamsOptions, mint_cursor, proto::*, read_page, + stored_append_record, stored_output_record, stream_hash, }; #[async_trait::async_trait] @@ -1650,6 +1650,130 @@ mod cases { ["1"] ); } + /// A store read that asks to arm when it finds nothing, as a reader that sleeps on progress + /// makes it. + async fn arming_read( + case: &dyn Case, + chain: &ChainId, + topic: &str, + after: &str, + ) -> StoreReadResponse { + case.store() + .read(StoreReadRequest { + chain: Some(chain.clone()), + topic: topic.to_string(), + after_position: after.to_string(), + max_records: 10, + arm: true, + ..Default::default() + }) + .await + .unwrap() + } + + pub(crate) async fn an_empty_read_that_asks_arms_and_the_next_append_reports_it( + case: &dyn Case, + ) { + let chain = case.chain(); + let read = arming_read(case, &chain, OUT, "").await; + assert!(read.records.is_empty()); + assert!(read.armed, "an empty read that asks arms"); + let landed = Producer::new(case, &chain, OUT, "p", 1) + .append(&["1"]) + .await; + assert!(landed.arm.is_some_and(|arm| arm > 0), "{landed:?}"); + } + + pub(crate) async fn an_append_with_nothing_armed_reports_zero(case: &dyn Case) { + let chain = case.chain(); + let mut producer = Producer::new(case, &chain, OUT, "p", 1); + assert_eq!(producer.append(&["1"]).await.arm, Some(0)); + // A read that doesn't ask never arms. + let read = read(case, &chain, OUT, BEGINNING).await.unwrap(); + let after = &read.records[0].cursor; + let empty = read_with(case, &chain, OUT, after, Duration::ZERO, &read.state) + .await + .unwrap(); + assert!(empty.records.is_empty()); + assert_eq!(producer.append(&["2"]).await.arm, Some(0)); + } + + pub(crate) async fn an_arm_that_finds_a_newer_record_stays_unarmed(case: &dyn Case) { + let chain = case.chain(); + let mut producer = Producer::new(case, &chain, OUT, "p", 1); + let first = producer.append(&["1"]).await; + producer.append(&["2"]).await; + let read = arming_read(case, &chain, OUT, &first.last_position).await; + assert_eq!(read.records.len(), 1); + assert!(!read.armed, "an answer with records never arms"); + assert_eq!(producer.append(&["3"]).await.arm, Some(0)); + } + + pub(crate) async fn a_disarm_of_an_older_generation_keeps_the_newer_arm(case: &dyn Case) { + let store = case.store(); + let chain = case.chain(); + let mut producer = Producer::new(case, &chain, OUT, "p", 1); + assert!(arming_read(case, &chain, OUT, "").await.armed); + let landed = producer.append(&["1"]).await; + let older = landed.arm.unwrap(); + assert!( + arming_read(case, &chain, OUT, &landed.last_position) + .await + .armed + ); + let landed = producer.append(&["2"]).await; + let newer = landed.arm.unwrap(); + assert!(newer > older, "a new arm gets a higher generation"); + store.disarm(&chain, OUT, older).await.unwrap(); + assert_eq!(producer.append(&["3"]).await.arm, Some(newer)); + store.disarm(&chain, OUT, newer).await.unwrap(); + assert_eq!(producer.append(&["4"]).await.arm, Some(0)); + // A disarm of a topic that holds no arm changes nothing. + store.disarm(&chain, OTHER, newer).await.unwrap(); + } + + pub(crate) async fn a_promotion_reports_each_topics_arm(case: &dyn Case) { + let store = case.store(); + let chain = case.chain(); + assert!(arming_read(case, &chain, OUT, "").await.armed); + store + .stage(staged(&chain, "run-1", "t1", &[(OUT, "1"), (OTHER, "2")])) + .await + .unwrap(); + let promoted = store + .promote(&stage_ref(&chain, "t1", &[OUT, OTHER])) + .await + .unwrap(); + assert_eq!(promoted.outcome(), PromoteOutcome::Promoted); + assert!(promoted.arms[OUT] > 0, "{promoted:?}"); + assert_eq!(promoted.arms[OTHER], 0, "{promoted:?}"); + } + + pub(crate) async fn newest_or_arm_arms_only_an_empty_topic(case: &dyn Case) { + let store = case.store(); + let chain = case.chain(); + let NewestOrArm::Armed(generation) = store.newest_or_arm(&chain, OUT).await.unwrap() else { + panic!("an empty topic arms"); + }; + let mut producer = Producer::new(case, &chain, OUT, "p", 1); + let landed = producer.append(&["1"]).await; + assert_eq!(landed.arm, Some(generation)); + store.disarm(&chain, OUT, generation).await.unwrap(); + assert_eq!( + store.newest_or_arm(&chain, OUT).await.unwrap(), + NewestOrArm::Newest(landed.last_position) + ); + assert_eq!(producer.append(&["2"]).await.arm, Some(0)); + } + + pub(crate) async fn a_closed_topic_never_arms(case: &dyn Case) { + let store = case.store(); + let chain = case.chain(); + store.close_topic(&chain, OUT, None).await.unwrap(); + let read = arming_read(case, &chain, OUT, "").await; + assert!(read.closed); + assert!(!read.armed, "a closed answer never arms"); + } } macro_rules! conformance { @@ -1715,13 +1839,30 @@ macro_rules! store_cases { }; } +/// The arm a reader sets for the append notifications, for a store that arms. +macro_rules! arm_cases { + ($case:expr) => { + conformance!($case; + an_empty_read_that_asks_arms_and_the_next_append_reports_it, + an_append_with_nothing_armed_reports_zero, + an_arm_that_finds_a_newer_record_stays_unarmed, + a_disarm_of_an_older_generation_keeps_the_newer_arm, + a_promotion_reports_each_topics_arm, + newest_or_arm_arms_only_an_empty_topic, + a_closed_topic_never_arms, + ); + }; +} + mod memory { store_cases!(super::memory()); + arm_cases!(super::memory()); conformance!(super::memory(); deleting_an_owner_drops_every_chain_of_it); } mod notifying_memory { store_cases!(super::notifying(super::memory())); + arm_cases!(super::notifying(super::memory())); conformance!(super::notifying(super::memory()); deleting_an_owner_drops_every_chain_of_it, );