diff --git a/crates/streams/Cargo.toml b/crates/streams/Cargo.toml index 88d22bdd6..8224a1dd6 100644 --- a/crates/streams/Cargo.toml +++ b/crates/streams/Cargo.toml @@ -24,7 +24,7 @@ prost-types = { workspace = true } serde_json = { workspace = true, optional = true } 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"] } tonic = { workspace = true } [dependencies.temporalio-common] diff --git a/crates/streams/protos/temporal/sdk/streams/v1/runtime.proto b/crates/streams/protos/temporal/sdk/streams/v1/runtime.proto index 2f559b76e..d6f287332 100644 --- a/crates/streams/protos/temporal/sdk/streams/v1/runtime.proto +++ b/crates/streams/protos/temporal/sdk/streams/v1/runtime.proto @@ -212,16 +212,18 @@ message CloseStreamAsRequest { 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. +// Method `CloseScope`, asynchronous. Flushes and frees the scope's buffers, closes its writers and +// frees the scope. A failed writer or buffer 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. + // The first failure, if any. A buffer's comes first. StreamFailure failure = 2; + // The buffers whose failure no `BufferPush` or `BufferFlush` gave lang yet. + int32 failed_buffers = 3; } // A failure that a response reports, as lang raises it from a failed call. @@ -240,13 +242,14 @@ message FreeScopeRequest { 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. +// Method `Shutdown`, asynchronous. Closes every scope, buffer and writer. Then only `Shutdown`, +// the static and the `Free` methods, and the pump calls run. A dropped one runs again. message ShutdownRequest { } message ShutdownResponse { int32 failed_writers = 1; + int32 failed_buffers = 2; } // Method `WorkflowAppend`, asynchronous. The append Local Activity of Workflow code. The attempts @@ -360,3 +363,86 @@ message WithdrawsRequest { message WithdrawsResponse { bool withdraws = 1; } + +// Method `OpenBuffer`, synchronous. Opens a buffered publisher of one Stream, which writes as the +// scope's Producer, one batch at a time. A failed batch fails the buffer, so a reader sees no gap. +message OpenBufferRequest { + uint64 scope_id = 1; + temporal.api.stream.v1.StreamRef stream = 2; + // A batch is due this long after its first value, or at once for a flush. Empty is 100 ms. + google.protobuf.Duration max_delay = 3; + // A batch is due at once with this many values. Zero is 1000. + int32 max_batch = 4; + // Set when lang has no codec, so the runtime writes each batch. Unset, lang's pump runs each + // batch through its codec with `BufferNextBatch` and `BufferWrite` until `done`. + bool encoded = 5; +} + +message OpenBufferResponse { + uint64 buffer_id = 1; +} + +// Method `BufferPush`, synchronous. Adds values to the buffer and returns at once. After a failed +// batch it fails with the batch's error. +message BufferPushRequest { + uint64 buffer_id = 1; + repeated temporal.api.stream.v1.AppendRecord records = 2; +} + +message BufferPushResponse { +} + +// Method `BufferNextBatch`, asynchronous. Waits for the next due batch and takes it, for lang's +// codec pump. A batch that no `BufferWrite` started yet comes again, for a pump that lang restarts. +message BufferNextBatchRequest { + uint64 buffer_id = 1; +} + +message BufferBatch { + uint64 batch_id = 1; + // The bodies before lang's codec. + repeated temporal.api.stream.v1.AppendRecord records = 2; + // The buffer is freed or failed, so the pump stops. + bool done = 3; +} + +// Method `BufferWrite`, asynchronous. Writes the batch in flight after lang's codec, or fails the +// buffer with the codec's error. A failed write fails the buffer. A stale batch fails the call. +message BufferWriteRequest { + uint64 buffer_id = 1; + uint64 batch_id = 2; + oneof body { + BufferRecords records = 3; + string codec_error = 4; + } +} + +message BufferRecords { + // The bodies after lang's codec, one for each value of the batch. Another count fails the call + // and the buffer. + repeated temporal.api.stream.v1.AppendRecord records = 1; +} + +message BufferWriteResponse { +} + +// Method `BufferFlush`, asynchronous. Waits until every value pushed before the call is written or +// failed. It is the backpressure point of a caller that pushes faster than the Store writes. +message BufferFlushRequest { + uint64 buffer_id = 1; +} + +message BufferFlushResponse { + // The offset of the newest Record that the buffer wrote. Zero if it wrote none. + int64 last_offset = 1; + // The values that a failed batch left unwritten. + int64 unwritten = 2; + // The failure of the buffer, if any. + StreamFailure failure = 3; +} + +// Method `FreeBuffer`, synchronous, with `FreeResponse`. Frees the buffer with no wait, and drops +// the values that no write took. A write in flight still completes. +message FreeBufferRequest { + uint64 buffer_id = 1; +} diff --git a/crates/streams/src/buffer.rs b/crates/streams/src/buffer.rs new file mode 100644 index 000000000..5ac3d1d96 --- /dev/null +++ b/crates/streams/src/buffer.rs @@ -0,0 +1,488 @@ +//! Buffered publishers. The runtime keeps the batching rules, and lang keeps only the codec, which +//! runs once per batch, since a remote codec costs a round trip per call. + +use crate::{ + Reason, StreamError, StreamResult, + ids::check_stream_id, + proto::{ + AppendRecord, BufferBatch, BufferFlushResponse, BufferWriteRequest, OpenBufferRequest, + StreamFailure, StreamRef, buffer_write_request::Body, + }, + runtime::Shared, + scope::Scope, + writer::Written, +}; +use parking_lot::Mutex; +use std::{ + collections::HashMap, + sync::{ + Arc, + atomic::{AtomicU64, Ordering}, + }, + time::Duration, +}; +use tokio::{sync::Notify, time::Instant}; + +const DEFAULT_DELAY: Duration = Duration::from_millis(100); +const LONGEST_DELAY: Duration = Duration::from_secs(365 * 24 * 60 * 60); +const DEFAULT_BATCH: usize = 1000; + +pub(crate) struct Buffer { + scope: Arc, + /// The runtime's own pump takes the batches, so lang's pump calls are refused. + encoded: bool, + stream: StreamRef, + max_delay: Duration, + max_batch: usize, + state: Mutex, + /// Wakes the pump and the flushes after each change of the state. + changed: Notify, +} + +#[derive(Default)] +struct State { + /// The values that no batch took yet. + pending: Vec, + /// When the oldest pending value arrived. + first_at: Option, + in_flight: Option, + last_batch: u64, + pushed: u64, + /// The values whose batch was written or failed. + settled: u64, + /// A batch is due at once while fewer values than this are settled, so a flush doesn't wait. + flush_to: u64, + error: Option, + /// Lang got the error from a push or a flush. + reported: bool, + last_offset: i64, + unwritten: u64, + /// Set when the scope's close starts to flush the buffer. Then a push fails. + closing: bool, + freed: bool, +} + +struct InFlight { + id: u64, + count: u64, + /// The values of a batch that lang's pump took, until a write starts. The next + /// `BufferNextBatch` offers them again, so a dropped pump doesn't strand the batch. + records: Option>, +} + +fn unknown_buffer() -> StreamError { + StreamError::new(Reason::UnknownWriter, "No buffer has the id") +} + +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 Buffer { + fn push(&self, records: Vec) -> StreamResult<()> { + let mut state = self.state.lock(); + if let Some(error) = &state.error { + let error = error.clone(); + state.reported = true; + return Err(error); + } + if state.freed || state.closing { + return Err(unknown_buffer()); + } + if records.is_empty() { + return Err(StreamError::invalid("A push takes one or more values")); + } + if state.pending.is_empty() { + state.first_at = Some(Instant::now()); + } + state.pushed += records.len() as u64; + state.pending.extend(records); + drop(state); + self.changed.notify_waiters(); + Ok(()) + } + + /// Waits for the next due batch and takes it. `None` once the buffer is freed or failed. + async fn next_batch(&self) -> Option<(u64, Vec)> { + loop { + let notified = self.changed.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + let due_at = { + let mut state = self.state.lock(); + if state.freed || state.error.is_some() { + return None; + } + if let Some(InFlight { + id, + records: Some(records), + .. + }) = &state.in_flight + { + return Some((*id, records.clone())); + } + match state.first_at { + Some(first_at) if state.in_flight.is_none() => { + let due_at = first_at + self.max_delay; + let full = state.pending.len() >= self.max_batch; + if full || state.flush_to > state.settled || due_at <= Instant::now() { + return Some(state.take(self.max_batch, !self.encoded)); + } + Some(due_at) + } + _ => None, + } + }; + match due_at { + Some(at) => { + tokio::select! { + _ = notified => {} + _ = tokio::time::sleep_until(at) => {} + } + } + None => notified.await, + } + } + } + + /// Starts the write of the batch in flight, once. Records of another count than the batch's + /// fail the call and the buffer, as a codec error does. + fn start_write(&self, batch_id: u64, records: Option) -> StreamResult<()> { + let mut state = self.state.lock(); + // A free already counted the batch as unwritten. + if state.freed { + return Err(unknown_buffer()); + } + let Some(in_flight) = state + .in_flight + .as_mut() + .filter(|in_flight| in_flight.id == batch_id && in_flight.records.is_some()) + else { + return Err(StreamError::invalid(format!( + "Batch {batch_id} isn't a batch in flight that nobody writes" + ))); + }; + in_flight.records = None; + if records.is_some_and(|count| count as u64 != in_flight.count) { + let error = StreamError::invalid(format!( + "Batch {batch_id} has {} values, and the codec must give one Record for each", + in_flight.count + )); + drop(state); + let written = Written { + error: Some(error.clone()), + ..Default::default() + }; + self.settle(batch_id, written); + return Err(error); + } + Ok(()) + } + + /// Ends the batch in flight with the outcome of its write. + fn settle(&self, batch_id: u64, written: Written) { + let mut state = self.state.lock(); + let count = match &state.in_flight { + Some(in_flight) if in_flight.id == batch_id => in_flight.count, + _ => return, + }; + state.in_flight = None; + state.settled += count; + state.last_offset = state.last_offset.max(written.last); + if let Some(error) = written.error { + // A failed batch fails the later values too, so a reader sees no gap. + let left = state.pending.len() as u64; + state.pending.clear(); + state.first_at = None; + state.settled += left; + state.unwritten += count.saturating_sub(written.records as u64) + left; + state.error = Some(error); + } + drop(state); + self.changed.notify_waiters(); + } + + /// Waits until every value pushed before the call is settled. + async fn flush(&self) -> BufferFlushResponse { + let target = { + let mut state = self.state.lock(); + state.flush_to = state.flush_to.max(state.pushed); + state.pushed + }; + self.changed.notify_waiters(); + loop { + let notified = self.changed.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + { + let mut state = self.state.lock(); + if state.settled >= target || state.freed { + let mut failure = state.error.as_ref().map(StreamError::failure); + if failure.is_none() && state.settled < target { + failure = Some(StreamError::failure(&StreamError::new( + Reason::UnknownWriter, + "The buffer was freed before it wrote every value", + ))); + } + state.reported |= failure.is_some(); + return BufferFlushResponse { + last_offset: state.last_offset, + unwritten: state.unwritten as i64, + failure, + }; + } + } + notified.await; + } + } + + fn free(&self) { + let mut state = self.state.lock(); + state.freed = true; + state.unwritten += state.pending.len() as u64; + state.pending.clear(); + // Lang's pump can't write a batch that it took once the buffer is gone. + if let Some(InFlight { + count, + records: Some(_), + .. + }) = state.in_flight + { + state.unwritten += count; + } + drop(state); + self.changed.notify_waiters(); + } +} + +impl State { + /// Takes the next batch. With `keep`, it keeps its values until lang's write starts. + fn take(&mut self, max: usize, keep: bool) -> (u64, Vec) { + let rest = self.pending.split_off(self.pending.len().min(max)); + let batch = std::mem::replace(&mut self.pending, rest); + // The values left behind waited since the batch's first value, so they're due at once. + if self.pending.is_empty() { + self.first_at = None; + } + self.last_batch += 1; + self.in_flight = Some(InFlight { + id: self.last_batch, + count: batch.len() as u64, + records: keep.then(|| batch.clone()), + }); + (self.last_batch, batch) + } +} + +/// Fails the batch in flight when the write task ends early, such as by a panic. After the write +/// settled the batch, the settle here finds no batch and does nothing. +struct SettleOnDrop { + buffer: Arc, + batch_id: u64, +} + +impl Drop for SettleOnDrop { + fn drop(&mut self) { + let error = StreamError::storage("The write of a batch ended before it answered"); + let written = Written { + error: Some(error), + ..Default::default() + }; + self.buffer.settle(self.batch_id, written); + } +} + +#[derive(Default)] +pub(crate) struct Buffers { + next: AtomicU64, + table: Mutex>>, +} + +impl Buffers { + fn get(&self, id: u64) -> StreamResult> { + self.table + .lock() + .get(&id) + .cloned() + .ok_or_else(unknown_buffer) + } + + fn lang_pumped(&self, id: u64) -> StreamResult> { + let buffer = self.get(id)?; + if buffer.encoded { + return Err(StreamError::invalid( + "The runtime writes the batches of an encoded buffer", + )); + } + Ok(buffer) + } + + pub(crate) fn free(&self, id: u64) -> StreamResult<()> { + let buffer = self.table.lock().remove(&id).ok_or_else(unknown_buffer)?; + buffer.scope.remove_buffer(id); + buffer.free(); + Ok(()) + } + + /// Frees every buffer, so that the runtime's pumps stop. + pub(crate) fn free_all(&self) { + for (_, buffer) in self.table.lock().drain() { + buffer.free(); + } + } +} + +impl Shared { + pub(crate) fn open_buffer(self: &Arc, request: OpenBufferRequest) -> StreamResult { + let scope = self.scopes.get(request.scope_id)?; + let stream = request + .stream + .ok_or_else(|| StreamError::invalid("The request has no Stream"))?; + check_stream_id(&stream.id)?; + let handle = match request.encoded { + true => Some(self.tokio_handle().ok_or_else(|| { + StreamError::invalid("An encoded buffer needs a Tokio runtime to write on") + })?), + false => None, + }; + let buffer = Arc::new(Buffer { + scope: scope.clone(), + encoded: request.encoded, + stream, + // A delay of more than a year counts as a year, so that `Instant` math can't overflow. + max_delay: (duration(request.max_delay).unwrap_or(DEFAULT_DELAY)).min(LONGEST_DELAY), + max_batch: match request.max_batch { + n if n > 0 => n as usize, + _ => DEFAULT_BATCH, + }, + state: Mutex::default(), + changed: Notify::new(), + }); + let id = self.buffers.next.fetch_add(1, Ordering::Relaxed) + 1; + self.buffers.table.lock().insert(id, buffer.clone()); + if let Err(error) = scope.add_buffer(id) { + self.buffers.free(id)?; + return Err(error); + } + if let Some(handle) = handle { + // The pump holds a weak reference, so it doesn't keep a dropped runtime alive. + let shared = Arc::downgrade(self); + handle.spawn(async move { + while let Some((batch_id, records)) = buffer.next_batch().await { + let guard = SettleOnDrop { + buffer: buffer.clone(), + batch_id, + }; + let Some(shared) = shared.upgrade() else { + return; + }; + let written = shared.write_batch(&guard.buffer, records).await; + guard.buffer.settle(batch_id, written); + } + }); + } + Ok(id) + } + + pub(crate) fn buffer_push(&self, id: u64, records: Vec) -> StreamResult<()> { + self.buffers.get(id)?.push(records) + } + + pub(crate) async fn buffer_next_batch(&self, id: u64) -> StreamResult { + let buffer = match self.buffers.lang_pumped(id) { + // A freed buffer is gone from the table, and its pump stops. + Err(error) if error.reason == Reason::UnknownWriter => { + return Ok(BufferBatch { + done: true, + ..Default::default() + }); + } + buffer => buffer?, + }; + Ok(match buffer.next_batch().await { + Some((batch_id, records)) => BufferBatch { + batch_id, + records, + done: false, + }, + None => BufferBatch { + done: true, + ..Default::default() + }, + }) + } + + pub(crate) async fn buffer_write( + self: &Arc, + request: BufferWriteRequest, + ) -> StreamResult<()> { + let buffer = self.buffers.lang_pumped(request.buffer_id)?; + let body = request + .body + .ok_or_else(|| StreamError::invalid("The write has no records"))?; + let handle = self + .tokio_handle() + .ok_or_else(|| StreamError::invalid("A buffer needs a Tokio runtime to write on"))?; + let batch_id = request.batch_id; + let count = match &body { + Body::Records(batch) => Some(batch.records.len()), + Body::CodecError(_) => None, + }; + buffer.start_write(batch_id, count)?; + let records = match body { + Body::Records(batch) => batch.records, + Body::CodecError(message) => { + let failed = format!("The codec failed on a batch: {message}"); + let error = Some(StreamError::invalid(failed)); + buffer.settle( + batch_id, + Written { + error, + ..Default::default() + }, + ); + return Ok(()); + } + }; + // The write runs on a task of its own, so a dropped call doesn't leave the batch in flight. + // Nobody may await that task, so the guard settles the batch if the task ends early, or if + // a runtime that shut down never runs it. + let (shared, guard) = (self.clone(), SettleOnDrop { buffer, batch_id }); + let write = handle.spawn(async move { + let written = shared.write_batch(&guard.buffer, records).await; + guard.buffer.settle(batch_id, written); + }); + let _ = write.await; + Ok(()) + } + + pub(crate) async fn buffer_flush(&self, id: u64) -> StreamResult { + Ok(self.buffers.get(id)?.flush().await) + } + + /// Flushes and frees a scope's buffer. Returns the failure that lang didn't get yet, also for + /// values that a free during the flush dropped. + pub(crate) async fn end_buffer(&self, id: u64) -> Option { + let buffer = self.buffers.get(id).ok()?; + let reported = { + let mut state = buffer.state.lock(); + state.closing = true; + state.reported + }; + let flushed = buffer.flush().await; + let _ = self.buffers.free(id); + flushed.failure.filter(|_| !reported) + } + + async fn write_batch(&self, buffer: &Buffer, records: Vec) -> Written { + match buffer.scope.writer(self, &buffer.stream).await { + Ok(writer) => self.writers.write_counted(writer, records).await, + Err(error) => Written { + error: Some(error), + ..Default::default() + }, + } + } +} diff --git a/crates/streams/src/lib.rs b/crates/streams/src/lib.rs index 667a622bf..568676bbe 100644 --- a/crates/streams/src/lib.rs +++ b/crates/streams/src/lib.rs @@ -6,6 +6,7 @@ //! client and calls it with an operation name and proto bytes. Core has no hook for it. A Workflow //! makes each call from a Local Activity. +mod buffer; pub mod digest; pub mod epoch; mod error; diff --git a/crates/streams/src/runtime.rs b/crates/streams/src/runtime.rs index 7db0b58b6..d4c68ae72 100644 --- a/crates/streams/src/runtime.rs +++ b/crates/streams/src/runtime.rs @@ -2,13 +2,16 @@ use crate::{ Reason, StreamError, StreamResult, + buffer::Buffers, memory::MemoryStore, metrics::Metrics, proto::{ - AppendRecordsRequest, CloseScopeRequest, CloseStreamRequest, CloseWriterRequest, - CloseWriterResponse, CloseWritersRequest, CloseWritersResponse, DescribeStreamRequest, - FreeReaderRequest, FreeResponse, FreeScopeRequest, FreeWriterRequest, FreeWriterResponse, - MemoryStreamConfig, OpenReaderResponse, OpenWriterRequest, OpenWriterResponse, + AppendRecordsRequest, BufferFlushRequest, BufferNextBatchRequest, BufferPushRequest, + BufferPushResponse, BufferWriteResponse, CloseScopeRequest, CloseStreamRequest, + CloseWriterRequest, CloseWriterResponse, CloseWritersRequest, CloseWritersResponse, + DescribeStreamRequest, FreeBufferRequest, FreeReaderRequest, FreeResponse, + FreeScopeRequest, FreeWriterRequest, FreeWriterResponse, MemoryStreamConfig, + OpenBufferResponse, OpenReaderResponse, OpenWriterRequest, OpenWriterResponse, ReadRecordsRequest, ShutdownRequest, WorkflowAppendRequest, WriteRecordsRequest, WriteRecordsResponse, }, @@ -22,7 +25,7 @@ use crate::{ use prost::Message; use std::{ sync::{ - Arc, + Arc, OnceLock, atomic::{AtomicBool, Ordering}, }, time::Duration, @@ -62,6 +65,10 @@ pub(crate) struct Shared { pub(crate) scopes: Scopes, pub(crate) pending: Pending, pub(crate) readers: Readers, + pub(crate) buffers: Buffers, + /// The Tokio runtime of the build or of the first asynchronous call, for the writes of an + /// encoded buffer. + handle: OnceLock, shut_down: AtomicBool, /// Set when a `Shutdown` ends every scope and writer. Then the `Free` methods do nothing. finished: AtomicBool, @@ -97,6 +104,19 @@ impl Shared { }) .await } + + /// The Tokio runtime to spawn on. The runtime's own comes first, since it lives as long as the + /// runtime. + pub(crate) fn tokio_handle(&self) -> Option { + (self.handle.get().cloned()).or_else(|| tokio::runtime::Handle::try_current().ok()) + } +} + +impl Drop for Shared { + fn drop(&mut self) { + // The pumps of encoded buffers hold no reference to the runtime, so this stops them. + self.buffers.free_all(); + } } /// The Stream runtime on one Store. @@ -138,6 +158,10 @@ impl StreamRuntime { scopes: Scopes::default(), pending: Pending::default(), readers: Readers::default(), + buffers: Buffers::default(), + handle: tokio::runtime::Handle::try_current() + .map(OnceLock::from) + .unwrap_or_default(), shut_down: AtomicBool::new(false), finished: AtomicBool::new(false), }); @@ -183,6 +207,12 @@ impl StreamRuntime { /// with [StreamError::code] and [StreamError::details]. pub async fn call(&self, method: &str, input: &[u8]) -> StreamResult> { self.check_open(method)?; + if let (None, Ok(handle)) = ( + self.shared.handle.get(), + tokio::runtime::Handle::try_current(), + ) { + let _ = self.shared.handle.set(handle); + } // 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; @@ -273,6 +303,24 @@ impl StreamRuntime { .encode_to_vec(), "WorkflowClose" => shared.workflow_close(decode(input)?).await?.encode_to_vec(), "ReadNext" => shared.read_next(decode(input)?).await?.encode_to_vec(), + "BufferNextBatch" => { + let request = decode::(input)?; + shared + .buffer_next_batch(request.buffer_id) + .await? + .encode_to_vec() + } + "BufferWrite" => { + shared.buffer_write(decode(input)?).await?; + BufferWriteResponse {}.encode_to_vec() + } + "BufferFlush" => { + let request = decode::(input)?; + shared + .buffer_flush(request.buffer_id) + .await? + .encode_to_vec() + } "Shutdown" => { decode::(input)?; shared.shut_down.store(true, Ordering::SeqCst); @@ -320,6 +368,20 @@ impl StreamRuntime { let reader_id = shared.open_reader(decode(input)?)?; Ok(OpenReaderResponse { reader_id }.encode_to_vec()) } + "OpenBuffer" => { + let buffer_id = shared.open_buffer(decode(input)?)?; + Ok(OpenBufferResponse { buffer_id }.encode_to_vec()) + } + "BufferPush" => { + let request = decode::(input)?; + shared.buffer_push(request.buffer_id, request.records)?; + Ok(BufferPushResponse {}.encode_to_vec()) + } + "FreeBuffer" => { + let request = decode::(input)?; + shared.buffers.free(request.buffer_id)?; + Ok(FreeResponse {}.encode_to_vec()) + } "FreeReader" => { let request = decode::(input)?; shared.readers.free(request.reader_id)?; @@ -334,14 +396,15 @@ impl StreamRuntime { } } - /// Refuses every method after `Shutdown`, except `Shutdown`, the `Free` methods and the static - /// ones. + /// Refuses every method after `Shutdown`, except `Shutdown`, the `Free` methods, the static + /// ones and the pump calls. fn check_open(&self, method: &str) -> StreamResult<()> { - // A dropped `Shutdown` leaves work that the next one does. + // A dropped `Shutdown` leaves work that the next one does. The shutdown flushes the + // buffers, so lang's codec pumps keep running. if !self.shared.is_shut_down() || method.starts_with("Free") || crate::sdk::is_static(method) - || method == "Shutdown" + || matches!(method, "Shutdown" | "BufferNextBatch" | "BufferWrite") { return Ok(()); } diff --git a/crates/streams/src/scope.rs b/crates/streams/src/scope.rs index 6f188658a..a4a7360d5 100644 --- a/crates/streams/src/scope.rs +++ b/crates/streams/src/scope.rs @@ -1,5 +1,5 @@ //! Producer scopes. A scope is the Producer of one Activity attempt or one client, and the writers -//! that write as it, one for each Stream. Lang keeps only the scope id. +//! and buffers that write as it, with one writer for each Stream. Lang keeps only the scope id. use crate::{ Reason, StreamError, StreamResult, derive_producer, ids, @@ -37,6 +37,9 @@ pub(crate) struct Scope { #[derive(Default)] struct ScopeWriters { by_stream: HashMap, + buffers: Vec, + /// Set when the scope's close or free starts. Then the scope takes no new buffer. + closing: bool, /// Set by the close or free of the scope. A writer that opens after it is released at once. ended: bool, } @@ -104,6 +107,33 @@ impl Scope { Ok(id) } + pub(crate) fn add_buffer(&self, id: u64) -> StreamResult<()> { + let mut writers = self.writers.lock(); + if writers.closing { + return Err(unknown_scope()); + } + writers.buffers.push(id); + Ok(()) + } + + pub(crate) fn remove_buffer(&self, id: u64) { + self.writers.lock().buffers.retain(|buffer| *buffer != id); + } + + /// Starts the close of the scope and returns its first buffer, which the caller ends. + fn first_buffer(&self) -> Option { + let mut writers = self.writers.lock(); + writers.closing = true; + writers.buffers.first().copied() + } + + /// Starts the close of the scope and returns its buffers, which the caller frees. + fn take_buffers(&self) -> Vec { + let mut writers = self.writers.lock(); + writers.closing = true; + std::mem::take(&mut writers.buffers) + } + /// Ends the scope and returns its writers, which the caller frees. fn end(&self) -> Vec { let mut writers = self.writers.lock(); @@ -359,6 +389,7 @@ impl Shared { Ok(self.end_scope(&scope, None).await) } + /// Flushes and frees the scope's buffers first, since their writes need its writers. async fn end_scope( &self, scope: &Scope, @@ -371,6 +402,14 @@ impl Shared { armed: true, }; let mut response = CloseScopeResponse::default(); + while let Some(buffer) = scope.first_buffer() { + let failed = self.end_buffer(buffer).await; + scope.remove_buffer(buffer); + if let Some(failure) = failed { + response.failed_buffers += 1; + response.failure.get_or_insert(failure); + } + } while let Some(writer) = scope.take_writer() { if let Err(error) = self.writers.close(writer, Holder::Scope).await { response.failed_writers += 1; @@ -389,6 +428,10 @@ impl Shared { /// Releases what the scope holds, with no wait. fn release(&self, scope: &Scope) { + for buffer in scope.take_buffers() { + // Lang may have freed the buffer already. + let _ = self.buffers.free(buffer); + } for writer in scope.end() { let _ = self.writers.free(writer, Holder::Scope); } @@ -400,7 +443,9 @@ impl Shared { // One scope at a time, so that a dropped shutdown leaves the rest for the next one. while let Some((id, scope)) = self.scopes.take_any() { let requeue = Some((id, scope.clone())); - response.failed_writers += self.end_scope(&scope, requeue).await.failed_writers; + let closed = self.end_scope(&scope, requeue).await; + response.failed_writers += closed.failed_writers; + response.failed_buffers += closed.failed_buffers; } response.failed_writers += self.writers.close_all().await as i32; response @@ -415,7 +460,8 @@ mod tests { memory::MemoryStore, proto::{ ActivityContext, AppendRecord, CloseScopeRequest, CloseScopeResponse, - CreateStreamRequest, FreeScopeRequest, ShutdownRequest, + CreateStreamRequest, FreeBufferRequest, FreeScopeRequest, OpenBufferRequest, + OpenBufferResponse, ShutdownRequest, }, testing::support::Flaky, }; @@ -641,6 +687,35 @@ mod tests { } } + #[tokio::test] + async fn a_freed_buffer_leaves_its_scope() { + let (_, ext) = flaky().await; + let open = OpenClientScopeRequest::default().encode_to_vec(); + let scope_id = + OpenScopeResponse::decode(ext.call_sync("OpenClientScope", &open).unwrap().as_slice()) + .unwrap() + .scope_id; + for _ in 0..3 { + let open = OpenBufferRequest { + scope_id, + stream: stream("a"), + ..Default::default() + }; + let opened = OpenBufferResponse::decode( + ext.call_sync("OpenBuffer", &open.encode_to_vec()) + .unwrap() + .as_slice(), + ) + .unwrap(); + let free = FreeBufferRequest { + buffer_id: opened.buffer_id, + }; + ext.call_sync("FreeBuffer", &free.encode_to_vec()).unwrap(); + } + let scope = ext.shared().scopes.get(scope_id).unwrap(); + assert!(scope.writers.lock().buffers.is_empty()); + } + #[tokio::test] async fn a_dropped_activity_scope_open_leaves_no_scope() { let (store, ext) = flaky().await; diff --git a/crates/streams/src/testing/buffer.rs b/crates/streams/src/testing/buffer.rs new file mode 100644 index 000000000..21f9e5b95 --- /dev/null +++ b/crates/streams/src/testing/buffer.rs @@ -0,0 +1,792 @@ +//! Buffered publishers, with and without lang's codec pump, through the runtime's byte call, as +//! lang sends them. +//! +//! Every Store runs the suite through its own runtime. + +use super::{ + case, + support::{Flaky, Needs, stores_for}, +}; +use crate::{ + MemoryStreamRuntime, StreamError, StreamRuntime, + proto::{ + AppendRecord, BufferBatch, BufferFlushRequest, BufferFlushResponse, BufferNextBatchRequest, + BufferPushRequest, BufferPushResponse, BufferRecords, BufferWriteRequest, + BufferWriteResponse, CloseScopeRequest, CloseScopeResponse, CloseStreamAsRequest, + CloseStreamRequest, CloseStreamResponse, CreateStreamRequest, CreateStreamResponse, + FreeBufferRequest, FreeResponse, OpenBufferRequest, OpenBufferResponse, + OpenClientScopeRequest, OpenScopeResponse, ReadRecordsRequest, ReadRecordsResponse, + ShutdownRequest, ShutdownResponse, StreamErrorReason, StreamRef, + buffer_write_request::Body, stream_record::Item, + }, +}; +use prost::Message; +use std::{ + sync::{ + Arc, + atomic::{AtomicI64, AtomicUsize, Ordering}, + }, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, +}; +use temporalio_common::protos::temporal::api::common::v1::Payload; + +static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + +#[derive(Clone)] +struct Case { + ext: Arc, + namespace: String, + stream: String, +} + +async fn case_on(ext: Arc) -> Case { + 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, + namespace: format!("ns-{now}-{n}"), + stream: format!("s-{now}-{n}"), + }; + let _: CreateStreamResponse = case + .call( + "CreateStream", + CreateStreamRequest { + namespace: case.namespace.clone(), + stream_id: case.stream.clone(), + ..Default::default() + }, + ) + .await + .unwrap(); + case +} + +async fn cases() -> Vec { + let mut cases = vec![]; + for ext in super::support::stream_runtimes() { + cases.push(case_on(ext).await); + } + cases +} + +fn millis(ms: u64) -> Option { + let duration = Duration::from_millis(ms); + Some(prost_types::Duration { + seconds: duration.as_secs() as i64, + nanos: duration.subsec_nanos() as i32, + }) +} + +fn records(texts: &[&str]) -> Vec { + texts + .iter() + .map(|text| AppendRecord { + body: Some(Payload { + data: text.as_bytes().to_vec(), + ..Default::default() + }), + }) + .collect() +} + +impl Case { + async fn call( + &self, + method: &str, + request: Req, + ) -> Result { + let bytes = self.ext.call(method, &request.encode_to_vec()).await?; + Ok(Resp::decode(bytes.as_slice()).unwrap()) + } + + fn call_sync( + &self, + method: &str, + request: Req, + ) -> Result { + let bytes = self.ext.call_sync(method, &request.encode_to_vec())?; + Ok(Resp::decode(bytes.as_slice()).unwrap()) + } + + fn stream_ref(&self) -> StreamRef { + StreamRef { + id: self.stream.clone(), + ..Default::default() + } + } + + fn scope(&self) -> u64 { + let opened: OpenScopeResponse = self + .call_sync( + "OpenClientScope", + OpenClientScopeRequest { + namespace: self.namespace.clone(), + identity: "me".into(), + ..Default::default() + }, + ) + .unwrap(); + opened.scope_id + } + + fn buffer(&self, scope_id: u64, encoded: bool, delay_ms: u64, max_batch: i32) -> u64 { + let opened: OpenBufferResponse = self + .call_sync( + "OpenBuffer", + OpenBufferRequest { + scope_id, + stream: Some(self.stream_ref()), + max_delay: millis(delay_ms), + max_batch, + encoded, + }, + ) + .unwrap(); + opened.buffer_id + } + + fn push(&self, buffer_id: u64, texts: &[&str]) -> Result<(), StreamError> { + let _: BufferPushResponse = self.call_sync( + "BufferPush", + BufferPushRequest { + buffer_id, + records: records(texts), + }, + )?; + Ok(()) + } + + async fn flush(&self, buffer_id: u64) -> BufferFlushResponse { + self.call("BufferFlush", BufferFlushRequest { buffer_id }) + .await + .unwrap() + } + + async fn next_batch(&self, buffer_id: u64) -> BufferBatch { + self.call("BufferNextBatch", BufferNextBatchRequest { buffer_id }) + .await + .unwrap() + } + + async fn write(&self, buffer_id: u64, batch_id: u64, body: Body) { + let _: BufferWriteResponse = self + .call( + "BufferWrite", + BufferWriteRequest { + buffer_id, + batch_id, + body: Some(body), + }, + ) + .await + .unwrap(); + } + + /// Lang's pump, with a codec that writes each body in upper case. Returns the batch count. + async fn pump(&self, buffer_id: u64, batches: Arc) { + loop { + let batch = self.next_batch(buffer_id).await; + if batch.done { + return; + } + batches.fetch_add(1, Ordering::SeqCst); + let mut records = batch.records; + for record in &mut records { + record.body.as_mut().unwrap().data.make_ascii_uppercase(); + } + let body = Body::Records(BufferRecords { records }); + self.write(buffer_id, batch.batch_id, body).await; + } + } + + async fn close_scope(&self, scope_id: u64) -> CloseScopeResponse { + self.call("CloseScope", CloseScopeRequest { scope_id }) + .await + .unwrap() + } + + async fn close_stream(&self) { + let other = self.scope(); + let _: CloseStreamResponse = self + .call( + "CloseStreamAs", + CloseStreamAsRequest { + scope_id: other, + close: Some(CloseStreamRequest { + stream: Some(self.stream_ref()), + ..Default::default() + }), + }, + ) + .await + .unwrap(); + } + + async fn texts(&self) -> Vec { + let read: ReadRecordsResponse = self + .call( + "ReadRecords", + ReadRecordsRequest { + namespace: self.namespace.clone(), + stream: Some(self.stream_ref()), + ..Default::default() + }, + ) + .await + .unwrap(); + read.records + .into_iter() + .filter_map(|record| match record.item { + Some(Item::Data(data)) => Some(String::from_utf8(data.body.unwrap().data).unwrap()), + _ => None, + }) + .collect() + } + + /// The texts once the Stream has `count` of them, within a few seconds. + async fn texts_when(&self, count: usize) -> Vec { + let until = Instant::now() + Duration::from_secs(5); + loop { + let texts = self.texts().await; + if texts.len() >= count || Instant::now() > until { + return texts; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + } +} + +/// An encoded buffer writes its values in one batch at most `max_delay` after the first one, and +/// a flush returns the last offset. +async fn an_encoded_buffer_writes_its_batch_after_the_delay() { + for case in cases().await { + let buffer = case.buffer(case.scope(), true, 200, 0); + case.push(buffer, &["a", "b"]).unwrap(); + case.push(buffer, &["c"]).unwrap(); + assert!(case.texts().await.is_empty()); + assert_eq!(case.texts_when(3).await, ["a", "b", "c"]); + case.push(buffer, &["d"]).unwrap(); + let flushed = case.flush(buffer).await; + assert!(flushed.last_offset > 0 && flushed.failure.is_none()); + assert_eq!(case.texts().await, ["a", "b", "c", "d"]); + } +} + +/// A full batch goes at once, with no wait for the delay. +async fn a_full_batch_goes_at_once() { + for case in cases().await { + let buffer = case.buffer(case.scope(), true, 60_000, 2); + case.push(buffer, &["1", "2", "3"]).unwrap(); + assert_eq!(case.texts_when(2).await[..2], ["1", "2"]); + assert_eq!(case.flush(buffer).await.unwritten, 0); + assert_eq!(case.texts().await, ["1", "2", "3"]); + } +} + +/// Lang's pump runs its codec once on each batch, and a flush makes the batch due at once. +async fn a_pump_runs_the_codec_once_per_batch() { + for case in cases().await { + let buffer = case.buffer(case.scope(), false, 60_000, 0); + let batches = Arc::new(AtomicUsize::new(0)); + let pump = tokio::spawn({ + let (case, batches) = (case.clone(), batches.clone()); + async move { case.pump(buffer, batches).await } + }); + case.push(buffer, &["a", "b", "c"]).unwrap(); + let flushed = case.flush(buffer).await; + assert!(flushed.last_offset > 0); + assert_eq!(case.texts().await, ["A", "B", "C"]); + assert_eq!(batches.load(Ordering::SeqCst), 1); + let _: FreeResponse = case + .call_sync("FreeBuffer", FreeBufferRequest { buffer_id: buffer }) + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), pump) + .await + .unwrap() + .unwrap(); + } +} + +/// After a failed batch, the buffer writes nothing more. Each later push fails, and a flush gives +/// the failure and the unwritten count. +async fn a_failed_batch_fails_every_later_call() { + for case in cases().await { + let buffer = case.buffer(case.scope(), true, 50, 0); + case.close_stream().await; + case.push(buffer, &["a", "b"]).unwrap(); + let flushed = case.flush(buffer).await; + assert_eq!(flushed.failure.unwrap().reason(), StreamErrorReason::Closed); + assert_eq!((flushed.unwritten, flushed.last_offset), (2, 0)); + let error = case.push(buffer, &["c"]).unwrap_err(); + assert_eq!(error.reason.token(), "stream_closed"); + } +} + +/// A codec failure fails the buffer with its message. +async fn a_codec_failure_fails_the_buffer() { + for case in cases().await { + let buffer = case.buffer(case.scope(), false, 60_000, 0); + case.push(buffer, &["a"]).unwrap(); + let flush = case.flush(buffer); + let pump = async { + let batch = case.next_batch(buffer).await; + let body = Body::CodecError("the key is gone".into()); + case.write(buffer, batch.batch_id, body).await; + assert!(case.next_batch(buffer).await.done); + }; + let (flushed, ()) = tokio::join!(flush, pump); + let failure = flushed.failure.unwrap(); + assert_eq!(failure.reason(), StreamErrorReason::Invalid); + assert!(failure.message.contains("the key is gone")); + assert!(case.texts().await.is_empty()); + } +} + +/// The close of a scope writes its buffers' values first. It reports a buffer failure that no +/// push or flush gave lang yet. +async fn a_scope_close_flushes_its_buffers() { + for case in cases().await { + let scope = case.scope(); + let buffer = case.buffer(scope, true, 60_000, 0); + case.push(buffer, &["a", "b"]).unwrap(); + let closed = case.close_scope(scope).await; + assert_eq!(closed, CloseScopeResponse::default()); + assert_eq!(case.texts().await, ["a", "b"]); + let error = case.push(buffer, &["c"]).unwrap_err(); + assert_eq!(error.reason.token(), "stream_unknown_writer"); + + let scope = case.scope(); + let buffer = case.buffer(scope, true, 60_000, 0); + case.push(buffer, &["late"]).unwrap(); + case.close_stream().await; + let closed = case.close_scope(scope).await; + assert_eq!(closed.failed_buffers, 1); + assert_eq!(closed.failure.unwrap().reason(), StreamErrorReason::Closed); + } +} + +/// A buffer failure that lang already got doesn't fail the scope's close again. +async fn a_reported_buffer_failure_doesnt_count_at_the_scope_close() { + for case in cases().await { + let scope = case.scope(); + let buffer = case.buffer(scope, true, 50, 0); + case.close_stream().await; + case.push(buffer, &["a"]).unwrap(); + assert!(case.flush(buffer).await.failure.is_some()); + assert_eq!(case.close_scope(scope).await.failed_buffers, 0); + } +} + +/// The shutdown writes every buffer's values. A second runtime on the same Store reads them, since +/// the first one refuses every call after it. +async fn the_shutdown_writes_every_buffer() { + for target in stores_for(Needs::CALLS).await { + let case = case_on(MemoryStreamRuntime::with_store(target.store.clone(), None)).await; + let buffer = case.buffer(case.scope(), true, 60_000, 0); + case.push(buffer, &["a"]).unwrap(); + let shut: ShutdownResponse = case.call("Shutdown", ShutdownRequest {}).await.unwrap(); + assert_eq!((shut.failed_writers, shut.failed_buffers), (0, 0)); + let reader = Case { + ext: MemoryStreamRuntime::with_store(target.store, None), + ..case + }; + assert_eq!(reader.texts().await, ["a"]); + } +} + +/// The runtime refuses a pump call on an encoded buffer, and a write of a batch that isn't in +/// flight or that a write already took. +async fn the_pump_calls_are_checked() { + for case in cases().await { + let scope = case.scope(); + let encoded = case.buffer(scope, true, 60_000, 0); + let next = case.call::<_, BufferBatch>( + "BufferNextBatch", + BufferNextBatchRequest { buffer_id: encoded }, + ); + let error = tokio::time::timeout(Duration::from_secs(5), next).await; + assert_eq!(error.unwrap().unwrap_err().reason.token(), "stream_invalid"); + let buffer = case.buffer(scope, false, 0, 0); + case.push(buffer, &["a"]).unwrap(); + let batch = case.next_batch(buffer).await; + let write = |batch_id| BufferWriteRequest { + buffer_id: buffer, + batch_id, + body: Some(Body::Records(BufferRecords { + records: records(&["a"]), + })), + }; + let wrong = case + .call::<_, BufferWriteResponse>("BufferWrite", write(batch.batch_id + 1)) + .await; + assert_eq!(wrong.unwrap_err().reason.token(), "stream_invalid"); + let _: BufferWriteResponse = case + .call("BufferWrite", write(batch.batch_id)) + .await + .unwrap(); + let again = case + .call::<_, BufferWriteResponse>("BufferWrite", write(batch.batch_id)) + .await; + assert_eq!(again.unwrap_err().reason.token(), "stream_invalid"); + assert_eq!(case.texts().await, ["a"]); + } +} + +/// The shutdown waits for lang's pump to write the values of its buffer. +async fn the_shutdown_waits_for_lang_s_pump() { + for target in stores_for(Needs::CALLS).await { + let case = case_on(MemoryStreamRuntime::with_store(target.store.clone(), None)).await; + let buffer = case.buffer(case.scope(), false, 60_000, 0); + let pump = tokio::spawn({ + let case = case.clone(); + async move { case.pump(buffer, Arc::default()).await } + }); + case.push(buffer, &["a"]).unwrap(); + let shutdown = case.call::<_, ShutdownResponse>("Shutdown", ShutdownRequest {}); + let shut = tokio::time::timeout(Duration::from_secs(5), shutdown).await; + assert_eq!(shut.unwrap().unwrap().failed_buffers, 0); + tokio::time::timeout(Duration::from_secs(5), pump) + .await + .unwrap() + .unwrap(); + let reader = Case { + ext: MemoryStreamRuntime::with_store(target.store, None), + ..case + }; + assert_eq!(reader.texts().await, ["A"]); + } +} + +/// A `BufferWrite` of the one value `a`, for batch `batch_id`. +fn write_a(buffer_id: u64, batch_id: u64) -> BufferWriteRequest { + BufferWriteRequest { + buffer_id, + batch_id, + body: Some(Body::Records(BufferRecords { + records: records(&["a"]), + })), + } +} + +/// A `BufferWrite` that a cancel drops while the Store is down still writes its batch once the +/// Store is back, so a flush doesn't wait for it forever. A freed buffer gives its pump `done`. +async fn a_dropped_write_still_settles_its_batch() { + for (store, case) in flaky_cases(-1).await { + let buffer = case.buffer(case.scope(), false, 0, 0); + case.push(buffer, &["a"]).unwrap(); + let batch = case.next_batch(buffer).await; + store.down.store(true, Ordering::SeqCst); + let write = + case.call::<_, BufferWriteResponse>("BufferWrite", write_a(buffer, batch.batch_id)); + assert!( + tokio::time::timeout(Duration::from_millis(200), write) + .await + .is_err() + ); + store.down.store(false, Ordering::SeqCst); + let flushed = tokio::time::timeout(Duration::from_secs(5), case.flush(buffer)).await; + assert!(flushed.unwrap().last_offset > 0); + assert_eq!(case.texts().await, ["a"]); + let _: FreeResponse = case + .call_sync("FreeBuffer", FreeBufferRequest { buffer_id: buffer }) + .unwrap(); + assert!(case.next_batch(buffer).await.done); + } +} + +/// A write that panics, with its `BufferWrite` dropped, still fails its batch, so a flush ends. +async fn a_panicking_write_fails_its_batch() { + for (store, case) in flaky_cases(-1).await { + let buffer = case.buffer(case.scope(), false, 0, 0); + case.push(buffer, &["a"]).unwrap(); + let batch = case.next_batch(buffer).await; + store.panic_next.store(true, Ordering::SeqCst); + let write = + case.call::<_, BufferWriteResponse>("BufferWrite", write_a(buffer, batch.batch_id)); + // The write's first poll spawns the write task, and the select drops the call at once. + tokio::select! { + biased; + _ = write => panic!("the write answered in its first poll"), + _ = std::future::ready(()) => {} + } + let flushed = tokio::time::timeout(Duration::from_secs(5), case.flush(buffer)).await; + let flushed = flushed.unwrap(); + assert_eq!( + flushed.failure.unwrap().reason(), + StreamErrorReason::Storage + ); + assert_eq!(flushed.unwritten, 1); + } +} + +/// The runtime's own pump fails a batch whose write panics, so a flush ends. +async fn an_encoded_pump_whose_write_panics_fails_its_batch() { + for (store, case) in flaky_cases(-1).await { + let buffer = case.buffer(case.scope(), true, 60_000, 0); + store.panic_next.store(true, Ordering::SeqCst); + case.push(buffer, &["a"]).unwrap(); + let flushed = tokio::time::timeout(Duration::from_secs(5), case.flush(buffer)).await; + let flushed = flushed.unwrap(); + assert_eq!( + flushed.failure.unwrap().reason(), + StreamErrorReason::Storage + ); + assert_eq!(flushed.unwritten, 1); + } +} + +/// A free still works while a shutdown runs, so lang can end a buffer whose pump is gone. +async fn a_free_ends_a_shutdown_whose_pump_is_gone() { + for (_, case) in flaky_cases(-1).await { + let buffer = case.buffer(case.scope(), false, 60_000, 0); + case.push(buffer, &["a"]).unwrap(); + let shutdown = tokio::spawn({ + let case = case.clone(); + async move { + case.call::<_, ShutdownResponse>("Shutdown", ShutdownRequest {}) + .await + } + }); + tokio::time::sleep(Duration::from_millis(100)).await; + let _: FreeResponse = case + .call_sync("FreeBuffer", FreeBufferRequest { buffer_id: buffer }) + .unwrap(); + let shut = tokio::time::timeout(Duration::from_secs(5), shutdown).await; + let shut = shut.unwrap().unwrap().unwrap(); + assert_eq!(shut.failed_buffers, 1); + } +} + +/// A runtime on each backend under test, on a fault-injecting Store over the backend's Store that +/// takes `appends` appends before it refuses the later ones, with the case's Stream created. +async fn flaky_cases(appends: i64) -> Vec<(Arc, Case)> { + let mut cases = vec![]; + for target in stores_for(Needs::CALLS).await { + let store = Arc::new(Flaky { + appends_left: AtomicI64::new(appends), + ..Flaky::new(target.store) + }); + let case = case_on(MemoryStreamRuntime::with_store(store.clone(), None)).await; + cases.push((store, case)); + } + cases +} + +/// A batch that the writer splits, and whose second part fails, counts only that part as +/// unwritten, and gives the first part's last offset. +async fn a_partly_written_batch_counts_only_the_rest() { + for (_, case) in flaky_cases(1).await { + let buffer = case.buffer(case.scope(), true, 60_000, 0); + let big = "x".repeat(1_400_000); + case.push(buffer, &[&big, &big, &big]).unwrap(); + let flushed = case.flush(buffer).await; + assert_eq!(flushed.failure.unwrap().reason(), StreamErrorReason::Closed); + assert_eq!(flushed.unwritten, 1); + assert!(flushed.last_offset > 0); + assert_eq!(case.texts().await.len(), 2); + } +} + +/// A dropped scope close frees the buffers that it didn't reach. +async fn a_dropped_scope_close_frees_its_buffers() { + for (store, case) in flaky_cases(-1).await { + let scope = case.scope(); + let (first, second) = ( + case.buffer(scope, true, 60_000, 0), + case.buffer(scope, true, 60_000, 0), + ); + case.push(first, &["a"]).unwrap(); + case.push(second, &["b"]).unwrap(); + // The first buffer's write retries while the Store is down, so the close waits on it. + store.down.store(true, Ordering::SeqCst); + let close = + case.call::<_, CloseScopeResponse>("CloseScope", CloseScopeRequest { scope_id: scope }); + assert!( + tokio::time::timeout(Duration::from_millis(300), close) + .await + .is_err() + ); + for buffer in [first, second] { + let error = case.push(buffer, &["c"]).unwrap_err(); + assert_eq!(error.reason.token(), "stream_unknown_writer"); + } + store.down.store(false, Ordering::SeqCst); + } +} + +/// A huge delay counts as a year, so the batch math doesn't overflow and a flush still writes. +async fn a_huge_delay_doesnt_overflow() { + for case in cases().await { + let opened: OpenBufferResponse = case + .call_sync( + "OpenBuffer", + OpenBufferRequest { + scope_id: case.scope(), + stream: Some(case.stream_ref()), + max_delay: Some(prost_types::Duration { + seconds: i64::MAX, + nanos: 0, + }), + encoded: true, + ..Default::default() + }, + ) + .unwrap(); + case.push(opened.buffer_id, &["a"]).unwrap(); + let flushed = tokio::time::timeout(Duration::from_secs(5), case.flush(opened.buffer_id)); + assert!(flushed.await.unwrap().last_offset > 0); + } +} + +/// A batch that a dropped pump took comes again on the next `BufferNextBatch`, with its id and +/// values. +async fn a_batch_that_a_dropped_pump_took_comes_again() { + for case in cases().await { + let buffer = case.buffer(case.scope(), false, 0, 0); + case.push(buffer, &["a"]).unwrap(); + let taken = case.next_batch(buffer).await; + let again = tokio::time::timeout(Duration::from_secs(5), case.next_batch(buffer)).await; + let again = again.unwrap(); + assert_eq!( + (again.batch_id, again.records), + (taken.batch_id, taken.records) + ); + let records = records(&["a"]); + let body = Body::Records(BufferRecords { records }); + case.write(buffer, taken.batch_id, body).await; + assert!(case.flush(buffer).await.last_offset > 0); + assert_eq!(case.texts().await, ["a"]); + } +} + +/// A write must give one Record for each value of its batch. Another count fails the buffer, so a +/// flush doesn't wait for that batch. +async fn a_write_with_another_count_fails_the_buffer() { + for case in cases().await { + let buffer = case.buffer(case.scope(), false, 0, 0); + case.push(buffer, &["a"]).unwrap(); + let batch = case.next_batch(buffer).await; + let write = |texts: &[&str]| BufferWriteRequest { + buffer_id: buffer, + batch_id: batch.batch_id, + body: Some(Body::Records(BufferRecords { + records: records(texts), + })), + }; + let error = case + .call::<_, BufferWriteResponse>("BufferWrite", write(&["a", "b"])) + .await + .unwrap_err(); + assert_eq!(error.reason.token(), "stream_invalid"); + let flushed = tokio::time::timeout(Duration::from_secs(5), case.flush(buffer)).await; + let flushed = flushed.unwrap(); + assert_eq!( + flushed.failure.unwrap().reason(), + StreamErrorReason::Invalid + ); + assert_eq!(flushed.unwritten, 1); + assert!(case.texts().await.is_empty()); + } +} + +/// A push fails once the scope's close flushes the buffer, so the close drops no value. +async fn a_push_during_the_scope_close_fails() { + for (store, case) in flaky_cases(-1).await { + let scope = case.scope(); + let buffer = case.buffer(scope, true, 60_000, 0); + case.push(buffer, &["a"]).unwrap(); + store.down.store(true, Ordering::SeqCst); + let close = tokio::spawn({ + let case = case.clone(); + async move { + case.call::<_, CloseScopeResponse>( + "CloseScope", + CloseScopeRequest { scope_id: scope }, + ) + .await + } + }); + tokio::time::sleep(Duration::from_millis(100)).await; + let error = case.push(buffer, &["b"]).unwrap_err(); + assert_eq!(error.reason.token(), "stream_unknown_writer"); + store.down.store(false, Ordering::SeqCst); + let closed = close.await.unwrap().unwrap(); + assert_eq!(closed, CloseScopeResponse::default()); + assert_eq!(case.texts().await, ["a"]); + } +} + +/// A dropped shutdown puts back the scope that it was closing, and the next one writes its +/// buffer's values. +async fn a_dropped_shutdown_writes_the_buffers_next_time() { + for (store, case) in flaky_cases(-1).await { + // One value a batch, so that `b` waits behind `a` while the Store is down. + let buffer = case.buffer(case.scope(), true, 60_000, 1); + store.down.store(true, Ordering::SeqCst); + case.push(buffer, &["a", "b"]).unwrap(); + let shutdown = case.call::<_, ShutdownResponse>("Shutdown", ShutdownRequest {}); + assert!( + tokio::time::timeout(Duration::from_millis(300), shutdown) + .await + .is_err() + ); + store.down.store(false, Ordering::SeqCst); + let shut: ShutdownResponse = case.call("Shutdown", ShutdownRequest {}).await.unwrap(); + assert_eq!((shut.failed_writers, shut.failed_buffers), (0, 0)); + let reader = Case { + ext: MemoryStreamRuntime::with_store(store, None), + ..case + }; + assert_eq!(reader.texts().await, ["a", "b"]); + } +} + +/// A free counts a batch that lang's pump took but didn't write as unwritten. +async fn a_free_counts_a_taken_batch_as_unwritten() { + for case in cases().await { + let buffer = case.buffer(case.scope(), false, 0, 0); + case.push(buffer, &["a", "b"]).unwrap(); + case.next_batch(buffer).await; + let flush = tokio::spawn({ + let case = case.clone(); + async move { case.flush(buffer).await } + }); + tokio::time::sleep(Duration::from_millis(50)).await; + let _: FreeResponse = case + .call_sync("FreeBuffer", FreeBufferRequest { buffer_id: buffer }) + .unwrap(); + let flushed = tokio::time::timeout(Duration::from_secs(5), flush).await; + let flushed = flushed.unwrap().unwrap(); + assert_eq!(flushed.unwritten, 2); + assert!(flushed.failure.is_some()); + } +} + +/// The cases of the suite, in file order. +pub const CASES: &[super::Case] = &[ + case!(an_encoded_buffer_writes_its_batch_after_the_delay), + case!(a_full_batch_goes_at_once), + case!(a_pump_runs_the_codec_once_per_batch), + case!(a_failed_batch_fails_every_later_call), + case!(a_codec_failure_fails_the_buffer), + case!(a_scope_close_flushes_its_buffers), + case!(a_reported_buffer_failure_doesnt_count_at_the_scope_close), + case!(the_shutdown_writes_every_buffer), + case!(the_pump_calls_are_checked), + case!(the_shutdown_waits_for_lang_s_pump), + case!(a_dropped_write_still_settles_its_batch), + case!(a_partly_written_batch_counts_only_the_rest), + case!(a_dropped_scope_close_frees_its_buffers), + case!(a_huge_delay_doesnt_overflow), + case!(a_batch_that_a_dropped_pump_took_comes_again), + case!(a_write_with_another_count_fails_the_buffer), + case!(a_push_during_the_scope_close_fails), + case!(a_dropped_shutdown_writes_the_buffers_next_time), + case!(a_free_counts_a_taken_batch_as_unwritten), + case!(a_panicking_write_fails_its_batch), + case!(a_free_ends_a_shutdown_whose_pump_is_gone), + case!(an_encoded_pump_whose_write_panics_fails_its_batch), +]; diff --git a/crates/streams/src/testing/mod.rs b/crates/streams/src/testing/mod.rs index e37ff7106..9678da858 100644 --- a/crates/streams/src/testing/mod.rs +++ b/crates/streams/src/testing/mod.rs @@ -3,6 +3,7 @@ pub(crate) mod support; +pub mod buffer; pub mod reader; pub mod scope; pub mod service; @@ -176,4 +177,9 @@ mod tests { async fn the_reader_suite_passes_on_the_memory_store() { run(vec![Arc::new(MemoryBackend)], reader::CASES).await; } + + #[tokio::test] + async fn the_buffer_suite_passes_on_the_memory_store() { + run(vec![Arc::new(MemoryBackend)], buffer::CASES).await; + } } diff --git a/crates/streams/src/testing/support.rs b/crates/streams/src/testing/support.rs index ba57b6cf4..a92e0d355 100644 --- a/crates/streams/src/testing/support.rs +++ b/crates/streams/src/testing/support.rs @@ -5,7 +5,7 @@ use super::{RawKeys, backends}; use crate::{ - StreamError, StreamResult, StreamRuntime, + Reason, StreamError, StreamResult, StreamRuntime, digest::batch_digest, epoch::withdraws, ids::local_activity_producer_id, @@ -415,6 +415,11 @@ pub(crate) struct Flaky { pub(crate) down: AtomicBool, /// The appends that it takes before it goes down. Below zero is never. pub(crate) down_after: AtomicI64, + /// The appends that it takes before it refuses every later one as `CLOSED`. Below zero is no + /// limit. + pub(crate) appends_left: AtomicI64, + /// Panics in the next append, as a bug in a Store would. + pub(crate) panic_next: AtomicBool, } impl Flaky { @@ -425,6 +430,8 @@ impl Flaky { lose_next: AtomicBool::new(false), down: AtomicBool::new(false), down_after: AtomicI64::new(-1), + appends_left: AtomicI64::new(-1), + panic_next: AtomicBool::new(false), } } } @@ -458,12 +465,20 @@ impl Store for Flaky { stream: &StreamRef, batch: &Batch<'_>, ) -> StreamResult { + assert!( + !self.panic_next.swap(false, Ordering::SeqCst), + "the Store panics" + ); if self.down_after.fetch_sub(1, Ordering::SeqCst) == 0 { self.down.store(true, Ordering::SeqCst); } if self.down.load(Ordering::SeqCst) { return Err(StreamError::storage("the Store is down")); } + if self.appends_left.fetch_sub(1, Ordering::SeqCst) == 0 { + self.appends_left.store(0, Ordering::SeqCst); + return Err(StreamError::new(Reason::Closed, "the Store refuses")); + } let appended = self.inner.append(namespace, stream, batch).await?; if self.lose_next.swap(false, Ordering::SeqCst) { return Err(StreamError::storage("the answer was lost")); diff --git a/crates/streams/src/writer.rs b/crates/streams/src/writer.rs index 5c8a9ae0f..836b8cf03 100644 --- a/crates/streams/src/writer.rs +++ b/crates/streams/src/writer.rs @@ -88,42 +88,46 @@ impl Writer { /// Appends `records` now and returns the offsets of the first and the last one. Records past /// the Store's batch limit go in more batches, in order. - async fn write(&self, mut records: Vec) -> StreamResult<(i64, i64)> { + async fn write(&self, mut records: Vec) -> Written { let _writing = self.writing.lock().await; + let mut written = Written::default(); // A close that started while this write waited has ended the writer. if self.state.lock().closed { - return Err(StreamError::new( + written.error = Some(StreamError::new( Reason::UnknownWriter, "The writer closed while the write waited", )); + return written; } - let (mut first, mut last) = (0, 0); - while !records.is_empty() { + while !records.is_empty() && written.error.is_none() { let count = self.batch_len(records.iter()); let rest = records.split_off(count); if let Some(error) = &self.state.lock().failed { - return Err(error.clone()); + written.error = Some(error.clone()); + break; } let sequence = next_sequence(); - let (batch_first, batch_last) = match self.send(records, sequence).await { - Ok(offsets) => offsets, + match self.send(records, sequence).await { + Ok((batch_first, batch_last)) => { + if written.first == 0 { + written.first = batch_first; + } + written.last = batch_last; + written.records += count; + } // The next write takes a higher `sequence`, so the Store refuses a late copy of // this batch after it. - Err(error) if error.reason == Reason::Storage => return Err(error), + Err(error) if error.reason == Reason::Storage => written.error = Some(error), Err(error) => { self.state.lock().failed = Some(error.clone()); - return Err(error); + written.error = Some(error); } - }; - if first == 0 { - first = batch_first; } - last = batch_last; records = rest; } let mut state = self.state.lock(); - state.last_offset = state.last_offset.max(last); - Ok((first, last)) + state.last_offset = state.last_offset.max(written.last); + written } /// Appends one batch and returns its first and last offset. After a Store failure it sends the @@ -176,6 +180,15 @@ fn check_sizes(producer: &Producer, records: &[AppendRecord]) -> StreamResult<() Ok(()) } +/// What a write wrote. A write that failed after some batches has their Records and offsets. +#[derive(Default)] +pub(crate) struct Written { + pub(crate) first: i64, + pub(crate) last: i64, + pub(crate) records: usize, + pub(crate) error: Option, +} + /// The namespace, Stream id, Producer id and epoch key of a writer. type WriterKey = (String, String, String, String); @@ -363,20 +376,38 @@ impl Writers { id: u64, records: Vec, ) -> StreamResult<(i64, i64)> { - let writer = self.get(id)?; + let written = self.write_counted(id, records).await; + match written.error { + Some(error) => Err(error), + None => Ok((written.first, written.last)), + } + } + + /// Writes as [Self::write] does, and also tells how many Records went in before a failure. + pub(crate) async fn write_counted(&self, id: u64, records: Vec) -> Written { + let refused = |error| Written { + error: Some(error), + ..Default::default() + }; + let writer = match self.get(id) { + Ok(writer) => writer, + Err(error) => return refused(error), + }; { let state = writer.state.lock(); if let Some(error) = &state.failed { - return Err(error.clone()); + return refused(error.clone()); } if state.closed { - return Err(StreamError::new( + return refused(StreamError::new( Reason::UnknownWriter, format!("Writer {id} is closed"), )); } } - check_sizes(&writer.producer, &records)?; + if let Err(error) = check_sizes(&writer.producer, &records) { + return refused(error); + } writer.write(records).await }