[Jetty] Made a reader read again after it subscribes.

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
@@ -362,6 +365,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 {

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

crates/streams/Cargo.toml

@@ -32,7 +32,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]

crates/streams/src/handle.rs

@@ -893,6 +893,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.

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<Option<NotifyError>>,
         polled: Mutex<Vec<(ChainId, String, i64)>>,
         poll_answer: Mutex<Option<Result<Option<NotifierPoll>, NotifyError>>>,
+        /// How long the poll takes before it answers.
+        poll_delay: Mutex<Duration>,
     }
 
     #[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();