[Jetty] Attached a Workflow's stream reader to the stream's notifier.

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;

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/<run id>/<subscription id>`, 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

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(&notification.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!(

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<K> {
+pub(crate) struct Remembered<K> {
     entries: Mutex<HashMap<K, Checked>>,
     limit: usize,
 }
 
 impl<K: std::hash::Hash + Eq + Clone> Remembered<K> {
-    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<Checked> {
+    pub(crate) fn get(&self, key: &K) -> Option<Checked> {
         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<K: std::hash::Hash + Eq + Clone> Remembered<K> {
             },
         );
     }
-
-    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<ChainId> {
+        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<ChainId> {
         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 {

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::{

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.