[Brook] Added buffered writers for Activities and clients.

CHANGELOG.md

@@ -19,3 +19,7 @@
   closes in the background after `after_complete` accepts. `DrainCloses` waits for them. A flush
   splits at the batch limit, drops oversize Records and Records for a Stream that someone else
   closed, and runs `WAIT` with `durable_writes`.
+- Added buffered writers for Activities and clients: `OpenWriter`, `Publish`, `FlushWriter` and
+  `CloseWriter`. A writer flushes every `batch_interval`, at once on `force_flush`, and refuses a
+  publish past its buffer limit. A writer of an Activity attempt above 1 first sends an empty
+  batch, so the new epoch withdraws the earlier attempt's Records.

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

@@ -111,6 +111,16 @@ message CloseWriterResponse {
     int64 last_offset = 1;
 }
 
+// Method `CloseWriters`, asynchronous. Flushes and drops every open writer, for example when the
+// plugin's client or Worker shuts down, so no flusher task stays behind.
+message CloseWritersRequest {
+}
+
+message CloseWritersResponse {
+    // The writers whose last flush failed.
+    int32 failed = 1;
+}
+
 // Method `DrainCloses`, asynchronous. Waits for the closes that `after_complete` started, for
 // example at Worker shutdown.
 message DrainClosesRequest {

crates/streams/src/extension.rs

@@ -5,12 +5,15 @@ use crate::{
     metrics::Metrics,
     proto::{
         AppendRecordsRequest, BufferWorkflowCallsRequest, BufferWorkflowCallsResponse,
-        CloseStreamRequest, DescribeStreamRequest, DrainClosesRequest, DrainClosesResponse,
-        ReadRecordsRequest, RedisStreamConfig,
+        CloseStreamRequest, CloseWriterRequest, CloseWriterResponse, CloseWritersRequest,
+        CloseWritersResponse, DescribeStreamRequest, DrainClosesRequest, DrainClosesResponse,
+        FlushWriterRequest, FlushWriterResponse, OpenWriterRequest, OpenWriterResponse,
+        PublishRequest, PublishResponse, ReadRecordsRequest, RedisStreamConfig,
     },
     redis::{RedisStore, RedisStoreOptions},
     service::StreamService,
     workflow::{FlushOptions, WorkflowHooks, WorkflowRuns},
+    writer::Writers,
 };
 use prost::Message;
 use std::{sync::Arc, time::Duration};
@@ -77,6 +80,7 @@ pub(crate) struct Shared {
     metrics: Metrics,
     service: OnceCell<Arc<StreamService>>,
     pub(crate) runs: WorkflowRuns,
+    writers: Writers,
 }
 
 impl Shared {
@@ -129,6 +133,7 @@ impl RedisStreamExtension {
             metrics: Metrics::new(&meter),
             service: OnceCell::new(),
             runs: WorkflowRuns::default(),
+            writers: Writers::default(),
         });
         let hooks = Arc::new(WorkflowHooks {
             shared: shared.clone(),
@@ -160,6 +165,15 @@ impl RedisStreamExtension {
                     .buffer(decode::<BufferWorkflowCallsRequest>(input)?)?;
                 BufferWorkflowCallsResponse {}.encode_to_vec()
             }
+            "Publish" => {
+                let request = decode::<PublishRequest>(input)?;
+                self.shared.writers.publish(
+                    request.writer_id,
+                    request.records,
+                    request.force_flush,
+                )?;
+                PublishResponse {}.encode_to_vec()
+            }
             _ => return Err(ExtensionError::unknown_method(method)),
         };
         Ok(output)
@@ -199,6 +213,29 @@ impl RedisStreamExtension {
                 .describe_stream(decode::<DescribeStreamRequest>(input)?)
                 .await?
                 .encode_to_vec(),
+            "OpenWriter" => {
+                let request = decode::<OpenWriterRequest>(input)?;
+                let writer_id = self.shared.writers.open(service, request).await?;
+                OpenWriterResponse { writer_id }.encode_to_vec()
+            }
+            "FlushWriter" => {
+                let request = decode::<FlushWriterRequest>(input)?;
+                let last_offset = self.shared.writers.flush(request.writer_id).await?;
+                FlushWriterResponse { last_offset }.encode_to_vec()
+            }
+            "CloseWriters" => {
+                decode::<CloseWritersRequest>(input)?;
+                let failed = self.shared.writers.close_all().await;
+                CloseWritersResponse {
+                    failed: failed as i32,
+                }
+                .encode_to_vec()
+            }
+            "CloseWriter" => {
+                let request = decode::<CloseWriterRequest>(input)?;
+                let last_offset = self.shared.writers.close(request.writer_id).await?;
+                CloseWriterResponse { last_offset }.encode_to_vec()
+            }
             _ => return Err(ExtensionError::unknown_method(method)),
         };
         Ok(output)

crates/streams/src/lib.rs

@@ -21,6 +21,7 @@ mod server {
 }
 mod service;
 mod workflow;
+mod writer;
 
 pub use error::{Reason, StreamError, StreamResult};
 pub use extension::{EXTENSION_NAME, ExtensionOptions, RedisStreamExtension};

crates/streams/src/metrics.rs

@@ -13,6 +13,7 @@ pub(crate) struct Metrics {
     flush_failures: Counter,
     fenced: Counter,
     superseded: Counter,
+    buffer_full: Counter,
     close_retries: Counter,
     records_dropped: Counter,
     blocking_reads_waiting: UpDownCounter,
@@ -54,6 +55,11 @@ impl Metrics {
                 "Superseded flags that the Store set",
                 "",
             )),
+            buffer_full: meter.counter(parameters(
+                "stream_buffer_full",
+                "Publishes refused because a writer's buffer was full",
+                "",
+            )),
             close_retries: meter.counter(parameters(
                 "stream_close_retries",
                 "Close retries after a Workflow Task completion",
@@ -88,6 +94,10 @@ impl Metrics {
         self.superseded.add(1, &self.attributes);
     }
 
+    pub(crate) fn buffer_full(&self) {
+        self.buffer_full.add(1, &self.attributes);
+    }
+
     pub(crate) fn close_retried(&self) {
         self.close_retries.add(1, &self.attributes);
     }