From fc0c2e2597318a887003b78f21bba2d59aa496ea Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 20 Oct 2026 20:18:44 -0700 Subject: [PATCH 1/3] Added buffered writers for Activities and clients. A publish from an Activity or a client must return at once, so the crate buffers it and appends one batch per interval. The writer keeps the batch sequence of its epoch, which keeps a resend after a Store failure from going in twice. --- CHANGELOG.md | 4 + crates/streams/src/extension.rs | 33 ++- crates/streams/src/lib.rs | 1 + crates/streams/src/metrics.rs | 10 + crates/streams/src/writer.rs | 270 +++++++++++++++++++++++++ crates/streams/tests/writer.rs | 348 ++++++++++++++++++++++++++++++++ 6 files changed, 664 insertions(+), 2 deletions(-) create mode 100644 crates/streams/src/writer.rs create mode 100644 crates/streams/tests/writer.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a1db8f..9e0fef2 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/crates/streams/src/extension.rs b/crates/streams/src/extension.rs index 4bc07c1..ac57a19 100644 --- a/crates/streams/src/extension.rs +++ b/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, 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}; @@ -76,6 +79,7 @@ pub(crate) struct Shared { metrics: Metrics, service: OnceCell>, pub(crate) runs: WorkflowRuns, + writers: Writers, } impl Shared { @@ -128,6 +132,7 @@ impl RedisStreamExtension { metrics: Metrics::new(&meter), service: OnceCell::new(), runs: WorkflowRuns::default(), + writers: Writers::default(), }); let hooks = Arc::new(WorkflowHooks { shared: shared.clone(), @@ -159,6 +164,15 @@ impl RedisStreamExtension { .buffer(decode::(input)?)?; BufferWorkflowCallsResponse {}.encode_to_vec() } + "Publish" => { + let request = decode::(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) @@ -198,6 +212,21 @@ impl RedisStreamExtension { .describe_stream(decode::(input)?) .await? .encode_to_vec(), + "OpenWriter" => { + let request = decode::(input)?; + let writer_id = self.shared.writers.open(service, request).await?; + OpenWriterResponse { writer_id }.encode_to_vec() + } + "FlushWriter" => { + let request = decode::(input)?; + let last_offset = self.shared.writers.flush(request.writer_id).await?; + FlushWriterResponse { last_offset }.encode_to_vec() + } + "CloseWriter" => { + let request = decode::(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) diff --git a/crates/streams/src/lib.rs b/crates/streams/src/lib.rs index c129ead..e683a35 100644 --- a/crates/streams/src/lib.rs +++ b/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}; diff --git a/crates/streams/src/metrics.rs b/crates/streams/src/metrics.rs index 079526d..9860726 100644 --- a/crates/streams/src/metrics.rs +++ b/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); } diff --git a/crates/streams/src/writer.rs b/crates/streams/src/writer.rs new file mode 100644 index 0000000..e7c7877 --- /dev/null +++ b/crates/streams/src/writer.rs @@ -0,0 +1,270 @@ +//! Buffered writers for Activities and clients. `publish` only buffers, and a background flusher +//! appends every `batch_interval`, so a Producer pays one Store call per batch, not per Record. + +use crate::{ + Reason, StreamError, StreamResult, + proto::{ + AppendRecord, AppendRecordsRequest, OpenWriterRequest, Producer, ProducerKind, StreamRef, + }, + service::StreamService, +}; +use parking_lot::Mutex; +use std::{ + collections::{HashMap, VecDeque}, + sync::{ + Arc, + atomic::{AtomicU64, Ordering}, + }, + time::Duration, +}; +use tokio::{sync::Notify, time::Instant}; + +const DEFAULT_INTERVAL: Duration = Duration::from_millis(100); +const DEFAULT_BATCH: usize = 1000; +const DEFAULT_RETRY: Duration = Duration::from_secs(10); +const DEFAULT_BUFFERED: usize = 10_000; +const FIRST_BACKOFF: Duration = Duration::from_millis(50); + +#[derive(Default)] +struct State { + buffer: VecDeque, + /// The last `sequence` the writer sent in its epoch. + sequence: i64, + last_offset: i64, + /// A flush that failed for good. Every later call fails with it, since its Records are lost. + failed: Option, + closed: bool, +} + +struct Writer { + namespace: String, + stream: StreamRef, + producer: Producer, + interval: Duration, + max_batch: usize, + max_retry: Duration, + max_buffered: usize, + service: Arc, + state: Mutex, + wake: Notify, + /// One flush at a time, so the batches of one epoch go out in `sequence` order. + flushing: tokio::sync::Mutex<()>, +} + +fn duration(value: Option) -> Option { + let value = value?; + let duration = Duration::new(value.seconds.max(0) as u64, value.nanos.max(0) as u32); + (!duration.is_zero()).then_some(duration) +} + +impl Writer { + /// Sends the buffer, one batch after another, until it is empty or a batch fails for good. + async fn flush(&self) -> StreamResult { + let _flushing = self.flushing.lock().await; + loop { + let (records, sequence) = { + let mut state = self.state.lock(); + if let Some(error) = &state.failed { + return Err(error.clone()); + } + if state.buffer.is_empty() { + return Ok(state.last_offset); + } + let count = state.buffer.len().min(self.max_batch); + let records: Vec<_> = state.buffer.drain(..count).collect(); + state.sequence += 1; + (records, state.sequence) + }; + match self.send(records, sequence).await { + Ok(last) => { + let mut state = self.state.lock(); + state.last_offset = state.last_offset.max(last); + } + Err(error) => { + self.state.lock().failed = Some(error.clone()); + return Err(error); + } + } + } + } + + /// Appends one batch, and sends it again with the same `sequence` after a Store failure, so + /// the dedupe key keeps it from going in twice. + async fn send(&self, records: Vec, sequence: i64) -> StreamResult { + let request = AppendRecordsRequest { + namespace: self.namespace.clone(), + stream: Some(self.stream.clone()), + producer: Some(self.producer.clone()), + sequence, + records, + ..Default::default() + }; + let until = Instant::now() + self.max_retry; + let mut backoff = FIRST_BACKOFF; + loop { + match self.service.append_records(request.clone()).await { + Ok(response) => return Ok(response.last_offset), + Err(error) + if error.reason == Reason::Storage && Instant::now() + backoff < until => + { + tokio::time::sleep(backoff).await; + backoff = (backoff * 2).min(Duration::from_secs(1)); + } + Err(error) => return Err(error), + } + } + } + + async fn run(self: Arc) { + loop { + let _ = tokio::time::timeout(self.interval, self.wake.notified()).await; + if self.state.lock().closed { + return; + } + // A failure stays in the state, and the next publish or flush gives it to lang. + let _ = self.flush().await; + } + } +} + +/// The open writers of one extension. +#[derive(Default)] +pub(crate) struct Writers { + next: AtomicU64, + open: Mutex>>, + /// The writer of each (namespace, Stream, Producer id, epoch), so a second open shares its + /// sequence instead of reusing it. + by_producer: Mutex>, +} + +impl Writers { + /// Opens a writer, or finds the one this Producer epoch has on the Stream. A writer of an + /// Activity attempt above 1 first sends an empty batch, so the new epoch withdraws the + /// earlier attempt's Records also when it publishes nothing. + pub(crate) async fn open( + &self, + service: &Arc, + request: OpenWriterRequest, + ) -> StreamResult { + let stream = request + .stream + .ok_or_else(|| StreamError::invalid("The writer has no Stream"))?; + crate::ids::check_stream_id(&stream.id)?; + let producer = request + .producer + .ok_or_else(|| StreamError::invalid("The writer has no Producer"))?; + crate::ids::check_producer(&producer)?; + if producer.kind == ProducerKind::Workflow as i32 { + return Err(StreamError::invalid("Workflow code doesn't open a writer")); + } + let epoch = producer.epoch.unwrap_or_default(); + let key = format!( + "{}\n{}\n{}\n{}", + request.namespace, + stream.id, + producer.id, + crate::epoch::epoch_key(&epoch) + ); + if let Some(id) = self.by_producer.lock().get(&key) { + return Ok(*id); + } + let writer = Arc::new(Writer { + namespace: request.namespace, + stream, + producer, + interval: duration(request.batch_interval).unwrap_or(DEFAULT_INTERVAL), + max_batch: match request.max_batch_size { + n if n > 0 => n as usize, + _ => DEFAULT_BATCH, + }, + max_retry: duration(request.max_retry_duration).unwrap_or(DEFAULT_RETRY), + max_buffered: match request.max_buffered_records { + n if n > 0 => n as usize, + _ => DEFAULT_BUFFERED, + }, + service: service.clone(), + state: Mutex::default(), + wake: Notify::new(), + flushing: tokio::sync::Mutex::new(()), + }); + let retry = matches!( + ProducerKind::try_from(writer.producer.kind).unwrap_or_default(), + ProducerKind::Activity | ProducerKind::LocalActivity | ProducerKind::StandaloneActivity + ) && epoch.attempt > 1; + if retry { + writer.state.lock().sequence = 1; + match writer.send(vec![], 1).await { + // A Stream with no key has nothing to withdraw yet. + Ok(_) => {} + Err(error) if error.reason == Reason::NotFound => {} + Err(error) => return Err(error), + } + } + let id = self.next.fetch_add(1, Ordering::Relaxed) + 1; + let mut by_producer = self.by_producer.lock(); + if let Some(id) = by_producer.get(&key) { + return Ok(*id); + } + by_producer.insert(key, id); + self.open.lock().insert(id, writer.clone()); + tokio::spawn(writer.run()); + Ok(id) + } + + fn get(&self, id: u64) -> StreamResult> { + self.open + .lock() + .get(&id) + .cloned() + .ok_or_else(|| StreamError::new(Reason::UnknownWriter, format!("No writer {id}"))) + } + + /// Buffers Records. It never waits, so lang may call it from any thread. + pub(crate) fn publish( + &self, + id: u64, + records: Vec, + force_flush: bool, + ) -> StreamResult<()> { + let writer = self.get(id)?; + { + let mut state = writer.state.lock(); + if let Some(error) = &state.failed { + return Err(error.clone()); + } + if state.buffer.len() + records.len() > writer.max_buffered { + writer.service.metrics.buffer_full(); + return Err(StreamError::new( + Reason::BufferFull, + format!( + "The writer has {} Records waiting, and takes at most {}", + state.buffer.len(), + writer.max_buffered + ), + )); + } + state.buffer.extend(records); + if !force_flush && state.buffer.len() < writer.max_batch { + return Ok(()); + } + } + writer.wake.notify_one(); + Ok(()) + } + + /// Waits until every Record published so far is in the Store. + pub(crate) async fn flush(&self, id: u64) -> StreamResult { + self.get(id)?.flush().await + } + + /// Flushes, then drops the writer. + pub(crate) async fn close(&self, id: u64) -> StreamResult { + let writer = self.get(id)?; + let result = writer.flush().await; + writer.state.lock().closed = true; + writer.wake.notify_one(); + self.open.lock().remove(&id); + self.by_producer.lock().retain(|_, open| *open != id); + result + } +} diff --git a/crates/streams/tests/writer.rs b/crates/streams/tests/writer.rs new file mode 100644 index 0000000..aab43e9 --- /dev/null +++ b/crates/streams/tests/writer.rs @@ -0,0 +1,348 @@ +//! The buffered writers of Activities and clients, through the byte calls, on a real Redis. + +mod support; + +use prost::Message; +use prost_types::Timestamp; +use std::{ + sync::Arc, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, +}; +use support::server_url; +use temporalio_common::protos::temporal::api::common::v1::Payload; +use temporalio_sdk_core::{ExtensionError, WorkerExtension}; +use temporalio_streams::{ + ExtensionOptions, RedisStreamExtension, + epoch::withdraws, + ids::{activity_producer_id, client_producer_id}, + proto::{ + AppendRecord, CloseStreamRequest, CloseWriterRequest, CloseWriterResponse, + CreateStreamRequest, CreateStreamResponse, FlushWriterRequest, FlushWriterResponse, + OpenWriterRequest, OpenWriterResponse, Producer, ProducerEpoch, ProducerKind, + PublishRequest, ReadRecordsRequest, ReadRecordsResponse, RedisStreamConfig, + StreamErrorDetails, StreamRecord, StreamRecordKind, StreamRef, + }, +}; +use tonic::Code; + +static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + +struct Case { + ext: Arc, + namespace: String, + stream: String, +} + +async fn cases() -> Vec { + let mut cases = vec![]; + for (variable, cluster) in [ + ("STREAMS_REDIS_URL", false), + ("STREAMS_REDIS_CLUSTER_URL", true), + ] { + let Some(url) = server_url(variable) else { + continue; + }; + let config = RedisStreamConfig { + urls: vec![url], + cluster, + key_prefix: "streams-test".to_string(), + ..Default::default() + }; + let n = NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos(); + let case = Case { + ext: RedisStreamExtension::new(ExtensionOptions::from_config(config).unwrap(), None), + namespace: format!("ns-{now}-{n}"), + stream: format!("s-{now}-{n}"), + }; + let created: CreateStreamResponse = case + .call( + "CreateStream", + CreateStreamRequest { + namespace: case.namespace.clone(), + stream_id: case.stream.clone(), + ..Default::default() + }, + ) + .await + .unwrap(); + assert!(created.created); + cases.push(case); + } + cases +} + +fn client() -> Producer { + Producer { + kind: ProducerKind::Client as i32, + id: client_producer_id("me", "p"), + epoch: Some(ProducerEpoch { + attempt: 1, + ..Default::default() + }), + run_id: String::new(), + } +} + +fn activity(attempt: i32) -> Producer { + Producer { + kind: ProducerKind::Activity as i32, + id: activity_producer_id("r1", "a", "s"), + epoch: Some(ProducerEpoch { + attempt_scheduled_time: Some(Timestamp { + seconds: 10 * attempt as i64, + nanos: 0, + }), + attempt, + ..Default::default() + }), + run_id: "r1".to_string(), + } +} + +fn records(texts: &[&str]) -> Vec { + texts + .iter() + .map(|text| AppendRecord { + topic: String::new(), + body: Some(Payload { + data: text.as_bytes().to_vec(), + ..Default::default() + }), + }) + .collect() +} + +fn reason(error: &ExtensionError) -> String { + StreamErrorDetails::decode(error.details.as_slice()) + .unwrap() + .reason +} + +fn millis(ms: u64) -> Option { + Some(prost_types::Duration { + seconds: (ms / 1000) as i64, + nanos: ((ms % 1000) * 1_000_000) as i32, + }) +} + +impl Case { + async fn call( + &self, + method: &str, + request: Req, + ) -> Result { + let bytes = self + .ext + .call_async(method.to_string(), request.encode_to_vec()) + .await?; + Ok(Resp::decode(bytes.as_slice()).unwrap()) + } + + async fn open(&self, producer: Producer, change: impl Fn(&mut OpenWriterRequest)) -> u64 { + let mut request = OpenWriterRequest { + namespace: self.namespace.clone(), + stream: Some(StreamRef { + id: self.stream.clone(), + }), + producer: Some(producer), + ..Default::default() + }; + change(&mut request); + let opened: OpenWriterResponse = self.call("OpenWriter", request).await.unwrap(); + opened.writer_id + } + + fn publish( + &self, + writer_id: u64, + texts: &[&str], + force_flush: bool, + ) -> Result<(), ExtensionError> { + let request = PublishRequest { + writer_id, + records: records(texts), + force_flush, + }; + self.ext + .call("Publish", &request.encode_to_vec()) + .map(|_| ()) + } + + async fn flush(&self, writer_id: u64) -> Result { + let flushed: FlushWriterResponse = self + .call("FlushWriter", FlushWriterRequest { writer_id }) + .await?; + Ok(flushed.last_offset) + } + + async fn read(&self, wait_ms: u64) -> ReadRecordsResponse { + self.call( + "ReadRecords", + ReadRecordsRequest { + namespace: self.namespace.clone(), + stream: Some(StreamRef { + id: self.stream.clone(), + }), + wait: millis(wait_ms), + ..Default::default() + }, + ) + .await + .unwrap() + } + + async fn visible(&self) -> Vec { + let mut kept: Vec = vec![]; + for record in self.read(0).await.records { + match StreamRecordKind::try_from(record.kind).unwrap() { + StreamRecordKind::Superseded => { + let s = record.supersession.unwrap(); + kept.retain(|old| !withdraws(&s, old)); + } + StreamRecordKind::Data => kept.push(record), + _ => {} + } + } + kept.into_iter() + .map(|record| String::from_utf8(record.body.unwrap().data).unwrap()) + .collect() + } +} + +#[tokio::test] +async fn a_publish_goes_out_with_the_next_batch_interval() { + for case in cases().await { + let writer = case + .open(client(), |r| r.batch_interval = millis(100)) + .await; + case.publish(writer, &["a", "b"], false).unwrap(); + assert!(case.read(0).await.records.is_empty()); + let read = case.read(3000).await; + assert_eq!(read.records.len(), 2); + } +} + +#[tokio::test] +async fn a_forced_flush_goes_out_at_once() { + for case in cases().await { + let writer = case + .open(client(), |r| r.batch_interval = millis(60_000)) + .await; + let started = Instant::now(); + case.publish(writer, &["a"], true).unwrap(); + assert_eq!(case.read(3000).await.records.len(), 1); + assert!(started.elapsed() < Duration::from_secs(2)); + } +} + +#[tokio::test] +async fn a_flush_waits_for_the_store_and_batches_keep_their_order() { + for case in cases().await { + let writer = case + .open(client(), |r| { + r.batch_interval = millis(60_000); + r.max_batch_size = 2; + }) + .await; + case.publish(writer, &["1", "2", "3"], false).unwrap(); + case.publish(writer, &["4", "5"], false).unwrap(); + let last = case.flush(writer).await.unwrap(); + let read = case.read(0).await; + assert_eq!(read.records.last().unwrap().offset, last); + assert_eq!(case.visible().await, ["1", "2", "3", "4", "5"]); + // Opening the same Producer epoch again gives the same writer, so it keeps the sequence. + assert_eq!(case.open(client(), |_| {}).await, writer); + } +} + +#[tokio::test] +async fn a_full_buffer_refuses_a_publish() { + for case in cases().await { + let writer = case + .open(client(), |r| { + r.batch_interval = millis(60_000); + r.max_buffered_records = 2; + }) + .await; + case.publish(writer, &["1", "2"], false).unwrap(); + let error = case.publish(writer, &["3"], false).unwrap_err(); + assert_eq!( + (error.code, reason(&error).as_str()), + (Code::ResourceExhausted, "stream_buffer_full") + ); + case.flush(writer).await.unwrap(); + case.publish(writer, &["3"], false).unwrap(); + } +} + +#[tokio::test] +async fn a_failed_flush_fails_the_flush_and_every_later_publish() { + for case in cases().await { + let writer = case + .open(client(), |r| r.batch_interval = millis(60_000)) + .await; + let close = CloseStreamRequest { + namespace: case.namespace.clone(), + stream: Some(StreamRef { + id: case.stream.clone(), + }), + producer: Some(client()), + ..Default::default() + }; + let _: temporalio_streams::proto::CloseStreamResponse = + case.call("CloseStream", close).await.unwrap(); + case.publish(writer, &["late"], false).unwrap(); + let error = case.flush(writer).await.unwrap_err(); + assert_eq!( + (error.code, reason(&error).as_str()), + (Code::FailedPrecondition, "stream_closed") + ); + assert_eq!( + reason(&case.publish(writer, &["x"], false).unwrap_err()), + "stream_closed" + ); + } +} + +#[tokio::test] +async fn an_activity_retry_withdraws_the_earlier_attempt_even_with_nothing_to_publish() { + for case in cases().await { + let first = case.open(activity(1), |_| {}).await; + case.publish(first, &["a1"], false).unwrap(); + case.flush(first).await.unwrap(); + let second = case.open(activity(2), |_| {}).await; + assert!(case.visible().await.is_empty()); + case.publish(second, &["a2"], false).unwrap(); + case.flush(second).await.unwrap(); + assert_eq!(case.visible().await, ["a2"]); + // A late flush of the earlier attempt is fenced. + case.publish(first, &["late"], false).unwrap(); + let error = case.flush(first).await.unwrap_err(); + assert_eq!(reason(&error), "stream_fenced"); + } +} + +#[tokio::test] +async fn a_closed_writer_is_gone_after_its_last_flush() { + for case in cases().await { + let writer = case + .open(client(), |r| r.batch_interval = millis(60_000)) + .await; + case.publish(writer, &["last"], false).unwrap(); + let closed: CloseWriterResponse = case + .call("CloseWriter", CloseWriterRequest { writer_id: writer }) + .await + .unwrap(); + assert!(closed.last_offset > 0); + assert_eq!(case.visible().await, ["last"]); + let error = case.publish(writer, &["x"], false).unwrap_err(); + assert_eq!( + (error.code, reason(&error).as_str()), + (Code::NotFound, "stream_unknown_writer") + ); + } +} From 05402e30eaeb346860731368619cb67755fe98ae Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 20 Oct 2026 20:32:47 -0700 Subject: [PATCH 2/3] Split writer batches by bytes and kept the sequence of a reopen. A writer batch over the Store's limit failed for good and lost the buffer, and a reopen of a closed writer sent a used sequence again, which dropped its Records as resends. --- crates/streams/src/writer.rs | 77 +++++++++++++++++++++++++++++----- crates/streams/tests/writer.rs | 41 ++++++++++++++++++ 2 files changed, 108 insertions(+), 10 deletions(-) diff --git a/crates/streams/src/writer.rs b/crates/streams/src/writer.rs index e7c7877..06ad99e 100644 --- a/crates/streams/src/writer.rs +++ b/crates/streams/src/writer.rs @@ -6,9 +6,11 @@ use crate::{ proto::{ AppendRecord, AppendRecordsRequest, OpenWriterRequest, Producer, ProducerKind, StreamRef, }, + redis::{MAX_BATCH_BYTES, MAX_RECORD_BYTES}, service::StreamService, }; use parking_lot::Mutex; +use prost::Message; use std::{ collections::{HashMap, VecDeque}, sync::{ @@ -70,7 +72,18 @@ impl Writer { if state.buffer.is_empty() { return Ok(state.last_offset); } - let count = state.buffer.len().min(self.max_batch); + // The Store refuses a batch over its limit, so a batch ends at the count or the + // bytes, whichever comes first. + let mut count = 0; + let mut bytes = 0; + for record in state.buffer.iter().take(self.max_batch) { + let size = stored_size(&self.producer, record); + if count > 0 && bytes + size > MAX_BATCH_BYTES { + break; + } + bytes += size; + count += 1; + } let records: Vec<_> = state.buffer.drain(..count).collect(); state.sequence += 1; (records, state.sequence) @@ -118,8 +131,12 @@ impl Writer { async fn run(self: Arc) { loop { let _ = tokio::time::timeout(self.interval, self.wake.notified()).await; - if self.state.lock().closed { - return; + { + let state = self.state.lock(); + // A writer that failed for good writes nothing more, so its task ends too. + if state.closed || state.failed.is_some() { + return; + } } // A failure stays in the state, and the next publish or flush gives it to lang. let _ = self.flush().await; @@ -127,6 +144,11 @@ impl Writer { } } +/// The size of a Record once stored, with a margin for the envelope's own fields. +fn stored_size(producer: &Producer, record: &AppendRecord) -> usize { + record.encoded_len() + producer.encoded_len() + 32 +} + /// The open writers of one extension. #[derive(Default)] pub(crate) struct Writers { @@ -135,6 +157,10 @@ pub(crate) struct Writers { /// The writer of each (namespace, Stream, Producer id, epoch), so a second open shares its /// sequence instead of reusing it. by_producer: Mutex>, + /// The last `sequence` of each closed writer, so a reopen of the same epoch goes on from + /// it instead of sending a used `sequence` again. + closed_sequences: Mutex>, + keys: Mutex>, } impl Writers { @@ -187,10 +213,18 @@ impl Writers { wake: Notify::new(), flushing: tokio::sync::Mutex::new(()), }); - let retry = matches!( - ProducerKind::try_from(writer.producer.kind).unwrap_or_default(), - ProducerKind::Activity | ProducerKind::LocalActivity | ProducerKind::StandaloneActivity - ) && epoch.attempt > 1; + let reopened = self.closed_sequences.lock().get(&key).copied(); + if let Some(sequence) = reopened { + writer.state.lock().sequence = sequence; + } + let retry = reopened.is_none() + && matches!( + ProducerKind::try_from(writer.producer.kind).unwrap_or_default(), + ProducerKind::Activity + | ProducerKind::LocalActivity + | ProducerKind::StandaloneActivity + ) + && epoch.attempt > 1; if retry { writer.state.lock().sequence = 1; match writer.send(vec![], 1).await { @@ -205,7 +239,8 @@ impl Writers { if let Some(id) = by_producer.get(&key) { return Ok(*id); } - by_producer.insert(key, id); + by_producer.insert(key.clone(), id); + self.keys.lock().insert(id, key); self.open.lock().insert(id, writer.clone()); tokio::spawn(writer.run()); Ok(id) @@ -232,6 +267,22 @@ impl Writers { if let Some(error) = &state.failed { return Err(error.clone()); } + if state.closed { + return Err(StreamError::new( + Reason::UnknownWriter, + format!("Writer {id} is closed"), + )); + } + if let Some(big) = records + .iter() + .map(|record| stored_size(&writer.producer, record)) + .find(|size| *size > MAX_RECORD_BYTES) + { + return Err(StreamError::new( + Reason::TooLarge, + format!("A Record has about {big} bytes, more than {MAX_RECORD_BYTES}"), + )); + } if state.buffer.len() + records.len() > writer.max_buffered { writer.service.metrics.buffer_full(); return Err(StreamError::new( @@ -260,11 +311,17 @@ impl Writers { /// Flushes, then drops the writer. pub(crate) async fn close(&self, id: u64) -> StreamResult { let writer = self.get(id)?; - let result = writer.flush().await; + // A publish checks the closed state under the same lock, so every Record it took is + // in the buffer that this flush sends. writer.state.lock().closed = true; + let result = writer.flush().await; writer.wake.notify_one(); self.open.lock().remove(&id); - self.by_producer.lock().retain(|_, open| *open != id); + if let Some(key) = self.keys.lock().remove(&id) { + self.by_producer.lock().remove(&key); + let sequence = writer.state.lock().sequence; + self.closed_sequences.lock().insert(key, sequence); + } result } } diff --git a/crates/streams/tests/writer.rs b/crates/streams/tests/writer.rs index aab43e9..4b77c39 100644 --- a/crates/streams/tests/writer.rs +++ b/crates/streams/tests/writer.rs @@ -346,3 +346,44 @@ async fn a_closed_writer_is_gone_after_its_last_flush() { ); } } + +#[tokio::test] +async fn a_flush_splits_big_records_by_bytes_and_refuses_an_oversize_one() { + for case in cases().await { + let writer = case + .open(client(), |r| r.batch_interval = millis(60_000)) + .await; + let big = "x".repeat(1_400_000); + case.publish(writer, &[&big, &big, &big, "end"], false) + .unwrap(); + case.flush(writer).await.unwrap(); + assert_eq!(case.visible().await.len(), 4); + let huge = "y".repeat(1_700_000); + let error = case.publish(writer, &[&huge], false).unwrap_err(); + assert_eq!( + (error.code, reason(&error).as_str()), + (Code::InvalidArgument, "stream_too_large") + ); + } +} + +#[tokio::test] +async fn a_reopen_of_a_closed_writer_goes_on_with_its_sequence() { + for case in cases().await { + let writer = case + .open(client(), |r| r.batch_interval = millis(60_000)) + .await; + case.publish(writer, &["1"], false).unwrap(); + let _: CloseWriterResponse = case + .call("CloseWriter", CloseWriterRequest { writer_id: writer }) + .await + .unwrap(); + let again = case + .open(client(), |r| r.batch_interval = millis(60_000)) + .await; + assert_ne!(again, writer); + case.publish(again, &["2"], false).unwrap(); + case.flush(again).await.unwrap(); + assert_eq!(case.visible().await, ["1", "2"]); + } +} From 17bf3f909d7e5c9aa87eb055f34e0e0127d5865f Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 20 Oct 2026 20:40:26 -0700 Subject: [PATCH 3/3] Added CloseWriters, which flushes and drops every open writer. A writer that lang never closed kept its flusher task until the process ended, so a shutdown now closes them all. --- .../temporal/sdk/streams/v1/extension.proto | 10 +++++++ crates/streams/src/extension.rs | 16 +++++++++--- crates/streams/src/writer.rs | 12 +++++++++ crates/streams/tests/writer.rs | 26 +++++++++++++++++++ 4 files changed, 60 insertions(+), 4 deletions(-) diff --git a/crates/streams/protos/temporal/sdk/streams/v1/extension.proto b/crates/streams/protos/temporal/sdk/streams/v1/extension.proto index b063acf..c4c0c00 100644 --- a/crates/streams/protos/temporal/sdk/streams/v1/extension.proto +++ b/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 { diff --git a/crates/streams/src/extension.rs b/crates/streams/src/extension.rs index 47f52d3..330868e 100644 --- a/crates/streams/src/extension.rs +++ b/crates/streams/src/extension.rs @@ -5,10 +5,10 @@ use crate::{ metrics::Metrics, proto::{ AppendRecordsRequest, BufferWorkflowCallsRequest, BufferWorkflowCallsResponse, - CloseStreamRequest, CloseWriterRequest, CloseWriterResponse, DescribeStreamRequest, - DrainClosesRequest, DrainClosesResponse, FlushWriterRequest, FlushWriterResponse, - OpenWriterRequest, OpenWriterResponse, PublishRequest, PublishResponse, ReadRecordsRequest, - RedisStreamConfig, + CloseStreamRequest, CloseWriterRequest, CloseWriterResponse, CloseWritersRequest, + CloseWritersResponse, DescribeStreamRequest, DrainClosesRequest, DrainClosesResponse, + FlushWriterRequest, FlushWriterResponse, OpenWriterRequest, OpenWriterResponse, + PublishRequest, PublishResponse, ReadRecordsRequest, RedisStreamConfig, }, redis::{RedisStore, RedisStoreOptions}, service::StreamService, @@ -223,6 +223,14 @@ impl RedisStreamExtension { let last_offset = self.shared.writers.flush(request.writer_id).await?; FlushWriterResponse { last_offset }.encode_to_vec() } + "CloseWriters" => { + decode::(input)?; + let failed = self.shared.writers.close_all().await; + CloseWritersResponse { + failed: failed as i32, + } + .encode_to_vec() + } "CloseWriter" => { let request = decode::(input)?; let last_offset = self.shared.writers.close(request.writer_id).await?; diff --git a/crates/streams/src/writer.rs b/crates/streams/src/writer.rs index 06ad99e..a342ca7 100644 --- a/crates/streams/src/writer.rs +++ b/crates/streams/src/writer.rs @@ -308,6 +308,18 @@ impl Writers { self.get(id)?.flush().await } + /// Closes every open writer. Answers how many failed their last flush. + pub(crate) async fn close_all(&self) -> usize { + let ids: Vec = self.open.lock().keys().copied().collect(); + let mut failed = 0; + for id in ids { + if self.close(id).await.is_err() { + failed += 1; + } + } + failed + } + /// Flushes, then drops the writer. pub(crate) async fn close(&self, id: u64) -> StreamResult { let writer = self.get(id)?; diff --git a/crates/streams/tests/writer.rs b/crates/streams/tests/writer.rs index 4b77c39..070ee66 100644 --- a/crates/streams/tests/writer.rs +++ b/crates/streams/tests/writer.rs @@ -387,3 +387,29 @@ async fn a_reopen_of_a_closed_writer_goes_on_with_its_sequence() { assert_eq!(case.visible().await, ["1", "2"]); } } + +#[tokio::test] +async fn closing_every_writer_flushes_them_all() { + for case in cases().await { + let a = case + .open(client(), |r| r.batch_interval = millis(60_000)) + .await; + let b = case + .open(activity(1), |r| r.batch_interval = millis(60_000)) + .await; + case.publish(a, &["a"], false).unwrap(); + case.publish(b, &["b"], false).unwrap(); + let closed: temporalio_streams::proto::CloseWritersResponse = case + .call( + "CloseWriters", + temporalio_streams::proto::CloseWritersRequest {}, + ) + .await + .unwrap(); + assert_eq!(closed.failed, 0); + let mut texts = case.visible().await; + texts.sort(); + assert_eq!(texts, ["a", "b"]); + assert!(case.publish(a, &["x"], false).is_err()); + } +}