[Brook K12] Added Producer scopes for Activity attempts and clients.

crates/streams/protos/temporal/sdk/streams/v1/runtime.proto

@@ -6,7 +6,9 @@ package temporal.sdk.streams.v1;
 
 import "google/protobuf/duration.proto";
 import "google/protobuf/timestamp.proto";
+import "temporal/api/enums/v1/stream.proto";
 import "temporal/api/stream/v1/message.proto";
+import "temporal/api/streamservice/v1/request_response.proto";
 
 // The settings of the memory Stream runtime, for tests. Each field has the number of the same
 // field in the config of the other Stores, such as Redis, so a backend can send that config here
@@ -160,3 +162,89 @@ message ScopedStreamIdRequest {
 message ScopedStreamIdResponse {
     string id = 1;
 }
+
+// Method `OpenActivityScope`, asynchronous. Opens the scope of one Activity attempt, which writes
+// as the Producer of `DeriveProducer`, with one writer for each Stream. `CloseScope` ends it.
+message OpenActivityScopeRequest {
+    string namespace = 1;
+    ActivityContext context = 2;
+    // The Streams in the attempt's arguments. From attempt 2, the open withdraws the earlier
+    // attempts' Records on each. A missing, closed or expired Stream doesn't fail it.
+    repeated temporal.api.stream.v1.StreamRef streams = 3;
+}
+
+// Method `OpenClientScope`, synchronous. Opens the scope of the client Producer `client:<name>` at
+// `epoch`, or with no name `client:<identity>:<key>` at epoch 1, with a new key for each scope.
+message OpenClientScopeRequest {
+    string namespace = 1;
+    string identity = 2;
+    string name = 3;
+    // 1 or more, with a name only.
+    int32 epoch = 4;
+}
+
+message OpenScopeResponse {
+    uint64 scope_id = 1;
+    temporal.api.stream.v1.Producer producer = 2;
+}
+
+// Method `Publish`, asynchronous. Appends Records as the scope's Producer and answers when the
+// Store has them. The first publish to a Stream opens the scope's writer, as `OpenWriter` does.
+message PublishRequest {
+    uint64 scope_id = 1;
+    temporal.api.stream.v1.StreamRef stream = 2;
+    // The bodies are in their stored form, after lang's converter and codec.
+    repeated temporal.api.stream.v1.AppendRecord records = 3;
+}
+
+message PublishResponse {
+    // The offsets of the first and the last Record. Zero when the request has no Record.
+    int64 first_offset = 1;
+    int64 last_offset = 2;
+}
+
+// Method `CloseStreamAs`, asynchronous, with `CloseStreamResponse`. Closes the scope's writer for
+// the Stream, then the Stream, as the scope's Producer. An Activity scope fails with NOT_OWNER.
+message CloseStreamAsRequest {
+    uint64 scope_id = 1;
+    // The runtime sets `namespace` and `producer` from the scope. Payloads are in their stored
+    // form.
+    temporal.api.streamservice.v1.CloseStreamRequest close = 2;
+}
+
+// Method `CloseScope`, asynchronous. Closes every writer of the scope and frees it. A failed
+// writer doesn't fail the call. The response counts it instead.
+message CloseScopeRequest {
+    uint64 scope_id = 1;
+}
+
+message CloseScopeResponse {
+    int32 failed_writers = 1;
+    // The first failure, if any.
+    StreamFailure failure = 2;
+}
+
+// A failure that a response reports, as lang raises it from a failed call.
+message StreamFailure {
+    temporal.api.enums.v1.StreamErrorReason reason = 1;
+    string message = 2;
+}
+
+// Method `FreeScope`, synchronous. Frees the scope and releases its writers with no wait, such as
+// from a finalizer. A write in flight still completes.
+message FreeScopeRequest {
+    uint64 scope_id = 1;
+}
+
+// The response of `FreeScope` and the other `Free` methods, which do nothing once a shutdown ended.
+message FreeResponse {
+}
+
+// Method `Shutdown`, asynchronous. Closes every scope and writer, for the backend's shutdown. Then
+// only `Shutdown`, the static and the `Free` methods run. A dropped one runs again.
+message ShutdownRequest {
+}
+
+message ShutdownResponse {
+    int32 failed_writers = 1;
+}

crates/streams/src/error.rs

@@ -1,6 +1,6 @@
 //! The error of every failed call, and its form for lang.
 
-use crate::proto::{StreamErrorDetails, StreamErrorReason};
+use crate::proto::{StreamErrorDetails, StreamErrorReason, StreamFailure};
 use prost::Message;
 use tonic::Code;
 
@@ -38,9 +38,10 @@ pub enum Reason {
     Invalid,
     /// The Store failed or did not answer.
     Storage,
-    /// No writer has the id.
+    /// No writer has the id. A scope, reader or buffer id that the runtime doesn't have gives it
+    /// too.
     UnknownWriter,
-    /// The Store lacks a capability that the call needs.
+    /// The Store lacks a capability that the call needs, or the runtime shut down.
     Unsupported,
 }
 
@@ -146,6 +147,14 @@ impl StreamError {
         }
         .encode_to_vec()
     }
+
+    /// The error as a response reports it, when a failed write doesn't fail the call.
+    pub fn failure(&self) -> StreamFailure {
+        StreamFailure {
+            reason: self.reason.proto() as i32,
+            message: self.message.clone(),
+        }
+    }
 }
 
 #[cfg(test)]

crates/streams/src/lib.rs

@@ -16,6 +16,7 @@ pub mod offset;
 pub mod protos;
 mod record;
 mod runtime;
+mod scope;
 mod sdk;
 mod service;
 pub mod store;

crates/streams/src/runtime.rs

@@ -1,21 +1,29 @@
 //! The Stream runtime, a plain object that lang's bridge builds per client and calls by name.
 
 use crate::{
-    StreamError, StreamResult,
+    Reason, StreamError, StreamResult,
     memory::MemoryStore,
     metrics::Metrics,
     proto::{
-        AppendRecordsRequest, CloseStreamRequest, CloseWriterRequest, CloseWriterResponse,
-        CloseWritersRequest, CloseWritersResponse, DescribeStreamRequest, FreeWriterRequest,
-        FreeWriterResponse, MemoryStreamConfig, OpenWriterRequest, OpenWriterResponse,
-        ReadRecordsRequest, WriteRecordsRequest, WriteRecordsResponse,
+        AppendRecordsRequest, CloseScopeRequest, CloseStreamRequest, CloseWriterRequest,
+        CloseWriterResponse, CloseWritersRequest, CloseWritersResponse, DescribeStreamRequest,
+        FreeResponse, FreeScopeRequest, FreeWriterRequest, FreeWriterResponse, MemoryStreamConfig,
+        OpenWriterRequest, OpenWriterResponse, ReadRecordsRequest, ShutdownRequest,
+        WriteRecordsRequest, WriteRecordsResponse,
     },
+    scope::Scopes,
     service::StreamService,
     store::{DEFAULT_RETENTION, DEFAULT_SAFETY_TTL, Store},
-    writer::Writers,
+    writer::{Holder, Writers},
 };
 use prost::Message;
-use std::{sync::Arc, time::Duration};
+use std::{
+    sync::{
+        Arc,
+        atomic::{AtomicBool, Ordering},
+    },
+    time::Duration,
+};
 use temporalio_common::telemetry::metrics::TemporalMeter;
 use tokio::sync::OnceCell;
 
@@ -47,10 +55,30 @@ pub(crate) struct Shared {
     metrics: Metrics,
     factory: Arc<dyn StoreFactory>,
     service: OnceCell<Arc<StreamService>>,
-    writers: Writers,
+    pub(crate) writers: Writers,
+    pub(crate) scopes: Scopes,
+    shut_down: AtomicBool,
+    /// Set when a `Shutdown` ends every scope and writer. Then the `Free` methods do nothing.
+    finished: AtomicBool,
 }
 
 impl Shared {
+    pub(crate) fn is_shut_down(&self) -> bool {
+        self.shut_down.load(Ordering::SeqCst)
+    }
+
+    fn is_finished(&self) -> bool {
+        self.finished.load(Ordering::SeqCst)
+    }
+
+    /// A free of something that a running shutdown already took has nothing left to do.
+    fn taken_by_shutdown(&self, freed: StreamResult<()>) -> StreamResult<()> {
+        match freed {
+            Err(error) if error.reason == Reason::UnknownWriter && self.is_shut_down() => Ok(()),
+            freed => freed,
+        }
+    }
+
     /// The service, which connects its Store on its first use.
     pub(crate) async fn service(&self) -> StreamResult<&Arc<StreamService>> {
         self.service
@@ -102,6 +130,9 @@ impl StreamRuntime {
             factory,
             service: OnceCell::new(),
             writers: Writers::default(),
+            scopes: Scopes::default(),
+            shut_down: AtomicBool::new(false),
+            finished: AtomicBool::new(false),
         });
         Arc::new(Self { shared })
     }
@@ -129,14 +160,25 @@ impl StreamRuntime {
     /// Releases one handle of a writer with no wait, as the `FreeWriter` method does. A bridge
     /// calls it from a finalizer, which has no async runtime to await on.
     pub fn free_writer(&self, writer_id: u64) -> StreamResult<()> {
-        self.shared.writers.free(writer_id)
+        if self.shared.is_finished() {
+            return Ok(());
+        }
+        let freed = self.shared.writers.free(writer_id, Holder::Lang);
+        self.shared.taken_by_shutdown(freed)
+    }
+
+    #[cfg(test)]
+    pub(crate) fn shared(&self) -> &Arc<Shared> {
+        &self.shared
     }
 
     /// Runs any method on encoded input and returns the encoded output. The bridge raises an error
     /// with [StreamError::code] and [StreamError::details].
     pub async fn call(&self, method: &str, input: &[u8]) -> StreamResult<Vec<u8>> {
+        self.check_open(method)?;
         // The methods of `call_sync` don't connect, so they work when the Store is down too.
         let service = || self.shared.service();
+        let shared = &self.shared;
         let output = match method {
             "CreateStream" => service()
                 .await?
@@ -165,7 +207,15 @@ impl StreamRuntime {
                 .encode_to_vec(),
             "OpenWriter" => {
                 let request = decode::<OpenWriterRequest>(input)?;
-                let writer_id = self.shared.writers.open(service().await?, request).await?;
+                let writers = &self.shared.writers;
+                let writer_id = writers
+                    .open(service().await?, request, Holder::Lang)
+                    .await?;
+                // A shutdown during the open has already closed every writer.
+                if shared.is_shut_down() {
+                    let _ = shared.writers.free(writer_id, Holder::Lang);
+                    return Err(shut_down());
+                }
                 OpenWriterResponse { writer_id }.encode_to_vec()
             }
             "WriteRecords" => {
@@ -183,6 +233,7 @@ impl StreamRuntime {
             }
             "CloseWriters" => {
                 decode::<CloseWritersRequest>(input)?;
+                // The writers of scopes stay, since another Worker of the client may use them.
                 let failed = self.shared.writers.close_all().await;
                 CloseWritersResponse {
                     failed: failed as i32,
@@ -191,9 +242,31 @@ impl StreamRuntime {
             }
             "CloseWriter" => {
                 let request = decode::<CloseWriterRequest>(input)?;
-                let last_offset = self.shared.writers.close(request.writer_id).await?;
+                let last_offset = (self.shared.writers)
+                    .close(request.writer_id, Holder::Lang)
+                    .await?;
                 CloseWriterResponse { last_offset }.encode_to_vec()
             }
+            "OpenActivityScope" => shared
+                .open_activity_scope(decode(input)?)
+                .await?
+                .encode_to_vec(),
+            "Publish" => shared.publish(decode(input)?).await?.encode_to_vec(),
+            "CloseStreamAs" => shared
+                .close_stream_as(decode(input)?)
+                .await?
+                .encode_to_vec(),
+            "CloseScope" => {
+                let request = decode::<CloseScopeRequest>(input)?;
+                shared.close_scope(request.scope_id).await?.encode_to_vec()
+            }
+            "Shutdown" => {
+                decode::<ShutdownRequest>(input)?;
+                shared.shut_down.store(true, Ordering::SeqCst);
+                let closed = shared.close_everything().await;
+                shared.finished.store(true, Ordering::SeqCst);
+                closed.encode_to_vec()
+            }
             _ => return self.call_sync(method, input),
         };
         Ok(output)
@@ -202,15 +275,52 @@ impl StreamRuntime {
     /// Runs a method that never waits, with no async runtime entered, so a finalizer can call it.
     /// `runtime.proto` marks these methods as synchronous or static.
     pub fn call_sync(&self, method: &str, input: &[u8]) -> StreamResult<Vec<u8>> {
+        // A finalizer can run after the shutdown, which already freed everything. Before that, a
+        // free still works, so that lang can end a buffer whose pump is gone.
+        if method.starts_with("Free") && self.shared.is_finished() {
+            return Ok(vec![]);
+        }
+        self.check_open(method)?;
+        let output = self.run_sync(method, input);
+        match output {
+            Err(error) if method.starts_with("Free") => {
+                self.shared.taken_by_shutdown(Err(error)).map(|()| vec![])
+            }
+            output => output,
+        }
+    }
+
+    fn run_sync(&self, method: &str, input: &[u8]) -> StreamResult<Vec<u8>> {
+        let shared = &self.shared;
         match method {
             "FreeWriter" => {
                 let request = decode::<FreeWriterRequest>(input)?;
                 self.free_writer(request.writer_id)?;
                 Ok(FreeWriterResponse {}.encode_to_vec())
             }
+            "OpenClientScope" => Ok(shared.open_client_scope(decode(input)?)?.encode_to_vec()),
+            "FreeScope" => {
+                let request = decode::<FreeScopeRequest>(input)?;
+                shared.free_scope(request.scope_id)?;
+                Ok(FreeResponse {}.encode_to_vec())
+            }
             _ => crate::call_static(method, input),
         }
     }
+
+    /// Refuses every method after `Shutdown`, except `Shutdown`, the `Free` methods and the static
+    /// ones.
+    fn check_open(&self, method: &str) -> StreamResult<()> {
+        // A dropped `Shutdown` leaves work that the next one does.
+        if !self.shared.is_shut_down()
+            || method.starts_with("Free")
+            || crate::sdk::is_static(method)
+            || method == "Shutdown"
+        {
+            return Ok(());
+        }
+        Err(shut_down())
+    }
 }
 
 /// The name of the memory runtime, for a bridge that selects a runtime by name.
@@ -265,6 +375,10 @@ fn duration(value: Option<prost_types::Duration>) -> Option<Duration> {
     (!duration.is_zero()).then_some(duration)
 }
 
+pub(crate) fn shut_down() -> StreamError {
+    StreamError::new(Reason::Unsupported, "The Stream runtime shut down")
+}
+
 pub(crate) fn decode<M: Message + Default>(input: &[u8]) -> StreamResult<M> {
     M::decode(input).map_err(|error| StreamError::invalid(format!("The request is bad: {error}")))
 }