From e2e81bdde8f731658b09a364363a50305993eab1 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 13 Oct 2026 18:40:45 -0700 Subject: [PATCH] Let a stream read arm the stream when it waits and finds nothing. The read takes arm and the answer says armed. Only the pass that waits out the read's own wait arms, and a done answer never does. An armed answer also carries wake_after, the counter of the read's position, so the reader waits for a higher counter and a notification that came first still counts. --- .../temporal/sdk/core/streams/streams.proto | 11 +++ crates/streams/src/handle.rs | 85 ++++++++++++++++++- crates/streams/src/reader.rs | 15 +++- 3 files changed, 107 insertions(+), 4 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 af1b96989b..4c74d12dab 100644 --- a/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto +++ b/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto @@ -159,6 +159,10 @@ message ReadRequest { uint32 max_records = 4; // The `state` of the previous answer of the same read. Empty on its first call. bytes state = 5; + // For a reader that sleeps on the stream's notifier. When the read waits its whole `wait` and + // finds nothing, the store arms the stream in the same step, so the next append notifies the + // notifier. A reader that doesn't wait on the notifier leaves it unset. + bool arm = 6; } // A later attempt of a producer started writing. Readers synthesize it, and no store holds it. @@ -193,6 +197,13 @@ message ReadResponse { // The owner's chain ended and the read delivered every record the stream retains. No later // call returns a record. bool done = 4; + // The read asked to arm, found nothing and armed the stream, so the next append notifies the + // stream's notifier. Never set with records or with `done`. + bool armed = 5; + // Set with `armed`: the notification counter of the read's position, zero for an empty + // stream. A notification with a higher counter names a record after the position, so the + // reader waits for a counter above this one. One that landed before the wait started counts. + int64 wake_after = 6; } message LatestRequest { diff --git a/crates/streams/src/handle.rs b/crates/streams/src/handle.rs index c26261b458..8c2df3f34b 100644 --- a/crates/streams/src/handle.rs +++ b/crates/streams/src/handle.rs @@ -246,7 +246,8 @@ impl Streams { /// the interval over. So many readers parked on a quiet stream don't load the server. Once /// the owner ended, or the store marked the stream closed, the read delivers what is left /// and answers `done`. A read pinned to a run ends when that run closes, even if the chain - /// continued as new. + /// continued as new. A read with `arm` that waits its whole `wait` and finds nothing arms the + /// stream and answers `armed`. pub async fn read(&self, request: ReadRequest) -> StreamResult { let stream = address(request.stream.as_ref())?; let mut state = ReadState::decode_from(&request.state)?; @@ -279,6 +280,9 @@ impl Streams { remaining.min(until_check) }; page_request.wait = wait.try_into().ok(); + // Only the pass that waits out the read's own wait arms, so a shorter pass before an + // owner check doesn't wake the caller for a gap shorter than the reader chose. + page_request.arm = request.arm && !state.owner_ended && remaining <= until_check; let (mut page, closed) = read_with(&*self.store, &target, &page_request, &mut state).await?; page_request.after = page.cursor.clone(); @@ -294,6 +298,8 @@ impl Streams { } if state.owner_ended { page.done = true; + page.armed = false; + page.wake_after = 0; break page; } if unix_ms() >= state.next_check_ms { @@ -1173,6 +1179,83 @@ mod tests { panic!("the read never ended"); } + #[tokio::test] + async fn only_the_pass_that_waits_out_the_read_arms() { + // Owner checks every 20 ms cut the wait into passes. One arm for the whole wait keeps a + // gap shorter than the reader's wait from costing a notification. + let setup = setup(quick()); + let page = setup + .streams + .read(ReadRequest { + arm: true, + ..read_after("", &[], Duration::from_millis(300)) + }) + .await + .unwrap(); + assert!(page.armed && page.records.is_empty() && !page.done); + assert_eq!(page.wake_after, 0, "an empty stream wakes on any counter"); + assert_eq!( + setup + .store + .newest_or_arm(&chain("run-1"), "out") + .await + .unwrap(), + crate::NewestOrArm::Armed(2), + "the read armed once" + ); + } + + #[tokio::test] + async fn an_armed_read_waits_for_a_counter_above_its_position() { + let setup = setup(StreamsOptions::default()); + setup.streams.append(append("p", 1, "1")).await.unwrap(); + let landed = setup.streams.append(append("p", 2, "2")).await.unwrap(); + let page = setup + .streams + .read(ReadRequest { + arm: true, + ..read_after(&landed.last_cursor, &[], Duration::from_millis(20)) + }) + .await + .unwrap(); + assert!(page.armed); + let at = |cursor: &str| { + crate::cursor_position( + cursor, + setup.store.name(), + &setup.streams.hash(&stream("out")), + ) + .unwrap() + .unwrap() + .to_string() + }; + assert_eq!( + Some(page.wake_after), + crate::progress_counter(&at(&landed.last_cursor)) + ); + let next = setup.streams.append(append("p", 3, "3")).await.unwrap(); + assert!( + crate::progress_counter(&at(&next.last_cursor)).unwrap() > page.wake_after, + "the next record's notification wakes the reader" + ); + } + + #[tokio::test] + async fn a_read_that_finds_its_owner_ended_answers_done_unarmed() { + let setup = setup(quick()); + setup.owners.end("wf", WorkflowExecutionStatus::Completed); + let page = setup + .streams + .read(ReadRequest { + arm: true, + ..read_after("", &[], Duration::from_secs(5)) + }) + .await + .unwrap(); + assert!(page.done); + assert!(!page.armed); + } + #[tokio::test] async fn a_read_ends_when_its_owner_chain_ends_and_marks_it_closed() { let setup = setup(quick()); diff --git a/crates/streams/src/reader.rs b/crates/streams/src/reader.rs index 7c3bc20e26..dd7c2ea99c 100644 --- a/crates/streams/src/reader.rs +++ b/crates/streams/src/reader.rs @@ -7,6 +7,7 @@ use crate::{ BEGINNING, END, StreamError, StreamResult, StreamStore, cursor_position, mint_cursor, + progress_counter, proto::{ ChainId, ReadRecord, ReadRequest, ReadResponse, StoreLatestRequest, StoreReadRequest, StreamRecord, StreamRecordKind, Supersession, read_record, @@ -218,16 +219,16 @@ pub(crate) async fn read_with( .read(StoreReadRequest { chain: Some(target.chain.clone()), topic: target.topic.clone(), - after_position: position, + after_position: position.clone(), wait: request.wait, max_records: match request.max_records { 0 => DEFAULT_MAX_RECORDS, max => max, }, - arm: false, + arm: request.arm, }) .await?; - let closed = page.closed; + let (closed, armed) = (page.closed, page.armed); let mut records = Vec::with_capacity(page.records.len()); let mut previous = start; for stored in page.records { @@ -272,6 +273,14 @@ pub(crate) async fn read_with( cursor: previous, state: Vec::new(), done: false, + armed, + // An armed read found nothing after `position`, so a notification counter above that + // position's counter can only come from a newer record. + wake_after: if armed { + progress_counter(&position).unwrap_or(0) + } else { + 0 + }, }, closed, ))