From 5aa933010b9bcc587cf4f49980be7a826b7920e1 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Wed, 14 Oct 2026 17:10:56 -0700 Subject: [PATCH] Read a stream once more after a subscribe or a poll starts its notifier. A notification for a stream with no notifier is dropped, so a record from before the subscribe woke no one. The subscribe answers the newest counter, and a client wait rereads once its poll is out. --- .../temporal/sdk/core/streams/streams.proto | 11 +- crates/sdk-core/CHANGELOG.md | 5 + crates/streams/Cargo.toml | 2 +- crates/streams/src/handle.rs | 22 ++++ crates/streams/src/subscribe.rs | 122 +++++++++++++++++- 5 files changed, 153 insertions(+), 9 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 2ca8ab6a90..df86d7b891 100644 --- a/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto +++ b/crates/protos/protos/local/temporal/sdk/core/streams/streams.proto @@ -40,8 +40,11 @@ service StreamService { rpc Unsubscribe(UnsubscribeRequest) returns (UnsubscribeResponse); // Waits on the stream's notifier, for a client reader whose read came back empty. Returns when // the notifier's counter is above `after_counter`, when the stream closed, or at the server's - // wait limit (`streamnotifier.pollMaxWait`). The reader then reads again. After the wait limit, - // that read finds records whose notification was lost, or comes back empty again. + // wait limit (`streamnotifier.pollMaxWait`). The poll starts the notifier, and Core reads the + // stream once more just after it sends the poll, since a record that landed before the notifier + // existed woke no one. A newer record there answers at once. The reader then reads again. + // After the wait limit, that read finds records whose notification was lost, or comes back + // empty again. rpc WaitForNotification(WaitForNotificationRequest) returns (WaitForNotificationResponse); // Runs a stream driver over the store's notification log until the caller cancels the call. // It answers only when the server takes no stream notifications, or when the store keeps no @@ -358,6 +361,10 @@ message SubscribeResponse { string chain_first_run_id = 1; // Whether this call sent the attach. False when this process attached the reader before. bool attached = 2; + // The notification counter of the stream's newest record, read after the attach landed, zero + // for an empty stream. A record that landed before the notifier existed woke no one, so a + // reader whose last read stopped below this counter reads once more before it sleeps. + int64 newest_counter = 3; } message UnsubscribeRequest { diff --git a/crates/sdk-core/CHANGELOG.md b/crates/sdk-core/CHANGELOG.md index 8462f668ff..beee0237cd 100644 --- a/crates/sdk-core/CHANGELOG.md +++ b/crates/sdk-core/CHANGELOG.md @@ -64,6 +64,11 @@ relevant information. `NotifyStreams` call per namespace, and acks the entries only once the call landed. It takes entries another driver left unacked for `claim_idle` (60 seconds). Each driver needs its own `consumer`. `StreamDriver` runs the same in Rust. +* Experimental: a reader reads first and subscribes only when a read comes back empty, since a + notification for a stream with no notifier is dropped. `SubscribeResponse.newest_counter` names + the newest record after the attach landed, so a reader whose read stopped below it reads once + more before it sleeps. `WaitForNotification` reads the stream once more just after it sends the + poll, and answers at once for a record that landed before the notifier existed. * Experimental: the new `FlushNotifications` stream call waits up to 10 seconds for the notifications still out. A process calls it before it stops. * Experimental: a server without the stream notifier answers the first notifier call with diff --git a/crates/streams/Cargo.toml b/crates/streams/Cargo.toml index 010105d424..45281988f5 100644 --- a/crates/streams/Cargo.toml +++ b/crates/streams/Cargo.toml @@ -31,7 +31,7 @@ redis = { version = "1.7", optional = true, default-features = false, features = ] } sha2 = "0.10" thiserror = { workspace = true } -tokio = { version = "1.47", default-features = false, features = ["rt", "sync", "time"] } +tokio = { version = "1.47", default-features = false, features = ["macros", "rt", "sync", "time"] } tracing = "0.1" [dependencies.temporalio-common] diff --git a/crates/streams/src/handle.rs b/crates/streams/src/handle.rs index 053c8deb26..0a74385d63 100644 --- a/crates/streams/src/handle.rs +++ b/crates/streams/src/handle.rs @@ -844,6 +844,28 @@ impl Streams { fn cursor(&self, stream: &StreamAddress, position: &str) -> String { mint_cursor(self.store.name(), &self.hash(stream), position) } + + /// The notification counter and the cursor of the newest record of `stream` on `chain`, or + /// zero and an empty cursor when it holds none. + pub(crate) async fn newest( + &self, + stream: &StreamAddress, + chain: &ChainId, + ) -> StreamResult<(i64, String)> { + let position = self + .store + .latest(StoreLatestRequest { + chain: Some(chain.clone()), + topic: stream.topic.clone(), + }) + .await? + .position; + if position.is_empty() { + return Ok((0, String::new())); + } + let counter = crate::progress_counter(&position).unwrap_or(0); + Ok((counter, self.cursor(stream, &position))) + } } /// The address of a stream this release keeps, checked. diff --git a/crates/streams/src/subscribe.rs b/crates/streams/src/subscribe.rs index d1f135684c..8326d9b001 100644 --- a/crates/streams/src/subscribe.rs +++ b/crates/streams/src/subscribe.rs @@ -11,11 +11,15 @@ use crate::{ WorkflowAttachment, handle::{Remembered, address}, proto::{ - StreamAddress, SubscribeRequest, SubscribeResponse, UnsubscribeRequest, + ChainId, StreamAddress, SubscribeRequest, SubscribeResponse, UnsubscribeRequest, UnsubscribeResponse, WaitForNotificationRequest, WaitForNotificationResponse, }, }; -use std::sync::Arc; +use std::{sync::Arc, time::Duration}; + +/// How long after it sends a poll Core reads the stream once more. The poll starts the stream's +/// notifier, and a record that lands later notifies it. +const REREAD_AFTER: Duration = Duration::from_millis(200); /// The longest subscription id the server takes. const MAX_SUBSCRIPTION_ID_BYTES: usize = 256; @@ -123,7 +127,12 @@ impl Subscriptions { ); let key = SubscriptionKey::new(stream, &request_id); if let Some(attached) = self.attached.get(&key) { + let chain = ChainId { + first_run_id: attached.first_run_id.clone(), + ..streams.chain(stream).await? + }; return Ok(SubscribeResponse { + newest_counter: streams.newest(stream, &chain).await?.0, chain_first_run_id: attached.first_run_id, attached: false, }); @@ -144,9 +153,12 @@ impl Subscriptions { .await .map_err(subscribe_error)?; self.attached.put(key, &chain.first_run_id, OWNER_RECHECK); + // Read after the attach landed, so a record after this one notifies the reader. + let newest_counter = streams.newest(stream, &chain).await?.0; Ok(SubscribeResponse { chain_first_run_id: chain.first_run_id, attached: true, + newest_counter, }) } @@ -193,10 +205,28 @@ impl Subscriptions { let stream = address(request.stream.as_ref())?; let client = self.client()?; let chain = streams.chain(stream).await?; - let polled = client - .poll(chain, stream.topic.clone(), request.after_counter) - .await - .map_err(subscribe_error)?; + let poll = client.poll(chain.clone(), stream.topic.clone(), request.after_counter); + tokio::pin!(poll); + // The poll starts the notifier, so a record from before it woke no one. One more read + // once the poll is out finds it. + let reread = async { + tokio::time::sleep(REREAD_AFTER).await; + streams.newest(stream, &chain).await + }; + let polled = tokio::select! { + polled = &mut poll => polled, + newest = reread => match newest { + Ok((counter, position)) if counter > request.after_counter => { + return Ok(WaitForNotificationResponse { + counter, + position, + ..Default::default() + }); + } + _ => poll.await, + }, + } + .map_err(subscribe_error)?; let Some(polled) = polled else { return Ok(WaitForNotificationResponse { timed_out: true, @@ -254,6 +284,8 @@ mod tests { answer: Mutex>, polled: Mutex>, poll_answer: Mutex, NotifyError>>>, + /// How long the poll takes before it answers. + poll_delay: Mutex, } #[async_trait::async_trait] @@ -293,6 +325,8 @@ mod tests { .lock() .unwrap() .push((chain, topic, after_counter)); + let delay = *self.poll_delay.lock().unwrap(); + tokio::time::sleep(delay).await; self.poll_answer .lock() .unwrap() @@ -596,6 +630,82 @@ mod tests { .map(|answer| WaitForNotificationResponse::decode(answer.as_slice()).unwrap()) } + /// Appends `count` records to the chain the owner names now. + async fn append_records(setup: &Setup, count: i64) { + use crate::StreamStore; + let first_run_id = setup.owner.0.lock().unwrap().clone(); + setup + .store + .append(crate::proto::StoreAppendRequest { + chain: Some(ChainId { + namespace: "ns".to_string(), + workflow_id: "owner".to_string(), + first_run_id, + }), + topic: "tokens".to_string(), + producer_id: "p".to_string(), + attempt: 1, + sequence: 1, + digest: vec![1], + records: (0..count).map(|_| b"record".to_vec()).collect(), + }) + .await + .unwrap(); + } + + #[tokio::test] + async fn a_subscription_names_the_newest_counter_it_found_after_the_attach() { + let setup = setup(); + let empty = subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + assert_eq!(empty.newest_counter, 0, "an empty stream"); + append_records(&setup, 2).await; + // A record from before the attach woke no one, so the reader reads once more. + let again = subscribe(&setup.service, subscribe_request("1")) + .await + .unwrap(); + assert!(!again.attached); + assert_eq!(again.newest_counter, 2); + let fresh = subscribe(&setup.service, subscribe_request("2")) + .await + .unwrap(); + assert!(fresh.attached); + assert_eq!(fresh.newest_counter, 2); + } + + #[tokio::test] + async fn a_wait_answers_at_once_for_a_record_from_before_the_poll() { + let setup = setup(); + append_records(&setup, 1).await; + *setup.server.poll_delay.lock().unwrap() = Duration::from_secs(30); + let started = tokio::time::Instant::now(); + let answer = wait(&setup, 0, Ok(None)).await.unwrap(); + assert!(started.elapsed() < Duration::from_secs(5)); + assert!(!answer.timed_out && !answer.closed); + assert_eq!(answer.counter, 1); + assert!(answer.position.starts_with("memory:"), "{answer:?}"); + } + + #[tokio::test] + async fn a_wait_with_nothing_newer_waits_for_the_poll() { + let setup = setup(); + append_records(&setup, 1).await; + *setup.server.poll_delay.lock().unwrap() = Duration::from_millis(400); + let answer = wait( + &setup, + 1, + Ok(Some(NotifierPoll { + counter: 2, + position: "cursor-2".to_string(), + ..Default::default() + })), + ) + .await + .unwrap(); + assert_eq!((answer.counter, answer.position.as_str()), (2, "cursor-2")); + } + #[tokio::test] async fn a_wait_asks_the_notifier_of_the_streams_chain_for_a_counter_above_the_read() { let setup = setup();