From 9c2a42791761e4cc97e3474d63c9b0708d84e1a9 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Tue, 20 Oct 2026 21:18:36 +0800 Subject: [PATCH 01/10] Add bounded delivery admission, completions, and lifecycle barriers --- Cargo.toml | 2 - src/completion.rs | 131 ++++++ src/delivery.rs | 86 ++++ src/error.rs | 18 + src/internal/delivery.rs | 893 +++++++++++++++++++++++++++++++++++++++ src/internal/mod.rs | 1 + src/lib.rs | 7 + src/log_queue/config.rs | 17 +- src/log_queue/mod.rs | 1 - src/log_queue/queue.rs | 788 +++++++++------------------------- src/log_queue/worker.rs | 34 -- src/logger.rs | 106 ++--- src/span.rs | 142 +++++-- src/test_utils.rs | 16 +- tests/http_flow.rs | 83 +++- 15 files changed, 1552 insertions(+), 773 deletions(-) create mode 100644 src/completion.rs create mode 100644 src/delivery.rs create mode 100644 src/internal/delivery.rs delete mode 100644 src/log_queue/worker.rs diff --git a/Cargo.toml b/Cargo.toml index db3ccc5..528e49d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,13 +14,11 @@ categories = ["api-bindings", "asynchronous", "development-tools"] anyhow = "1" base64 = "0.23" bon = "3" -arc-swap = "1" async-trait = "0.1" backoff = { version = "0.4", features = ["tokio"] } bytes = "1" futures = "0.3" chrono = { version = "0.4", features = ["serde"] } -crossbeam = "0.8" indexmap = "2" reqwest = { version = "0.13", default-features = false, features = ["json", "multipart", "query", "stream", "rustls"] } serde = { version = "1", features = ["derive"] } diff --git a/src/completion.rs b/src/completion.rs new file mode 100644 index 0000000..89e7ccb --- /dev/null +++ b/src/completion.rs @@ -0,0 +1,131 @@ +//! Observation of work submitted when an SDK method is called. +use crate::ExecutionError; +use std::{ + future::{Future, IntoFuture}, + ops::Deref, + pin::Pin, + task::{Context, Poll}, +}; +use tokio::sync::oneshot; + +/// A synchronously recorded issuance result and an eventual operation result. +/// +/// Inspect `&*completion` to check admission without waiting. Awaiting returns +/// that error, or the accepted operation's result. Dropping this token does not +/// cancel accepted work. No work is started by polling the token. +/// +/// ```no_run +/// use braintrust_sdk_rust::{BraintrustClient, SpanHandle, SpanLog}; +/// use std::time::Duration; +/// +/// # async fn example(client: &BraintrustClient, span: &SpanHandle) +/// # -> Result<(), Box> { +/// let completion = span.log(SpanLog::default()); // admission happens now +/// if let Err(error) = &*completion { +/// eprintln!("submission rejected: {error}"); +/// } +/// completion.await?; // observes admission failure or full batch delivery +/// let report = client.flush().await?; +/// client.shutdown(Duration::from_secs(10)).await?; +/// # Ok(()) +/// # } +/// ``` +pub struct Completion { + issued: Result<(), E>, + receiver: oneshot::Receiver>, +} +impl Completion { + pub(crate) fn pair() -> (oneshot::Sender>, Self) { + let (sender, receiver) = oneshot::channel(); + ( + sender, + Self { + issued: Ok(()), + receiver, + }, + ) + } + pub(crate) fn record_issuance(&mut self, result: Result<(), E>) { + self.issued = result; + } +} +impl Deref for Completion { + type Target = Result<(), E>; + fn deref(&self) -> &Self::Target { + &self.issued + } +} +impl std::fmt::Debug for Completion +where + E: std::fmt::Debug, +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Completion") + .field("issued", &self.issued) + .finish_non_exhaustive() + } +} +/// The future produced by consuming a [`Completion`]. +pub struct CompletionFuture { + issued: Option>, + receiver: oneshot::Receiver>, +} +// Neither field is structurally pinned; the channel receiver is Unpin. +impl Unpin for CompletionFuture {} +impl> IntoFuture for Completion { + type Output = Result; + type IntoFuture = CompletionFuture; + fn into_future(self) -> Self::IntoFuture { + CompletionFuture { + issued: Some(self.issued), + receiver: self.receiver, + } + } +} +impl> Future for CompletionFuture { + type Output = Result; + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.get_mut(); + if let Some(Err(error)) = this.issued.take() { + return Poll::Ready(Err(error)); + } + match Pin::new(&mut this.receiver).poll(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(Ok(result)) => Poll::Ready(result), + Poll::Ready(Err(_)) => Poll::Ready(Err(ExecutionError::WorkerLost.into())), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + #[tokio::test] + async fn issuance_error_needs_neither_clone_nor_producer() { + #[derive(Debug, PartialEq)] + enum Error { + Admission, + Execution, + } + impl From for Error { + fn from(_: ExecutionError) -> Self { + Self::Execution + } + } + let (sender, mut token) = Completion::<(), Error>::pair(); + token.record_issuance(Err(Error::Admission)); + assert_eq!(&*token, &Err(Error::Admission)); + drop(sender); + assert_eq!(token.await, Err(Error::Admission)); + } + #[tokio::test] + async fn successful_issuance_stays_successful_after_operation_failure() { + let (sender, token) = Completion::<(), ExecutionError>::pair(); + sender.send(Err(ExecutionError::TimedOut)).unwrap(); + assert!(token.is_ok()); + assert_eq!(token.await, Err(ExecutionError::TimedOut)); + let (sender, token) = Completion::<(), ExecutionError>::pair(); + drop(sender); + assert_eq!(token.await, Err(ExecutionError::WorkerLost)); + } +} diff --git a/src/delivery.rs b/src/delivery.rs new file mode 100644 index 0000000..d706d65 --- /dev/null +++ b/src/delivery.rs @@ -0,0 +1,86 @@ +//! Outcomes of queued ingestion and its lifecycle operations. +use std::sync::Arc; +use thiserror::Error; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] +pub enum EnqueueError { + #[error("delivery admission capacity is full")] + Full, + #[error("delivery queue is closed")] + Closed, +} +#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] +pub enum ExecutionError { + #[error("delivery worker stopped before completing accepted work")] + WorkerLost, + #[error("delivery shutdown timed out")] + TimedOut, +} +/// A terminal preparation, encoding or upload failure shared by a logical batch. +#[derive(Debug, Error)] +#[error("log delivery failed: {source}")] +pub struct DeliveryError { + #[source] + pub(crate) source: anyhow::Error, +} +impl DeliveryError { + /// Inspect the original preparation, encoding or transport error. + pub fn cause(&self) -> &anyhow::Error { + &self.source + } +} +#[derive(Debug, Clone, Error)] +pub enum WriteError { + #[error(transparent)] + Admission(#[from] EnqueueError), + #[error(transparent)] + Execution(#[from] ExecutionError), + #[error(transparent)] + Delivery(Arc), +} +/// Acknowledgment that the logical batch containing this write succeeded. +/// Backend query visibility is separate from delivery acknowledgment. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DeliveryReceipt; +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct FlushReport { + /// Accepted operations delivered through this flush's cutoff, before merging. + pub delivered: u64, +} +#[derive(Debug, Clone, Error)] +#[error("{failed} accepted writes failed: {cause}")] +pub struct FlushFailure { + pub failed: u64, + pub cause: Arc, +} +#[derive(Debug, Clone, Error)] +pub enum FlushError { + #[error(transparent)] + Admission(#[from] EnqueueError), + #[error(transparent)] + Execution(#[from] ExecutionError), + #[error(transparent)] + Delivery(FlushFailure), +} +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct ShutdownReport { + pub delivered: u64, +} +#[derive(Debug, Clone, Error)] +pub enum ShutdownError { + #[error(transparent)] + Admission(#[from] EnqueueError), + #[error(transparent)] + Execution(#[from] ExecutionError), + #[error(transparent)] + Delivery(FlushFailure), +} +impl From for ShutdownError { + fn from(error: FlushError) -> Self { + match error { + FlushError::Admission(error) => Self::Admission(error), + FlushError::Execution(error) => Self::Execution(error), + FlushError::Delivery(error) => Self::Delivery(error), + } + } +} diff --git a/src/error.rs b/src/error.rs index 9b6729e..7444f6a 100644 --- a/src/error.rs +++ b/src/error.rs @@ -19,3 +19,21 @@ pub enum BraintrustError { } pub type Result = std::result::Result; + +// Compatibility for resource APIs that still return BraintrustError while the +// public operation surfaces migrate to typed completion errors. +impl From for BraintrustError { + fn from(error: crate::FlushError) -> Self { + if let crate::FlushError::Delivery(failure) = &error { + if let Some(Self::Api { status, message }) = + failure.cause.cause().downcast_ref::() + { + return Self::Api { + status: *status, + message: message.clone(), + }; + } + } + Self::Background(error.to_string()) + } +} diff --git a/src/internal/delivery.rs b/src/internal/delivery.rs new file mode 100644 index 0000000..cf42b56 --- /dev/null +++ b/src/internal/delivery.rs @@ -0,0 +1,893 @@ +//! Atomic admission, fixed flush cutoffs and active-batch outcome ownership. +use crate::{ + Completion, DeliveryError, DeliveryReceipt, EnqueueError, ExecutionError, FlushError, + FlushFailure, FlushReport, ShutdownError, ShutdownReport, WriteError, +}; +use futures::{ + future::{BoxFuture, Shared}, + FutureExt, +}; +use std::{ + collections::VecDeque, + sync::{Arc, Mutex, MutexGuard}, + time::Duration, +}; +use tokio::{ + runtime::Handle, + sync::{oneshot, Notify, Semaphore}, + task::AbortHandle, +}; +use tracing::warn; + +type WriteSender = oneshot::Sender>; +type FlushSender = oneshot::Sender>; +#[derive(Clone)] +enum BatchError { + Delivery(Arc), + Execution(ExecutionError), +} +struct Accepted { + command: C, + sender: WriteSender, +} +struct Batch { + commands: Vec, + completion: Arc, +} +struct BatchCompletion { + count: usize, + senders: Mutex>>, +} +fn lock(value: &Mutex) -> MutexGuard<'_, T> { + value + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} +impl BatchCompletion { + fn finish(&self, result: &Result<(), BatchError>) { + if let Some(writes) = lock(&self.senders).take() { + let write_result = match result { + Ok(()) => Ok(DeliveryReceipt), + Err(BatchError::Delivery(error)) => Err(WriteError::Delivery(error.clone())), + Err(BatchError::Execution(error)) => Err(WriteError::Execution(*error)), + }; + for sender in writes { + let _ = sender.send(write_result.clone()); + } + } + } +} +/// One bounded summary per active flush, rather than cloning every covered +/// batch into each observer. Its cutoff is the last batch present at admission. +struct Barrier { + through: Arc, + sender: Option, + delivered: u64, + failure: Option, + shutdown: bool, +} +impl Barrier { + fn record(&mut self, count: u64, result: &Result<(), BatchError>) { + match result { + Ok(()) => self.delivered += count, + Err(BatchError::Delivery(error)) => match &mut self.failure { + Some(failure) => failure.failed += count, + None => { + self.failure = Some(FlushFailure { + failed: count, + cause: error.clone(), + }) + } + }, + Err(BatchError::Execution(error)) => { + if let Some(sender) = self.sender.take() { + let _ = sender.send(Err((*error).into())); + } + } + } + } + fn finish(&mut self) { + if let Some(sender) = self.sender.take() { + let result = match self.failure.take() { + Some(failure) => Err(FlushError::Delivery(failure)), + None => Ok(FlushReport { + delivered: self.delivered, + }), + }; + let _ = sender.send(result); + } + } +} +struct State { + pending: Vec>, + ready: VecDeque>, + active: Option>, + barriers: Vec, + outstanding: usize, + accepting: bool, + terminal: Option, + depth_warned: bool, +} +impl State { + fn seal(&mut self) { + if self.pending.is_empty() { + return; + } + let (commands, writes): (Vec<_>, Vec<_>) = std::mem::take(&mut self.pending) + .into_iter() + .map(|item| (item.command, item.sender)) + .unzip(); + let count = commands.len(); + self.ready.push_back(Batch { + commands, + completion: Arc::new(BatchCompletion { + count, + senders: Mutex::new(Some(writes)), + }), + }); + } + fn barrier(&mut self, sender: FlushSender, shutdown: bool) { + self.seal(); + if let Some(through) = self + .ready + .back() + .map(|batch| batch.completion.clone()) + .or_else(|| self.active.clone()) + { + self.barriers.push(Barrier { + through, + sender: Some(sender), + delivered: 0, + failure: None, + shutdown, + }); + } else { + let _ = sender.send(Ok(FlushReport::default())); + } + } + fn finish(&mut self, batch: &Arc, result: Result<(), BatchError>) { + if self + .active + .as_ref() + .is_some_and(|active| Arc::ptr_eq(active, batch)) + { + self.outstanding -= batch.count; + self.active = None; + if self.outstanding == 0 { + self.depth_warned = false; + } + batch.finish(&result); + // Every registered barrier covers this batch: completed cutoffs + // are removed here, before admitting any subsequent flush. + self.barriers.retain_mut(|barrier| { + barrier.record(batch.count as u64, &result); + if Arc::ptr_eq(&barrier.through, batch) { + barrier.finish(); + false + } else { + true + } + }); + } + } +} +struct Inner { + state: Mutex>, + wake: Notify, + capacity: usize, + control_limit: usize, + shutdown_controls: Arc, + abort: Mutex>, +} +impl Inner { + fn stop(&self, error: ExecutionError) { + let mut state = lock(&self.state); + state.accepting = false; + if state.terminal.is_some() { + return; + } + state.terminal = Some(error); + for accepted in state.pending.drain(..) { + let _ = accepted.sender.send(Err(error.into())); + } + let result = Err(BatchError::Execution(error)); + if let Some(active) = state.active.take() { + active.finish(&result); + } + for batch in state.ready.drain(..) { + batch.completion.finish(&result); + } + for mut barrier in state.barriers.drain(..) { + barrier.record(0, &result); + } + state.outstanding = 0; + self.wake.notify_one(); + } + fn close(&self) { + lock(&self.state).accepting = false; + self.wake.notify_one(); + } +} + +/// Producers own no completed-outcome registry. A batch owns its write senders +/// until all its physical requests finish; barriers summarize only active work. +pub(crate) struct WriteService { + inner: Arc>, + owners: Arc>, + runtime: Handle, + exit: Shared>>, +} +impl Clone for WriteService { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + owners: self.owners.clone(), + runtime: self.runtime.clone(), + exit: self.exit.clone(), + } + } +} +struct ProducerOwner { + inner: std::sync::Weak>, +} +impl Drop for ProducerOwner { + fn drop(&mut self) { + if let Some(inner) = self.inner.upgrade() { + inner.close(); + } + } +} +impl WriteService { + pub fn new(capacity: usize, runtime: Handle, process: F) -> Self + where + F: Fn(Vec) -> BoxFuture<'static, anyhow::Result<()>> + Send + 'static, + { + let control_limit = capacity.clamp(1, 1024); + let inner = Arc::new(Inner { + state: Mutex::new(State { + pending: Vec::new(), + ready: VecDeque::new(), + active: None, + barriers: Vec::new(), + outstanding: 0, + accepting: true, + terminal: None, + depth_warned: false, + }), + wake: Notify::new(), + capacity, + control_limit, + shutdown_controls: Arc::new(Semaphore::new(control_limit)), + abort: Mutex::new(None), + }); + let (sender, receiver) = oneshot::channel(); + let exit = async move { receiver.await.unwrap_or(Err(ExecutionError::WorkerLost)) } + .boxed() + .shared(); + // Construct outside the future so cancellation before its first poll + // still rejects pending work and notifies every accepted observer. + let guard = WorkerGuard { + inner: inner.clone(), + exit: Some(sender), + }; + let task = runtime.spawn(async move { run(guard, process).await }); + *lock(&inner.abort) = Some(task.abort_handle()); + let owners = Arc::new(ProducerOwner { + inner: Arc::downgrade(&inner), + }); + Self { + inner, + owners, + runtime, + exit, + } + } + + #[cfg(test)] + pub fn enqueue(&self, command: C) -> Completion { + self.enqueue_with_rejection(command, |_, _| {}) + } + + pub fn enqueue_with_rejection( + &self, + command: C, + rejected: impl FnOnce(C, EnqueueError), + ) -> Completion { + let (sender, mut completion) = Completion::pair(); + let mut state = lock(&self.inner.state); + let error = if !state.accepting { + Some(EnqueueError::Closed) + } else if state.outstanding >= self.inner.capacity { + Some(EnqueueError::Full) + } else { + None + }; + if let Some(error) = error { + completion.record_issuance(Err(error.into())); + drop(state); + rejected(command, error); + } else { + // Sender registration and publication have one linearization point. + state.pending.push(Accepted { command, sender }); + state.outstanding += 1; + if !state.depth_warned + && state.outstanding >= self.inner.capacity - self.inner.capacity / 5 + { + state.depth_warned = true; + warn!( + outstanding = state.outstanding, + capacity = self.inner.capacity, + "delivery queue has reached 80% of capacity" + ); + } + self.inner.wake.notify_one(); + } + completion + } + + pub fn flush(&self) -> Completion { + let (sender, mut completion) = Completion::pair(); + let mut state = lock(&self.inner.state); + if !state.accepting { + completion.record_issuance(Err(EnqueueError::Closed.into())); + } else if state + .barriers + .iter() + .filter(|barrier| !barrier.shutdown) + .count() + >= self.inner.control_limit + { + completion.record_issuance(Err(EnqueueError::Full.into())); + } else { + state.barrier(sender, false); + self.inner.wake.notify_one(); + } + completion + } + + pub fn shutdown(&self, timeout: Duration) -> Completion { + let (sender, mut completion) = Completion::pair(); + let Ok(permit) = self.inner.shutdown_controls.clone().try_acquire_owned() else { + completion.record_issuance(Err(EnqueueError::Full.into())); + return completion; + }; + let (barrier_sender, barrier) = Completion::pair(); + { + let mut state = lock(&self.inner.state); + if let Some(error) = state.terminal { + completion.record_issuance(Err(error.into())); + return completion; + } + state.accepting = false; + state.barrier(barrier_sender, true); + } + let inner = self.inner.clone(); + let exit = self.exit.clone(); + self.inner.wake.notify_one(); + self.runtime.spawn(async move { + let _permit = permit; + let work = async { + let report = barrier.await; + exit.await.map_err(ShutdownError::Execution)?; + report + .map(|report| ShutdownReport { + delivered: report.delivered, + }) + .map_err(Into::into) + }; + let result = match tokio::time::timeout(timeout, work).await { + Ok(result) => result, + Err(_) => { + inner.stop(ExecutionError::TimedOut); + if let Some(abort) = lock(&inner.abort).as_ref() { + abort.abort(); + } + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + } + }; + let _ = sender.send(result); + }); + completion + } + + #[cfg(test)] + pub fn is_empty(&self) -> bool { + lock(&self.inner.state).outstanding == 0 + } + #[cfg(test)] + fn outstanding(&self) -> usize { + lock(&self.inner.state).outstanding + } +} +struct WorkerGuard { + inner: Arc>, + exit: Option>>, +} +impl Drop for WorkerGuard { + fn drop(&mut self) { + if let Some(exit) = self.exit.take() { + let error = lock(&self.inner.state) + .terminal + .unwrap_or(ExecutionError::WorkerLost); + self.inner.stop(error); + let _ = exit.send(Err(error)); + } + } +} +async fn run(mut guard: WorkerGuard, process: F) +where + C: Send + 'static, + F: Fn(Vec) -> BoxFuture<'static, anyhow::Result<()>>, +{ + loop { + let wake = guard.inner.wake.notified(); + let next = { + let mut state = lock(&guard.inner.state); + state.seal(); + match state.ready.pop_front() { + Some(batch) => { + state.active = Some(batch.completion.clone()); + Some(batch) + } + None if !state.accepting => break, + None => None, + } + }; + if let Some(batch) = next { + let result = process(batch.commands) + .await + .map_err(|source| BatchError::Delivery(Arc::new(DeliveryError { source }))); + if let Err(BatchError::Delivery(error)) = &result { + warn!(%error, "delivery batch failed"); + } + lock(&guard.inner.state).finish(&batch.completion, result); + } else { + wake.await; + } + } + if let Some(exit) = guard.exit.take() { + let _ = exit.send(Ok(())); + } +} +#[cfg(test)] +mod tests { + use super::*; + use std::{ + future::IntoFuture, + sync::atomic::{AtomicBool, Ordering}, + }; + use tokio::sync::mpsc; + + fn gated( + capacity: usize, + ) -> ( + WriteService, + mpsc::UnboundedReceiver>, + Arc, + ) { + let (entered, receiver) = mpsc::unbounded_channel(); + let gate = Arc::new(Semaphore::new(0)); + let process_gate = gate.clone(); + let service = WriteService::new(capacity, Handle::current(), move |commands| { + let gate = process_gate.clone(); + let entered = entered.clone(); + Box::pin(async move { + entered.send(commands).unwrap(); + gate.acquire().await.unwrap().forget(); + Ok(()) + }) + }); + (service, receiver, gate) + } + + #[tokio::test] + async fn concurrent_flushes_observe_fixed_cutoffs_including_active_work() { + let (service, mut entered, gate) = gated(8); + let one = service.enqueue(1); + let first = service.flush(); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + let concurrent = service.flush(); + let two = service.enqueue(2); + let mut second = Box::pin(service.flush().into_future()); + let three = service.enqueue(3); // beyond second's cutoff + gate.add_permits(1); + assert_eq!(first.await.unwrap().delivered, 1); + assert_eq!(concurrent.await.unwrap().delivered, 1); + one.await.unwrap(); + assert_eq!(entered.recv().await.unwrap(), vec![2]); + assert!(futures::poll!(second.as_mut()).is_pending()); + gate.add_permits(1); + assert_eq!(second.await.unwrap().delivered, 2); + two.await.unwrap(); + assert_eq!(entered.recv().await.unwrap(), vec![3]); + // The second flush finished without waiting for the later arrival. + gate.add_permits(1); + three.await.unwrap(); + assert_eq!(service.flush().await.unwrap().delivered, 0); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + } + + #[tokio::test] + async fn a_new_flush_does_not_inherit_a_failed_prefix_from_an_older_active_flush() { + let (entered, mut receiver) = mpsc::unbounded_channel(); + let gate = Arc::new(Semaphore::new(0)); + let process_gate = gate.clone(); + let service = WriteService::new(4, Handle::current(), move |commands: Vec| { + let entered = entered.clone(); + let gate = process_gate.clone(); + Box::pin(async move { + entered.send(commands.clone()).unwrap(); + gate.acquire().await.unwrap().forget(); + if commands == [1] { + anyhow::bail!("first batch failed"); + } + Ok(()) + }) + }); + let first = service.enqueue(1); + assert_eq!(receiver.recv().await.unwrap(), vec![1]); + let second = service.enqueue(2); + let covering = service.flush(); + gate.add_permits(1); + assert!(matches!(first.await, Err(WriteError::Delivery(_)))); + assert_eq!(receiver.recv().await.unwrap(), vec![2]); + let fresh = service.flush(); + gate.add_permits(1); + second.await.unwrap(); + let Err(FlushError::Delivery(failure)) = covering.await else { + panic!("expected captured failure"); + }; + assert_eq!(failure.failed, 1); + assert_eq!(fresh.await.unwrap().delivered, 1); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + } + + #[tokio::test] + async fn graceful_shutdown_reports_its_failed_batch_after_notifying_all_writes() { + let service = WriteService::new(4, Handle::current(), |_: Vec| { + Box::pin(async { + anyhow::bail!("terminal failure"); + }) + }); + let first = service.enqueue(1); + let second = service.enqueue(2); + let shutdown = service.shutdown(Duration::from_secs(1)); + let Err(ShutdownError::Delivery(failure)) = shutdown.await else { + panic!("expected delivery failure"); + }; + assert_eq!(failure.failed, 2); + for completion in [first, second] { + assert!(matches!(completion.await, Err(WriteError::Delivery(_)))); + } + assert!(service.is_empty()); + } + + #[tokio::test] + async fn capacity_includes_in_flight_work_and_rejection_is_observable_twice() { + let (service, mut entered, gate) = gated(1); + let first = service.enqueue(1); + entered.recv().await.unwrap(); + let rejected = service.enqueue(2); + assert!(matches!( + &*rejected, + Err(WriteError::Admission(EnqueueError::Full)) + )); + assert!(matches!( + rejected.await, + Err(WriteError::Admission(EnqueueError::Full)) + )); + assert_eq!(service.outstanding(), 1); + gate.add_permits(1); + first.await.unwrap(); + assert!(service.is_empty()); + gate.add_permits(1); + service.enqueue(3).await.unwrap(); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + } + + #[tokio::test] + async fn dropped_observer_does_not_cancel_automatic_delivery_or_fill_the_queue() { + let (service, mut entered, gate) = gated(2); + drop(service.enqueue(1)); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + let flush = service.flush(); + drop(flush); // lifecycle observers also do not cancel accepted work + gate.add_permits(1); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + assert!(service.is_empty()); + } + + #[tokio::test] + async fn every_failed_batch_notifies_all_members_releases_errors_and_allows_fresh_work() { + let fail = Arc::new(AtomicBool::new(true)); + let process_fail = fail.clone(); + let service = WriteService::new(4, Handle::current(), move |_: Vec| { + let fail = process_fail.load(Ordering::Relaxed); + Box::pin(async move { + if fail { + anyhow::bail!("terminal failure"); + } + Ok(()) + }) + }); + for _ in 0..100 { + let first = service.enqueue(1); + let second = service.enqueue(2); + let flush = service.flush(); + let Err(FlushError::Delivery(failure)) = flush.await else { + panic!("expected failure"); + }; + assert_eq!(failure.failed, 2); + let weak = Arc::downgrade(&failure.cause); + for completion in [first, second] { + let Err(WriteError::Delivery(error)) = completion.await else { + panic!("expected failure"); + }; + assert!(Arc::ptr_eq(&error, &failure.cause)); + } + drop(failure); + assert!(weak.upgrade().is_none(), "SDK must release terminal errors"); + assert_eq!(service.outstanding(), 0); + assert_eq!(service.flush().await.unwrap().delivered, 0); + } + fail.store(false, Ordering::Relaxed); + let fresh = service.enqueue(3); + assert_eq!(service.flush().await.unwrap().delivered, 1); + fresh.await.unwrap(); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + } + + #[tokio::test] + async fn shutdown_stops_admission_immediately_drains_and_is_shared_by_clones() { + let (service, mut entered, gate) = gated(4); + let write = service.enqueue(1); + let shutdown = service.shutdown(Duration::from_secs(1)); + let concurrent = service.clone().shutdown(Duration::from_secs(1)); + assert!(shutdown.is_ok()); + assert!(matches!( + service.enqueue(2).await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); + entered.recv().await.unwrap(); + gate.add_permits(1); + write.await.unwrap(); + assert_eq!(shutdown.await.unwrap().delivered, 1); + assert_eq!(concurrent.await.unwrap().delivered, 1); + assert!(service.is_empty()); + } + + #[tokio::test] + async fn shutdown_timeout_fails_active_work_even_when_flush_slots_are_full() { + let (service, mut entered, _) = gated(1); + let write = service.enqueue(1); + entered.recv().await.unwrap(); + let flush = service.flush(); + let rejected = service.flush(); + assert!(matches!( + &*rejected, + Err(FlushError::Admission(EnqueueError::Full)) + )); + let shutdown = service.shutdown(Duration::from_millis(10)); + assert!( + shutdown.is_ok(), + "shutdown must not compete with flush slots" + ); + assert!(matches!( + shutdown.await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + write.await, + Err(WriteError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + flush.await, + Err(FlushError::Execution(ExecutionError::TimedOut)) + )); + assert!(service.is_empty()); + assert!(matches!( + service.enqueue(2).await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); + } + + #[tokio::test] + async fn shutdown_timeout_also_retires_batches_queued_behind_active_work() { + let (service, mut entered, _) = gated(4); + let first = service.enqueue(1); + entered.recv().await.unwrap(); + let second = service.enqueue(2); + let third = service.enqueue(3); + let flush = service.flush(); + let shutdown = service.shutdown(Duration::from_millis(10)); + assert!(matches!( + shutdown.await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); + for completion in [first, second, third] { + assert!(matches!( + completion.await, + Err(WriteError::Execution(ExecutionError::TimedOut)) + )); + } + assert!(matches!( + flush.await, + Err(FlushError::Execution(ExecutionError::TimedOut)) + )); + assert_eq!(service.outstanding(), 0); + assert!(lock(&service.inner.state).barriers.is_empty()); + assert!(lock(&service.inner.state).ready.is_empty()); + } + + #[tokio::test] + async fn very_large_write_capacity_does_not_overflow_lifecycle_control_limits() { + let (service, _, _) = gated(usize::MAX); + assert_eq!(service.inner.control_limit, 1024); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + } + + #[tokio::test] + async fn worker_cancellation_before_first_poll_completes_every_accepted_sender() { + let (service, _, _) = gated(4); + let first = service.enqueue(1); + let second = service.enqueue(2); + let flush = service.flush(); + lock(&service.inner.abort).as_ref().unwrap().abort(); + for completion in [first, second] { + assert!(matches!( + completion.await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + } + assert!(matches!( + flush.await, + Err(FlushError::Execution(ExecutionError::WorkerLost)) + )); + assert!(service.is_empty()); + assert!(matches!( + service.enqueue(3).await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); + } + + #[tokio::test] + async fn worker_panic_notifies_active_and_pending_members() { + let (entered, mut receiver) = mpsc::unbounded_channel(); + let gate = Arc::new(Semaphore::new(0)); + let process_gate = gate.clone(); + let service = WriteService::new(4, Handle::current(), move |_: Vec| { + let entered = entered.clone(); + let gate = process_gate.clone(); + Box::pin(async move { + entered.send(()).unwrap(); + gate.acquire().await.unwrap().forget(); + panic!("test processor failed"); + }) + }); + let active = service.enqueue(1); + receiver.recv().await.unwrap(); + let pending = service.enqueue(2); + let flush = service.flush(); + gate.add_permits(1); + for completion in [active, pending] { + assert!(matches!( + completion.await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + } + assert!(matches!( + flush.await, + Err(FlushError::Execution(ExecutionError::WorkerLost)) + )); + assert!(service.is_empty()); + } + + #[tokio::test] + async fn dropping_the_last_producers_drains_work_and_stops_the_worker() { + let (service, mut entered, gate) = gated(4); + let write = service.enqueue(1); + let clone = service.clone(); + let exit = service.exit.clone(); + std::thread::scope(|scope| { + scope.spawn(move || drop(service)); + scope.spawn(move || drop(clone)); + }); + entered.recv().await.unwrap(); + gate.add_permits(1); + write.await.unwrap(); + tokio::time::timeout(Duration::from_secs(1), exit) + .await + .unwrap() + .unwrap(); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn fast_workers_and_racing_flushes_cannot_lose_registered_completions() { + let seen = Arc::new(Mutex::new(std::collections::HashSet::new())); + let processed = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let process_seen = seen.clone(); + let process_count = processed.clone(); + let service = WriteService::new(1024, Handle::current(), move |commands: Vec| { + lock(&process_seen).extend(commands.iter().copied()); + process_count.fetch_add(commands.len(), Ordering::Relaxed); + Box::pin(async { Ok(()) }) + }); + let (writes, flushes) = std::thread::scope(|scope| { + let producers: Vec<_> = (0..4) + .map(|producer| { + let service = &service; + scope.spawn(move || { + (0..50) + .map(|item| service.enqueue(producer * 50 + item)) + .collect::>() + }) + }) + .collect(); + let flushers: Vec<_> = (0..2) + .map(|_| { + let service = &service; + scope.spawn(move || (0..50).map(|_| service.flush()).collect::>()) + }) + .collect(); + ( + producers + .into_iter() + .flat_map(|producer| producer.join().unwrap()) + .collect::>(), + flushers + .into_iter() + .flat_map(|flusher| flusher.join().unwrap()) + .collect::>(), + ) + }); + tokio::time::timeout(Duration::from_secs(5), async { + for write in writes { + write.await.unwrap(); + } + for flush in flushes { + flush.await.unwrap(); + } + }) + .await + .unwrap(); + assert_eq!(lock(&seen).len(), 200); + assert_eq!(processed.load(Ordering::Relaxed), 200); + assert!(service.is_empty()); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn racing_producers_cannot_overbook_capacity() { + let (service, _entered, gate) = gated(8); + let completions = std::thread::scope(|scope| { + let threads: Vec<_> = (0..32) + .map(|number| { + let service = &service; + scope.spawn(move || service.enqueue(number)) + }) + .collect(); + threads + .into_iter() + .map(|thread| thread.join().unwrap()) + .collect::>() + }); + assert_eq!( + completions + .iter() + .filter(|completion| completion.is_ok()) + .count(), + 8 + ); + assert_eq!(service.outstanding(), 8); + gate.add_permits(8); + let mut delivered = 0; + for completion in completions { + if completion.await.is_ok() { + delivered += 1; + } + } + assert_eq!(delivered, 8); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + assert!(service.is_empty()); + } +} diff --git a/src/internal/mod.rs b/src/internal/mod.rs index 0b285bc..a10d05f 100644 --- a/src/internal/mod.rs +++ b/src/internal/mod.rs @@ -1,4 +1,5 @@ pub(crate) mod batching; +pub(crate) mod delivery; pub(crate) mod ingest; pub(crate) mod object_store; pub(crate) mod retry; diff --git a/src/lib.rs b/src/lib.rs index c252d60..eb51164 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,5 +1,7 @@ pub mod api; +mod completion; mod dataset; +mod delivery; mod error; pub mod eval; mod experiments; @@ -17,10 +19,15 @@ mod stream; pub(crate) mod test_utils; mod types; +pub use completion::{Completion, CompletionFuture}; pub use dataset::{ DatasetBuilder, DatasetBuilderError, DatasetInsert, DatasetInsertBuilder, DatasetInsertBuilderError, DatasetIterator, DatasetRecord, DatasetSummary, }; +pub use delivery::{ + DeliveryError, DeliveryReceipt, EnqueueError, ExecutionError, FlushError, FlushFailure, + FlushReport, ShutdownError, ShutdownReport, WriteError, +}; pub use error::{BraintrustError, Result}; pub use experiments::{ BaseExperimentInfo, Experiment, ExperimentBuilder, ExperimentBuilderError, ExperimentLog, diff --git a/src/log_queue/config.rs b/src/log_queue/config.rs index a5d1a7b..b371107 100644 --- a/src/log_queue/config.rs +++ b/src/log_queue/config.rs @@ -122,13 +122,12 @@ pub struct LogQueueConfig { #[builder(default = default_num_retries())] num_retries: usize, - /// Sync flush mode (BRAINTRUST_SYNC_FLUSH, default: false) + /// Legacy compatibility value. Delivery is always automatic; ignored. #[builder(default = default_sync_flush())] sync_flush: bool, - /// Return preparation and HTTP failures to an explicit flush caller. - /// Checked clients also disable automatic flushing so no failure can be - /// consumed before that caller observes it. + /// Legacy compatibility value. Flush reports covered active failures and + /// delivery remains automatic; ignored. #[builder(default = false)] checked_delivery: bool, @@ -144,11 +143,9 @@ pub struct LogQueueConfig { /// Debug: save all payloads (BRAINTRUST_ALL_PUBLISH_PAYLOADS_DIR) all_publish_payloads_dir: Option, - /// When true, the queue drops new items once it reaches `queue_max_size`. - /// Defaults to false (unbounded), matching the TypeScript SDK's default behavior where - /// data loss via queue overflow is opt-in. Set to true to apply a hard cap and protect - /// against runaway memory growth in high-throughput scenarios. - #[builder(default = false)] + /// Legacy compatibility value. Admission is always bounded by queue_max_size, + /// including in-flight work; ignored. + #[builder(default = true)] enforce_queue_size_limit: bool, } @@ -220,7 +217,7 @@ mod tests { assert_eq!(config.batch_max_bytes(), DEFAULT_BATCH_MAX_BYTES); assert_eq!(config.num_retries(), DEFAULT_NUM_RETRIES); assert!(!config.sync_flush()); - assert!(!config.enforce_queue_size_limit()); // unbounded by default (matches TS SDK) + assert!(config.enforce_queue_size_limit()); } #[test] diff --git a/src/log_queue/mod.rs b/src/log_queue/mod.rs index a5dd237..48a4b3b 100644 --- a/src/log_queue/mod.rs +++ b/src/log_queue/mod.rs @@ -4,7 +4,6 @@ mod http; mod merge; mod queue; mod row_key; -pub(crate) mod worker; // Public API exports pub use config::LogQueueConfig; diff --git a/src/log_queue/queue.rs b/src/log_queue/queue.rs index b168884..fc6b59b 100644 --- a/src/log_queue/queue.rs +++ b/src/log_queue/queue.rs @@ -6,20 +6,20 @@ use super::config::LogQueueConfig; use super::http::{fetch_version_info, save_payload_debug, send_batch_with_retry}; use super::merge::merge_row_into; use super::row_key::RowKey; -use super::worker::LogCommand; use crate::error::{BraintrustError, Result}; +use crate::internal::delivery::WriteService; use crate::logger::LoginState; use crate::span_components::{project_logs_identifier, ProjectLogsIdentifier}; use crate::types::{LogDestination, Logs3Row, ParentSpanInfo, SpanObjectType, SpanPayload}; -use arc_swap::ArcSwap; -use crossbeam::channel::{unbounded, Receiver, Sender, TrySendError}; +use crate::{ + Completion, DeliveryReceipt, FlushError, FlushReport, ShutdownError, ShutdownReport, WriteError, +}; use futures::future::join_all; use indexmap::IndexMap; -use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; -use tokio::sync::{mpsc, oneshot, Mutex as TokioMutex, OnceCell}; -use tokio::task::JoinHandle; +use tokio::sync::{Mutex as TokioMutex, OnceCell}; use tracing::warn; use url::Url; @@ -35,77 +35,18 @@ struct SubmitCommand { parent_info: Option, } -struct CmdChannel { - sender: Sender, - receiver: Receiver, - queued_count: AtomicUsize, -} - -fn make_cmd_channel() -> CmdChannel { - let (sender, receiver) = unbounded(); - CmdChannel { - sender, - receiver, - queued_count: AtomicUsize::new(0), - } -} - -fn reserve_queue_slot( - queued_count: &AtomicUsize, - queue_max_size: usize, - enforce_queue_size_limit: bool, -) -> bool { - let mut current = queued_count.load(Ordering::Relaxed); - loop { - if enforce_queue_size_limit && current >= queue_max_size { - return false; - } - - match queued_count.compare_exchange_weak( - current, - current + 1, - Ordering::Relaxed, - Ordering::Relaxed, - ) { - Ok(_) => return true, - Err(next) => current = next, - } - } -} - -/// All state shared between `LogQueue` and the background worker. -/// -/// Always accessed through `Arc`. Fields no longer need their own `Arc` -/// wrappers — the outer `Arc` provides shared ownership and the field types provide -/// their own interior mutability where needed. -pub(super) struct LogQueueCore { - /// Lock-free queue channel (sender, receiver) pair, atomically swappable. - /// Holds unprocessed `SubmitCommand`s; preparation (project registration, row - /// building) happens during `flush_internal`. - channel: ArcSwap, +/// Preparation and protocol orchestration, independent of admission/lifecycle. +struct LogQueueCore { dropped_count: AtomicUsize, last_drop_log_time: AtomicU64, - queue_depth_warning_emitted: AtomicBool, config: LogQueueConfig, login_state: LoginState, client: reqwest::Client, - /// Client-configured data-plane URL. Used as a fallback for flushing when the - /// shared login state has no URL (e.g. per-request-credential callers that - /// never log in). api_url: Url, - /// Client-configured control-plane URL. Fallback for project registration - /// when the shared login state has no URL. app_url: Url, - /// Cache of (org_id:project_name) → project_id to avoid redundant HTTP calls. project_cache: TokioMutex>, - /// Handle to the currently in-progress flush task, if any. - /// Using JoinHandle allows checking completion status and awaiting the result. - flush_handle: TokioMutex>>>, - /// Lazily cached (max_request_size, can_use_overflow) from GET /version. - /// Fetched once on first flush and reused thereafter. version_info: OnceCell<(usize, bool)>, } - impl LogQueueCore { fn new( config: LogQueueConfig, @@ -114,131 +55,19 @@ impl LogQueueCore { api_url: Url, app_url: Url, ) -> Arc { - let channel = make_cmd_channel(); Arc::new(Self { - channel: ArcSwap::from_pointee(channel), dropped_count: AtomicUsize::new(0), last_drop_log_time: AtomicU64::new(0), - queue_depth_warning_emitted: AtomicBool::new(false), config, login_state, client, api_url, app_url, project_cache: TokioMutex::new(IndexMap::new()), - flush_handle: TokioMutex::new(None), version_info: OnceCell::new(), }) } - /// Push a submit command directly to the lock-free queue. - /// - /// This is called by `LogQueue::submit` so that ALL drops (queue-full and - /// channel-closed) are handled uniformly here — with drop count tracking, - /// throttled warnings, and optional debug dumps. - fn push(self: &Arc, cmd: SubmitCommand) { - let channel = self.channel.load(); - if !reserve_queue_slot( - &channel.queued_count, - self.config.queue_max_size(), - self.config.enforce_queue_size_limit(), - ) { - self.dropped_count.fetch_add(1, Ordering::Relaxed); - self.log_drop_warning(); - self.dump_dropped_row_if_configured(cmd); - return; - } - - match channel.sender.try_send(cmd) { - Ok(()) => self.log_queue_depth_warning(channel.queued_count.load(Ordering::Relaxed)), - Err(TrySendError::Full(_)) => unreachable!("the log queue is unbounded"), - Err(TrySendError::Disconnected(_)) => { - channel.queued_count.fetch_sub(1, Ordering::Relaxed); - self.dropped_count.fetch_add(1, Ordering::Relaxed); - self.log_drop_warning(); - } - } - } - - /// Trigger a background flush if none is in progress. - /// Non-blocking — spawns a tokio task if needed. - /// No-op if no tokio runtime is available (e.g. in sync test contexts). - async fn trigger_flush_background(self: &Arc) { - let mut guard = self.flush_handle.lock().await; - - // Check if there's an existing flush in progress - if let Some(handle) = guard.as_ref() { - if !handle.is_finished() { - // Flush already in progress - return; - } - } - - // Consume the finished handle to surface any panic or error from the task. - if let Some(handle) = guard.take() { - if let Err(err) = handle.await { - warn!("background flush task failed: {:?}", err); - } - } - - // No flush in progress, spawn a new one - let Ok(runtime_handle) = tokio::runtime::Handle::try_current() else { - // No runtime available - return; - }; - - let core = self.clone(); - let join_handle = runtime_handle.spawn(async move { core.flush_internal().await }); - - *guard = Some(join_handle); - } - - /// Flush all rows in the lock-free queue and wait for completion. - /// - /// If a flush is already in progress, waits for it to complete and then - /// flushes any remaining items (matches TypeScript SDK behavior). - /// - /// Called by the background worker in response to `LogCommand::Flush` / - /// `LogCommand::TriggerFlush`. Does NOT drain the worker's mpsc channel — - /// callers that need that should use `LogQueue::flush_all()`. - pub(super) async fn flush(self: &Arc) -> std::result::Result<(), anyhow::Error> { - loop { - // Trigger a flush (no-op if one already running) - self.trigger_flush_background().await; - - // Wait for the current flush to complete, if any - let handle = { - let mut guard = self.flush_handle.lock().await; - guard.take() - }; - - if let Some(handle) = handle { - // Wait for the flush task to complete and propagate any errors - match handle.await { - Ok(Ok(())) => { - // Flush succeeded - } - Ok(Err(e)) => { - // Flush failed - return Err(e); - } - Err(e) => { - // Task panicked or was cancelled - return Err(anyhow::anyhow!("flush task failed: {}", e)); - } - } - } - - // Check if queue is now empty - if self.is_empty() { - return Ok(()); - } - - // Queue is non-empty — items arrived while the flush was running. - // Loop back to the top, which calls trigger_flush_background() again. - } - } - /// Get the cached (max_request_size, can_use_overflow) from GET /version. /// Fetches lazily on first call and caches the result. async fn get_version_info( @@ -255,14 +84,16 @@ impl LogQueueCore { .await } - /// Internal flush implementation: drain → prepare → chunk → merge → batch → send. + /// Process one detached logical batch: prepare → merge → encode → upload. /// - /// Commands are drained atomically. Each command is prepared (project - /// registration, row building) before batching. Rows are then processed in + /// The write service owns membership and the fixed cutoff. Each command is + /// prepared (project registration, row building) before batching. Rows are then processed in /// sequential chunks of `min(batch_max_items, flush_chunk_size)`. /// Within each chunk, batches are sent concurrently. - async fn flush_internal(self: &Arc) -> std::result::Result<(), anyhow::Error> { - let cmds = self.drain_all(); + async fn process( + self: &Arc, + cmds: Vec, + ) -> std::result::Result<(), anyhow::Error> { if cmds.is_empty() { return Ok(()); } @@ -280,11 +111,7 @@ impl LogQueueCore { Ok(url) => url, Err(e) => { warn!(error = %e, "Invalid API URL"); - return if self.config.checked_delivery() { - Err(e.into()) - } else { - Ok(()) - }; + return Err(e.into()); } }; @@ -320,11 +147,7 @@ impl LogQueueCore { } if groups.is_empty() { - return if self.config.checked_delivery() { - first_error.map_or(Ok(()), Err) - } else { - Ok(()) - }; + return first_error.map_or(Ok(()), Err); } // Version info reflects a server-wide limit, so it is fetched once and @@ -423,11 +246,7 @@ impl LogQueueCore { } } } - if self.config.checked_delivery() { - first_error.map_or(Ok(()), Err) - } else { - Ok(()) - } + first_error.map_or(Ok(()), Err) } /// Prepare a `Logs3Row` from a submit command's token, payload, and parent info. @@ -726,39 +545,6 @@ impl LogQueueCore { Ok(register_response.project.id) } - /// Drain all commands from the queue with precise flush boundaries. - /// - /// Atomically swaps in a new channel using lock-free CAS, then drains the old channel. - fn drain_all(&self) -> Vec { - // A new channel starts a new fill cycle, so it may emit the high-water warning. - self.queue_depth_warning_emitted - .store(false, Ordering::Relaxed); - let old_channel = self.channel.swap(Arc::new(make_cmd_channel())); - let mut items = Vec::new(); - while let Ok(item) = old_channel.receiver.try_recv() { - old_channel.queued_count.fetch_sub(1, Ordering::Relaxed); - items.push(item); - } - items - } - - /// Check if the queue is approximately empty. - pub fn is_empty(&self) -> bool { - self.channel.load().receiver.is_empty() - } - - /// Get the approximate number of items in the queue (tests only). - #[cfg(test)] - pub fn len(&self) -> usize { - self.channel.load().receiver.len() - } - - /// Get the total number of dropped items (tests only). - #[cfg(test)] - pub fn dropped_count(&self) -> usize { - self.dropped_count.load(Ordering::Relaxed) - } - /// Log a drop warning with throttling. /// /// Reports the number of items dropped since the last warning was emitted, @@ -790,23 +576,6 @@ impl LogQueueCore { } } - /// Warn once per fill cycle when the queue reaches 80% of its configured maximum. - fn log_queue_depth_warning(&self, queue_size: usize) { - let queue_max_size = self.config.queue_max_size(); - let warning_threshold = queue_max_size.saturating_mul(4).div_ceil(5); - - if queue_size >= warning_threshold - && !self - .queue_depth_warning_emitted - .swap(true, Ordering::Relaxed) - { - warn!( - queue_size, - queue_max_size, "Log queue has reached 80% of the configured maximum depth" - ); - } - } - /// If debug directories are configured, serialize the dropped command's payload /// and write it to disk. Spawns a background task — best-effort, non-blocking. fn dump_dropped_row_if_configured(&self, cmd: SubmitCommand) { @@ -942,405 +711,226 @@ impl LogQueueCore { } } -/// Lock-free log queue with batching and HTTP dispatch. -/// -/// This is the complete log processing system: -/// - Owns the background worker for flush coordination -/// - Lock-free queue for accumulating submit commands (bounded or unbounded) -/// - Row preparation (project registration, destination resolution) deferred to flush time -/// - Automatic batching and serialization via spawn_blocking -/// - HTTP dispatch with retry logic (concurrent per-chunk batches) -/// - Atomic drain for precise flush boundaries -/// - Matches TypeScript SDK `activeFlush` pattern: only one flush runs at a time -/// - Lazy version fetch for dynamic maxRequestSize and overflow support +/// Compatibility adapter connecting span preparation to the bounded write service. #[derive(Clone)] pub struct LogQueue { - /// All shared state, also held by the background worker via Arc. core: Arc, - /// Channel to the background worker (for flush commands only). - worker_sender: mpsc::Sender, - /// Handle to the background worker task. - #[allow(dead_code)] - worker_handle: Arc>, + writes: WriteService, } - impl LogQueue { - /// Create a new log queue that spawns and owns the background worker. - /// - /// The worker handles flush coordination; row preparation (project registration, - /// destination resolution) happens inside `flush_internal` at flush time. - #[allow(clippy::too_many_arguments)] pub fn new( config: LogQueueConfig, login_state: LoginState, client: reqwest::Client, - worker_queue_size: usize, + _worker_queue_size: usize, api_url: Url, app_url: Url, ) -> Self { - let core = LogQueueCore::new( - config, - login_state.clone(), - client.clone(), - api_url, - app_url, + let capacity = config.queue_max_size(); + let core = LogQueueCore::new(config, login_state, client, api_url, app_url); + let processor = core.clone(); + let writes = WriteService::new( + capacity, + tokio::runtime::Handle::current(), + move |commands| { + let processor = processor.clone(); + Box::pin(async move { processor.process(commands).await }) + }, ); - - let (worker_sender, worker_receiver) = mpsc::channel(worker_queue_size.max(32)); - - let worker_handle = Arc::new(tokio::spawn(super::worker::run_worker( - worker_receiver, - core.clone(), - ))); - - Self { - core, - worker_sender, - worker_handle, - } + Self { core, writes } } - #[cfg(test)] - pub(crate) fn enforces_queue_size_limit(&self) -> bool { - self.core.config.enforce_queue_size_limit() - } - - /// Submit a span payload for processing. - /// - /// The command is pushed directly into the lock-free queue. All drops - /// (queue-full and channel-closed) go through `LogQueueCore::push`, which - /// handles drop counting, throttled warnings, and optional debug dumps - /// uniformly. Row preparation (project registration, destination resolution) - /// is deferred to flush time. pub fn submit( &self, token: impl Into, payload: SpanPayload, parent_info: Option, - ) { - self.core.push(SubmitCommand { + ) -> Completion { + let command = SubmitCommand { token: token.into(), payload, parent_info, - }); - // Signal the worker to flush. try_send is non-blocking and best-effort: - // if the channel is full the worker already has a flush pending. - if !self.core.config.sync_flush() { - let _ = self.worker_sender.try_send(LogCommand::TriggerFlush); - } - } - - /// Flush all pending data through the worker, waiting for completion. - /// - /// Sends a `Flush` command through the worker channel, which ensures any - /// previously submitted rows (still in the worker's mpsc queue) are processed - /// before the flush begins. This is the correct flush path for external callers. - pub async fn flush_all(&self) -> Result<()> { - let (tx, rx) = oneshot::channel(); - self.worker_sender - .send(LogCommand::Flush(tx)) - .await - .map_err(|_| BraintrustError::ChannelClosed)?; - rx.await - .map_err(|_| BraintrustError::ChannelClosed)? - .map_err(|error| { - error.downcast::().unwrap_or_else(|error| { - error - .downcast::() - .map(BraintrustError::Http) - .unwrap_or_else(|error| BraintrustError::Background(error.to_string())) - }) + }; + self.writes + .enqueue_with_rejection(command, |command, error| { + self.core.dropped_count.fetch_add(1, Ordering::Relaxed); + if matches!(error, crate::EnqueueError::Full) { + self.core.log_drop_warning(); + } else { + warn!(%error, "span submission rejected"); + } + self.core.dump_dropped_row_if_configured(command); }) } - /// Return true if the lock-free row queue is approximately empty. - /// - /// Note: this inspects only the lock-free queue, not the worker's mpsc backlog, - /// so it may briefly report empty while a submit is in flight to the worker. + pub fn flush(&self) -> Completion { + self.writes.flush() + } + pub fn shutdown( + &self, + timeout: std::time::Duration, + ) -> Completion { + self.writes.shutdown(timeout) + } + #[cfg(test)] pub(crate) fn is_empty(&self) -> bool { - self.core.is_empty() + self.writes.is_empty() } - - /// Trigger a non-blocking background flush via the worker. pub async fn trigger_flush_command(&self) -> Result<()> { - self.worker_sender - .send(LogCommand::TriggerFlush) - .await - .map_err(|_| BraintrustError::ChannelClosed) - } -} - -impl Drop for LogQueue { - fn drop(&mut self) { - // Skip if the lock-free queue is visibly empty. Credentials travel with each - // queued row, so a flush works regardless of shared login state. - // - // Note: `is_empty()` only inspects the lock-free row queue, not the worker's - // mpsc channel. Items submitted but not yet processed by the worker would be - // missed by this guard. In practice the window is very small (a race between - // Drop and the worker task), and when it does occur the worker's own "channel - // closed" handler performs a final flush. Callers that need a hard guarantee - // should call `flush_all()` before dropping. - if self.core.is_empty() { - return; - } - - // Use flush_all() rather than core.flush() directly. Routing through the - // worker ensures any Submit commands ahead of this Drop in the mpsc queue are - // processed before the flush begins. - // - // block_in_place panics on current-thread runtimes (e.g. #[tokio::test] - // default). In that case we skip the flush here — callers on a current-thread - // runtime should call flush_all() explicitly before dropping. - if let Ok(handle) = tokio::runtime::Handle::try_current() { - // `block_in_place` panics on a current-thread runtime. In that case, - // skip the synchronous flush; dropping `worker_sender` closes the - // command channel and the worker performs a final flush. - if handle.runtime_flavor() != tokio::runtime::RuntimeFlavor::MultiThread { - return; - } - tokio::task::block_in_place(|| { - handle.block_on(async { - let _ = self.flush_all().await; - }); - }); - } else { - let rt = tokio::runtime::Runtime::new().unwrap(); - rt.block_on(async { - let _ = self.flush_all().await; - }); + let token = self.flush(); + if let Err(error) = &*token { + return Err(error.clone().into()); } + Ok(()) } } #[cfg(test)] mod tests { use super::*; - use crate::log_queue::config::LogQueueConfig; - - fn make_test_cmd(id: &str) -> SubmitCommand { - SubmitCommand { - token: "test-token".to_string(), - payload: SpanPayload { - row_id: id.to_string(), - span_id: id.to_string(), - is_merge: false, - span_components: None, - org_id: "org".to_string(), - org_name: Some("test-org".to_string()), - project_name: None, - input: None, - output: None, - expected: None, - error: None, - scores: None, - metadata: None, - metrics: None, - tags: None, - context: None, - span_attributes: None, - extra: std::collections::HashMap::new(), - }, - parent_info: Some(ParentSpanInfo::Experiment { - object_id: "exp-test".to_string(), - }), + use serde_json::{json, Value}; + use wiremock::{ + matchers::{header, method, path}, + Mock, MockServer, ResponseTemplate, + }; + + fn payload(id: &str, merge: bool) -> SpanPayload { + SpanPayload { + row_id: id.into(), + span_id: id.into(), + is_merge: merge, + span_components: None, + org_id: "org".into(), + org_name: Some("tenant".into()), + project_name: None, + input: Some(json!({"value": if merge { 2 } else { 1 }})), + output: None, + expected: None, + error: None, + scores: None, + metadata: None, + metrics: None, + tags: None, + context: None, + span_attributes: None, + extra: std::collections::HashMap::new(), } } - - fn make_queue(max_size: usize) -> LogQueue { - // No login state set - prevents background flushes from making real HTTP calls - // which would hang tests waiting for connections to non-existent servers. - let config = LogQueueConfig::builder() - .queue_max_size(max_size) - .enforce_queue_size_limit(true) // Enable bounding so drop tests work correctly - .sync_flush(true) // Disable auto-flush to avoid background HTTP calls - .build(); - let login_state = LoginState::new(); + fn queue(server: &MockServer) -> LogQueue { LogQueue::new( - config, - login_state, + LogQueueConfig::builder().num_retries(0).build(), + LoginState::new(), reqwest::Client::new(), - 256, - Url::parse("http://localhost").unwrap(), - Url::parse("http://localhost").unwrap(), + 32, + Url::parse(&server.uri()).unwrap(), + Url::parse(&server.uri()).unwrap(), ) } - - #[tokio::test] - async fn test_drop_does_not_panic_on_current_thread_runtime() { - let config = LogQueueConfig::builder() - .queue_max_size(10) - .enforce_queue_size_limit(true) - .sync_flush(true) - .build(); - let login_state = LoginState::new(); - login_state.set( - "test-key".to_string(), - "org".to_string(), - "test-org".to_string(), - "http://127.0.0.1:9".to_string(), - "http://127.0.0.1:9".to_string(), - ); - let queue = LogQueue::new( - config, - login_state, - reqwest::Client::new(), - 256, - Url::parse("http://localhost").unwrap(), - Url::parse("http://localhost").unwrap(), - ); - queue.core.push(make_test_cmd("1")); - // Logged in + non-empty queue: Drop takes the flush path. `#[tokio::test]` - // runs on a current-thread runtime, where this previously panicked in - // `block_in_place`. - drop(queue); - } - - #[tokio::test] - async fn test_push_queues_item() { - let queue = make_queue(10); - assert!(queue.core.is_empty()); - queue.core.push(make_test_cmd("1")); - assert_eq!(queue.core.len(), 1); - } - - #[tokio::test] - async fn test_queue_depth_warning_emits_at_eighty_percent_once_per_fill_cycle() { - let queue = make_queue(10); - - for id in 1..8 { - queue.core.push(make_test_cmd(&id.to_string())); - } - assert!(!queue - .core - .queue_depth_warning_emitted - .load(Ordering::Relaxed)); - - queue.core.push(make_test_cmd("8")); - assert!(queue - .core - .queue_depth_warning_emitted - .load(Ordering::Relaxed)); - - queue.core.push(make_test_cmd("9")); - assert!(queue - .core - .queue_depth_warning_emitted - .load(Ordering::Relaxed)); - - queue.core.drain_all(); - assert!(!queue - .core - .queue_depth_warning_emitted - .load(Ordering::Relaxed)); - } - - #[tokio::test] - async fn test_push_drops_when_full() { - let queue = make_queue(2); - queue.core.push(make_test_cmd("1")); - queue.core.push(make_test_cmd("2")); - assert_eq!(queue.core.len(), 2); - - // This push should be dropped since the queue is full - queue.core.push(make_test_cmd("3")); - assert_eq!(queue.core.len(), 2); - // The throttled warning fires immediately on the first drop (last_drop_log_time - // starts at 0, which is well beyond the throttle window), resetting the counter. - assert_eq!(queue.core.dropped_count(), 0); - } - - #[tokio::test] - async fn test_drop_count_accumulates_within_throttle_window() { - // With a very long logging period the counter accumulates across multiple drops - // without resetting (the throttle hasn't expired yet after the first warning fires). - let config = LogQueueConfig::builder() - .queue_max_size(1) - .enforce_queue_size_limit(true) - .sync_flush(true) - .build(); - let login_state = LoginState::new(); - let queue = LogQueue::new( - config, - login_state, - reqwest::Client::new(), - 256, - Url::parse("http://localhost").unwrap(), - Url::parse("http://localhost").unwrap(), - ); - - queue.core.push(make_test_cmd("1")); // fills queue - - // First drop: warning fires (last_drop_log_time=0 is always past throttle), - // resets counter to 0 and records timestamp. - queue.core.push(make_test_cmd("2")); // dropped, count reset to 0 after warning - assert_eq!(queue.core.dropped_count(), 0); - - // Subsequent drops within the same throttle window accumulate without reset. - queue.core.push(make_test_cmd("3")); // dropped, counter now 1 - queue.core.push(make_test_cmd("4")); // dropped, counter now 2 - assert_eq!(queue.core.dropped_count(), 2); + fn parent() -> Option { + Some(ParentSpanInfo::Experiment { + object_id: "experiment".into(), + }) } #[tokio::test] - async fn test_unbounded_queue_never_drops() { - let config = LogQueueConfig::builder() - .enforce_queue_size_limit(false) - .sync_flush(true) - .build(); - let login_state = LoginState::new(); - let queue = LogQueue::new( - config, - login_state, - reqwest::Client::new(), - 256, - Url::parse("http://localhost").unwrap(), - Url::parse("http://localhost").unwrap(), + async fn late_route_failure_fails_merged_members_and_successful_siblings_then_retires() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/version")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({}))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .and(header("authorization", "Bearer good")) + .respond_with(ResponseTemplate::new(200)) + .expect(2) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .and(header("authorization", "Bearer bad")) + .respond_with(ResponseTemplate::new(400)) + .expect(1) + .mount(&server) + .await; + let queue = queue(&server); + let first = queue.submit("good", payload("merged", false), parent()); + let update = queue.submit("good", payload("merged", true), parent()); + let last = queue.submit("bad", payload("other", false), parent()); + let Err(FlushError::Delivery(failure)) = queue.flush().await else { + panic!("expected failed flush"); + }; + assert_eq!( + failure.failed, 3, + "count accepted operations before merging" ); - - // Push many more items than the default queue_max_size - for i in 0..20_000 { - queue.core.push(make_test_cmd(&i.to_string())); + for token in [first, update, last] { + let Err(WriteError::Delivery(error)) = token.await else { + panic!("expected batch error"); + }; + assert!(Arc::ptr_eq(&error, &failure.cause)); } - assert_eq!(queue.core.len(), 20_000); - assert_eq!(queue.core.dropped_count(), 0); - } - - #[tokio::test] - async fn test_drain_all_returns_all_items() { - let queue = make_queue(10); - queue.core.push(make_test_cmd("1")); - queue.core.push(make_test_cmd("2")); - queue.core.push(make_test_cmd("3")); - - let drained = queue.core.drain_all(); - assert_eq!(drained.len(), 3); - assert!(queue.core.is_empty()); - } - - #[tokio::test] - async fn test_drain_all_atomic_boundary() { - // Items pushed after drain starts should go to the new channel, not the drained batch - let queue = make_queue(10); - queue.core.push(make_test_cmd("1")); - queue.core.push(make_test_cmd("2")); - - let drained = queue.core.drain_all(); - assert_eq!(drained.len(), 2); - - // Push after drain goes to the new channel - queue.core.push(make_test_cmd("3")); - assert_eq!(queue.core.len(), 1); + assert!(queue.is_empty()); + assert_eq!(queue.flush().await.unwrap().delivered, 0); + let requests = server.received_requests().await.unwrap(); + let uploads: Vec<_> = requests + .iter() + .filter(|req| req.url.path() == "/logs3") + .collect(); + assert_eq!(uploads.len(), 2); + let first_body: Value = serde_json::from_slice(&uploads[0].body).unwrap(); + assert_eq!(first_body["rows"].as_array().unwrap().len(), 1); + assert_eq!(first_body["rows"][0]["input"]["value"], 2); + let fresh = queue.submit("good", payload("fresh", false), parent()); + assert_eq!(queue.flush().await.unwrap().delivered, 1); + fresh.await.unwrap(); + queue + .shutdown(std::time::Duration::from_secs(1)) + .await + .unwrap(); } #[tokio::test] - async fn test_push_no_auto_flush_in_sync_mode() { - // make_queue uses sync_flush=true, so push() should not trigger background flush - let queue = make_queue(10); - - queue.core.push(make_test_cmd("1")); - // No flush should be in progress - in sync mode user calls flush() explicitly - assert!(queue.core.flush_handle.try_lock().unwrap().is_none()); - assert_eq!(queue.core.len(), 1); + async fn rejected_split_sibling_fails_every_member_despite_an_earlier_committed_request() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/version")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({}))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(|request: &wiremock::Request| { + let body: Value = serde_json::from_slice(&request.body).unwrap(); + let rows = body["rows"].as_array().unwrap(); + ResponseTemplate::new(if rows.len() > 1 { + 413 + } else if rows[0]["id"] == "late-failure" { + 400 + } else { + 200 + }) + }) + .expect(3) + .mount(&server) + .await; + let queue = queue(&server); + let first = queue.submit("token", payload("committed", false), parent()); + let last = queue.submit("token", payload("late-failure", false), parent()); + let Err(FlushError::Delivery(failure)) = queue.flush().await else { + panic!("expected failure"); + }; + assert_eq!(failure.failed, 2); + for token in [first, last] { + assert!(matches!(token.await, Err(WriteError::Delivery(_)))); + } + assert!(queue.is_empty()); + queue.flush().await.unwrap(); + queue + .shutdown(std::time::Duration::from_secs(1)) + .await + .unwrap(); } } diff --git a/src/log_queue/worker.rs b/src/log_queue/worker.rs deleted file mode 100644 index 45b281e..0000000 --- a/src/log_queue/worker.rs +++ /dev/null @@ -1,34 +0,0 @@ -use std::sync::Arc; -use tokio::sync::{mpsc, oneshot}; -use tracing::warn; - -use super::queue::LogQueueCore; - -/// Commands sent through the worker channel (flush coordination only). -pub(super) enum LogCommand { - Flush(oneshot::Sender>), - TriggerFlush, -} - -pub(super) async fn run_worker(mut receiver: mpsc::Receiver, core: Arc) { - loop { - match receiver.recv().await { - Some(LogCommand::Flush(response)) => { - let result = core.flush().await; - let _ = response.send(result); - } - Some(LogCommand::TriggerFlush) => { - if let Err(e) = core.flush().await { - warn!(error = %e, "background flush failed"); - } - } - None => { - // Channel closed — flush any remaining items and exit. - if let Err(e) = core.flush().await { - warn!(error = %e, "final flush failed"); - } - break; - } - } - } -} diff --git a/src/logger.rs b/src/logger.rs index ae76973..fc468e2 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -202,7 +202,7 @@ impl BraintrustClientBuilder { batch_max_items: None, batch_max_bytes: None, queue_max_size: None, - enforce_queue_size_limit: false, + enforce_queue_size_limit: true, checked_delivery: false, environment: None, span_origin: None, @@ -247,17 +247,16 @@ impl BraintrustClientBuilder { self } - /// Enable or disable enforcement of the queue size limit. - /// - /// When enabled, the queue drops newly arriving events once it reaches - /// `queue_max_size`. Defaults to false, leaving the queue unbounded. + /// Legacy compatibility setting. Delivery admission is always bounded by + /// `with_queue_max_size`, including in-flight writes. This value is ignored. pub fn enforce_queue_size_limit(mut self, enforce: bool) -> Self { self.enforce_queue_size_limit = enforce; self } - /// Make `flush` report preparation and HTTP delivery failures. Automatic - /// flushing is disabled so the explicit caller observes every result. + /// Legacy compatibility setting. Delivery is always automatic and flush + /// reports failures in covered active work. This value is ignored. Await a + /// write completion to observe that write's result independently of flush. pub fn checked_delivery(mut self, checked: bool) -> Self { self.checked_delivery = checked; self @@ -374,7 +373,6 @@ impl BraintrustClientBuilder { .maybe_queue_max_size(self.queue_max_size) .enforce_queue_size_limit(self.enforce_queue_size_limit) .checked_delivery(self.checked_delivery) - .maybe_sync_flush(self.checked_delivery.then_some(true)) .build(); // LogQueue owns the background worker. It is given the client-configured @@ -732,7 +730,10 @@ impl BraintrustClient { extra: event_data.extra, }; - self.submit_payload(token, payload, None); + let completion = self.submit_payload(token, payload, None); + if let Err(error) = &*completion { + return Err(BraintrustError::Background(error.to_string())); + } Ok(()) } @@ -909,24 +910,29 @@ impl BraintrustClient { ExperimentBuilder::new(submitter, token, org_id) } - /// Submit a span payload for logging (fire-and-forget). - /// - /// Returns immediately after queuing. HTTP submission happens in the background. - /// Errors are logged as warnings but not propagated to callers. - /// Submit a span payload synchronously (fire-and-forget). - /// The payload is queued internally and processed by a background worker. + /// Submit immediately; the token observes admission and full batch delivery. pub(crate) fn submit_payload( &self, token: impl Into, payload: SpanPayload, parent_info: Option, - ) { + ) -> crate::Completion { self.inner.queue.submit(token, payload, parent_info) } - /// Flush all pending log events. - pub async fn flush(&self) -> Result<()> { - self.inner.queue.flush_all().await + /// Stop accepting writes, drain accepted work and stop the delivery worker. + /// A timeout fails remaining completions and discards active SDK state. + pub fn shutdown( + &self, + timeout: std::time::Duration, + ) -> crate::Completion { + self.inner.queue.shutdown(timeout) + } + + /// Capture a cutoff immediately and observe delivery of the covered batch. + /// New writes belong to later work; completed failures are never replayed. + pub fn flush(&self) -> crate::Completion { + self.inner.queue.flush() } /// Trigger a non-blocking background flush. @@ -936,53 +942,19 @@ impl BraintrustClient { } } -impl Drop for BraintrustClient { - fn drop(&mut self) { - // Only flush on the last client reference, and only when there is pending - // data. `strong_count == 1` means no other `BraintrustClient` clones are - // alive. We check queue emptiness rather than login state so that - // per-request-credential (skip_login / multi-tenant) callers, which never - // populate the shared login state, still get their pending spans flushed. - if Arc::strong_count(&self.inner) != 1 || self.inner.queue.is_empty() { - return; - } - - // Flush via flush_all(), which routes through the background worker's mpsc channel. - // This ensures any Submit commands still sitting in the worker's backlog are processed - // before the flush begins — resolving the race in LogQueue::Drop where is_empty() only - // inspects the lock-free crossbeam queue, not the worker channel. - // After this flush completes, LogQueue::Drop will find an empty queue and exit early. - // - // block_in_place panics on current-thread runtimes (e.g. #[tokio::test] default). - // In that case we skip the flush here; LogQueue::Drop provides a best-effort fallback - // for the lock-free queue, and tests should always call flush() explicitly. - let inner = self.inner.clone(); - if let Ok(handle) = tokio::runtime::Handle::try_current() { - if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread { - tokio::task::block_in_place(|| { - handle.block_on(async move { - let _ = inner.queue.flush_all().await; - }); - }); - } - } else { - // No async runtime available — spin up a temporary one. - let rt = tokio::runtime::Runtime::new().unwrap(); - rt.block_on(async move { - let _ = inner.queue.flush_all().await; - }); - } - } -} - #[async_trait::async_trait] impl SpanSubmitter for BraintrustClient { - fn submit(&self, token: String, payload: SpanPayload, parent_info: Option) { + fn submit( + &self, + token: String, + payload: SpanPayload, + parent_info: Option, + ) -> crate::Completion { self.submit_payload(token, payload, parent_info) } - async fn flush(&self) -> Result<()> { - BraintrustClient::flush(self).await + fn flush(&self) -> crate::Completion { + BraintrustClient::flush(self) } async fn trigger_flush(&self) -> Result<()> { @@ -1530,17 +1502,23 @@ mod tests { } #[tokio::test] - async fn builder_configures_queue_size_limit_enforcement() { + async fn builder_capacity_is_enforced_even_with_the_legacy_unbounded_setting() { let client = BraintrustClient::builder() .app_url("https://example.com") .api_url("https://example.com") .skip_login(true) - .enforce_queue_size_limit(true) + .with_queue_max_size(1) + .enforce_queue_size_limit(false) .build() .await .expect("client"); - assert!(client.inner.queue.enforces_queue_size_limit()); + let span = client.span_builder_with_credentials("token", "org").build(); + assert!(span.log(crate::SpanLog::default()).is_ok()); + assert!(matches!( + &*span.log(crate::SpanLog::default()), + Err(crate::WriteError::Admission(crate::EnqueueError::Full)) + )); } #[tokio::test] diff --git a/src/span.rs b/src/span.rs index 30a88dc..beb4022 100644 --- a/src/span.rs +++ b/src/span.rs @@ -249,13 +249,16 @@ impl SpanLogBuilder { #[async_trait] pub(crate) trait SpanSubmitter: Send + Sync { - /// Submit a span payload for queuing (fire-and-forget). - /// The payload is queued internally and processed by a background worker. - /// Dropping is handled internally if the queue is full. - fn submit(&self, token: String, payload: SpanPayload, parent_info: Option); + /// Attempt admission immediately and observe the logical batch outcome. + fn submit( + &self, + token: String, + payload: SpanPayload, + parent_info: Option, + ) -> crate::Completion; /// Flush all queued span data through to the server. - async fn flush(&self) -> Result<()>; + fn flush(&self) -> crate::Completion; /// Trigger a non-blocking background flush. async fn trigger_flush(&self) -> Result<()>; @@ -572,19 +575,25 @@ impl SpanHandle { /// Log event data to this span. All fields are optional. /// Multiple calls will merge data (later values overwrite earlier ones). - /// Each call synchronously queues the current span snapshot for background processing. - pub fn log(&self, event: SpanLog) { - let payload = { - let mut inner = self - .inner - .lock() - .expect("span mutex should not be poisoned"); - apply_span_log_to_data(&mut inner, event); - current_span_payload(&mut inner) - }; - - self.submitter - .submit(self.token.clone(), payload, self.parent_info.clone()); + /// Each call attempts admission immediately. Await the returned token for full + /// delivery; ignoring it leaves accepted work running automatically. + pub fn log( + &self, + event: SpanLog, + ) -> crate::Completion { + let mut inner = self + .inner + .lock() + .expect("span mutex should not be poisoned"); + apply_span_log_to_data(&mut inner, event); + let payload = current_span_payload(&mut inner); + let completion = + self.submitter + .submit(self.token.clone(), payload, self.parent_info.clone()); + if completion.is_ok() { + inner.has_flushed = true; + } + completion } /// Flush all queued span data through to Braintrust. @@ -592,8 +601,8 @@ impl SpanHandle { /// This does not create a new queued span update. Call [`SpanHandle::log`] or /// [`SpanHandle::end`] first to submit changes, then call `flush()` when you need /// to ensure the background queue has been processed. - pub async fn flush(&self) -> Result<()> { - self.submitter.flush().await + pub fn flush(&self) -> crate::Completion { + self.submitter.flush() } /// Mark the span as ended with the current timestamp. @@ -609,21 +618,23 @@ impl SpanHandle { /// Calling this multiple times is idempotent: once an end time is set, /// subsequent calls return the previously-set value. pub fn end_with_time(&self, end_time: f64) -> f64 { - let (end_time, payload) = { - let mut inner = self - .inner - .lock() - .expect("span mutex should not be poisoned"); - if let Some(existing) = inner.end_time { - return existing; - } - inner.end_time = Some(end_time); - (end_time, current_span_payload(&mut inner)) - }; - - self.submitter - .submit(self.token.clone(), payload, self.parent_info.clone()); - end_time + let mut inner = self + .inner + .lock() + .expect("span mutex should not be poisoned"); + let chosen_time = *inner.end_time.get_or_insert(end_time); + if inner.end_submitted { + return chosen_time; + } + let payload = current_span_payload(&mut inner); + let completion = + self.submitter + .submit(self.token.clone(), payload, self.parent_info.clone()); + if completion.is_ok() { + inner.has_flushed = true; + inner.end_submitted = true; + } + chosen_time } /// Trigger a non-blocking background flush. @@ -758,9 +769,7 @@ fn current_span_payload(inner: &mut SpanData) -> SpanPayload { inner.metrics.insert("end".to_string(), end); } - let payload: SpanPayload = inner.clone().into(); - inner.has_flushed = true; - payload + inner.clone().into() } #[derive(Clone, Default)] @@ -768,6 +777,7 @@ struct SpanData { row_id: String, span_id: String, has_flushed: bool, + end_submitted: bool, org_id: String, org_name: Option, project_name: Option, @@ -1060,6 +1070,62 @@ mod tests { use crate::ProjectLogsIdentifier; use serde_json::json; + #[test] + fn rejected_admission_preserves_first_write_and_retryable_final_timestamp() { + use std::sync::atomic::{AtomicBool, Ordering}; + struct Submitter { + reject: AtomicBool, + accepted: Mutex>, + } + #[async_trait] + impl SpanSubmitter for Submitter { + fn submit( + &self, + _: String, + payload: SpanPayload, + _: Option, + ) -> crate::Completion { + let (sender, mut completion) = crate::Completion::pair(); + if self.reject.load(Ordering::Relaxed) { + completion.record_issuance(Err(crate::EnqueueError::Full.into())); + } else { + self.accepted.lock().unwrap().push(payload); + let _ = sender.send(Ok(crate::DeliveryReceipt)); + } + completion + } + fn flush(&self) -> crate::Completion { + let (sender, completion) = crate::Completion::pair(); + let _ = sender.send(Ok(crate::FlushReport::default())); + completion + } + async fn trigger_flush(&self) -> Result<()> { + Ok(()) + } + } + let submitter = Arc::new(Submitter { + reject: AtomicBool::new(true), + accepted: Mutex::new(Vec::new()), + }); + let span = SpanBuilder::new(submitter.clone(), "token", "org").build(); + assert!(span.log(SpanLog::default()).is_err()); + submitter.reject.store(false, Ordering::Relaxed); + assert!(span.log(SpanLog::default()).is_ok()); + assert!( + !submitter.accepted.lock().unwrap()[0].is_merge, + "rejected writes do not count as a first admission" + ); + submitter.reject.store(true, Ordering::Relaxed); + assert_eq!(span.end_with_time(10.0), 10.0); + assert_eq!(submitter.accepted.lock().unwrap().len(), 1); + submitter.reject.store(false, Ordering::Relaxed); + assert_eq!(span.end_with_time(20.0), 10.0); + assert_eq!(span.end_with_time(30.0), 10.0); + let accepted = submitter.accepted.lock().unwrap(); + assert_eq!(accepted.len(), 2); + assert_eq!(accepted[1].metrics.as_ref().unwrap()["end"], 10.0); + } + #[tokio::test] async fn span_logs_input_and_output() { let (span, collector) = build_test_span(); diff --git a/src/test_utils.rs b/src/test_utils.rs index 0075a6b..b9718d4 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -58,16 +58,26 @@ impl MockBraintrustClient { #[async_trait] impl SpanSubmitter for MockBraintrustClient { - fn submit(&self, token: String, payload: SpanPayload, parent_info: Option) { + fn submit( + &self, + token: String, + payload: SpanPayload, + parent_info: Option, + ) -> crate::Completion { self.collector.push(CapturedSpan { token, payload, parent: parent_info, }); + let (sender, completion) = crate::Completion::pair(); + let _ = sender.send(Ok(crate::DeliveryReceipt)); + completion } - async fn flush(&self) -> Result<()> { - Ok(()) + fn flush(&self) -> crate::Completion { + let (sender, completion) = crate::Completion::pair(); + let _ = sender.send(Ok(crate::FlushReport::default())); + completion } async fn trigger_flush(&self) -> Result<()> { diff --git a/tests/http_flow.rs b/tests/http_flow.rs index 28b1d65..fb11bd9 100644 --- a/tests/http_flow.rs +++ b/tests/http_flow.rs @@ -63,9 +63,8 @@ async fn submits_with_bearer_token() { } #[tokio::test] -async fn flush_is_fire_and_forget() { - // flush() should succeed even when project registration would fail. - // Errors are logged but not propagated (resilient design). +async fn preparation_failure_reaches_write_and_covering_flush() { + // Preparation failure terminates the logical batch, then is forgotten. let server = MockServer::start().await; Mock::given(method("POST")) @@ -87,28 +86,26 @@ async fn flush_is_fire_and_forget() { .project_name("demo") .build(); - span.log( + let completion = span.log( SpanLog::builder() .input(Value::String("input".into())) .build() .expect("build"), ); - // flush() waits for queued work to be processed. Background submission errors - // are still swallowed by the queue and should not bubble out here. - let result = span.flush().await; - assert!( - result.is_ok(), - "flush() should not surface background submission errors" - ); - - // Drain pending items before drop so the Drop impl's synchronous flush path - // is a no-op (queue is empty). The 500 from project registration is swallowed. - let _ = client.flush().await; + assert!(span.flush().await.is_err()); + assert!(matches!( + completion.await, + Err(braintrust_sdk_rust::WriteError::Delivery(_)) + )); + client + .flush() + .await + .expect("old failures must not poison a later flush"); } #[tokio::test] -async fn checked_flush_reports_backend_rejection() { +async fn automatic_delivery_reports_backend_rejection_through_completion() { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/api/project/register")) @@ -125,7 +122,6 @@ async fn checked_flush_reports_backend_rejection() { let client = BraintrustClient::builder() .skip_login(true) - .checked_delivery(true) .api_url(server.uri()) .app_url(server.uri()) .build() @@ -135,9 +131,52 @@ async fn checked_flush_reports_backend_rejection() { .span_builder_with_credentials("secret-token", "org") .project_name("demo") .build(); - span.log(SpanLog::builder().input(json!("input")).build().unwrap()); - assert!(matches!( - client.flush().await, - Err(BraintrustError::Api { status: 401, .. }) - )); + let completion = span.log(SpanLog::builder().input(json!("input")).build().unwrap()); + assert!( + matches!(completion.await, Err(braintrust_sdk_rust::WriteError::Delivery(error)) + if matches!(error.cause().downcast_ref::(), Some(BraintrustError::Api { status: 401, .. }))) + ); + client + .flush() + .await + .expect("completed background failure is retired"); +} + +#[tokio::test] +async fn submission_and_flush_can_be_issued_without_a_call_site_runtime() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/project/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"project":{"id":"proj"}}))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .skip_login(true) + .api_url(server.uri()) + .app_url(server.uri()) + .build() + .await + .unwrap(); + let span = client + .span_builder_with_credentials("token", "org") + .project_name("demo") + .build(); + let (write, flush) = std::thread::spawn(move || { + let write = span.log(SpanLog::builder().input(json!("input")).build().unwrap()); + let flush = client.flush(); + assert!(write.is_ok()); + assert!(flush.is_ok()); + (write, flush) + }) + .join() + .unwrap(); + write.await.unwrap(); + assert_eq!(flush.await.unwrap().delivered, 1); } From 9ade6fc3076d9f596b827d6092420809fea53bbe Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Thu, 29 Oct 2026 19:21:48 +0800 Subject: [PATCH 02/10] Keep delivery outcomes extensible --- src/delivery.rs | 44 +++++++++++++++++++++++++++++++++++----- src/internal/delivery.rs | 2 +- src/span.rs | 2 +- src/test_utils.rs | 2 +- tests/http_flow.rs | 2 +- 5 files changed, 43 insertions(+), 9 deletions(-) diff --git a/src/delivery.rs b/src/delivery.rs index d706d65..2a67f04 100644 --- a/src/delivery.rs +++ b/src/delivery.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use thiserror::Error; #[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] +#[non_exhaustive] pub enum EnqueueError { #[error("delivery admission capacity is full")] Full, @@ -10,6 +11,7 @@ pub enum EnqueueError { Closed, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] +#[non_exhaustive] pub enum ExecutionError { #[error("delivery worker stopped before completing accepted work")] WorkerLost, @@ -30,6 +32,7 @@ impl DeliveryError { } } #[derive(Debug, Clone, Error)] +#[non_exhaustive] pub enum WriteError { #[error(transparent)] Admission(#[from] EnqueueError), @@ -41,19 +44,43 @@ pub enum WriteError { /// Acknowledgment that the logical batch containing this write succeeded. /// Backend query visibility is separate from delivery acknowledgment. #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct DeliveryReceipt; +pub struct DeliveryReceipt { + pub(crate) _private: (), +} +impl DeliveryReceipt { + pub(crate) fn new() -> Self { + Self { _private: () } + } +} #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub struct FlushReport { /// Accepted operations delivered through this flush's cutoff, before merging. - pub delivered: u64, + pub(crate) delivered: u64, +} +impl FlushReport { + /// Number of accepted operations delivered through the flush cutoff. + pub fn delivered(&self) -> u64 { + self.delivered + } } #[derive(Debug, Clone, Error)] #[error("{failed} accepted writes failed: {cause}")] pub struct FlushFailure { - pub failed: u64, - pub cause: Arc, + pub(crate) failed: u64, + pub(crate) cause: Arc, +} +impl FlushFailure { + /// Number of accepted writes that failed. + pub fn failed(&self) -> u64 { + self.failed + } + /// Shared cause of the failed writes. + pub fn cause(&self) -> &Arc { + &self.cause + } } #[derive(Debug, Clone, Error)] +#[non_exhaustive] pub enum FlushError { #[error(transparent)] Admission(#[from] EnqueueError), @@ -64,9 +91,16 @@ pub enum FlushError { } #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub struct ShutdownReport { - pub delivered: u64, + pub(crate) delivered: u64, +} +impl ShutdownReport { + /// Number of accepted operations delivered through the shutdown cutoff. + pub fn delivered(&self) -> u64 { + self.delivered + } } #[derive(Debug, Clone, Error)] +#[non_exhaustive] pub enum ShutdownError { #[error(transparent)] Admission(#[from] EnqueueError), diff --git a/src/internal/delivery.rs b/src/internal/delivery.rs index cf42b56..be07675 100644 --- a/src/internal/delivery.rs +++ b/src/internal/delivery.rs @@ -47,7 +47,7 @@ impl BatchCompletion { fn finish(&self, result: &Result<(), BatchError>) { if let Some(writes) = lock(&self.senders).take() { let write_result = match result { - Ok(()) => Ok(DeliveryReceipt), + Ok(()) => Ok(DeliveryReceipt::new()), Err(BatchError::Delivery(error)) => Err(WriteError::Delivery(error.clone())), Err(BatchError::Execution(error)) => Err(WriteError::Execution(*error)), }; diff --git a/src/span.rs b/src/span.rs index beb4022..7fad7f9 100644 --- a/src/span.rs +++ b/src/span.rs @@ -1090,7 +1090,7 @@ mod tests { completion.record_issuance(Err(crate::EnqueueError::Full.into())); } else { self.accepted.lock().unwrap().push(payload); - let _ = sender.send(Ok(crate::DeliveryReceipt)); + let _ = sender.send(Ok(crate::DeliveryReceipt::new())); } completion } diff --git a/src/test_utils.rs b/src/test_utils.rs index b9718d4..ec1dcd5 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -70,7 +70,7 @@ impl SpanSubmitter for MockBraintrustClient { parent: parent_info, }); let (sender, completion) = crate::Completion::pair(); - let _ = sender.send(Ok(crate::DeliveryReceipt)); + let _ = sender.send(Ok(crate::DeliveryReceipt::new())); completion } diff --git a/tests/http_flow.rs b/tests/http_flow.rs index fb11bd9..4b22843 100644 --- a/tests/http_flow.rs +++ b/tests/http_flow.rs @@ -178,5 +178,5 @@ async fn submission_and_flush_can_be_issued_without_a_call_site_runtime() { .join() .unwrap(); write.await.unwrap(); - assert_eq!(flush.await.unwrap().delivered, 1); + assert_eq!(flush.await.unwrap().delivered(), 1); } From 79f28750d64749ae5995774704ffdf468cfef7bf Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Thu, 29 Oct 2026 19:23:46 +0800 Subject: [PATCH 03/10] Expose completion for explicit span updates --- src/logger.rs | 20 ++++++++++++++++++-- tests/http_flow.rs | 39 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/src/logger.rs b/src/logger.rs index fc468e2..ca88697 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -610,6 +610,7 @@ impl BraintrustClient { .org_id() .ok_or_else(|| BraintrustError::InvalidConfig("Not logged in".into()))?; self.update_span_internal(api_key, org_id, login_state.org_name(), exported, event) + .map(|_| ()) } /// Update an existing span using explicit credentials instead of shared login state. @@ -628,6 +629,21 @@ impl BraintrustClient { let org_id = org_id.into(); self.update_span_internal(token, org_id, None, exported, event) + .map(|_| ()) + } + + /// Update an existing span and retain its eventual delivery result. + /// + /// Validation and admission are performed before this method returns. The + /// returned completion observes delivery even if it finishes before a flush. + pub fn update_span_with_credentials_tracked( + &self, + token: impl Into, + org_id: impl Into, + exported: &str, + event: crate::span::SpanLog, + ) -> Result> { + self.update_span_internal(token.into(), org_id.into(), None, exported, event) } fn update_span_internal( @@ -637,7 +653,7 @@ impl BraintrustClient { org_name: Option, exported: &str, event: crate::span::SpanLog, - ) -> Result<()> { + ) -> Result> { let components = SpanComponents::parse(exported)?; let row_id = components.row_id.clone().ok_or_else(|| { BraintrustError::InvalidConfig("Exported span must have a row_id".into()) @@ -734,7 +750,7 @@ impl BraintrustClient { if let Err(error) = &*completion { return Err(BraintrustError::Background(error.to_string())); } - Ok(()) + Ok(completion) } /// Perform login synchronously. diff --git a/tests/http_flow.rs b/tests/http_flow.rs index 4b22843..82457dd 100644 --- a/tests/http_flow.rs +++ b/tests/http_flow.rs @@ -142,6 +142,45 @@ async fn automatic_delivery_reports_backend_rejection_through_completion() { .expect("completed background failure is retired"); } +#[tokio::test] +async fn tracked_stateless_update_reports_automatic_delivery_failure() { + use braintrust_sdk_rust::{SpanComponents, SpanObjectType, WriteError}; + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(401)) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .skip_login(true) + .api_url(server.uri()) + .app_url(server.uri()) + .build() + .await + .unwrap(); + let exported = SpanComponents { + object_type: SpanObjectType::ProjectLogs, + object_id: Some("project".into()), + compute_object_metadata_args: None, + row_id: Some("row".into()), + span_id: Some("span".into()), + root_span_id: Some("root".into()), + span_parents: None, + propagated_event: None, + } + .to_str(); + let completion = client + .update_span_with_credentials_tracked( + "token", + "org", + &exported, + SpanLog::builder().input(json!("input")).build().unwrap(), + ) + .unwrap(); + assert!(matches!(completion.await, Err(WriteError::Delivery(_)))); + client.flush().await.expect("completed failure is retired"); +} + #[tokio::test] async fn submission_and_flush_can_be_issued_without_a_call_site_runtime() { let server = MockServer::start().await; From df2333a2942b5c5ed8751624adb97b1b198f2a94 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Thu, 29 Oct 2026 19:43:26 +0800 Subject: [PATCH 04/10] Expose completion submission result --- src/completion.rs | 16 +++++++++++++--- src/logger.rs | 2 +- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/src/completion.rs b/src/completion.rs index 89e7ccb..b9e6e81 100644 --- a/src/completion.rs +++ b/src/completion.rs @@ -10,7 +10,7 @@ use tokio::sync::oneshot; /// A synchronously recorded issuance result and an eventual operation result. /// -/// Inspect `&*completion` to check admission without waiting. Awaiting returns +/// Inspect [`Completion::submission`] to check admission without waiting. Awaiting returns /// that error, or the accepted operation's result. Dropping this token does not /// cancel accepted work. No work is started by polling the token. /// @@ -21,7 +21,7 @@ use tokio::sync::oneshot; /// # async fn example(client: &BraintrustClient, span: &SpanHandle) /// # -> Result<(), Box> { /// let completion = span.log(SpanLog::default()); // admission happens now -/// if let Err(error) = &*completion { +/// if let Err(error) = completion.submission() { /// eprintln!("submission rejected: {error}"); /// } /// completion.await?; // observes admission failure or full batch delivery @@ -35,6 +35,14 @@ pub struct Completion { receiver: oneshot::Receiver>, } impl Completion { + /// Borrow the immediate submission result without waiting for delivery. + /// + /// An accepted submission may still fail during delivery. Await this + /// completion to observe the final result. + pub fn submission(&self) -> &Result<(), E> { + &self.issued + } + pub(crate) fn pair() -> (oneshot::Sender>, Self) { let (sender, receiver) = oneshot::channel(); ( @@ -114,7 +122,8 @@ mod tests { } let (sender, mut token) = Completion::<(), Error>::pair(); token.record_issuance(Err(Error::Admission)); - assert_eq!(&*token, &Err(Error::Admission)); + assert_eq!(token.submission(), &Err(Error::Admission)); + assert_eq!(&*token, token.submission()); drop(sender); assert_eq!(token.await, Err(Error::Admission)); } @@ -122,6 +131,7 @@ mod tests { async fn successful_issuance_stays_successful_after_operation_failure() { let (sender, token) = Completion::<(), ExecutionError>::pair(); sender.send(Err(ExecutionError::TimedOut)).unwrap(); + assert!(token.submission().is_ok()); assert!(token.is_ok()); assert_eq!(token.await, Err(ExecutionError::TimedOut)); let (sender, token) = Completion::<(), ExecutionError>::pair(); diff --git a/src/logger.rs b/src/logger.rs index ca88697..493e23e 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -747,7 +747,7 @@ impl BraintrustClient { }; let completion = self.submit_payload(token, payload, None); - if let Err(error) = &*completion { + if let Err(error) = completion.submission() { return Err(BraintrustError::Background(error.to_string())); } Ok(completion) From e6b22898a78737f1de5d7b9aa6a860bf1c3b09dc Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Thu, 29 Oct 2026 20:30:54 +0800 Subject: [PATCH 05/10] Use lock-free queue snapshots for delivery --- Cargo.toml | 3 + src/internal/delivery.rs | 643 +++++++++++++++++++++++++++++++++------ src/log_queue/queue.rs | 39 ++- 3 files changed, 575 insertions(+), 110 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 528e49d..b31c851 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,10 +15,13 @@ anyhow = "1" base64 = "0.23" bon = "3" async-trait = "0.1" +arc-swap = "1" backoff = { version = "0.4", features = ["tokio"] } bytes = "1" futures = "0.3" chrono = { version = "0.4", features = ["serde"] } +crossbeam-utils = "0.8" +sdd = "4.8" indexmap = "2" reqwest = { version = "0.13", default-features = false, features = ["json", "multipart", "query", "stream", "rustls"] } serde = { version = "1", features = ["derive"] } diff --git a/src/internal/delivery.rs b/src/internal/delivery.rs index be07675..192f504 100644 --- a/src/internal/delivery.rs +++ b/src/internal/delivery.rs @@ -1,26 +1,65 @@ -//! Atomic admission, fixed flush cutoffs and active-batch outcome ownership. +//! Lock-free admission, atomic queue snapshots and batch outcome ownership. use crate::{ Completion, DeliveryError, DeliveryReceipt, EnqueueError, ExecutionError, FlushError, FlushFailure, FlushReport, ShutdownError, ShutdownReport, WriteError, }; +use arc_swap::ArcSwap; +use crossbeam_utils::atomic::AtomicCell; use futures::{ future::{BoxFuture, Shared}, + task::AtomicWaker, FutureExt, }; use std::{ collections::VecDeque, - sync::{Arc, Mutex, MutexGuard}, + sync::{ + atomic::{AtomicBool, AtomicU8, AtomicUsize, Ordering}, + Arc, Mutex, MutexGuard, + }, time::Duration, }; use tokio::{ runtime::Handle, - sync::{oneshot, Notify, Semaphore}, + sync::{oneshot, Semaphore}, task::AbortHandle, }; use tracing::warn; type WriteSender = oneshot::Sender>; type FlushSender = oneshot::Sender>; +const CLOSED: usize = 1 << (usize::BITS - 1); + +/// One worker consumes coalesced notifications; producers never lock a wait list. +struct Wake { + ready: AtomicBool, + waker: AtomicWaker, +} +impl Wake { + fn new() -> Self { + Self { + ready: AtomicBool::new(false), + waker: AtomicWaker::new(), + } + } + fn notify_one(&self) { + self.ready.store(true, Ordering::Release); + self.waker.wake(); + } + async fn notified(&self) { + futures::future::poll_fn(|cx| { + if self.ready.swap(false, Ordering::AcqRel) { + return std::task::Poll::Ready(()); + } + self.waker.register(cx.waker()); + if self.ready.swap(false, Ordering::AcqRel) { + std::task::Poll::Ready(()) + } else { + std::task::Poll::Pending + } + }) + .await + } +} #[derive(Clone)] enum BatchError { Delivery(Arc), @@ -30,14 +69,56 @@ struct Accepted { command: C, sender: WriteSender, } +type IngressItem = AtomicCell>>>; +struct QueueEpoch { + queue: sdd::Queue>, + publishers: AtomicUsize, + active: AtomicBool, +} +impl QueueEpoch { + fn new() -> Self { + Self { + queue: sdd::Queue::new(), + publishers: AtomicUsize::new(0), + active: AtomicBool::new(false), + } + } + fn enter(&self) -> bool { + let mut current = self.publishers.load(Ordering::SeqCst); + loop { + if current & CLOSED != 0 { + return false; + } + match self.publishers.compare_exchange_weak( + current, + current + 1, + Ordering::SeqCst, + Ordering::SeqCst, + ) { + Ok(_) => return true, + Err(next) => current = next, + } + } + } + fn close(&self) { + self.publishers.fetch_or(CLOSED, Ordering::SeqCst); + } + fn settled(&self) -> bool { + self.publishers.load(Ordering::SeqCst) & !CLOSED == 0 + } +} struct Batch { - commands: Vec, + epoch: Arc>, completion: Arc, } struct BatchCompletion { - count: usize, + count: AtomicUsize, senders: Mutex>>, } +struct ActiveBatch { + epoch: Arc>, + completion: Arc, +} fn lock(value: &Mutex) -> MutexGuard<'_, T> { value .lock() @@ -99,68 +180,36 @@ impl Barrier { } } struct State { - pending: Vec>, ready: VecDeque>, - active: Option>, + active: Option>, barriers: Vec, - outstanding: usize, accepting: bool, terminal: Option, - depth_warned: bool, + shutdown_pending: Vec, } impl State { - fn seal(&mut self) { - if self.pending.is_empty() { - return; - } - let (commands, writes): (Vec<_>, Vec<_>) = std::mem::take(&mut self.pending) - .into_iter() - .map(|item| (item.command, item.sender)) - .unzip(); - let count = commands.len(); - self.ready.push_back(Batch { - commands, - completion: Arc::new(BatchCompletion { - count, - senders: Mutex::new(Some(writes)), - }), + fn barrier(&mut self, sender: FlushSender, shutdown: bool, through: Arc) { + self.barriers.push(Barrier { + through, + sender: Some(sender), + delivered: 0, + failure: None, + shutdown, }); } - fn barrier(&mut self, sender: FlushSender, shutdown: bool) { - self.seal(); - if let Some(through) = self - .ready - .back() - .map(|batch| batch.completion.clone()) - .or_else(|| self.active.clone()) - { - self.barriers.push(Barrier { - through, - sender: Some(sender), - delivered: 0, - failure: None, - shutdown, - }); - } else { - let _ = sender.send(Ok(FlushReport::default())); - } - } - fn finish(&mut self, batch: &Arc, result: Result<(), BatchError>) { + fn finish(&mut self, batch: &Arc, result: Result<(), BatchError>) -> usize { if self .active .as_ref() - .is_some_and(|active| Arc::ptr_eq(active, batch)) + .is_some_and(|active| Arc::ptr_eq(&active.completion, batch)) { - self.outstanding -= batch.count; self.active = None; - if self.outstanding == 0 { - self.depth_warned = false; - } batch.finish(&result); // Every registered barrier covers this batch: completed cutoffs // are removed here, before admitting any subsequent flush. + let count = batch.count.load(Ordering::Acquire); self.barriers.retain_mut(|barrier| { - barrier.record(batch.count as u64, &result); + barrier.record(count as u64, &result); if Arc::ptr_eq(&barrier.through, batch) { barrier.finish(); false @@ -168,49 +217,210 @@ impl State { true } }); + return count; } + 0 } } struct Inner { state: Mutex>, - wake: Notify, + ingress: ArcSwap>, + outstanding: AtomicUsize, + publishers: AtomicUsize, + terminal: AtomicU8, + depth_warned: AtomicBool, + wake: Wake, capacity: usize, control_limit: usize, shutdown_controls: Arc, abort: Mutex>, + #[cfg(test)] + after_take: Mutex, std::sync::mpsc::Receiver<()>)>>, } impl Inner { + fn begin_publication(&self) -> Result<(), EnqueueError> { + let mut current = self.publishers.load(Ordering::SeqCst); + loop { + if current & CLOSED != 0 { + return Err(EnqueueError::Closed); + } + match self.publishers.compare_exchange_weak( + current, + current + 1, + Ordering::SeqCst, + Ordering::SeqCst, + ) { + Ok(_) => return Ok(()), + Err(next) => current = next, + } + } + } + + fn end_publication(&self) { + self.publishers.fetch_sub(1, Ordering::SeqCst); + self.wake.notify_one(); + } + + fn reserve(&self) -> Result { + let mut current = self.outstanding.load(Ordering::Relaxed); + loop { + if current >= self.capacity { + return Err(EnqueueError::Full); + } + match self.outstanding.compare_exchange_weak( + current, + current + 1, + Ordering::AcqRel, + Ordering::Relaxed, + ) { + Ok(_) => return Ok(current + 1), + Err(next) => current = next, + } + } + } + + fn release(&self, count: usize) { + if count != 0 && self.outstanding.fetch_sub(count, Ordering::AcqRel) == count { + self.depth_warned.store(false, Ordering::Relaxed); + } + } + + fn detach(&self, state: &mut State) -> Arc { + let epoch = self.ingress.swap(Arc::new(QueueEpoch::new())); + epoch.close(); + let completion = Arc::new(BatchCompletion { + count: AtomicUsize::new(0), + senders: Mutex::new(None), + }); + state.ready.push_back(Batch { + epoch, + completion: completion.clone(), + }); + completion + } + + fn register_shutdowns(&self, state: &mut State) -> bool { + if state.shutdown_pending.is_empty() { + return true; + } + if self.publishers.load(Ordering::SeqCst) & !CLOSED != 0 { + return false; + } + let through = self.detach(state); + for sender in std::mem::take(&mut state.shutdown_pending) { + state.barrier(sender, true, through.clone()); + } + true + } + + fn fail_epoch(&self, epoch: &QueueEpoch, error: ExecutionError) { + while let Some(entry) = epoch.queue.pop() { + let accepted = entry.take().expect("each queue entry is taken once"); + let _ = accepted.sender.send(Err(error.into())); + self.release(1); + } + } + + fn publish(&self, accepted: Accepted) + where + C: 'static, + { + let mut accepted = Some(accepted); + loop { + let epoch = self.ingress.load_full(); + if epoch.enter() { + self.publish_to_epoch(&epoch, accepted.take().unwrap()); + return; + } + let terminal = self.terminal.load(Ordering::SeqCst); + if terminal != 0 { + let error = if terminal == 2 { + ExecutionError::TimedOut + } else { + ExecutionError::WorkerLost + }; + let _ = accepted.take().unwrap().sender.send(Err(error.into())); + self.release(1); + self.end_publication(); + return; + } + } + } + + fn publish_to_epoch(&self, epoch: &QueueEpoch, accepted: Accepted) + where + C: 'static, + { + epoch.active.store(true, Ordering::Release); + epoch.queue.push(AtomicCell::new(Some(Box::new(accepted)))); + epoch.publishers.fetch_sub(1, Ordering::SeqCst); + let terminal = self.terminal.load(Ordering::SeqCst); + self.end_publication(); + if terminal != 0 { + self.fail_epoch( + epoch, + if terminal == 2 { + ExecutionError::TimedOut + } else { + ExecutionError::WorkerLost + }, + ); + } + } + fn stop(&self, error: ExecutionError) { + let code = match error { + ExecutionError::WorkerLost => 1, + ExecutionError::TimedOut => 2, + }; + let error = + match self + .terminal + .compare_exchange(0, code, Ordering::SeqCst, Ordering::SeqCst) + { + Ok(_) => error, + Err(2) => ExecutionError::TimedOut, + Err(_) => ExecutionError::WorkerLost, + }; + self.publishers.fetch_or(CLOSED, Ordering::SeqCst); let mut state = lock(&self.state); state.accepting = false; if state.terminal.is_some() { return; } state.terminal = Some(error); - for accepted in state.pending.drain(..) { - let _ = accepted.sender.send(Err(error.into())); - } let result = Err(BatchError::Execution(error)); if let Some(active) = state.active.take() { - active.finish(&result); + active.completion.finish(&result); + self.release(active.completion.count.load(Ordering::Acquire)); + self.fail_epoch(&active.epoch, error); } for batch in state.ready.drain(..) { - batch.completion.finish(&result); + batch.epoch.close(); + self.fail_epoch(&batch.epoch, error); } for mut barrier in state.barriers.drain(..) { barrier.record(0, &result); } - state.outstanding = 0; + for sender in state.shutdown_pending.drain(..) { + let _ = sender.send(Err(error.into())); + } + drop(state); + let current = self.ingress.load_full(); + current.close(); + self.fail_epoch(¤t, error); self.wake.notify_one(); } fn close(&self) { + self.publishers.fetch_or(CLOSED, Ordering::SeqCst); lock(&self.state).accepting = false; self.wake.notify_one(); } } -/// Producers own no completed-outcome registry. A batch owns its write senders -/// until all its physical requests finish; barriers summarize only active work. +/// Generic delivery coordination for an injected batch processor. Producers own +/// no completed-outcome registry: a batch owns its write senders until all its +/// physical requests finish, and barriers summarize only active work. pub(crate) struct WriteService { inner: Arc>, owners: Arc>, @@ -242,23 +452,32 @@ impl WriteService { where F: Fn(Vec) -> BoxFuture<'static, anyhow::Result<()>> + Send + 'static, { + assert!( + IngressItem::::is_lock_free(), + "delivery queue requires lock-free pointer-sized atomics" + ); let control_limit = capacity.clamp(1, 1024); let inner = Arc::new(Inner { state: Mutex::new(State { - pending: Vec::new(), ready: VecDeque::new(), active: None, barriers: Vec::new(), - outstanding: 0, accepting: true, terminal: None, - depth_warned: false, + shutdown_pending: Vec::new(), }), - wake: Notify::new(), + ingress: ArcSwap::from_pointee(QueueEpoch::new()), + outstanding: AtomicUsize::new(0), + publishers: AtomicUsize::new(0), + terminal: AtomicU8::new(0), + depth_warned: AtomicBool::new(false), + wake: Wake::new(), capacity, control_limit, shutdown_controls: Arc::new(Semaphore::new(control_limit)), abort: Mutex::new(None), + #[cfg(test)] + after_take: Mutex::new(None), }); let (sender, receiver) = oneshot::channel(); let exit = async move { receiver.await.unwrap_or(Err(ExecutionError::WorkerLost)) } @@ -294,33 +513,33 @@ impl WriteService { rejected: impl FnOnce(C, EnqueueError), ) -> Completion { let (sender, mut completion) = Completion::pair(); - let mut state = lock(&self.inner.state); - let error = if !state.accepting { - Some(EnqueueError::Closed) - } else if state.outstanding >= self.inner.capacity { - Some(EnqueueError::Full) - } else { - None - }; - if let Some(error) = error { + let reservation = + self.inner + .begin_publication() + .and_then(|()| match self.inner.reserve() { + Ok(depth) => Ok(depth), + Err(error) => { + self.inner.end_publication(); + Err(error) + } + }); + if let Err(error) = reservation { completion.record_issuance(Err(error.into())); - drop(state); rejected(command, error); } else { - // Sender registration and publication have one linearization point. - state.pending.push(Accepted { command, sender }); - state.outstanding += 1; - if !state.depth_warned - && state.outstanding >= self.inner.capacity - self.inner.capacity / 5 + let depth = reservation.expect("checked above"); + // Publication is the linearization point: completed submissions + // appear in FIFO order before a later flush snapshots the queue. + self.inner.publish(Accepted { command, sender }); + if depth >= self.inner.capacity - self.inner.capacity / 5 + && !self.inner.depth_warned.swap(true, Ordering::Relaxed) { - state.depth_warned = true; warn!( - outstanding = state.outstanding, + outstanding = depth, capacity = self.inner.capacity, "delivery queue has reached 80% of capacity" ); } - self.inner.wake.notify_one(); } completion } @@ -339,7 +558,8 @@ impl WriteService { { completion.record_issuance(Err(EnqueueError::Full.into())); } else { - state.barrier(sender, false); + let through = self.inner.detach(&mut state); + state.barrier(sender, false, through); self.inner.wake.notify_one(); } completion @@ -358,8 +578,10 @@ impl WriteService { completion.record_issuance(Err(error.into())); return completion; } + self.inner.publishers.fetch_or(CLOSED, Ordering::SeqCst); state.accepting = false; - state.barrier(barrier_sender, true); + state.shutdown_pending.push(barrier_sender); + self.inner.register_shutdowns(&mut state); } let inner = self.inner.clone(); let exit = self.exit.clone(); @@ -392,11 +614,11 @@ impl WriteService { #[cfg(test)] pub fn is_empty(&self) -> bool { - lock(&self.inner.state).outstanding == 0 + self.inner.outstanding.load(Ordering::Acquire) == 0 } #[cfg(test)] fn outstanding(&self) -> usize { - lock(&self.inner.state).outstanding + self.inner.outstanding.load(Ordering::Acquire) } } struct WorkerGuard { @@ -423,24 +645,107 @@ where let wake = guard.inner.wake.notified(); let next = { let mut state = lock(&guard.inner.state); - state.seal(); - match state.ready.pop_front() { - Some(batch) => { - state.active = Some(batch.completion.clone()); - Some(batch) + if state.terminal.is_some() { + break; + } else if !guard.inner.register_shutdowns(&mut state) { + None + } else { + if state.ready.is_empty() + && guard.inner.ingress.load().active.load(Ordering::Acquire) + { + guard.inner.detach(&mut state); + } + let settled = state + .ready + .front() + .is_some_and(|batch| batch.epoch.settled()); + match if settled { + state.ready.pop_front() + } else { + None + } { + Some(batch) => { + state.active = Some(ActiveBatch { + epoch: batch.epoch.clone(), + completion: batch.completion.clone(), + }); + Some(batch) + } + None if state.ready.is_empty() + && !state.accepting + && guard.inner.publishers.load(Ordering::SeqCst) & !CLOSED == 0 => + { + break + } + None => None, } - None if !state.accepting => break, - None => None, } }; if let Some(batch) = next { - let result = process(batch.commands) - .await - .map_err(|source| BatchError::Delivery(Arc::new(DeliveryError { source }))); + // The generation has been detached. Drain it without holding the + // coordinator lock, so concurrent flushes can establish cutoffs. + let mut commands = Vec::new(); + let mut senders = Vec::new(); + while let Some(entry) = batch.epoch.queue.pop() { + let accepted = entry.take().expect("each queue entry is taken once"); + commands.push(accepted.command); + senders.push(accepted.sender); + #[cfg(test)] + if let Some((entered, resume)) = lock(&guard.inner.after_take).take() { + entered.send(()).unwrap(); + resume.recv().unwrap(); + } + } + let mut senders = Some(senders); + let terminal = { + let state = lock(&guard.inner.state); + if let Some(error) = state.terminal { + Some(error) + } else { + batch + .completion + .count + .store(commands.len(), Ordering::Release); + *lock(&batch.completion.senders) = senders.take(); + None + } + }; + if let Some(error) = terminal { + // stop() drains the rest of the epoch; these entries were + // already taken by this worker and remain its responsibility. + for sender in senders.take().unwrap_or_default() { + let _ = sender.send(Err(error.into())); + } + guard.inner.release(commands.len()); + break; + } + let result = if commands.is_empty() { + Ok(()) + } else { + process(commands) + .await + .map_err(|source| BatchError::Delivery(Arc::new(DeliveryError { source }))) + }; if let Err(BatchError::Delivery(error)) = &result { warn!(%error, "delivery batch failed"); } - lock(&guard.inner.state).finish(&batch.completion, result); + loop { + let wake = guard.inner.wake.notified(); + let finished = { + let mut state = lock(&guard.inner.state); + if guard.inner.register_shutdowns(&mut state) { + let count = state.finish(&batch.completion, result.clone()); + guard.inner.release(count); + true + } else { + false + } + }; + if finished { + break; + } + wake.await; + } } else { wake.await; } @@ -458,6 +763,152 @@ mod tests { }; use tokio::sync::mpsc; + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn admission_does_not_wait_for_coordinator_lock() { + let (service, mut entered, gate) = gated(4); + let coordinator = lock(&service.inner.state); + let producer = service.clone(); + let (sent, received) = std::sync::mpsc::channel(); + let thread = std::thread::spawn(move || { + let completion = producer.enqueue(7); + sent.send(completion.submission().is_ok()).unwrap(); + }); + let admitted = received.recv_timeout(Duration::from_secs(2)); + drop(coordinator); + thread.join().unwrap(); + assert!(admitted.unwrap()); + assert_eq!(entered.recv().await.unwrap(), vec![7]); + gate.add_permits(1); + service.shutdown(Duration::from_secs(2)).await.unwrap(); + } + + #[tokio::test] + async fn shutdown_includes_a_publisher_already_in_progress() { + let (service, mut entered, gate) = gated(4); + let first = service.enqueue(1); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + + service.inner.begin_publication().unwrap(); + service.inner.reserve().unwrap(); + let (sender, second) = Completion::pair(); + let shutdown = service.shutdown(Duration::from_secs(2)); + let mut shutdown = Box::pin(shutdown.into_future()); + assert!( + tokio::time::timeout(Duration::from_millis(20), &mut shutdown) + .await + .is_err() + ); + service.inner.publish(Accepted { command: 2, sender }); + gate.add_permits(1); + assert!(first.await.is_ok()); + assert_eq!(entered.recv().await.unwrap(), vec![2]); + gate.add_permits(1); + assert!(second.await.is_ok()); + assert_eq!(shutdown.await.unwrap().delivered(), 2); + } + + #[tokio::test] + async fn flush_swaps_generation_and_waits_for_its_late_publisher() { + let (service, mut entered, gate) = gated(8); + let first = service.enqueue(1); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + + service.inner.begin_publication().unwrap(); + service.inner.reserve().unwrap(); + let old_epoch = service.inner.ingress.load_full(); + assert!(old_epoch.enter()); + let (sender, before) = Completion::pair(); + let mut flush = Box::pin(service.flush().into_future()); + assert!(!Arc::ptr_eq(&old_epoch, &service.inner.ingress.load_full())); + let after = service.enqueue(3); + service + .inner + .publish_to_epoch(&old_epoch, Accepted { command: 2, sender }); + + gate.add_permits(1); + first.await.unwrap(); + assert_eq!(entered.recv().await.unwrap(), vec![2]); + assert!(futures::poll!(flush.as_mut()).is_pending()); + gate.add_permits(1); + before.await.unwrap(); + assert_eq!(flush.await.unwrap().delivered(), 2); + assert_eq!(entered.recv().await.unwrap(), vec![3]); + gate.add_permits(1); + after.await.unwrap(); + service.shutdown(Duration::from_secs(2)).await.unwrap(); + } + + #[tokio::test] + async fn flush_detaches_without_draining_the_queue() { + let (service, mut entered, gate) = gated(8); + let writes: Vec<_> = (0..4).map(|item| service.enqueue(item)).collect(); + let old_epoch = service.inner.ingress.load_full(); + let flush = service.flush(); + assert!(!Arc::ptr_eq(&old_epoch, &service.inner.ingress.load_full())); + // On the current-thread runtime, the worker has not run yet. Flush + // must not have taken the first item out of the detached queue. + assert!(old_epoch.queue.peek(&sdd::Guard::new()).is_some()); + assert_eq!(entered.recv().await.unwrap(), vec![0, 1, 2, 3]); + gate.add_permits(1); + assert_eq!(flush.await.unwrap().delivered(), 4); + for write in writes { + write.await.unwrap(); + } + service.shutdown(Duration::from_secs(2)).await.unwrap(); + } + + #[tokio::test] + async fn worker_failure_completes_a_late_publisher() { + let (service, mut entered, _gate) = gated(4); + let first = service.enqueue(1); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + service.inner.begin_publication().unwrap(); + service.inner.reserve().unwrap(); + let old_epoch = service.inner.ingress.load_full(); + assert!(old_epoch.enter()); + let (sender, late) = Completion::pair(); + lock(&service.inner.abort).as_ref().unwrap().abort(); + assert!(matches!( + first.await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + service + .inner + .publish_to_epoch(&old_epoch, Accepted { command: 2, sender }); + assert!(matches!( + late.await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + assert!(service.is_empty()); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn timeout_fails_items_taken_and_not_yet_taken_by_worker() { + let (service, _entered, _gate) = gated(4); + let (taken, wait_taken) = std::sync::mpsc::channel(); + let (resume, wait_resume) = std::sync::mpsc::channel(); + *lock(&service.inner.after_take) = Some((taken, wait_resume)); + let coordinator = lock(&service.inner.state); + let first = service.enqueue(1); + let second = service.enqueue(2); + drop(coordinator); + wait_taken.recv_timeout(Duration::from_secs(2)).unwrap(); + let shutdown = service.shutdown(Duration::from_millis(50)); + assert!(matches!( + shutdown.await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); + resume.send(()).unwrap(); + for completion in [first, second] { + assert!(matches!( + completion.await, + Err(WriteError::Execution(ExecutionError::TimedOut)) + )); + } + let _ = service.exit.clone().await; + assert!(service.is_empty()); + } + fn gated( capacity: usize, ) -> ( diff --git a/src/log_queue/queue.rs b/src/log_queue/queue.rs index fc6b59b..eb0f549 100644 --- a/src/log_queue/queue.rs +++ b/src/log_queue/queue.rs @@ -36,9 +36,7 @@ struct SubmitCommand { } /// Preparation and protocol orchestration, independent of admission/lifecycle. -struct LogQueueCore { - dropped_count: AtomicUsize, - last_drop_log_time: AtomicU64, +struct SpanBatchProcessor { config: LogQueueConfig, login_state: LoginState, client: reqwest::Client, @@ -47,7 +45,7 @@ struct LogQueueCore { project_cache: TokioMutex>, version_info: OnceCell<(usize, bool)>, } -impl LogQueueCore { +impl SpanBatchProcessor { fn new( config: LogQueueConfig, login_state: LoginState, @@ -56,8 +54,6 @@ impl LogQueueCore { app_url: Url, ) -> Arc { Arc::new(Self { - dropped_count: AtomicUsize::new(0), - last_drop_log_time: AtomicU64::new(0), config, login_state, client, @@ -544,7 +540,16 @@ impl LogQueueCore { } Ok(register_response.project.id) } +} +/// Reports rejected submissions without coupling admission to batch processing. +struct RejectionReporter { + dropped_count: AtomicUsize, + last_drop_log_time: AtomicU64, + config: LogQueueConfig, +} + +impl RejectionReporter { /// Log a drop warning with throttling. /// /// Reports the number of items dropped since the last warning was emitted, @@ -714,7 +719,7 @@ impl LogQueueCore { /// Compatibility adapter connecting span preparation to the bounded write service. #[derive(Clone)] pub struct LogQueue { - core: Arc, + rejections: Arc, writes: WriteService, } impl LogQueue { @@ -727,8 +732,12 @@ impl LogQueue { app_url: Url, ) -> Self { let capacity = config.queue_max_size(); - let core = LogQueueCore::new(config, login_state, client, api_url, app_url); - let processor = core.clone(); + let rejections = Arc::new(RejectionReporter { + dropped_count: AtomicUsize::new(0), + last_drop_log_time: AtomicU64::new(0), + config: config.clone(), + }); + let processor = SpanBatchProcessor::new(config, login_state, client, api_url, app_url); let writes = WriteService::new( capacity, tokio::runtime::Handle::current(), @@ -737,7 +746,7 @@ impl LogQueue { Box::pin(async move { processor.process(commands).await }) }, ); - Self { core, writes } + Self { rejections, writes } } pub fn submit( @@ -753,13 +762,15 @@ impl LogQueue { }; self.writes .enqueue_with_rejection(command, |command, error| { - self.core.dropped_count.fetch_add(1, Ordering::Relaxed); + self.rejections + .dropped_count + .fetch_add(1, Ordering::Relaxed); if matches!(error, crate::EnqueueError::Full) { - self.core.log_drop_warning(); + self.rejections.log_drop_warning(); } else { warn!(%error, "span submission rejected"); } - self.core.dump_dropped_row_if_configured(command); + self.rejections.dump_dropped_row_if_configured(command); }) } @@ -778,7 +789,7 @@ impl LogQueue { } pub async fn trigger_flush_command(&self) -> Result<()> { let token = self.flush(); - if let Err(error) = &*token { + if let Err(error) = token.submission() { return Err(error.clone().into()); } Ok(()) From eedead3c399a69894b984ada696d7ed2ea1d0473 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Thu, 29 Oct 2026 20:35:50 +0800 Subject: [PATCH 06/10] Reduce queue generation read contention --- src/internal/delivery.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/internal/delivery.rs b/src/internal/delivery.rs index 192f504..b86acf6 100644 --- a/src/internal/delivery.rs +++ b/src/internal/delivery.rs @@ -327,7 +327,10 @@ impl Inner { { let mut accepted = Some(accepted); loop { - let epoch = self.ingress.load_full(); + // Keep the generation borrowed only for publication; unlike an + // owned load this avoids contending on its Arc reference count on + // every submission. A concurrent swap keeps the guard alive. + let epoch = self.ingress.load(); if epoch.enter() { self.publish_to_epoch(&epoch, accepted.take().unwrap()); return; From 4c5e67331a71ca2258f5169c0dbb28b064b9a152 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Thu, 29 Oct 2026 20:51:52 +0800 Subject: [PATCH 07/10] Simplify delivery batch ownership and lifecycle state --- src/internal/delivery.rs | 350 +++++++++++++++++---------------------- src/log_queue/queue.rs | 4 +- 2 files changed, 155 insertions(+), 199 deletions(-) diff --git a/src/internal/delivery.rs b/src/internal/delivery.rs index b86acf6..35859ed 100644 --- a/src/internal/delivery.rs +++ b/src/internal/delivery.rs @@ -1,4 +1,6 @@ -//! Lock-free admission, atomic queue snapshots and batch outcome ownership. +//! Producers publish into a lock-free FIFO generation. Flush swaps that generation +//! into the worker's ready list; one worker drains and processes batches in order. +//! The coordinator owns active write senders and flush cutoffs under one mutex. use crate::{ Completion, DeliveryError, DeliveryReceipt, EnqueueError, ExecutionError, FlushError, FlushFailure, FlushReport, ShutdownError, ShutdownReport, WriteError, @@ -60,11 +62,6 @@ impl Wake { .await } } -#[derive(Clone)] -enum BatchError { - Delivery(Arc), - Execution(ExecutionError), -} struct Accepted { command: C, sender: WriteSender, @@ -73,14 +70,12 @@ type IngressItem = AtomicCell>>>; struct QueueEpoch { queue: sdd::Queue>, publishers: AtomicUsize, - active: AtomicBool, } impl QueueEpoch { fn new() -> Self { Self { queue: sdd::Queue::new(), publishers: AtomicUsize::new(0), - active: AtomicBool::new(false), } } fn enter(&self) -> bool { @@ -109,49 +104,33 @@ impl QueueEpoch { } struct Batch { epoch: Arc>, - completion: Arc, -} -struct BatchCompletion { - count: AtomicUsize, - senders: Mutex>>, -} -struct ActiveBatch { - epoch: Arc>, - completion: Arc, + senders: Vec, } fn lock(value: &Mutex) -> MutexGuard<'_, T> { value .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) } -impl BatchCompletion { - fn finish(&self, result: &Result<(), BatchError>) { - if let Some(writes) = lock(&self.senders).take() { - let write_result = match result { - Ok(()) => Ok(DeliveryReceipt::new()), - Err(BatchError::Delivery(error)) => Err(WriteError::Delivery(error.clone())), - Err(BatchError::Execution(error)) => Err(WriteError::Execution(*error)), - }; - for sender in writes { - let _ = sender.send(write_result.clone()); - } - } +fn finish_writes(writes: Vec, result: Result<(), WriteError>) { + let write_result = result.map(|()| DeliveryReceipt::new()); + for sender in writes { + let _ = sender.send(write_result.clone()); } } /// One bounded summary per active flush, rather than cloning every covered /// batch into each observer. Its cutoff is the last batch present at admission. -struct Barrier { - through: Arc, +struct Barrier { + through: Arc>, sender: Option, delivered: u64, failure: Option, shutdown: bool, } -impl Barrier { - fn record(&mut self, count: u64, result: &Result<(), BatchError>) { +impl Barrier { + fn record(&mut self, count: u64, result: &Result<(), Arc>) { match result { Ok(()) => self.delivered += count, - Err(BatchError::Delivery(error)) => match &mut self.failure { + Err(error) => match &mut self.failure { Some(failure) => failure.failed += count, None => { self.failure = Some(FlushFailure { @@ -160,11 +139,6 @@ impl Barrier { }) } }, - Err(BatchError::Execution(error)) => { - if let Some(sender) = self.sender.take() { - let _ = sender.send(Err((*error).into())); - } - } } } fn finish(&mut self) { @@ -180,15 +154,13 @@ impl Barrier { } } struct State { - ready: VecDeque>, - active: Option>, - barriers: Vec, - accepting: bool, - terminal: Option, + ready: VecDeque>>, + active: Option>, + barriers: Vec>, shutdown_pending: Vec, } impl State { - fn barrier(&mut self, sender: FlushSender, shutdown: bool, through: Arc) { + fn barrier(&mut self, sender: FlushSender, shutdown: bool, through: Arc>) { self.barriers.push(Barrier { through, sender: Some(sender), @@ -197,29 +169,22 @@ impl State { shutdown, }); } - fn finish(&mut self, batch: &Arc, result: Result<(), BatchError>) -> usize { - if self - .active - .as_ref() - .is_some_and(|active| Arc::ptr_eq(&active.completion, batch)) - { - self.active = None; - batch.finish(&result); - // Every registered barrier covers this batch: completed cutoffs - // are removed here, before admitting any subsequent flush. - let count = batch.count.load(Ordering::Acquire); - self.barriers.retain_mut(|barrier| { - barrier.record(count as u64, &result); - if Arc::ptr_eq(&barrier.through, batch) { - barrier.finish(); - false - } else { - true - } - }); - return count; - } - 0 + fn finish(&mut self, result: &Result<(), Arc>) -> usize { + let active = self.active.take().expect("active batch"); + let count = active.senders.len(); + finish_writes(active.senders, result.clone().map_err(WriteError::Delivery)); + // Every registered barrier covers this batch: completed cutoffs + // are removed here, before admitting any subsequent flush. + self.barriers.retain_mut(|barrier| { + barrier.record(count as u64, result); + if Arc::ptr_eq(&barrier.through, &active.epoch) { + barrier.finish(); + false + } else { + true + } + }); + count } } struct Inner { @@ -238,6 +203,14 @@ struct Inner { after_take: Mutex, std::sync::mpsc::Receiver<()>)>>, } impl Inner { + fn terminal_error(&self) -> Option { + match self.terminal.load(Ordering::SeqCst) { + 0 => None, + 2 => Some(ExecutionError::TimedOut), + _ => Some(ExecutionError::WorkerLost), + } + } + fn begin_publication(&self) -> Result<(), EnqueueError> { let mut current = self.publishers.load(Ordering::SeqCst); loop { @@ -285,18 +258,11 @@ impl Inner { } } - fn detach(&self, state: &mut State) -> Arc { + fn detach(&self, state: &mut State) -> Arc> { let epoch = self.ingress.swap(Arc::new(QueueEpoch::new())); epoch.close(); - let completion = Arc::new(BatchCompletion { - count: AtomicUsize::new(0), - senders: Mutex::new(None), - }); - state.ready.push_back(Batch { - epoch, - completion: completion.clone(), - }); - completion + state.ready.push_back(epoch.clone()); + epoch } fn register_shutdowns(&self, state: &mut State) -> bool { @@ -335,13 +301,7 @@ impl Inner { self.publish_to_epoch(&epoch, accepted.take().unwrap()); return; } - let terminal = self.terminal.load(Ordering::SeqCst); - if terminal != 0 { - let error = if terminal == 2 { - ExecutionError::TimedOut - } else { - ExecutionError::WorkerLost - }; + if let Some(error) = self.terminal_error() { let _ = accepted.take().unwrap().sender.send(Err(error.into())); self.release(1); self.end_publication(); @@ -354,56 +314,41 @@ impl Inner { where C: 'static, { - epoch.active.store(true, Ordering::Release); epoch.queue.push(AtomicCell::new(Some(Box::new(accepted)))); epoch.publishers.fetch_sub(1, Ordering::SeqCst); - let terminal = self.terminal.load(Ordering::SeqCst); + let terminal = self.terminal_error(); self.end_publication(); - if terminal != 0 { - self.fail_epoch( - epoch, - if terminal == 2 { - ExecutionError::TimedOut - } else { - ExecutionError::WorkerLost - }, - ); + if let Some(error) = terminal { + self.fail_epoch(epoch, error); } } - fn stop(&self, error: ExecutionError) { + fn stop(&self, error: ExecutionError) -> ExecutionError { let code = match error { ExecutionError::WorkerLost => 1, ExecutionError::TimedOut => 2, }; - let error = - match self - .terminal - .compare_exchange(0, code, Ordering::SeqCst, Ordering::SeqCst) - { - Ok(_) => error, - Err(2) => ExecutionError::TimedOut, - Err(_) => ExecutionError::WorkerLost, - }; - self.publishers.fetch_or(CLOSED, Ordering::SeqCst); let mut state = lock(&self.state); - state.accepting = false; - if state.terminal.is_some() { - return; + if let Some(error) = self.terminal_error() { + return error; } - state.terminal = Some(error); - let result = Err(BatchError::Execution(error)); + // Serialize the first terminal transition with batch completion so a + // returning processor cannot report success after shutdown has failed it. + self.terminal.store(code, Ordering::SeqCst); + self.publishers.fetch_or(CLOSED, Ordering::SeqCst); if let Some(active) = state.active.take() { - active.completion.finish(&result); - self.release(active.completion.count.load(Ordering::Acquire)); + let count = active.senders.len(); + finish_writes(active.senders, Err(error.into())); + self.release(count); self.fail_epoch(&active.epoch, error); } - for batch in state.ready.drain(..) { - batch.epoch.close(); - self.fail_epoch(&batch.epoch, error); + for epoch in state.ready.drain(..) { + self.fail_epoch(&epoch, error); } for mut barrier in state.barriers.drain(..) { - barrier.record(0, &result); + if let Some(sender) = barrier.sender.take() { + let _ = sender.send(Err(error.into())); + } } for sender in state.shutdown_pending.drain(..) { let _ = sender.send(Err(error.into())); @@ -413,10 +358,10 @@ impl Inner { current.close(); self.fail_epoch(¤t, error); self.wake.notify_one(); + error } fn close(&self) { self.publishers.fetch_or(CLOSED, Ordering::SeqCst); - lock(&self.state).accepting = false; self.wake.notify_one(); } } @@ -465,8 +410,6 @@ impl WriteService { ready: VecDeque::new(), active: None, barriers: Vec::new(), - accepting: true, - terminal: None, shutdown_pending: Vec::new(), }), ingress: ArcSwap::from_pointee(QueueEpoch::new()), @@ -550,7 +493,9 @@ impl WriteService { pub fn flush(&self) -> Completion { let (sender, mut completion) = Completion::pair(); let mut state = lock(&self.inner.state); - if !state.accepting { + if let Some(error) = self.inner.terminal_error() { + completion.record_issuance(Err(error.into())); + } else if self.inner.publishers.load(Ordering::SeqCst) & CLOSED != 0 { completion.record_issuance(Err(EnqueueError::Closed.into())); } else if state .barriers @@ -577,12 +522,11 @@ impl WriteService { let (barrier_sender, barrier) = Completion::pair(); { let mut state = lock(&self.inner.state); - if let Some(error) = state.terminal { + if let Some(error) = self.inner.terminal_error() { completion.record_issuance(Err(error.into())); return completion; } self.inner.publishers.fetch_or(CLOSED, Ordering::SeqCst); - state.accepting = false; state.shutdown_pending.push(barrier_sender); self.inner.register_shutdowns(&mut state); } @@ -631,10 +575,7 @@ struct WorkerGuard { impl Drop for WorkerGuard { fn drop(&mut self) { if let Some(exit) = self.exit.take() { - let error = lock(&self.inner.state) - .terminal - .unwrap_or(ExecutionError::WorkerLost); - self.inner.stop(error); + let error = self.inner.stop(ExecutionError::WorkerLost); let _ = exit.send(Err(error)); } } @@ -648,48 +589,37 @@ where let wake = guard.inner.wake.notified(); let next = { let mut state = lock(&guard.inner.state); - if state.terminal.is_some() { + if guard.inner.terminal_error().is_some() { break; } else if !guard.inner.register_shutdowns(&mut state) { None } else { - if state.ready.is_empty() - && guard.inner.ingress.load().active.load(Ordering::Acquire) - { + // Check closure before queue contents: if the last producer + // finishes after this snapshot, its wake triggers another pass. + let closed = guard.inner.publishers.load(Ordering::SeqCst) == CLOSED; + if state.ready.is_empty() && !guard.inner.ingress.load().queue.is_empty() { guard.inner.detach(&mut state); } - let settled = state - .ready - .front() - .is_some_and(|batch| batch.epoch.settled()); - match if settled { - state.ready.pop_front() + if state.ready.front().is_some_and(|epoch| epoch.settled()) { + let epoch = state.ready.pop_front().expect("checked above"); + state.active = Some(Batch { + epoch: epoch.clone(), + senders: Vec::new(), + }); + Some(epoch) + } else if state.ready.is_empty() && closed { + break; } else { None - } { - Some(batch) => { - state.active = Some(ActiveBatch { - epoch: batch.epoch.clone(), - completion: batch.completion.clone(), - }); - Some(batch) - } - None if state.ready.is_empty() - && !state.accepting - && guard.inner.publishers.load(Ordering::SeqCst) & !CLOSED == 0 => - { - break - } - None => None, } } }; - if let Some(batch) = next { + if let Some(epoch) = next { // The generation has been detached. Drain it without holding the // coordinator lock, so concurrent flushes can establish cutoffs. let mut commands = Vec::new(); let mut senders = Vec::new(); - while let Some(entry) = batch.epoch.queue.pop() { + while let Some(entry) = epoch.queue.pop() { let accepted = entry.take().expect("each queue entry is taken once"); commands.push(accepted.command); senders.push(accepted.sender); @@ -699,45 +629,36 @@ where resume.recv().unwrap(); } } - let mut senders = Some(senders); - let terminal = { - let state = lock(&guard.inner.state); - if let Some(error) = state.terminal { - Some(error) - } else { - batch - .completion - .count - .store(commands.len(), Ordering::Release); - *lock(&batch.completion.senders) = senders.take(); - None - } - }; - if let Some(error) = terminal { - // stop() drains the rest of the epoch; these entries were - // already taken by this worker and remain its responsibility. - for sender in senders.take().unwrap_or_default() { - let _ = sender.send(Err(error.into())); + { + let mut state = lock(&guard.inner.state); + if let Some(error) = guard.inner.terminal_error() { + drop(state); + // stop() drains the rest of the epoch; these entries were + // already taken and remain the worker's responsibility. + finish_writes(senders, Err(error.into())); + guard.inner.release(commands.len()); + break; } - guard.inner.release(commands.len()); - break; + state.active.as_mut().expect("active batch").senders = senders; } let result = if commands.is_empty() { Ok(()) } else { process(commands) .await - .map_err(|source| BatchError::Delivery(Arc::new(DeliveryError { source }))) + .map_err(|source| Arc::new(DeliveryError { source })) }; - if let Err(BatchError::Delivery(error)) = &result { + if let Err(error) = &result { warn!(%error, "delivery batch failed"); } loop { let wake = guard.inner.wake.notified(); let finished = { let mut state = lock(&guard.inner.state); - if guard.inner.register_shutdowns(&mut state) { - let count = state.finish(&batch.completion, result.clone()); + if guard.inner.terminal_error().is_some() { + break; + } else if guard.inner.register_shutdowns(&mut state) { + let count = state.finish(&result); guard.inner.release(count); true } else { @@ -754,7 +675,8 @@ where } } if let Some(exit) = guard.exit.take() { - let _ = exit.send(Ok(())); + let result = guard.inner.terminal_error().map_or(Ok(()), Err); + let _ = exit.send(result); } } #[cfg(test)] @@ -912,6 +834,38 @@ mod tests { assert!(service.is_empty()); } + #[tokio::test] + async fn terminal_stop_is_preserved_when_the_processor_returns_success() { + let (service, mut entered, gate) = gated(4); + let active = service.enqueue(1); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + let pending = service.enqueue(2); + let flush = service.flush(); + + service.inner.stop(ExecutionError::TimedOut); + gate.add_permits(1); // The processor returns Ok after its batch was failed. + for write in [active, pending] { + assert!(matches!( + write.await, + Err(WriteError::Execution(ExecutionError::TimedOut)) + )); + } + assert!(matches!( + flush.await, + Err(FlushError::Execution(ExecutionError::TimedOut)) + )); + assert_eq!(service.exit.clone().await, Err(ExecutionError::TimedOut)); + assert!(service.is_empty()); + assert!(matches!( + service.flush().submission(), + Err(FlushError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + service.shutdown(Duration::from_secs(1)).await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); + } + fn gated( capacity: usize, ) -> ( @@ -1236,23 +1190,25 @@ mod tests { assert!(service.is_empty()); } - #[tokio::test] + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn dropping_the_last_producers_drains_work_and_stops_the_worker() { - let (service, mut entered, gate) = gated(4); - let write = service.enqueue(1); - let clone = service.clone(); - let exit = service.exit.clone(); - std::thread::scope(|scope| { - scope.spawn(move || drop(service)); - scope.spawn(move || drop(clone)); - }); - entered.recv().await.unwrap(); - gate.add_permits(1); - write.await.unwrap(); - tokio::time::timeout(Duration::from_secs(1), exit) - .await - .unwrap() - .unwrap(); + for _ in 0..64 { + let (service, mut entered, gate) = gated(4); + let write = service.enqueue(1); + let clone = service.clone(); + let exit = service.exit.clone(); + std::thread::scope(|scope| { + scope.spawn(move || drop(service)); + scope.spawn(move || drop(clone)); + }); + entered.recv().await.unwrap(); + gate.add_permits(1); + write.await.unwrap(); + tokio::time::timeout(Duration::from_secs(1), exit) + .await + .unwrap() + .unwrap(); + } } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/src/log_queue/queue.rs b/src/log_queue/queue.rs index eb0f549..de69ff0 100644 --- a/src/log_queue/queue.rs +++ b/src/log_queue/queue.rs @@ -27,8 +27,8 @@ use crate::api::projects::{ProjectRegisterRequest, ProjectRegisterResponse}; /// A raw span payload queued for deferred preparation and flushing. /// -/// Fields are private — construction and destructuring happen only within this -/// module. `worker.rs` never touches the fields directly. +/// Only the span batch processor interprets these fields; the delivery worker +/// treats commands as opaque values. struct SubmitCommand { token: String, payload: SpanPayload, From 1ce34fb8337cbcb48e06e5f31670aaba5771d173 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Thu, 29 Oct 2026 22:03:00 +0800 Subject: [PATCH 08/10] Flatten delivery control flow and separate lifecycle tests --- src/internal/delivery.rs | 1248 ++++++++------------------------ src/internal/delivery/tests.rs | 619 ++++++++++++++++ 2 files changed, 918 insertions(+), 949 deletions(-) create mode 100644 src/internal/delivery/tests.rs diff --git a/src/internal/delivery.rs b/src/internal/delivery.rs index 35859ed..5c3c478 100644 --- a/src/internal/delivery.rs +++ b/src/internal/delivery.rs @@ -29,8 +29,217 @@ use tracing::warn; type WriteSender = oneshot::Sender>; type FlushSender = oneshot::Sender>; +/// Generic delivery coordination for an injected batch processor. Producers own +/// no completed-outcome registry: a batch owns its write senders until all its +/// physical requests finish, and barriers summarize only active work. +pub(crate) struct WriteService { + inner: Arc>, + owners: Arc>, + abort: AbortHandle, + runtime: Handle, + exit: Shared>>, +} +impl Clone for WriteService { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + owners: self.owners.clone(), + abort: self.abort.clone(), + runtime: self.runtime.clone(), + exit: self.exit.clone(), + } + } +} +struct ProducerOwner { + inner: std::sync::Weak>, +} +impl Drop for ProducerOwner { + fn drop(&mut self) { + if let Some(inner) = self.inner.upgrade() { + inner.close(); + } + } +} +impl WriteService { + pub fn new(capacity: usize, runtime: Handle, process: F) -> Self + where + F: Fn(Vec) -> BoxFuture<'static, anyhow::Result<()>> + Send + 'static, + { + assert!( + IngressItem::::is_lock_free(), + "delivery queue requires lock-free pointer-sized atomics" + ); + let control_limit = capacity.clamp(1, 1024); + let inner = Arc::new(Inner { + state: Mutex::new(State { + ready: VecDeque::new(), + active: None, + barriers: VecDeque::new(), + shutdown_pending: Vec::new(), + }), + ingress: ArcSwap::from_pointee(QueueEpoch::new()), + outstanding: AtomicUsize::new(0), + publishers: AtomicUsize::new(0), + terminal: AtomicU8::new(0), + depth_warned: AtomicBool::new(false), + wake: Wake::new(), + capacity, + control_limit, + shutdown_controls: Arc::new(Semaphore::new(control_limit)), + #[cfg(test)] + after_take: Mutex::new(None), + }); + let (sender, receiver) = oneshot::channel(); + let exit = async move { receiver.await.unwrap_or(Err(ExecutionError::WorkerLost)) } + .boxed() + .shared(); + // Construct outside the future so cancellation before its first poll + // still rejects pending work and notifies every accepted observer. + let guard = WorkerGuard { + inner: inner.clone(), + exit: Some(sender), + }; + let task = runtime.spawn(async move { run(guard, process).await }); + let owners = Arc::new(ProducerOwner { + inner: Arc::downgrade(&inner), + }); + Self { + inner, + owners, + abort: task.abort_handle(), + runtime, + exit, + } + } + + #[cfg(test)] + pub fn enqueue(&self, command: C) -> Completion { + self.enqueue_with_rejection(command, |_, _| {}) + } + + pub fn enqueue_with_rejection( + &self, + command: C, + rejected: impl FnOnce(C, EnqueueError), + ) -> Completion { + let (sender, mut completion) = Completion::pair(); + let reservation = self.inner.begin_publication().and_then(|()| { + self.inner + .reserve() + .inspect_err(|_| self.inner.end_publication()) + }); + match reservation { + Err(error) => { + completion.record_issuance(Err(error.into())); + rejected(command, error); + } + Ok(depth) => { + // Publication is the linearization point: completed submissions + // appear in FIFO order before a later flush snapshots the queue. + self.inner.publish(Accepted { command, sender }); + if depth >= self.inner.capacity - self.inner.capacity / 5 + && !self.inner.depth_warned.swap(true, Ordering::Relaxed) + { + warn!( + outstanding = depth, + capacity = self.inner.capacity, + "delivery queue has reached 80% of capacity" + ); + } + } + } + completion + } + + pub fn flush(&self) -> Completion { + let (sender, mut completion) = Completion::pair(); + let mut state = lock(&self.inner.state); + if let Some(error) = self.inner.terminal_error() { + completion.record_issuance(Err(error.into())); + } else if self.inner.publishers.load(Ordering::SeqCst) & CLOSED != 0 { + completion.record_issuance(Err(EnqueueError::Closed.into())); + } else if state + .barriers + .iter() + .filter(|barrier| !barrier.shutdown) + .count() + >= self.inner.control_limit + { + completion.record_issuance(Err(EnqueueError::Full.into())); + } else { + let through = self.inner.detach(&mut state); + state.barrier(sender, false, through); + self.inner.wake.notify_one(); + } + completion + } + + pub fn shutdown(&self, timeout: Duration) -> Completion { + let (sender, mut completion) = Completion::pair(); + let Ok(permit) = self.inner.shutdown_controls.clone().try_acquire_owned() else { + completion.record_issuance(Err(EnqueueError::Full.into())); + return completion; + }; + let (barrier_sender, barrier) = Completion::pair(); + { + let mut state = lock(&self.inner.state); + if let Some(error) = self.inner.terminal_error() { + completion.record_issuance(Err(error.into())); + return completion; + } + self.inner.publishers.fetch_or(CLOSED, Ordering::SeqCst); + state.shutdown_pending.push(barrier_sender); + self.inner.register_shutdowns(&mut state); + } + let inner = self.inner.clone(); + let abort = self.abort.clone(); + let exit = self.exit.clone(); + self.inner.wake.notify_one(); + self.runtime.spawn(async move { + let _permit = permit; + let work = async { + let report = barrier.await; + exit.await.map_err(ShutdownError::Execution)?; + report + .map(|report| ShutdownReport { + delivered: report.delivered, + }) + .map_err(Into::into) + }; + let result = match tokio::time::timeout(timeout, work).await { + Ok(result) => result, + Err(_) => { + inner.stop(ExecutionError::TimedOut); + abort.abort(); + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + } + }; + let _ = sender.send(result); + }); + completion + } + + #[cfg(test)] + pub fn is_empty(&self) -> bool { + self.inner.outstanding.load(Ordering::Acquire) == 0 + } + #[cfg(test)] + fn outstanding(&self) -> usize { + self.inner.outstanding.load(Ordering::Acquire) + } +} + const CLOSED: usize = 1 << (usize::BITS - 1); +#[allow(deprecated, reason = "try_update requires Rust 1.95")] +fn enter(publishers: &AtomicUsize) -> bool { + publishers + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |count| { + (count & CLOSED == 0).then(|| count + 1) + }) + .is_ok() +} + /// One worker consumes coalesced notifications; producers never lock a wait list. struct Wake { ready: AtomicBool, @@ -79,21 +288,7 @@ impl QueueEpoch { } } fn enter(&self) -> bool { - let mut current = self.publishers.load(Ordering::SeqCst); - loop { - if current & CLOSED != 0 { - return false; - } - match self.publishers.compare_exchange_weak( - current, - current + 1, - Ordering::SeqCst, - Ordering::SeqCst, - ) { - Ok(_) => return true, - Err(next) => current = next, - } - } + enter(&self.publishers) } fn close(&self) { self.publishers.fetch_or(CLOSED, Ordering::SeqCst); @@ -121,7 +316,7 @@ fn finish_writes(writes: Vec, result: Result<(), WriteError>) { /// batch into each observer. Its cutoff is the last batch present at admission. struct Barrier { through: Arc>, - sender: Option, + sender: FlushSender, delivered: u64, failure: Option, shutdown: bool, @@ -141,29 +336,27 @@ impl Barrier { }, } } - fn finish(&mut self) { - if let Some(sender) = self.sender.take() { - let result = match self.failure.take() { - Some(failure) => Err(FlushError::Delivery(failure)), - None => Ok(FlushReport { - delivered: self.delivered, - }), - }; - let _ = sender.send(result); - } + fn finish(self) { + let result = match self.failure { + Some(failure) => Err(FlushError::Delivery(failure)), + None => Ok(FlushReport { + delivered: self.delivered, + }), + }; + let _ = self.sender.send(result); } } struct State { ready: VecDeque>>, active: Option>, - barriers: Vec>, + barriers: VecDeque>, shutdown_pending: Vec, } impl State { fn barrier(&mut self, sender: FlushSender, shutdown: bool, through: Arc>) { - self.barriers.push(Barrier { + self.barriers.push_back(Barrier { through, - sender: Some(sender), + sender, delivered: 0, failure: None, shutdown, @@ -175,15 +368,18 @@ impl State { finish_writes(active.senders, result.clone().map_err(WriteError::Delivery)); // Every registered barrier covers this batch: completed cutoffs // are removed here, before admitting any subsequent flush. - self.barriers.retain_mut(|barrier| { + for barrier in &mut self.barriers { barrier.record(count as u64, result); - if Arc::ptr_eq(&barrier.through, &active.epoch) { - barrier.finish(); - false - } else { - true - } - }); + } + // Cutoffs follow admission order, so completed observers leave from + // the front and their senders never need an intermediate empty state. + while self + .barriers + .front() + .is_some_and(|barrier| Arc::ptr_eq(&barrier.through, &active.epoch)) + { + self.barriers.pop_front().expect("checked above").finish(); + } count } } @@ -198,7 +394,6 @@ struct Inner { capacity: usize, control_limit: usize, shutdown_controls: Arc, - abort: Mutex>, #[cfg(test)] after_take: Mutex, std::sync::mpsc::Receiver<()>)>>, } @@ -212,21 +407,9 @@ impl Inner { } fn begin_publication(&self) -> Result<(), EnqueueError> { - let mut current = self.publishers.load(Ordering::SeqCst); - loop { - if current & CLOSED != 0 { - return Err(EnqueueError::Closed); - } - match self.publishers.compare_exchange_weak( - current, - current + 1, - Ordering::SeqCst, - Ordering::SeqCst, - ) { - Ok(_) => return Ok(()), - Err(next) => current = next, - } - } + enter(&self.publishers) + .then_some(()) + .ok_or(EnqueueError::Closed) } fn end_publication(&self) { @@ -234,22 +417,14 @@ impl Inner { self.wake.notify_one(); } + #[allow(deprecated, reason = "try_update requires Rust 1.95")] fn reserve(&self) -> Result { - let mut current = self.outstanding.load(Ordering::Relaxed); - loop { - if current >= self.capacity { - return Err(EnqueueError::Full); - } - match self.outstanding.compare_exchange_weak( - current, - current + 1, - Ordering::AcqRel, - Ordering::Relaxed, - ) { - Ok(_) => return Ok(current + 1), - Err(next) => current = next, - } - } + self.outstanding + .fetch_update(Ordering::AcqRel, Ordering::Relaxed, |depth| { + (depth < self.capacity).then(|| depth + 1) + }) + .map(|previous| previous + 1) + .map_err(|_| EnqueueError::Full) } fn release(&self, count: usize) { @@ -291,18 +466,17 @@ impl Inner { where C: 'static, { - let mut accepted = Some(accepted); loop { // Keep the generation borrowed only for publication; unlike an // owned load this avoids contending on its Arc reference count on // every submission. A concurrent swap keeps the guard alive. let epoch = self.ingress.load(); if epoch.enter() { - self.publish_to_epoch(&epoch, accepted.take().unwrap()); + self.publish_to_epoch(&epoch, accepted); return; } if let Some(error) = self.terminal_error() { - let _ = accepted.take().unwrap().sender.send(Err(error.into())); + let _ = accepted.sender.send(Err(error.into())); self.release(1); self.end_publication(); return; @@ -345,10 +519,8 @@ impl Inner { for epoch in state.ready.drain(..) { self.fail_epoch(&epoch, error); } - for mut barrier in state.barriers.drain(..) { - if let Some(sender) = barrier.sender.take() { - let _ = sender.send(Err(error.into())); - } + for barrier in state.barriers.drain(..) { + let _ = barrier.sender.send(Err(error.into())); } for sender in state.shutdown_pending.drain(..) { let _ = sender.send(Err(error.into())); @@ -366,208 +538,6 @@ impl Inner { } } -/// Generic delivery coordination for an injected batch processor. Producers own -/// no completed-outcome registry: a batch owns its write senders until all its -/// physical requests finish, and barriers summarize only active work. -pub(crate) struct WriteService { - inner: Arc>, - owners: Arc>, - runtime: Handle, - exit: Shared>>, -} -impl Clone for WriteService { - fn clone(&self) -> Self { - Self { - inner: self.inner.clone(), - owners: self.owners.clone(), - runtime: self.runtime.clone(), - exit: self.exit.clone(), - } - } -} -struct ProducerOwner { - inner: std::sync::Weak>, -} -impl Drop for ProducerOwner { - fn drop(&mut self) { - if let Some(inner) = self.inner.upgrade() { - inner.close(); - } - } -} -impl WriteService { - pub fn new(capacity: usize, runtime: Handle, process: F) -> Self - where - F: Fn(Vec) -> BoxFuture<'static, anyhow::Result<()>> + Send + 'static, - { - assert!( - IngressItem::::is_lock_free(), - "delivery queue requires lock-free pointer-sized atomics" - ); - let control_limit = capacity.clamp(1, 1024); - let inner = Arc::new(Inner { - state: Mutex::new(State { - ready: VecDeque::new(), - active: None, - barriers: Vec::new(), - shutdown_pending: Vec::new(), - }), - ingress: ArcSwap::from_pointee(QueueEpoch::new()), - outstanding: AtomicUsize::new(0), - publishers: AtomicUsize::new(0), - terminal: AtomicU8::new(0), - depth_warned: AtomicBool::new(false), - wake: Wake::new(), - capacity, - control_limit, - shutdown_controls: Arc::new(Semaphore::new(control_limit)), - abort: Mutex::new(None), - #[cfg(test)] - after_take: Mutex::new(None), - }); - let (sender, receiver) = oneshot::channel(); - let exit = async move { receiver.await.unwrap_or(Err(ExecutionError::WorkerLost)) } - .boxed() - .shared(); - // Construct outside the future so cancellation before its first poll - // still rejects pending work and notifies every accepted observer. - let guard = WorkerGuard { - inner: inner.clone(), - exit: Some(sender), - }; - let task = runtime.spawn(async move { run(guard, process).await }); - *lock(&inner.abort) = Some(task.abort_handle()); - let owners = Arc::new(ProducerOwner { - inner: Arc::downgrade(&inner), - }); - Self { - inner, - owners, - runtime, - exit, - } - } - - #[cfg(test)] - pub fn enqueue(&self, command: C) -> Completion { - self.enqueue_with_rejection(command, |_, _| {}) - } - - pub fn enqueue_with_rejection( - &self, - command: C, - rejected: impl FnOnce(C, EnqueueError), - ) -> Completion { - let (sender, mut completion) = Completion::pair(); - let reservation = - self.inner - .begin_publication() - .and_then(|()| match self.inner.reserve() { - Ok(depth) => Ok(depth), - Err(error) => { - self.inner.end_publication(); - Err(error) - } - }); - if let Err(error) = reservation { - completion.record_issuance(Err(error.into())); - rejected(command, error); - } else { - let depth = reservation.expect("checked above"); - // Publication is the linearization point: completed submissions - // appear in FIFO order before a later flush snapshots the queue. - self.inner.publish(Accepted { command, sender }); - if depth >= self.inner.capacity - self.inner.capacity / 5 - && !self.inner.depth_warned.swap(true, Ordering::Relaxed) - { - warn!( - outstanding = depth, - capacity = self.inner.capacity, - "delivery queue has reached 80% of capacity" - ); - } - } - completion - } - - pub fn flush(&self) -> Completion { - let (sender, mut completion) = Completion::pair(); - let mut state = lock(&self.inner.state); - if let Some(error) = self.inner.terminal_error() { - completion.record_issuance(Err(error.into())); - } else if self.inner.publishers.load(Ordering::SeqCst) & CLOSED != 0 { - completion.record_issuance(Err(EnqueueError::Closed.into())); - } else if state - .barriers - .iter() - .filter(|barrier| !barrier.shutdown) - .count() - >= self.inner.control_limit - { - completion.record_issuance(Err(EnqueueError::Full.into())); - } else { - let through = self.inner.detach(&mut state); - state.barrier(sender, false, through); - self.inner.wake.notify_one(); - } - completion - } - - pub fn shutdown(&self, timeout: Duration) -> Completion { - let (sender, mut completion) = Completion::pair(); - let Ok(permit) = self.inner.shutdown_controls.clone().try_acquire_owned() else { - completion.record_issuance(Err(EnqueueError::Full.into())); - return completion; - }; - let (barrier_sender, barrier) = Completion::pair(); - { - let mut state = lock(&self.inner.state); - if let Some(error) = self.inner.terminal_error() { - completion.record_issuance(Err(error.into())); - return completion; - } - self.inner.publishers.fetch_or(CLOSED, Ordering::SeqCst); - state.shutdown_pending.push(barrier_sender); - self.inner.register_shutdowns(&mut state); - } - let inner = self.inner.clone(); - let exit = self.exit.clone(); - self.inner.wake.notify_one(); - self.runtime.spawn(async move { - let _permit = permit; - let work = async { - let report = barrier.await; - exit.await.map_err(ShutdownError::Execution)?; - report - .map(|report| ShutdownReport { - delivered: report.delivered, - }) - .map_err(Into::into) - }; - let result = match tokio::time::timeout(timeout, work).await { - Ok(result) => result, - Err(_) => { - inner.stop(ExecutionError::TimedOut); - if let Some(abort) = lock(&inner.abort).as_ref() { - abort.abort(); - } - Err(ShutdownError::Execution(ExecutionError::TimedOut)) - } - }; - let _ = sender.send(result); - }); - completion - } - - #[cfg(test)] - pub fn is_empty(&self) -> bool { - self.inner.outstanding.load(Ordering::Acquire) == 0 - } - #[cfg(test)] - fn outstanding(&self) -> usize { - self.inner.outstanding.load(Ordering::Acquire) - } -} struct WorkerGuard { inner: Arc>, exit: Option>>, @@ -614,63 +584,63 @@ where } } }; - if let Some(epoch) = next { - // The generation has been detached. Drain it without holding the - // coordinator lock, so concurrent flushes can establish cutoffs. - let mut commands = Vec::new(); - let mut senders = Vec::new(); - while let Some(entry) = epoch.queue.pop() { - let accepted = entry.take().expect("each queue entry is taken once"); - commands.push(accepted.command); - senders.push(accepted.sender); - #[cfg(test)] - if let Some((entered, resume)) = lock(&guard.inner.after_take).take() { - entered.send(()).unwrap(); - resume.recv().unwrap(); - } + let Some(epoch) = next else { + wake.await; + continue; + }; + // The generation has been detached. Drain it without holding the + // coordinator lock, so concurrent flushes can establish cutoffs. + let mut commands = Vec::new(); + let mut senders = Vec::new(); + while let Some(entry) = epoch.queue.pop() { + let accepted = entry.take().expect("each queue entry is taken once"); + commands.push(accepted.command); + senders.push(accepted.sender); + #[cfg(test)] + if let Some((entered, resume)) = lock(&guard.inner.after_take).take() { + entered.send(()).unwrap(); + resume.recv().unwrap(); + } + } + { + let mut state = lock(&guard.inner.state); + if let Some(error) = guard.inner.terminal_error() { + drop(state); + // stop() drains the rest of the epoch; these entries were + // already taken and remain the worker's responsibility. + finish_writes(senders, Err(error.into())); + guard.inner.release(commands.len()); + break; } - { + state.active.as_mut().expect("active batch").senders = senders; + } + let result = if commands.is_empty() { + Ok(()) + } else { + process(commands) + .await + .map_err(|source| Arc::new(DeliveryError { source })) + }; + if let Err(error) = &result { + warn!(%error, "delivery batch failed"); + } + loop { + let wake = guard.inner.wake.notified(); + let finished = { let mut state = lock(&guard.inner.state); - if let Some(error) = guard.inner.terminal_error() { - drop(state); - // stop() drains the rest of the epoch; these entries were - // already taken and remain the worker's responsibility. - finish_writes(senders, Err(error.into())); - guard.inner.release(commands.len()); + if guard.inner.terminal_error().is_some() { break; + } else if guard.inner.register_shutdowns(&mut state) { + let count = state.finish(&result); + guard.inner.release(count); + true + } else { + false } - state.active.as_mut().expect("active batch").senders = senders; - } - let result = if commands.is_empty() { - Ok(()) - } else { - process(commands) - .await - .map_err(|source| Arc::new(DeliveryError { source })) }; - if let Err(error) = &result { - warn!(%error, "delivery batch failed"); - } - loop { - let wake = guard.inner.wake.notified(); - let finished = { - let mut state = lock(&guard.inner.state); - if guard.inner.terminal_error().is_some() { - break; - } else if guard.inner.register_shutdowns(&mut state) { - let count = state.finish(&result); - guard.inner.release(count); - true - } else { - false - } - }; - if finished { - break; - } - wake.await; + if finished { + break; } - } else { wake.await; } } @@ -680,624 +650,4 @@ where } } #[cfg(test)] -mod tests { - use super::*; - use std::{ - future::IntoFuture, - sync::atomic::{AtomicBool, Ordering}, - }; - use tokio::sync::mpsc; - - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn admission_does_not_wait_for_coordinator_lock() { - let (service, mut entered, gate) = gated(4); - let coordinator = lock(&service.inner.state); - let producer = service.clone(); - let (sent, received) = std::sync::mpsc::channel(); - let thread = std::thread::spawn(move || { - let completion = producer.enqueue(7); - sent.send(completion.submission().is_ok()).unwrap(); - }); - let admitted = received.recv_timeout(Duration::from_secs(2)); - drop(coordinator); - thread.join().unwrap(); - assert!(admitted.unwrap()); - assert_eq!(entered.recv().await.unwrap(), vec![7]); - gate.add_permits(1); - service.shutdown(Duration::from_secs(2)).await.unwrap(); - } - - #[tokio::test] - async fn shutdown_includes_a_publisher_already_in_progress() { - let (service, mut entered, gate) = gated(4); - let first = service.enqueue(1); - assert_eq!(entered.recv().await.unwrap(), vec![1]); - - service.inner.begin_publication().unwrap(); - service.inner.reserve().unwrap(); - let (sender, second) = Completion::pair(); - let shutdown = service.shutdown(Duration::from_secs(2)); - let mut shutdown = Box::pin(shutdown.into_future()); - assert!( - tokio::time::timeout(Duration::from_millis(20), &mut shutdown) - .await - .is_err() - ); - service.inner.publish(Accepted { command: 2, sender }); - gate.add_permits(1); - assert!(first.await.is_ok()); - assert_eq!(entered.recv().await.unwrap(), vec![2]); - gate.add_permits(1); - assert!(second.await.is_ok()); - assert_eq!(shutdown.await.unwrap().delivered(), 2); - } - - #[tokio::test] - async fn flush_swaps_generation_and_waits_for_its_late_publisher() { - let (service, mut entered, gate) = gated(8); - let first = service.enqueue(1); - assert_eq!(entered.recv().await.unwrap(), vec![1]); - - service.inner.begin_publication().unwrap(); - service.inner.reserve().unwrap(); - let old_epoch = service.inner.ingress.load_full(); - assert!(old_epoch.enter()); - let (sender, before) = Completion::pair(); - let mut flush = Box::pin(service.flush().into_future()); - assert!(!Arc::ptr_eq(&old_epoch, &service.inner.ingress.load_full())); - let after = service.enqueue(3); - service - .inner - .publish_to_epoch(&old_epoch, Accepted { command: 2, sender }); - - gate.add_permits(1); - first.await.unwrap(); - assert_eq!(entered.recv().await.unwrap(), vec![2]); - assert!(futures::poll!(flush.as_mut()).is_pending()); - gate.add_permits(1); - before.await.unwrap(); - assert_eq!(flush.await.unwrap().delivered(), 2); - assert_eq!(entered.recv().await.unwrap(), vec![3]); - gate.add_permits(1); - after.await.unwrap(); - service.shutdown(Duration::from_secs(2)).await.unwrap(); - } - - #[tokio::test] - async fn flush_detaches_without_draining_the_queue() { - let (service, mut entered, gate) = gated(8); - let writes: Vec<_> = (0..4).map(|item| service.enqueue(item)).collect(); - let old_epoch = service.inner.ingress.load_full(); - let flush = service.flush(); - assert!(!Arc::ptr_eq(&old_epoch, &service.inner.ingress.load_full())); - // On the current-thread runtime, the worker has not run yet. Flush - // must not have taken the first item out of the detached queue. - assert!(old_epoch.queue.peek(&sdd::Guard::new()).is_some()); - assert_eq!(entered.recv().await.unwrap(), vec![0, 1, 2, 3]); - gate.add_permits(1); - assert_eq!(flush.await.unwrap().delivered(), 4); - for write in writes { - write.await.unwrap(); - } - service.shutdown(Duration::from_secs(2)).await.unwrap(); - } - - #[tokio::test] - async fn worker_failure_completes_a_late_publisher() { - let (service, mut entered, _gate) = gated(4); - let first = service.enqueue(1); - assert_eq!(entered.recv().await.unwrap(), vec![1]); - service.inner.begin_publication().unwrap(); - service.inner.reserve().unwrap(); - let old_epoch = service.inner.ingress.load_full(); - assert!(old_epoch.enter()); - let (sender, late) = Completion::pair(); - lock(&service.inner.abort).as_ref().unwrap().abort(); - assert!(matches!( - first.await, - Err(WriteError::Execution(ExecutionError::WorkerLost)) - )); - service - .inner - .publish_to_epoch(&old_epoch, Accepted { command: 2, sender }); - assert!(matches!( - late.await, - Err(WriteError::Execution(ExecutionError::WorkerLost)) - )); - assert!(service.is_empty()); - } - - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn timeout_fails_items_taken_and_not_yet_taken_by_worker() { - let (service, _entered, _gate) = gated(4); - let (taken, wait_taken) = std::sync::mpsc::channel(); - let (resume, wait_resume) = std::sync::mpsc::channel(); - *lock(&service.inner.after_take) = Some((taken, wait_resume)); - let coordinator = lock(&service.inner.state); - let first = service.enqueue(1); - let second = service.enqueue(2); - drop(coordinator); - wait_taken.recv_timeout(Duration::from_secs(2)).unwrap(); - let shutdown = service.shutdown(Duration::from_millis(50)); - assert!(matches!( - shutdown.await, - Err(ShutdownError::Execution(ExecutionError::TimedOut)) - )); - resume.send(()).unwrap(); - for completion in [first, second] { - assert!(matches!( - completion.await, - Err(WriteError::Execution(ExecutionError::TimedOut)) - )); - } - let _ = service.exit.clone().await; - assert!(service.is_empty()); - } - - #[tokio::test] - async fn terminal_stop_is_preserved_when_the_processor_returns_success() { - let (service, mut entered, gate) = gated(4); - let active = service.enqueue(1); - assert_eq!(entered.recv().await.unwrap(), vec![1]); - let pending = service.enqueue(2); - let flush = service.flush(); - - service.inner.stop(ExecutionError::TimedOut); - gate.add_permits(1); // The processor returns Ok after its batch was failed. - for write in [active, pending] { - assert!(matches!( - write.await, - Err(WriteError::Execution(ExecutionError::TimedOut)) - )); - } - assert!(matches!( - flush.await, - Err(FlushError::Execution(ExecutionError::TimedOut)) - )); - assert_eq!(service.exit.clone().await, Err(ExecutionError::TimedOut)); - assert!(service.is_empty()); - assert!(matches!( - service.flush().submission(), - Err(FlushError::Execution(ExecutionError::TimedOut)) - )); - assert!(matches!( - service.shutdown(Duration::from_secs(1)).await, - Err(ShutdownError::Execution(ExecutionError::TimedOut)) - )); - } - - fn gated( - capacity: usize, - ) -> ( - WriteService, - mpsc::UnboundedReceiver>, - Arc, - ) { - let (entered, receiver) = mpsc::unbounded_channel(); - let gate = Arc::new(Semaphore::new(0)); - let process_gate = gate.clone(); - let service = WriteService::new(capacity, Handle::current(), move |commands| { - let gate = process_gate.clone(); - let entered = entered.clone(); - Box::pin(async move { - entered.send(commands).unwrap(); - gate.acquire().await.unwrap().forget(); - Ok(()) - }) - }); - (service, receiver, gate) - } - - #[tokio::test] - async fn concurrent_flushes_observe_fixed_cutoffs_including_active_work() { - let (service, mut entered, gate) = gated(8); - let one = service.enqueue(1); - let first = service.flush(); - assert_eq!(entered.recv().await.unwrap(), vec![1]); - let concurrent = service.flush(); - let two = service.enqueue(2); - let mut second = Box::pin(service.flush().into_future()); - let three = service.enqueue(3); // beyond second's cutoff - gate.add_permits(1); - assert_eq!(first.await.unwrap().delivered, 1); - assert_eq!(concurrent.await.unwrap().delivered, 1); - one.await.unwrap(); - assert_eq!(entered.recv().await.unwrap(), vec![2]); - assert!(futures::poll!(second.as_mut()).is_pending()); - gate.add_permits(1); - assert_eq!(second.await.unwrap().delivered, 2); - two.await.unwrap(); - assert_eq!(entered.recv().await.unwrap(), vec![3]); - // The second flush finished without waiting for the later arrival. - gate.add_permits(1); - three.await.unwrap(); - assert_eq!(service.flush().await.unwrap().delivered, 0); - service.shutdown(Duration::from_secs(1)).await.unwrap(); - } - - #[tokio::test] - async fn a_new_flush_does_not_inherit_a_failed_prefix_from_an_older_active_flush() { - let (entered, mut receiver) = mpsc::unbounded_channel(); - let gate = Arc::new(Semaphore::new(0)); - let process_gate = gate.clone(); - let service = WriteService::new(4, Handle::current(), move |commands: Vec| { - let entered = entered.clone(); - let gate = process_gate.clone(); - Box::pin(async move { - entered.send(commands.clone()).unwrap(); - gate.acquire().await.unwrap().forget(); - if commands == [1] { - anyhow::bail!("first batch failed"); - } - Ok(()) - }) - }); - let first = service.enqueue(1); - assert_eq!(receiver.recv().await.unwrap(), vec![1]); - let second = service.enqueue(2); - let covering = service.flush(); - gate.add_permits(1); - assert!(matches!(first.await, Err(WriteError::Delivery(_)))); - assert_eq!(receiver.recv().await.unwrap(), vec![2]); - let fresh = service.flush(); - gate.add_permits(1); - second.await.unwrap(); - let Err(FlushError::Delivery(failure)) = covering.await else { - panic!("expected captured failure"); - }; - assert_eq!(failure.failed, 1); - assert_eq!(fresh.await.unwrap().delivered, 1); - service.shutdown(Duration::from_secs(1)).await.unwrap(); - } - - #[tokio::test] - async fn graceful_shutdown_reports_its_failed_batch_after_notifying_all_writes() { - let service = WriteService::new(4, Handle::current(), |_: Vec| { - Box::pin(async { - anyhow::bail!("terminal failure"); - }) - }); - let first = service.enqueue(1); - let second = service.enqueue(2); - let shutdown = service.shutdown(Duration::from_secs(1)); - let Err(ShutdownError::Delivery(failure)) = shutdown.await else { - panic!("expected delivery failure"); - }; - assert_eq!(failure.failed, 2); - for completion in [first, second] { - assert!(matches!(completion.await, Err(WriteError::Delivery(_)))); - } - assert!(service.is_empty()); - } - - #[tokio::test] - async fn capacity_includes_in_flight_work_and_rejection_is_observable_twice() { - let (service, mut entered, gate) = gated(1); - let first = service.enqueue(1); - entered.recv().await.unwrap(); - let rejected = service.enqueue(2); - assert!(matches!( - &*rejected, - Err(WriteError::Admission(EnqueueError::Full)) - )); - assert!(matches!( - rejected.await, - Err(WriteError::Admission(EnqueueError::Full)) - )); - assert_eq!(service.outstanding(), 1); - gate.add_permits(1); - first.await.unwrap(); - assert!(service.is_empty()); - gate.add_permits(1); - service.enqueue(3).await.unwrap(); - service.shutdown(Duration::from_secs(1)).await.unwrap(); - } - - #[tokio::test] - async fn dropped_observer_does_not_cancel_automatic_delivery_or_fill_the_queue() { - let (service, mut entered, gate) = gated(2); - drop(service.enqueue(1)); - assert_eq!(entered.recv().await.unwrap(), vec![1]); - let flush = service.flush(); - drop(flush); // lifecycle observers also do not cancel accepted work - gate.add_permits(1); - service.shutdown(Duration::from_secs(1)).await.unwrap(); - assert!(service.is_empty()); - } - - #[tokio::test] - async fn every_failed_batch_notifies_all_members_releases_errors_and_allows_fresh_work() { - let fail = Arc::new(AtomicBool::new(true)); - let process_fail = fail.clone(); - let service = WriteService::new(4, Handle::current(), move |_: Vec| { - let fail = process_fail.load(Ordering::Relaxed); - Box::pin(async move { - if fail { - anyhow::bail!("terminal failure"); - } - Ok(()) - }) - }); - for _ in 0..100 { - let first = service.enqueue(1); - let second = service.enqueue(2); - let flush = service.flush(); - let Err(FlushError::Delivery(failure)) = flush.await else { - panic!("expected failure"); - }; - assert_eq!(failure.failed, 2); - let weak = Arc::downgrade(&failure.cause); - for completion in [first, second] { - let Err(WriteError::Delivery(error)) = completion.await else { - panic!("expected failure"); - }; - assert!(Arc::ptr_eq(&error, &failure.cause)); - } - drop(failure); - assert!(weak.upgrade().is_none(), "SDK must release terminal errors"); - assert_eq!(service.outstanding(), 0); - assert_eq!(service.flush().await.unwrap().delivered, 0); - } - fail.store(false, Ordering::Relaxed); - let fresh = service.enqueue(3); - assert_eq!(service.flush().await.unwrap().delivered, 1); - fresh.await.unwrap(); - service.shutdown(Duration::from_secs(1)).await.unwrap(); - } - - #[tokio::test] - async fn shutdown_stops_admission_immediately_drains_and_is_shared_by_clones() { - let (service, mut entered, gate) = gated(4); - let write = service.enqueue(1); - let shutdown = service.shutdown(Duration::from_secs(1)); - let concurrent = service.clone().shutdown(Duration::from_secs(1)); - assert!(shutdown.is_ok()); - assert!(matches!( - service.enqueue(2).await, - Err(WriteError::Admission(EnqueueError::Closed)) - )); - entered.recv().await.unwrap(); - gate.add_permits(1); - write.await.unwrap(); - assert_eq!(shutdown.await.unwrap().delivered, 1); - assert_eq!(concurrent.await.unwrap().delivered, 1); - assert!(service.is_empty()); - } - - #[tokio::test] - async fn shutdown_timeout_fails_active_work_even_when_flush_slots_are_full() { - let (service, mut entered, _) = gated(1); - let write = service.enqueue(1); - entered.recv().await.unwrap(); - let flush = service.flush(); - let rejected = service.flush(); - assert!(matches!( - &*rejected, - Err(FlushError::Admission(EnqueueError::Full)) - )); - let shutdown = service.shutdown(Duration::from_millis(10)); - assert!( - shutdown.is_ok(), - "shutdown must not compete with flush slots" - ); - assert!(matches!( - shutdown.await, - Err(ShutdownError::Execution(ExecutionError::TimedOut)) - )); - assert!(matches!( - write.await, - Err(WriteError::Execution(ExecutionError::TimedOut)) - )); - assert!(matches!( - flush.await, - Err(FlushError::Execution(ExecutionError::TimedOut)) - )); - assert!(service.is_empty()); - assert!(matches!( - service.enqueue(2).await, - Err(WriteError::Admission(EnqueueError::Closed)) - )); - } - - #[tokio::test] - async fn shutdown_timeout_also_retires_batches_queued_behind_active_work() { - let (service, mut entered, _) = gated(4); - let first = service.enqueue(1); - entered.recv().await.unwrap(); - let second = service.enqueue(2); - let third = service.enqueue(3); - let flush = service.flush(); - let shutdown = service.shutdown(Duration::from_millis(10)); - assert!(matches!( - shutdown.await, - Err(ShutdownError::Execution(ExecutionError::TimedOut)) - )); - for completion in [first, second, third] { - assert!(matches!( - completion.await, - Err(WriteError::Execution(ExecutionError::TimedOut)) - )); - } - assert!(matches!( - flush.await, - Err(FlushError::Execution(ExecutionError::TimedOut)) - )); - assert_eq!(service.outstanding(), 0); - assert!(lock(&service.inner.state).barriers.is_empty()); - assert!(lock(&service.inner.state).ready.is_empty()); - } - - #[tokio::test] - async fn very_large_write_capacity_does_not_overflow_lifecycle_control_limits() { - let (service, _, _) = gated(usize::MAX); - assert_eq!(service.inner.control_limit, 1024); - service.shutdown(Duration::from_secs(1)).await.unwrap(); - } - - #[tokio::test] - async fn worker_cancellation_before_first_poll_completes_every_accepted_sender() { - let (service, _, _) = gated(4); - let first = service.enqueue(1); - let second = service.enqueue(2); - let flush = service.flush(); - lock(&service.inner.abort).as_ref().unwrap().abort(); - for completion in [first, second] { - assert!(matches!( - completion.await, - Err(WriteError::Execution(ExecutionError::WorkerLost)) - )); - } - assert!(matches!( - flush.await, - Err(FlushError::Execution(ExecutionError::WorkerLost)) - )); - assert!(service.is_empty()); - assert!(matches!( - service.enqueue(3).await, - Err(WriteError::Admission(EnqueueError::Closed)) - )); - } - - #[tokio::test] - async fn worker_panic_notifies_active_and_pending_members() { - let (entered, mut receiver) = mpsc::unbounded_channel(); - let gate = Arc::new(Semaphore::new(0)); - let process_gate = gate.clone(); - let service = WriteService::new(4, Handle::current(), move |_: Vec| { - let entered = entered.clone(); - let gate = process_gate.clone(); - Box::pin(async move { - entered.send(()).unwrap(); - gate.acquire().await.unwrap().forget(); - panic!("test processor failed"); - }) - }); - let active = service.enqueue(1); - receiver.recv().await.unwrap(); - let pending = service.enqueue(2); - let flush = service.flush(); - gate.add_permits(1); - for completion in [active, pending] { - assert!(matches!( - completion.await, - Err(WriteError::Execution(ExecutionError::WorkerLost)) - )); - } - assert!(matches!( - flush.await, - Err(FlushError::Execution(ExecutionError::WorkerLost)) - )); - assert!(service.is_empty()); - } - - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn dropping_the_last_producers_drains_work_and_stops_the_worker() { - for _ in 0..64 { - let (service, mut entered, gate) = gated(4); - let write = service.enqueue(1); - let clone = service.clone(); - let exit = service.exit.clone(); - std::thread::scope(|scope| { - scope.spawn(move || drop(service)); - scope.spawn(move || drop(clone)); - }); - entered.recv().await.unwrap(); - gate.add_permits(1); - write.await.unwrap(); - tokio::time::timeout(Duration::from_secs(1), exit) - .await - .unwrap() - .unwrap(); - } - } - - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn fast_workers_and_racing_flushes_cannot_lose_registered_completions() { - let seen = Arc::new(Mutex::new(std::collections::HashSet::new())); - let processed = Arc::new(std::sync::atomic::AtomicUsize::new(0)); - let process_seen = seen.clone(); - let process_count = processed.clone(); - let service = WriteService::new(1024, Handle::current(), move |commands: Vec| { - lock(&process_seen).extend(commands.iter().copied()); - process_count.fetch_add(commands.len(), Ordering::Relaxed); - Box::pin(async { Ok(()) }) - }); - let (writes, flushes) = std::thread::scope(|scope| { - let producers: Vec<_> = (0..4) - .map(|producer| { - let service = &service; - scope.spawn(move || { - (0..50) - .map(|item| service.enqueue(producer * 50 + item)) - .collect::>() - }) - }) - .collect(); - let flushers: Vec<_> = (0..2) - .map(|_| { - let service = &service; - scope.spawn(move || (0..50).map(|_| service.flush()).collect::>()) - }) - .collect(); - ( - producers - .into_iter() - .flat_map(|producer| producer.join().unwrap()) - .collect::>(), - flushers - .into_iter() - .flat_map(|flusher| flusher.join().unwrap()) - .collect::>(), - ) - }); - tokio::time::timeout(Duration::from_secs(5), async { - for write in writes { - write.await.unwrap(); - } - for flush in flushes { - flush.await.unwrap(); - } - }) - .await - .unwrap(); - assert_eq!(lock(&seen).len(), 200); - assert_eq!(processed.load(Ordering::Relaxed), 200); - assert!(service.is_empty()); - service.shutdown(Duration::from_secs(1)).await.unwrap(); - } - - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn racing_producers_cannot_overbook_capacity() { - let (service, _entered, gate) = gated(8); - let completions = std::thread::scope(|scope| { - let threads: Vec<_> = (0..32) - .map(|number| { - let service = &service; - scope.spawn(move || service.enqueue(number)) - }) - .collect(); - threads - .into_iter() - .map(|thread| thread.join().unwrap()) - .collect::>() - }); - assert_eq!( - completions - .iter() - .filter(|completion| completion.is_ok()) - .count(), - 8 - ); - assert_eq!(service.outstanding(), 8); - gate.add_permits(8); - let mut delivered = 0; - for completion in completions { - if completion.await.is_ok() { - delivered += 1; - } - } - assert_eq!(delivered, 8); - service.shutdown(Duration::from_secs(1)).await.unwrap(); - assert!(service.is_empty()); - } -} +mod tests; diff --git a/src/internal/delivery/tests.rs b/src/internal/delivery/tests.rs new file mode 100644 index 0000000..61f5d92 --- /dev/null +++ b/src/internal/delivery/tests.rs @@ -0,0 +1,619 @@ +use super::*; +use std::{ + future::IntoFuture, + sync::atomic::{AtomicBool, Ordering}, +}; +use tokio::sync::mpsc; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn admission_does_not_wait_for_coordinator_lock() { + let (service, mut entered, gate) = gated(4); + let coordinator = lock(&service.inner.state); + let producer = service.clone(); + let (sent, received) = std::sync::mpsc::channel(); + let thread = std::thread::spawn(move || { + let completion = producer.enqueue(7); + sent.send(completion.submission().is_ok()).unwrap(); + }); + let admitted = received.recv_timeout(Duration::from_secs(2)); + drop(coordinator); + thread.join().unwrap(); + assert!(admitted.unwrap()); + assert_eq!(entered.recv().await.unwrap(), vec![7]); + gate.add_permits(1); + service.shutdown(Duration::from_secs(2)).await.unwrap(); +} + +#[tokio::test] +async fn shutdown_includes_a_publisher_already_in_progress() { + let (service, mut entered, gate) = gated(4); + let first = service.enqueue(1); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + + service.inner.begin_publication().unwrap(); + service.inner.reserve().unwrap(); + let (sender, second) = Completion::pair(); + let shutdown = service.shutdown(Duration::from_secs(2)); + let mut shutdown = Box::pin(shutdown.into_future()); + assert!( + tokio::time::timeout(Duration::from_millis(20), &mut shutdown) + .await + .is_err() + ); + service.inner.publish(Accepted { command: 2, sender }); + gate.add_permits(1); + assert!(first.await.is_ok()); + assert_eq!(entered.recv().await.unwrap(), vec![2]); + gate.add_permits(1); + assert!(second.await.is_ok()); + assert_eq!(shutdown.await.unwrap().delivered(), 2); +} + +#[tokio::test] +async fn flush_swaps_generation_and_waits_for_its_late_publisher() { + let (service, mut entered, gate) = gated(8); + let first = service.enqueue(1); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + + service.inner.begin_publication().unwrap(); + service.inner.reserve().unwrap(); + let old_epoch = service.inner.ingress.load_full(); + assert!(old_epoch.enter()); + let (sender, before) = Completion::pair(); + let mut flush = Box::pin(service.flush().into_future()); + assert!(!Arc::ptr_eq(&old_epoch, &service.inner.ingress.load_full())); + let after = service.enqueue(3); + service + .inner + .publish_to_epoch(&old_epoch, Accepted { command: 2, sender }); + + gate.add_permits(1); + first.await.unwrap(); + assert_eq!(entered.recv().await.unwrap(), vec![2]); + assert!(futures::poll!(flush.as_mut()).is_pending()); + gate.add_permits(1); + before.await.unwrap(); + assert_eq!(flush.await.unwrap().delivered(), 2); + assert_eq!(entered.recv().await.unwrap(), vec![3]); + gate.add_permits(1); + after.await.unwrap(); + service.shutdown(Duration::from_secs(2)).await.unwrap(); +} + +#[tokio::test] +async fn flush_detaches_without_draining_the_queue() { + let (service, mut entered, gate) = gated(8); + let writes: Vec<_> = (0..4).map(|item| service.enqueue(item)).collect(); + let old_epoch = service.inner.ingress.load_full(); + let flush = service.flush(); + assert!(!Arc::ptr_eq(&old_epoch, &service.inner.ingress.load_full())); + // On the current-thread runtime, the worker has not run yet. Flush + // must not have taken the first item out of the detached queue. + assert!(old_epoch.queue.peek(&sdd::Guard::new()).is_some()); + assert_eq!(entered.recv().await.unwrap(), vec![0, 1, 2, 3]); + gate.add_permits(1); + assert_eq!(flush.await.unwrap().delivered(), 4); + for write in writes { + write.await.unwrap(); + } + service.shutdown(Duration::from_secs(2)).await.unwrap(); +} + +#[tokio::test] +async fn worker_failure_completes_a_late_publisher() { + let (service, mut entered, _gate) = gated(4); + let first = service.enqueue(1); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + service.inner.begin_publication().unwrap(); + service.inner.reserve().unwrap(); + let old_epoch = service.inner.ingress.load_full(); + assert!(old_epoch.enter()); + let (sender, late) = Completion::pair(); + service.abort.abort(); + assert!(matches!( + first.await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + service + .inner + .publish_to_epoch(&old_epoch, Accepted { command: 2, sender }); + assert!(matches!( + late.await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + assert!(service.is_empty()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn timeout_fails_items_taken_and_not_yet_taken_by_worker() { + let (service, _entered, _gate) = gated(4); + let (taken, wait_taken) = std::sync::mpsc::channel(); + let (resume, wait_resume) = std::sync::mpsc::channel(); + *lock(&service.inner.after_take) = Some((taken, wait_resume)); + let coordinator = lock(&service.inner.state); + let first = service.enqueue(1); + let second = service.enqueue(2); + drop(coordinator); + wait_taken.recv_timeout(Duration::from_secs(2)).unwrap(); + let shutdown = service.shutdown(Duration::from_millis(50)); + assert!(matches!( + shutdown.await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); + resume.send(()).unwrap(); + for completion in [first, second] { + assert!(matches!( + completion.await, + Err(WriteError::Execution(ExecutionError::TimedOut)) + )); + } + let _ = service.exit.clone().await; + assert!(service.is_empty()); +} + +#[tokio::test] +async fn terminal_stop_is_preserved_when_the_processor_returns_success() { + let (service, mut entered, gate) = gated(4); + let active = service.enqueue(1); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + let pending = service.enqueue(2); + let flush = service.flush(); + + service.inner.stop(ExecutionError::TimedOut); + gate.add_permits(1); // The processor returns Ok after its batch was failed. + for write in [active, pending] { + assert!(matches!( + write.await, + Err(WriteError::Execution(ExecutionError::TimedOut)) + )); + } + assert!(matches!( + flush.await, + Err(FlushError::Execution(ExecutionError::TimedOut)) + )); + assert_eq!(service.exit.clone().await, Err(ExecutionError::TimedOut)); + assert!(service.is_empty()); + assert!(matches!( + service.flush().submission(), + Err(FlushError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + service.shutdown(Duration::from_secs(1)).await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); +} + +fn gated( + capacity: usize, +) -> ( + WriteService, + mpsc::UnboundedReceiver>, + Arc, +) { + let (entered, receiver) = mpsc::unbounded_channel(); + let gate = Arc::new(Semaphore::new(0)); + let process_gate = gate.clone(); + let service = WriteService::new(capacity, Handle::current(), move |commands| { + let gate = process_gate.clone(); + let entered = entered.clone(); + Box::pin(async move { + entered.send(commands).unwrap(); + gate.acquire().await.unwrap().forget(); + Ok(()) + }) + }); + (service, receiver, gate) +} + +#[tokio::test] +async fn concurrent_flushes_observe_fixed_cutoffs_including_active_work() { + let (service, mut entered, gate) = gated(8); + let one = service.enqueue(1); + let first = service.flush(); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + let concurrent = service.flush(); + let two = service.enqueue(2); + let mut second = Box::pin(service.flush().into_future()); + let three = service.enqueue(3); // beyond second's cutoff + gate.add_permits(1); + assert_eq!(first.await.unwrap().delivered, 1); + assert_eq!(concurrent.await.unwrap().delivered, 1); + one.await.unwrap(); + assert_eq!(entered.recv().await.unwrap(), vec![2]); + assert!(futures::poll!(second.as_mut()).is_pending()); + gate.add_permits(1); + assert_eq!(second.await.unwrap().delivered, 2); + two.await.unwrap(); + assert_eq!(entered.recv().await.unwrap(), vec![3]); + // The second flush finished without waiting for the later arrival. + gate.add_permits(1); + three.await.unwrap(); + assert_eq!(service.flush().await.unwrap().delivered, 0); + service.shutdown(Duration::from_secs(1)).await.unwrap(); +} + +#[tokio::test] +async fn a_new_flush_does_not_inherit_a_failed_prefix_from_an_older_active_flush() { + let (entered, mut receiver) = mpsc::unbounded_channel(); + let gate = Arc::new(Semaphore::new(0)); + let process_gate = gate.clone(); + let service = WriteService::new(4, Handle::current(), move |commands: Vec| { + let entered = entered.clone(); + let gate = process_gate.clone(); + Box::pin(async move { + entered.send(commands.clone()).unwrap(); + gate.acquire().await.unwrap().forget(); + if commands == [1] { + anyhow::bail!("first batch failed"); + } + Ok(()) + }) + }); + let first = service.enqueue(1); + assert_eq!(receiver.recv().await.unwrap(), vec![1]); + let second = service.enqueue(2); + let covering = service.flush(); + gate.add_permits(1); + assert!(matches!(first.await, Err(WriteError::Delivery(_)))); + assert_eq!(receiver.recv().await.unwrap(), vec![2]); + let fresh = service.flush(); + gate.add_permits(1); + second.await.unwrap(); + let Err(FlushError::Delivery(failure)) = covering.await else { + panic!("expected captured failure"); + }; + assert_eq!(failure.failed, 1); + assert_eq!(fresh.await.unwrap().delivered, 1); + service.shutdown(Duration::from_secs(1)).await.unwrap(); +} + +#[tokio::test] +async fn graceful_shutdown_reports_its_failed_batch_after_notifying_all_writes() { + let service = WriteService::new(4, Handle::current(), |_: Vec| { + Box::pin(async { + anyhow::bail!("terminal failure"); + }) + }); + let first = service.enqueue(1); + let second = service.enqueue(2); + let shutdown = service.shutdown(Duration::from_secs(1)); + let Err(ShutdownError::Delivery(failure)) = shutdown.await else { + panic!("expected delivery failure"); + }; + assert_eq!(failure.failed, 2); + for completion in [first, second] { + assert!(matches!(completion.await, Err(WriteError::Delivery(_)))); + } + assert!(service.is_empty()); +} + +#[tokio::test] +async fn capacity_includes_in_flight_work_and_rejection_is_observable_twice() { + let (service, mut entered, gate) = gated(1); + let first = service.enqueue(1); + entered.recv().await.unwrap(); + let rejected = service.enqueue(2); + assert!(matches!( + &*rejected, + Err(WriteError::Admission(EnqueueError::Full)) + )); + assert!(matches!( + rejected.await, + Err(WriteError::Admission(EnqueueError::Full)) + )); + assert_eq!(service.outstanding(), 1); + gate.add_permits(1); + first.await.unwrap(); + assert!(service.is_empty()); + gate.add_permits(1); + service.enqueue(3).await.unwrap(); + service.shutdown(Duration::from_secs(1)).await.unwrap(); +} + +#[tokio::test] +async fn dropped_observer_does_not_cancel_automatic_delivery_or_fill_the_queue() { + let (service, mut entered, gate) = gated(2); + drop(service.enqueue(1)); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + let flush = service.flush(); + drop(flush); // lifecycle observers also do not cancel accepted work + gate.add_permits(1); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + assert!(service.is_empty()); +} + +#[tokio::test] +async fn every_failed_batch_notifies_all_members_releases_errors_and_allows_fresh_work() { + let fail = Arc::new(AtomicBool::new(true)); + let process_fail = fail.clone(); + let service = WriteService::new(4, Handle::current(), move |_: Vec| { + let fail = process_fail.load(Ordering::Relaxed); + Box::pin(async move { + if fail { + anyhow::bail!("terminal failure"); + } + Ok(()) + }) + }); + for _ in 0..100 { + let first = service.enqueue(1); + let second = service.enqueue(2); + let flush = service.flush(); + let Err(FlushError::Delivery(failure)) = flush.await else { + panic!("expected failure"); + }; + assert_eq!(failure.failed, 2); + let weak = Arc::downgrade(&failure.cause); + for completion in [first, second] { + let Err(WriteError::Delivery(error)) = completion.await else { + panic!("expected failure"); + }; + assert!(Arc::ptr_eq(&error, &failure.cause)); + } + drop(failure); + assert!(weak.upgrade().is_none(), "SDK must release terminal errors"); + assert_eq!(service.outstanding(), 0); + assert_eq!(service.flush().await.unwrap().delivered, 0); + } + fail.store(false, Ordering::Relaxed); + let fresh = service.enqueue(3); + assert_eq!(service.flush().await.unwrap().delivered, 1); + fresh.await.unwrap(); + service.shutdown(Duration::from_secs(1)).await.unwrap(); +} + +#[tokio::test] +async fn shutdown_stops_admission_immediately_drains_and_is_shared_by_clones() { + let (service, mut entered, gate) = gated(4); + let write = service.enqueue(1); + let shutdown = service.shutdown(Duration::from_secs(1)); + let concurrent = service.clone().shutdown(Duration::from_secs(1)); + assert!(shutdown.is_ok()); + assert!(matches!( + service.enqueue(2).await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); + entered.recv().await.unwrap(); + gate.add_permits(1); + write.await.unwrap(); + assert_eq!(shutdown.await.unwrap().delivered, 1); + assert_eq!(concurrent.await.unwrap().delivered, 1); + assert!(service.is_empty()); +} + +#[tokio::test] +async fn shutdown_timeout_fails_active_work_even_when_flush_slots_are_full() { + let (service, mut entered, _) = gated(1); + let write = service.enqueue(1); + entered.recv().await.unwrap(); + let flush = service.flush(); + let rejected = service.flush(); + assert!(matches!( + &*rejected, + Err(FlushError::Admission(EnqueueError::Full)) + )); + let shutdown = service.shutdown(Duration::from_millis(10)); + assert!( + shutdown.is_ok(), + "shutdown must not compete with flush slots" + ); + assert!(matches!( + shutdown.await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + write.await, + Err(WriteError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + flush.await, + Err(FlushError::Execution(ExecutionError::TimedOut)) + )); + assert!(service.is_empty()); + assert!(matches!( + service.enqueue(2).await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); +} + +#[tokio::test] +async fn shutdown_timeout_also_retires_batches_queued_behind_active_work() { + let (service, mut entered, _) = gated(4); + let first = service.enqueue(1); + entered.recv().await.unwrap(); + let second = service.enqueue(2); + let third = service.enqueue(3); + let flush = service.flush(); + let shutdown = service.shutdown(Duration::from_millis(10)); + assert!(matches!( + shutdown.await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); + for completion in [first, second, third] { + assert!(matches!( + completion.await, + Err(WriteError::Execution(ExecutionError::TimedOut)) + )); + } + assert!(matches!( + flush.await, + Err(FlushError::Execution(ExecutionError::TimedOut)) + )); + assert_eq!(service.outstanding(), 0); + assert!(lock(&service.inner.state).barriers.is_empty()); + assert!(lock(&service.inner.state).ready.is_empty()); +} + +#[tokio::test] +async fn very_large_write_capacity_does_not_overflow_lifecycle_control_limits() { + let (service, _, _) = gated(usize::MAX); + assert_eq!(service.inner.control_limit, 1024); + service.shutdown(Duration::from_secs(1)).await.unwrap(); +} + +#[tokio::test] +async fn worker_cancellation_before_first_poll_completes_every_accepted_sender() { + let (service, _, _) = gated(4); + let first = service.enqueue(1); + let second = service.enqueue(2); + let flush = service.flush(); + service.abort.abort(); + for completion in [first, second] { + assert!(matches!( + completion.await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + } + assert!(matches!( + flush.await, + Err(FlushError::Execution(ExecutionError::WorkerLost)) + )); + assert!(service.is_empty()); + assert!(matches!( + service.enqueue(3).await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); +} + +#[tokio::test] +async fn worker_panic_notifies_active_and_pending_members() { + let (entered, mut receiver) = mpsc::unbounded_channel(); + let gate = Arc::new(Semaphore::new(0)); + let process_gate = gate.clone(); + let service = WriteService::new(4, Handle::current(), move |_: Vec| { + let entered = entered.clone(); + let gate = process_gate.clone(); + Box::pin(async move { + entered.send(()).unwrap(); + gate.acquire().await.unwrap().forget(); + panic!("test processor failed"); + }) + }); + let active = service.enqueue(1); + receiver.recv().await.unwrap(); + let pending = service.enqueue(2); + let flush = service.flush(); + gate.add_permits(1); + for completion in [active, pending] { + assert!(matches!( + completion.await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + } + assert!(matches!( + flush.await, + Err(FlushError::Execution(ExecutionError::WorkerLost)) + )); + assert!(service.is_empty()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn dropping_the_last_producers_drains_work_and_stops_the_worker() { + for _ in 0..64 { + let (service, mut entered, gate) = gated(4); + let write = service.enqueue(1); + let clone = service.clone(); + let exit = service.exit.clone(); + std::thread::scope(|scope| { + scope.spawn(move || drop(service)); + scope.spawn(move || drop(clone)); + }); + entered.recv().await.unwrap(); + gate.add_permits(1); + write.await.unwrap(); + tokio::time::timeout(Duration::from_secs(1), exit) + .await + .unwrap() + .unwrap(); + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn fast_workers_and_racing_flushes_cannot_lose_registered_completions() { + let seen = Arc::new(Mutex::new(std::collections::HashSet::new())); + let processed = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let process_seen = seen.clone(); + let process_count = processed.clone(); + let service = WriteService::new(1024, Handle::current(), move |commands: Vec| { + lock(&process_seen).extend(commands.iter().copied()); + process_count.fetch_add(commands.len(), Ordering::Relaxed); + Box::pin(async { Ok(()) }) + }); + let (writes, flushes) = std::thread::scope(|scope| { + let producers: Vec<_> = (0..4) + .map(|producer| { + let service = &service; + scope.spawn(move || { + (0..50) + .map(|item| service.enqueue(producer * 50 + item)) + .collect::>() + }) + }) + .collect(); + let flushers: Vec<_> = (0..2) + .map(|_| { + let service = &service; + scope.spawn(move || (0..50).map(|_| service.flush()).collect::>()) + }) + .collect(); + ( + producers + .into_iter() + .flat_map(|producer| producer.join().unwrap()) + .collect::>(), + flushers + .into_iter() + .flat_map(|flusher| flusher.join().unwrap()) + .collect::>(), + ) + }); + tokio::time::timeout(Duration::from_secs(5), async { + for write in writes { + write.await.unwrap(); + } + for flush in flushes { + flush.await.unwrap(); + } + }) + .await + .unwrap(); + assert_eq!(lock(&seen).len(), 200); + assert_eq!(processed.load(Ordering::Relaxed), 200); + assert!(service.is_empty()); + service.shutdown(Duration::from_secs(1)).await.unwrap(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn racing_producers_cannot_overbook_capacity() { + let (service, _entered, gate) = gated(8); + let completions = std::thread::scope(|scope| { + let threads: Vec<_> = (0..32) + .map(|number| { + let service = &service; + scope.spawn(move || service.enqueue(number)) + }) + .collect(); + threads + .into_iter() + .map(|thread| thread.join().unwrap()) + .collect::>() + }); + assert_eq!( + completions + .iter() + .filter(|completion| completion.is_ok()) + .count(), + 8 + ); + assert_eq!(service.outstanding(), 8); + gate.add_permits(8); + let mut delivered = 0; + for completion in completions { + if completion.await.is_ok() { + delivered += 1; + } + } + assert_eq!(delivered, 8); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + assert!(service.is_empty()); +} From 342784af6f9a484735507a7babf519e74d2dfd71 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Thu, 29 Oct 2026 22:21:41 +0800 Subject: [PATCH 09/10] Consolidate delivery producer ownership and admission --- src/internal/delivery.rs | 132 ++++++++++++++------------------- src/internal/delivery/tests.rs | 59 +++++++++------ 2 files changed, 92 insertions(+), 99 deletions(-) diff --git a/src/internal/delivery.rs b/src/internal/delivery.rs index 5c3c478..699aee4 100644 --- a/src/internal/delivery.rs +++ b/src/internal/delivery.rs @@ -33,31 +33,28 @@ type FlushSender = oneshot::Sender>; /// no completed-outcome registry: a batch owns its write senders until all its /// physical requests finish, and barriers summarize only active work. pub(crate) struct WriteService { - inner: Arc>, - owners: Arc>, - abort: AbortHandle, - runtime: Handle, - exit: Shared>>, + producer: Arc>, } impl Clone for WriteService { fn clone(&self) -> Self { Self { - inner: self.inner.clone(), - owners: self.owners.clone(), - abort: self.abort.clone(), - runtime: self.runtime.clone(), - exit: self.exit.clone(), + producer: self.producer.clone(), } } } -struct ProducerOwner { - inner: std::sync::Weak>, +/// Only producer handles retain this owner. The worker retains `Inner` alone, +/// so dropping the last producer closes admission while queued work can drain. +struct Producer { + inner: Arc>, + abort: AbortHandle, + runtime: Handle, + exit: Shared>>, + control_limit: usize, + shutdown_controls: Arc, } -impl Drop for ProducerOwner { +impl Drop for Producer { fn drop(&mut self) { - if let Some(inner) = self.inner.upgrade() { - inner.close(); - } + self.inner.close(); } } impl WriteService { @@ -84,8 +81,6 @@ impl WriteService { depth_warned: AtomicBool::new(false), wake: Wake::new(), capacity, - control_limit, - shutdown_controls: Arc::new(Semaphore::new(control_limit)), #[cfg(test)] after_take: Mutex::new(None), }); @@ -100,15 +95,15 @@ impl WriteService { exit: Some(sender), }; let task = runtime.spawn(async move { run(guard, process).await }); - let owners = Arc::new(ProducerOwner { - inner: Arc::downgrade(&inner), - }); Self { - inner, - owners, - abort: task.abort_handle(), - runtime, - exit, + producer: Arc::new(Producer { + inner, + abort: task.abort_handle(), + runtime, + exit, + control_limit, + shutdown_controls: Arc::new(Semaphore::new(control_limit)), + }), } } @@ -122,13 +117,9 @@ impl WriteService { command: C, rejected: impl FnOnce(C, EnqueueError), ) -> Completion { + let producer = &self.producer; let (sender, mut completion) = Completion::pair(); - let reservation = self.inner.begin_publication().and_then(|()| { - self.inner - .reserve() - .inspect_err(|_| self.inner.end_publication()) - }); - match reservation { + match producer.inner.admit() { Err(error) => { completion.record_issuance(Err(error.into())); rejected(command, error); @@ -136,13 +127,13 @@ impl WriteService { Ok(depth) => { // Publication is the linearization point: completed submissions // appear in FIFO order before a later flush snapshots the queue. - self.inner.publish(Accepted { command, sender }); - if depth >= self.inner.capacity - self.inner.capacity / 5 - && !self.inner.depth_warned.swap(true, Ordering::Relaxed) + producer.inner.publish(Accepted { command, sender }); + if depth >= producer.inner.capacity - producer.inner.capacity / 5 + && !producer.inner.depth_warned.swap(true, Ordering::Relaxed) { warn!( outstanding = depth, - capacity = self.inner.capacity, + capacity = producer.inner.capacity, "delivery queue has reached 80% of capacity" ); } @@ -152,50 +143,45 @@ impl WriteService { } pub fn flush(&self) -> Completion { + let producer = &self.producer; let (sender, mut completion) = Completion::pair(); - let mut state = lock(&self.inner.state); - if let Some(error) = self.inner.terminal_error() { + let mut state = lock(&producer.inner.state); + if let Some(error) = producer.inner.terminal_error() { completion.record_issuance(Err(error.into())); - } else if self.inner.publishers.load(Ordering::SeqCst) & CLOSED != 0 { + } else if producer.inner.publishers.load(Ordering::SeqCst) & CLOSED != 0 { completion.record_issuance(Err(EnqueueError::Closed.into())); - } else if state - .barriers - .iter() - .filter(|barrier| !barrier.shutdown) - .count() - >= self.inner.control_limit - { + } else if state.barriers.len() >= producer.control_limit { completion.record_issuance(Err(EnqueueError::Full.into())); } else { - let through = self.inner.detach(&mut state); - state.barrier(sender, false, through); - self.inner.wake.notify_one(); + let through = producer.inner.detach(&mut state); + state.barrier(sender, through); + producer.inner.wake.notify_one(); } completion } pub fn shutdown(&self, timeout: Duration) -> Completion { + let producer = &self.producer; let (sender, mut completion) = Completion::pair(); - let Ok(permit) = self.inner.shutdown_controls.clone().try_acquire_owned() else { + let Ok(permit) = producer.shutdown_controls.clone().try_acquire_owned() else { completion.record_issuance(Err(EnqueueError::Full.into())); return completion; }; let (barrier_sender, barrier) = Completion::pair(); { - let mut state = lock(&self.inner.state); - if let Some(error) = self.inner.terminal_error() { + let mut state = lock(&producer.inner.state); + if let Some(error) = producer.inner.terminal_error() { completion.record_issuance(Err(error.into())); return completion; } - self.inner.publishers.fetch_or(CLOSED, Ordering::SeqCst); + producer.inner.close(); state.shutdown_pending.push(barrier_sender); - self.inner.register_shutdowns(&mut state); + producer.inner.register_shutdowns(&mut state); } - let inner = self.inner.clone(); - let abort = self.abort.clone(); - let exit = self.exit.clone(); - self.inner.wake.notify_one(); - self.runtime.spawn(async move { + let inner = producer.inner.clone(); + let abort = producer.abort.clone(); + let exit = producer.exit.clone(); + producer.runtime.spawn(async move { let _permit = permit; let work = async { let report = barrier.await; @@ -221,11 +207,11 @@ impl WriteService { #[cfg(test)] pub fn is_empty(&self) -> bool { - self.inner.outstanding.load(Ordering::Acquire) == 0 + self.producer.inner.outstanding.load(Ordering::Acquire) == 0 } #[cfg(test)] fn outstanding(&self) -> usize { - self.inner.outstanding.load(Ordering::Acquire) + self.producer.inner.outstanding.load(Ordering::Acquire) } } @@ -319,7 +305,6 @@ struct Barrier { sender: FlushSender, delivered: u64, failure: Option, - shutdown: bool, } impl Barrier { fn record(&mut self, count: u64, result: &Result<(), Arc>) { @@ -353,13 +338,12 @@ struct State { shutdown_pending: Vec, } impl State { - fn barrier(&mut self, sender: FlushSender, shutdown: bool, through: Arc>) { + fn barrier(&mut self, sender: FlushSender, through: Arc>) { self.barriers.push_back(Barrier { through, sender, delivered: 0, failure: None, - shutdown, }); } fn finish(&mut self, result: &Result<(), Arc>) -> usize { @@ -392,8 +376,6 @@ struct Inner { depth_warned: AtomicBool, wake: Wake, capacity: usize, - control_limit: usize, - shutdown_controls: Arc, #[cfg(test)] after_take: Mutex, std::sync::mpsc::Receiver<()>)>>, } @@ -406,25 +388,25 @@ impl Inner { } } - fn begin_publication(&self) -> Result<(), EnqueueError> { - enter(&self.publishers) - .then_some(()) - .ok_or(EnqueueError::Closed) - } - fn end_publication(&self) { self.publishers.fetch_sub(1, Ordering::SeqCst); self.wake.notify_one(); } #[allow(deprecated, reason = "try_update requires Rust 1.95")] - fn reserve(&self) -> Result { + fn admit(&self) -> Result { + if !enter(&self.publishers) { + return Err(EnqueueError::Closed); + } self.outstanding .fetch_update(Ordering::AcqRel, Ordering::Relaxed, |depth| { (depth < self.capacity).then(|| depth + 1) }) .map(|previous| previous + 1) - .map_err(|_| EnqueueError::Full) + .map_err(|_| { + self.end_publication(); + EnqueueError::Full + }) } fn release(&self, count: usize) { @@ -449,7 +431,7 @@ impl Inner { } let through = self.detach(state); for sender in std::mem::take(&mut state.shutdown_pending) { - state.barrier(sender, true, through.clone()); + state.barrier(sender, through.clone()); } true } diff --git a/src/internal/delivery/tests.rs b/src/internal/delivery/tests.rs index 61f5d92..197515f 100644 --- a/src/internal/delivery/tests.rs +++ b/src/internal/delivery/tests.rs @@ -8,7 +8,7 @@ use tokio::sync::mpsc; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn admission_does_not_wait_for_coordinator_lock() { let (service, mut entered, gate) = gated(4); - let coordinator = lock(&service.inner.state); + let coordinator = lock(&service.producer.inner.state); let producer = service.clone(); let (sent, received) = std::sync::mpsc::channel(); let thread = std::thread::spawn(move || { @@ -30,8 +30,7 @@ async fn shutdown_includes_a_publisher_already_in_progress() { let first = service.enqueue(1); assert_eq!(entered.recv().await.unwrap(), vec![1]); - service.inner.begin_publication().unwrap(); - service.inner.reserve().unwrap(); + service.producer.inner.admit().unwrap(); let (sender, second) = Completion::pair(); let shutdown = service.shutdown(Duration::from_secs(2)); let mut shutdown = Box::pin(shutdown.into_future()); @@ -40,7 +39,10 @@ async fn shutdown_includes_a_publisher_already_in_progress() { .await .is_err() ); - service.inner.publish(Accepted { command: 2, sender }); + service + .producer + .inner + .publish(Accepted { command: 2, sender }); gate.add_permits(1); assert!(first.await.is_ok()); assert_eq!(entered.recv().await.unwrap(), vec![2]); @@ -55,15 +57,18 @@ async fn flush_swaps_generation_and_waits_for_its_late_publisher() { let first = service.enqueue(1); assert_eq!(entered.recv().await.unwrap(), vec![1]); - service.inner.begin_publication().unwrap(); - service.inner.reserve().unwrap(); - let old_epoch = service.inner.ingress.load_full(); + service.producer.inner.admit().unwrap(); + let old_epoch = service.producer.inner.ingress.load_full(); assert!(old_epoch.enter()); let (sender, before) = Completion::pair(); let mut flush = Box::pin(service.flush().into_future()); - assert!(!Arc::ptr_eq(&old_epoch, &service.inner.ingress.load_full())); + assert!(!Arc::ptr_eq( + &old_epoch, + &service.producer.inner.ingress.load_full() + )); let after = service.enqueue(3); service + .producer .inner .publish_to_epoch(&old_epoch, Accepted { command: 2, sender }); @@ -84,9 +89,12 @@ async fn flush_swaps_generation_and_waits_for_its_late_publisher() { async fn flush_detaches_without_draining_the_queue() { let (service, mut entered, gate) = gated(8); let writes: Vec<_> = (0..4).map(|item| service.enqueue(item)).collect(); - let old_epoch = service.inner.ingress.load_full(); + let old_epoch = service.producer.inner.ingress.load_full(); let flush = service.flush(); - assert!(!Arc::ptr_eq(&old_epoch, &service.inner.ingress.load_full())); + assert!(!Arc::ptr_eq( + &old_epoch, + &service.producer.inner.ingress.load_full() + )); // On the current-thread runtime, the worker has not run yet. Flush // must not have taken the first item out of the detached queue. assert!(old_epoch.queue.peek(&sdd::Guard::new()).is_some()); @@ -104,17 +112,17 @@ async fn worker_failure_completes_a_late_publisher() { let (service, mut entered, _gate) = gated(4); let first = service.enqueue(1); assert_eq!(entered.recv().await.unwrap(), vec![1]); - service.inner.begin_publication().unwrap(); - service.inner.reserve().unwrap(); - let old_epoch = service.inner.ingress.load_full(); + service.producer.inner.admit().unwrap(); + let old_epoch = service.producer.inner.ingress.load_full(); assert!(old_epoch.enter()); let (sender, late) = Completion::pair(); - service.abort.abort(); + service.producer.abort.abort(); assert!(matches!( first.await, Err(WriteError::Execution(ExecutionError::WorkerLost)) )); service + .producer .inner .publish_to_epoch(&old_epoch, Accepted { command: 2, sender }); assert!(matches!( @@ -129,8 +137,8 @@ async fn timeout_fails_items_taken_and_not_yet_taken_by_worker() { let (service, _entered, _gate) = gated(4); let (taken, wait_taken) = std::sync::mpsc::channel(); let (resume, wait_resume) = std::sync::mpsc::channel(); - *lock(&service.inner.after_take) = Some((taken, wait_resume)); - let coordinator = lock(&service.inner.state); + *lock(&service.producer.inner.after_take) = Some((taken, wait_resume)); + let coordinator = lock(&service.producer.inner.state); let first = service.enqueue(1); let second = service.enqueue(2); drop(coordinator); @@ -147,7 +155,7 @@ async fn timeout_fails_items_taken_and_not_yet_taken_by_worker() { Err(WriteError::Execution(ExecutionError::TimedOut)) )); } - let _ = service.exit.clone().await; + let _ = service.producer.exit.clone().await; assert!(service.is_empty()); } @@ -159,7 +167,7 @@ async fn terminal_stop_is_preserved_when_the_processor_returns_success() { let pending = service.enqueue(2); let flush = service.flush(); - service.inner.stop(ExecutionError::TimedOut); + service.producer.inner.stop(ExecutionError::TimedOut); gate.add_permits(1); // The processor returns Ok after its batch was failed. for write in [active, pending] { assert!(matches!( @@ -171,7 +179,10 @@ async fn terminal_stop_is_preserved_when_the_processor_returns_success() { flush.await, Err(FlushError::Execution(ExecutionError::TimedOut)) )); - assert_eq!(service.exit.clone().await, Err(ExecutionError::TimedOut)); + assert_eq!( + service.producer.exit.clone().await, + Err(ExecutionError::TimedOut) + ); assert!(service.is_empty()); assert!(matches!( service.flush().submission(), @@ -440,14 +451,14 @@ async fn shutdown_timeout_also_retires_batches_queued_behind_active_work() { Err(FlushError::Execution(ExecutionError::TimedOut)) )); assert_eq!(service.outstanding(), 0); - assert!(lock(&service.inner.state).barriers.is_empty()); - assert!(lock(&service.inner.state).ready.is_empty()); + assert!(lock(&service.producer.inner.state).barriers.is_empty()); + assert!(lock(&service.producer.inner.state).ready.is_empty()); } #[tokio::test] async fn very_large_write_capacity_does_not_overflow_lifecycle_control_limits() { let (service, _, _) = gated(usize::MAX); - assert_eq!(service.inner.control_limit, 1024); + assert_eq!(service.producer.control_limit, 1024); service.shutdown(Duration::from_secs(1)).await.unwrap(); } @@ -457,7 +468,7 @@ async fn worker_cancellation_before_first_poll_completes_every_accepted_sender() let first = service.enqueue(1); let second = service.enqueue(2); let flush = service.flush(); - service.abort.abort(); + service.producer.abort.abort(); for completion in [first, second] { assert!(matches!( completion.await, @@ -513,7 +524,7 @@ async fn dropping_the_last_producers_drains_work_and_stops_the_worker() { let (service, mut entered, gate) = gated(4); let write = service.enqueue(1); let clone = service.clone(); - let exit = service.exit.clone(); + let exit = service.producer.exit.clone(); std::thread::scope(|scope| { scope.spawn(move || drop(service)); scope.spawn(move || drop(clone)); From 1e4f901fe44b028f9a3ea59a204d7651cc38acfc Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Thu, 29 Oct 2026 22:35:02 +0800 Subject: [PATCH 10/10] Unify delivery shutdown with generation cutoffs --- src/internal/delivery.rs | 863 ++++++++++++++++++++++++++++----- src/internal/delivery/tests.rs | 630 ------------------------ 2 files changed, 735 insertions(+), 758 deletions(-) delete mode 100644 src/internal/delivery/tests.rs diff --git a/src/internal/delivery.rs b/src/internal/delivery.rs index 699aee4..3d346bb 100644 --- a/src/internal/delivery.rs +++ b/src/internal/delivery.rs @@ -5,7 +5,7 @@ use crate::{ Completion, DeliveryError, DeliveryReceipt, EnqueueError, ExecutionError, FlushError, FlushFailure, FlushReport, ShutdownError, ShutdownReport, WriteError, }; -use arc_swap::ArcSwap; +use arc_swap::{ArcSwap, Guard}; use crossbeam_utils::atomic::AtomicCell; use futures::{ future::{BoxFuture, Shared}, @@ -54,7 +54,7 @@ struct Producer { } impl Drop for Producer { fn drop(&mut self) { - self.inner.close(); + self.inner.close(&mut lock(&self.inner.state)); } } impl WriteService { @@ -72,11 +72,9 @@ impl WriteService { ready: VecDeque::new(), active: None, barriers: VecDeque::new(), - shutdown_pending: Vec::new(), }), - ingress: ArcSwap::from_pointee(QueueEpoch::new()), + ingress: ArcSwap::from_pointee(QueueEpoch::new(false)), outstanding: AtomicUsize::new(0), - publishers: AtomicUsize::new(0), terminal: AtomicU8::new(0), depth_warned: AtomicBool::new(false), wake: Wake::new(), @@ -124,10 +122,10 @@ impl WriteService { completion.record_issuance(Err(error.into())); rejected(command, error); } - Ok(depth) => { - // Publication is the linearization point: completed submissions - // appear in FIFO order before a later flush snapshots the queue. - producer.inner.publish(Accepted { command, sender }); + Ok((epoch, depth)) => { + // Admission fixes the snapshot membership; publication keeps + // FIFO order within that generation. + producer.inner.publish(&epoch, Accepted { command, sender }); if depth >= producer.inner.capacity - producer.inner.capacity / 5 && !producer.inner.depth_warned.swap(true, Ordering::Relaxed) { @@ -148,12 +146,12 @@ impl WriteService { let mut state = lock(&producer.inner.state); if let Some(error) = producer.inner.terminal_error() { completion.record_issuance(Err(error.into())); - } else if producer.inner.publishers.load(Ordering::SeqCst) & CLOSED != 0 { + } else if producer.inner.ingress.load().closed() { completion.record_issuance(Err(EnqueueError::Closed.into())); } else if state.barriers.len() >= producer.control_limit { completion.record_issuance(Err(EnqueueError::Full.into())); } else { - let through = producer.inner.detach(&mut state); + let through = producer.inner.detach(&mut state, false); state.barrier(sender, through); producer.inner.wake.notify_one(); } @@ -174,9 +172,17 @@ impl WriteService { completion.record_issuance(Err(error.into())); return completion; } - producer.inner.close(); - state.shutdown_pending.push(barrier_sender); - producer.inner.register_shutdowns(&mut state); + producer.inner.close(&mut state); + let through = state + .ready + .back() + .cloned() + .or_else(|| state.active.as_ref().map(|batch| batch.epoch.clone())); + if let Some(through) = through { + state.barrier(barrier_sender, through); + } else { + let _ = barrier_sender.send(Ok(FlushReport::default())); + } } let inner = producer.inner.clone(); let abort = producer.abort.clone(); @@ -262,15 +268,16 @@ struct Accepted { sender: WriteSender, } type IngressItem = AtomicCell>>>; +type EpochGuard = Guard>>; struct QueueEpoch { queue: sdd::Queue>, publishers: AtomicUsize, } impl QueueEpoch { - fn new() -> Self { + fn new(closed: bool) -> Self { Self { queue: sdd::Queue::new(), - publishers: AtomicUsize::new(0), + publishers: AtomicUsize::new(if closed { CLOSED } else { 0 }), } } fn enter(&self) -> bool { @@ -279,6 +286,9 @@ impl QueueEpoch { fn close(&self) { self.publishers.fetch_or(CLOSED, Ordering::SeqCst); } + fn closed(&self) -> bool { + self.publishers.load(Ordering::SeqCst) & CLOSED != 0 + } fn settled(&self) -> bool { self.publishers.load(Ordering::SeqCst) & !CLOSED == 0 } @@ -335,7 +345,6 @@ struct State { ready: VecDeque>>, active: Option>, barriers: VecDeque>, - shutdown_pending: Vec, } impl State { fn barrier(&mut self, sender: FlushSender, through: Arc>) { @@ -371,7 +380,6 @@ struct Inner { state: Mutex>, ingress: ArcSwap>, outstanding: AtomicUsize, - publishers: AtomicUsize, terminal: AtomicU8, depth_warned: AtomicBool, wake: Wake, @@ -388,25 +396,36 @@ impl Inner { } } - fn end_publication(&self) { - self.publishers.fetch_sub(1, Ordering::SeqCst); - self.wake.notify_one(); - } - #[allow(deprecated, reason = "try_update requires Rust 1.95")] - fn admit(&self) -> Result { - if !enter(&self.publishers) { - return Err(EnqueueError::Closed); + fn admit(&self) -> Result<(EpochGuard, usize), EnqueueError> { + loop { + // Enrollment binds this write to one flush/shutdown cutoff before + // admission succeeds. A closed current generation rejects new work; + // a stale generation means a concurrent flush advanced the queue. + let epoch = self.ingress.load(); + if !epoch.enter() { + if Arc::ptr_eq(&epoch, &self.ingress.load()) { + return Err(EnqueueError::Closed); + } + continue; + } + match self + .outstanding + .fetch_update(Ordering::AcqRel, Ordering::Relaxed, |depth| { + (depth < self.capacity).then(|| depth + 1) + }) { + Ok(previous) => return Ok((epoch, previous + 1)), + Err(_) => { + self.end_publication(&epoch); + return Err(EnqueueError::Full); + } + } } - self.outstanding - .fetch_update(Ordering::AcqRel, Ordering::Relaxed, |depth| { - (depth < self.capacity).then(|| depth + 1) - }) - .map(|previous| previous + 1) - .map_err(|_| { - self.end_publication(); - EnqueueError::Full - }) + } + + fn end_publication(&self, epoch: &QueueEpoch) { + epoch.publishers.fetch_sub(1, Ordering::SeqCst); + self.wake.notify_one(); } fn release(&self, count: usize) { @@ -415,27 +434,13 @@ impl Inner { } } - fn detach(&self, state: &mut State) -> Arc> { - let epoch = self.ingress.swap(Arc::new(QueueEpoch::new())); + fn detach(&self, state: &mut State, closing: bool) -> Arc> { + let epoch = self.ingress.swap(Arc::new(QueueEpoch::new(closing))); epoch.close(); state.ready.push_back(epoch.clone()); epoch } - fn register_shutdowns(&self, state: &mut State) -> bool { - if state.shutdown_pending.is_empty() { - return true; - } - if self.publishers.load(Ordering::SeqCst) & !CLOSED != 0 { - return false; - } - let through = self.detach(state); - for sender in std::mem::take(&mut state.shutdown_pending) { - state.barrier(sender, through.clone()); - } - true - } - fn fail_epoch(&self, epoch: &QueueEpoch, error: ExecutionError) { while let Some(entry) = epoch.queue.pop() { let accepted = entry.take().expect("each queue entry is taken once"); @@ -444,37 +449,15 @@ impl Inner { } } - fn publish(&self, accepted: Accepted) - where - C: 'static, - { - loop { - // Keep the generation borrowed only for publication; unlike an - // owned load this avoids contending on its Arc reference count on - // every submission. A concurrent swap keeps the guard alive. - let epoch = self.ingress.load(); - if epoch.enter() { - self.publish_to_epoch(&epoch, accepted); - return; - } - if let Some(error) = self.terminal_error() { - let _ = accepted.sender.send(Err(error.into())); - self.release(1); - self.end_publication(); - return; - } - } - } - - fn publish_to_epoch(&self, epoch: &QueueEpoch, accepted: Accepted) + fn publish(&self, epoch: &QueueEpoch, accepted: Accepted) where C: 'static, { epoch.queue.push(AtomicCell::new(Some(Box::new(accepted)))); - epoch.publishers.fetch_sub(1, Ordering::SeqCst); - let terminal = self.terminal_error(); - self.end_publication(); - if let Some(error) = terminal { + // Retire publication before checking terminal state: either stop() + // observes the queued item or this publisher performs its cleanup. + self.end_publication(epoch); + if let Some(error) = self.terminal_error() { self.fail_epoch(epoch, error); } } @@ -491,7 +474,7 @@ impl Inner { // Serialize the first terminal transition with batch completion so a // returning processor cannot report success after shutdown has failed it. self.terminal.store(code, Ordering::SeqCst); - self.publishers.fetch_or(CLOSED, Ordering::SeqCst); + self.close(&mut state); if let Some(active) = state.active.take() { let count = active.senders.len(); finish_writes(active.senders, Err(error.into())); @@ -504,18 +487,12 @@ impl Inner { for barrier in state.barriers.drain(..) { let _ = barrier.sender.send(Err(error.into())); } - for sender in state.shutdown_pending.drain(..) { - let _ = sender.send(Err(error.into())); - } - drop(state); - let current = self.ingress.load_full(); - current.close(); - self.fail_epoch(¤t, error); - self.wake.notify_one(); error } - fn close(&self) { - self.publishers.fetch_or(CLOSED, Ordering::SeqCst); + fn close(&self, state: &mut State) { + if !self.ingress.load().closed() { + self.detach(state, true); + } self.wake.notify_one(); } } @@ -543,27 +520,21 @@ where let mut state = lock(&guard.inner.state); if guard.inner.terminal_error().is_some() { break; - } else if !guard.inner.register_shutdowns(&mut state) { - None + } + if state.ready.is_empty() && !guard.inner.ingress.load().queue.is_empty() { + guard.inner.detach(&mut state, false); + } + if state.ready.front().is_some_and(|epoch| epoch.settled()) { + let epoch = state.ready.pop_front().expect("checked above"); + state.active = Some(Batch { + epoch: epoch.clone(), + senders: Vec::new(), + }); + Some(epoch) + } else if state.ready.is_empty() && guard.inner.ingress.load().closed() { + break; } else { - // Check closure before queue contents: if the last producer - // finishes after this snapshot, its wake triggers another pass. - let closed = guard.inner.publishers.load(Ordering::SeqCst) == CLOSED; - if state.ready.is_empty() && !guard.inner.ingress.load().queue.is_empty() { - guard.inner.detach(&mut state); - } - if state.ready.front().is_some_and(|epoch| epoch.settled()) { - let epoch = state.ready.pop_front().expect("checked above"); - state.active = Some(Batch { - epoch: epoch.clone(), - senders: Vec::new(), - }); - Some(epoch) - } else if state.ready.is_empty() && closed { - break; - } else { - None - } + None } }; let Some(epoch) = next else { @@ -606,25 +577,12 @@ where if let Err(error) = &result { warn!(%error, "delivery batch failed"); } - loop { - let wake = guard.inner.wake.notified(); - let finished = { - let mut state = lock(&guard.inner.state); - if guard.inner.terminal_error().is_some() { - break; - } else if guard.inner.register_shutdowns(&mut state) { - let count = state.finish(&result); - guard.inner.release(count); - true - } else { - false - } - }; - if finished { - break; - } - wake.await; + let mut state = lock(&guard.inner.state); + if guard.inner.terminal_error().is_some() { + break; } + let count = state.finish(&result); + guard.inner.release(count); } if let Some(exit) = guard.exit.take() { let result = guard.inner.terminal_error().map_or(Ok(()), Err); @@ -632,4 +590,653 @@ where } } #[cfg(test)] -mod tests; +mod tests { + use super::*; + use std::{ + future::IntoFuture, + sync::atomic::{AtomicBool, Ordering}, + }; + use tokio::sync::mpsc; + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn admission_does_not_wait_for_coordinator_lock() { + let (service, mut entered, gate) = gated(4); + let coordinator = lock(&service.producer.inner.state); + let producer = service.clone(); + let (sent, received) = std::sync::mpsc::channel(); + let thread = std::thread::spawn(move || { + let completion = producer.enqueue(7); + sent.send(completion.submission().is_ok()).unwrap(); + }); + let admitted = received.recv_timeout(Duration::from_secs(2)); + drop(coordinator); + thread.join().unwrap(); + assert!(admitted.unwrap()); + assert_eq!(entered.recv().await.unwrap(), vec![7]); + gate.add_permits(1); + service.shutdown(Duration::from_secs(2)).await.unwrap(); + } + + #[tokio::test] + async fn shutdown_includes_a_publisher_already_in_progress() { + let (service, mut entered, gate) = gated(4); + let first = service.enqueue(1); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + + let (epoch, _) = service.producer.inner.admit().unwrap(); + let (sender, second) = Completion::pair(); + let shutdown = service.shutdown(Duration::from_secs(2)); + let mut shutdown = Box::pin(shutdown.into_future()); + assert!( + tokio::time::timeout(Duration::from_millis(20), &mut shutdown) + .await + .is_err() + ); + service + .producer + .inner + .publish(&epoch, Accepted { command: 2, sender }); + gate.add_permits(1); + assert!(first.await.is_ok()); + assert_eq!(entered.recv().await.unwrap(), vec![2]); + gate.add_permits(1); + assert!(second.await.is_ok()); + assert_eq!(shutdown.await.unwrap().delivered(), 2); + } + + #[tokio::test] + async fn flush_swaps_generation_and_waits_for_its_late_publisher() { + let (service, mut entered, gate) = gated(8); + let first = service.enqueue(1); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + + let (old_epoch, _) = service.producer.inner.admit().unwrap(); + let (sender, before) = Completion::pair(); + let mut flush = Box::pin(service.flush().into_future()); + assert!(!Arc::ptr_eq( + &old_epoch, + &service.producer.inner.ingress.load_full() + )); + let after = service.enqueue(3); + service + .producer + .inner + .publish(&old_epoch, Accepted { command: 2, sender }); + + gate.add_permits(1); + first.await.unwrap(); + assert_eq!(entered.recv().await.unwrap(), vec![2]); + assert!(futures::poll!(flush.as_mut()).is_pending()); + gate.add_permits(1); + before.await.unwrap(); + assert_eq!(flush.await.unwrap().delivered(), 2); + assert_eq!(entered.recv().await.unwrap(), vec![3]); + gate.add_permits(1); + after.await.unwrap(); + service.shutdown(Duration::from_secs(2)).await.unwrap(); + } + + #[tokio::test] + async fn flush_detaches_without_draining_the_queue() { + let (service, mut entered, gate) = gated(8); + let writes: Vec<_> = (0..4).map(|item| service.enqueue(item)).collect(); + let old_epoch = service.producer.inner.ingress.load_full(); + let flush = service.flush(); + assert!(!Arc::ptr_eq( + &old_epoch, + &service.producer.inner.ingress.load_full() + )); + // On the current-thread runtime, the worker has not run yet. Flush + // must not have taken the first item out of the detached queue. + assert!(old_epoch.queue.peek(&sdd::Guard::new()).is_some()); + assert_eq!(entered.recv().await.unwrap(), vec![0, 1, 2, 3]); + gate.add_permits(1); + assert_eq!(flush.await.unwrap().delivered(), 4); + for write in writes { + write.await.unwrap(); + } + service.shutdown(Duration::from_secs(2)).await.unwrap(); + } + + #[tokio::test] + async fn worker_failure_completes_a_late_publisher() { + let (service, mut entered, _gate) = gated(4); + let first = service.enqueue(1); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + let (old_epoch, _) = service.producer.inner.admit().unwrap(); + let (sender, late) = Completion::pair(); + service.producer.abort.abort(); + assert!(matches!( + first.await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + service + .producer + .inner + .publish(&old_epoch, Accepted { command: 2, sender }); + assert!(matches!( + late.await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + assert!(service.is_empty()); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn timeout_fails_items_taken_and_not_yet_taken_by_worker() { + let (service, _entered, _gate) = gated(4); + let (taken, wait_taken) = std::sync::mpsc::channel(); + let (resume, wait_resume) = std::sync::mpsc::channel(); + *lock(&service.producer.inner.after_take) = Some((taken, wait_resume)); + let coordinator = lock(&service.producer.inner.state); + let first = service.enqueue(1); + let second = service.enqueue(2); + drop(coordinator); + wait_taken.recv_timeout(Duration::from_secs(2)).unwrap(); + let shutdown = service.shutdown(Duration::from_millis(50)); + assert!(matches!( + shutdown.await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); + resume.send(()).unwrap(); + for completion in [first, second] { + assert!(matches!( + completion.await, + Err(WriteError::Execution(ExecutionError::TimedOut)) + )); + } + let _ = service.producer.exit.clone().await; + assert!(service.is_empty()); + } + + #[tokio::test] + async fn terminal_stop_is_preserved_when_the_processor_returns_success() { + let (service, mut entered, gate) = gated(4); + let active = service.enqueue(1); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + let pending = service.enqueue(2); + let flush = service.flush(); + + service.producer.inner.stop(ExecutionError::TimedOut); + gate.add_permits(1); // The processor returns Ok after its batch was failed. + for write in [active, pending] { + assert!(matches!( + write.await, + Err(WriteError::Execution(ExecutionError::TimedOut)) + )); + } + assert!(matches!( + flush.await, + Err(FlushError::Execution(ExecutionError::TimedOut)) + )); + assert_eq!( + service.producer.exit.clone().await, + Err(ExecutionError::TimedOut) + ); + assert!(service.is_empty()); + assert!(matches!( + service.flush().submission(), + Err(FlushError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + service.shutdown(Duration::from_secs(1)).await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); + } + + #[tokio::test] + async fn repeated_shutdown_after_worker_exit_succeeds_without_new_work() { + let (service, _entered, _gate) = gated(4); + assert_eq!( + service + .shutdown(Duration::from_secs(1)) + .await + .unwrap() + .delivered(), + 0 + ); + assert_eq!( + service + .shutdown(Duration::from_millis(50)) + .await + .unwrap() + .delivered(), + 0 + ); + assert!(service.is_empty()); + } + + fn gated( + capacity: usize, + ) -> ( + WriteService, + mpsc::UnboundedReceiver>, + Arc, + ) { + let (entered, receiver) = mpsc::unbounded_channel(); + let gate = Arc::new(Semaphore::new(0)); + let process_gate = gate.clone(); + let service = WriteService::new(capacity, Handle::current(), move |commands| { + let gate = process_gate.clone(); + let entered = entered.clone(); + Box::pin(async move { + entered.send(commands).unwrap(); + gate.acquire().await.unwrap().forget(); + Ok(()) + }) + }); + (service, receiver, gate) + } + + #[tokio::test] + async fn concurrent_flushes_observe_fixed_cutoffs_including_active_work() { + let (service, mut entered, gate) = gated(8); + let one = service.enqueue(1); + let first = service.flush(); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + let concurrent = service.flush(); + let two = service.enqueue(2); + let mut second = Box::pin(service.flush().into_future()); + let three = service.enqueue(3); // beyond second's cutoff + gate.add_permits(1); + assert_eq!(first.await.unwrap().delivered, 1); + assert_eq!(concurrent.await.unwrap().delivered, 1); + one.await.unwrap(); + assert_eq!(entered.recv().await.unwrap(), vec![2]); + assert!(futures::poll!(second.as_mut()).is_pending()); + gate.add_permits(1); + assert_eq!(second.await.unwrap().delivered, 2); + two.await.unwrap(); + assert_eq!(entered.recv().await.unwrap(), vec![3]); + // The second flush finished without waiting for the later arrival. + gate.add_permits(1); + three.await.unwrap(); + assert_eq!(service.flush().await.unwrap().delivered, 0); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + } + + #[tokio::test] + async fn a_new_flush_does_not_inherit_a_failed_prefix_from_an_older_active_flush() { + let (entered, mut receiver) = mpsc::unbounded_channel(); + let gate = Arc::new(Semaphore::new(0)); + let process_gate = gate.clone(); + let service = WriteService::new(4, Handle::current(), move |commands: Vec| { + let entered = entered.clone(); + let gate = process_gate.clone(); + Box::pin(async move { + entered.send(commands.clone()).unwrap(); + gate.acquire().await.unwrap().forget(); + if commands == [1] { + anyhow::bail!("first batch failed"); + } + Ok(()) + }) + }); + let first = service.enqueue(1); + assert_eq!(receiver.recv().await.unwrap(), vec![1]); + let second = service.enqueue(2); + let covering = service.flush(); + gate.add_permits(1); + assert!(matches!(first.await, Err(WriteError::Delivery(_)))); + assert_eq!(receiver.recv().await.unwrap(), vec![2]); + let fresh = service.flush(); + gate.add_permits(1); + second.await.unwrap(); + let Err(FlushError::Delivery(failure)) = covering.await else { + panic!("expected captured failure"); + }; + assert_eq!(failure.failed, 1); + assert_eq!(fresh.await.unwrap().delivered, 1); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + } + + #[tokio::test] + async fn graceful_shutdown_reports_its_failed_batch_after_notifying_all_writes() { + let service = WriteService::new(4, Handle::current(), |_: Vec| { + Box::pin(async { + anyhow::bail!("terminal failure"); + }) + }); + let first = service.enqueue(1); + let second = service.enqueue(2); + let shutdown = service.shutdown(Duration::from_secs(1)); + let Err(ShutdownError::Delivery(failure)) = shutdown.await else { + panic!("expected delivery failure"); + }; + assert_eq!(failure.failed, 2); + for completion in [first, second] { + assert!(matches!(completion.await, Err(WriteError::Delivery(_)))); + } + assert!(service.is_empty()); + } + + #[tokio::test] + async fn capacity_includes_in_flight_work_and_rejection_is_observable_twice() { + let (service, mut entered, gate) = gated(1); + let first = service.enqueue(1); + entered.recv().await.unwrap(); + let rejected = service.enqueue(2); + assert!(matches!( + &*rejected, + Err(WriteError::Admission(EnqueueError::Full)) + )); + assert!(matches!( + rejected.await, + Err(WriteError::Admission(EnqueueError::Full)) + )); + assert_eq!(service.outstanding(), 1); + gate.add_permits(1); + first.await.unwrap(); + assert!(service.is_empty()); + gate.add_permits(1); + service.enqueue(3).await.unwrap(); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + } + + #[tokio::test] + async fn dropped_observer_does_not_cancel_automatic_delivery_or_fill_the_queue() { + let (service, mut entered, gate) = gated(2); + drop(service.enqueue(1)); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + let flush = service.flush(); + drop(flush); // lifecycle observers also do not cancel accepted work + gate.add_permits(1); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + assert!(service.is_empty()); + } + + #[tokio::test] + async fn every_failed_batch_notifies_all_members_releases_errors_and_allows_fresh_work() { + let fail = Arc::new(AtomicBool::new(true)); + let process_fail = fail.clone(); + let service = WriteService::new(4, Handle::current(), move |_: Vec| { + let fail = process_fail.load(Ordering::Relaxed); + Box::pin(async move { + if fail { + anyhow::bail!("terminal failure"); + } + Ok(()) + }) + }); + for _ in 0..100 { + let first = service.enqueue(1); + let second = service.enqueue(2); + let flush = service.flush(); + let Err(FlushError::Delivery(failure)) = flush.await else { + panic!("expected failure"); + }; + assert_eq!(failure.failed, 2); + let weak = Arc::downgrade(&failure.cause); + for completion in [first, second] { + let Err(WriteError::Delivery(error)) = completion.await else { + panic!("expected failure"); + }; + assert!(Arc::ptr_eq(&error, &failure.cause)); + } + drop(failure); + assert!(weak.upgrade().is_none(), "SDK must release terminal errors"); + assert_eq!(service.outstanding(), 0); + assert_eq!(service.flush().await.unwrap().delivered, 0); + } + fail.store(false, Ordering::Relaxed); + let fresh = service.enqueue(3); + assert_eq!(service.flush().await.unwrap().delivered, 1); + fresh.await.unwrap(); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + } + + #[tokio::test] + async fn shutdown_stops_admission_immediately_drains_and_is_shared_by_clones() { + let (service, mut entered, gate) = gated(4); + let write = service.enqueue(1); + let shutdown = service.shutdown(Duration::from_secs(1)); + let concurrent = service.clone().shutdown(Duration::from_secs(1)); + assert!(shutdown.is_ok()); + assert!(matches!( + service.enqueue(2).await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); + entered.recv().await.unwrap(); + gate.add_permits(1); + write.await.unwrap(); + assert_eq!(shutdown.await.unwrap().delivered, 1); + assert_eq!(concurrent.await.unwrap().delivered, 1); + assert!(service.is_empty()); + } + + #[tokio::test] + async fn shutdown_timeout_fails_active_work_even_when_flush_slots_are_full() { + let (service, mut entered, _) = gated(1); + let write = service.enqueue(1); + entered.recv().await.unwrap(); + let flush = service.flush(); + let rejected = service.flush(); + assert!(matches!( + &*rejected, + Err(FlushError::Admission(EnqueueError::Full)) + )); + let shutdown = service.shutdown(Duration::from_millis(10)); + assert!( + shutdown.is_ok(), + "shutdown must not compete with flush slots" + ); + assert!(matches!( + shutdown.await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + write.await, + Err(WriteError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + flush.await, + Err(FlushError::Execution(ExecutionError::TimedOut)) + )); + assert!(service.is_empty()); + assert!(matches!( + service.enqueue(2).await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); + } + + #[tokio::test] + async fn shutdown_timeout_also_retires_batches_queued_behind_active_work() { + let (service, mut entered, _) = gated(4); + let first = service.enqueue(1); + entered.recv().await.unwrap(); + let second = service.enqueue(2); + let third = service.enqueue(3); + let flush = service.flush(); + let shutdown = service.shutdown(Duration::from_millis(10)); + assert!(matches!( + shutdown.await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); + for completion in [first, second, third] { + assert!(matches!( + completion.await, + Err(WriteError::Execution(ExecutionError::TimedOut)) + )); + } + assert!(matches!( + flush.await, + Err(FlushError::Execution(ExecutionError::TimedOut)) + )); + assert_eq!(service.outstanding(), 0); + assert!(lock(&service.producer.inner.state).barriers.is_empty()); + assert!(lock(&service.producer.inner.state).ready.is_empty()); + } + + #[tokio::test] + async fn very_large_write_capacity_does_not_overflow_lifecycle_control_limits() { + let (service, _, _) = gated(usize::MAX); + assert_eq!(service.producer.control_limit, 1024); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + } + + #[tokio::test] + async fn worker_cancellation_before_first_poll_completes_every_accepted_sender() { + let (service, _, _) = gated(4); + let first = service.enqueue(1); + let second = service.enqueue(2); + let flush = service.flush(); + service.producer.abort.abort(); + for completion in [first, second] { + assert!(matches!( + completion.await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + } + assert!(matches!( + flush.await, + Err(FlushError::Execution(ExecutionError::WorkerLost)) + )); + assert!(service.is_empty()); + assert!(matches!( + service.enqueue(3).await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); + } + + #[tokio::test] + async fn worker_panic_notifies_active_and_pending_members() { + let (entered, mut receiver) = mpsc::unbounded_channel(); + let gate = Arc::new(Semaphore::new(0)); + let process_gate = gate.clone(); + let service = WriteService::new(4, Handle::current(), move |_: Vec| { + let entered = entered.clone(); + let gate = process_gate.clone(); + Box::pin(async move { + entered.send(()).unwrap(); + gate.acquire().await.unwrap().forget(); + panic!("test processor failed"); + }) + }); + let active = service.enqueue(1); + receiver.recv().await.unwrap(); + let pending = service.enqueue(2); + let flush = service.flush(); + gate.add_permits(1); + for completion in [active, pending] { + assert!(matches!( + completion.await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + } + assert!(matches!( + flush.await, + Err(FlushError::Execution(ExecutionError::WorkerLost)) + )); + assert!(service.is_empty()); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn dropping_the_last_producers_drains_work_and_stops_the_worker() { + for _ in 0..64 { + let (service, mut entered, gate) = gated(4); + let write = service.enqueue(1); + let clone = service.clone(); + let exit = service.producer.exit.clone(); + std::thread::scope(|scope| { + scope.spawn(move || drop(service)); + scope.spawn(move || drop(clone)); + }); + entered.recv().await.unwrap(); + gate.add_permits(1); + write.await.unwrap(); + tokio::time::timeout(Duration::from_secs(1), exit) + .await + .unwrap() + .unwrap(); + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn fast_workers_and_racing_flushes_cannot_lose_registered_completions() { + let seen = Arc::new(Mutex::new(std::collections::HashSet::new())); + let processed = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let process_seen = seen.clone(); + let process_count = processed.clone(); + let service = WriteService::new(1024, Handle::current(), move |commands: Vec| { + lock(&process_seen).extend(commands.iter().copied()); + process_count.fetch_add(commands.len(), Ordering::Relaxed); + Box::pin(async { Ok(()) }) + }); + let (writes, flushes) = std::thread::scope(|scope| { + let producers: Vec<_> = (0..4) + .map(|producer| { + let service = &service; + scope.spawn(move || { + (0..50) + .map(|item| service.enqueue(producer * 50 + item)) + .collect::>() + }) + }) + .collect(); + let flushers: Vec<_> = (0..2) + .map(|_| { + let service = &service; + scope.spawn(move || (0..50).map(|_| service.flush()).collect::>()) + }) + .collect(); + ( + producers + .into_iter() + .flat_map(|producer| producer.join().unwrap()) + .collect::>(), + flushers + .into_iter() + .flat_map(|flusher| flusher.join().unwrap()) + .collect::>(), + ) + }); + tokio::time::timeout(Duration::from_secs(5), async { + for write in writes { + write.await.unwrap(); + } + for flush in flushes { + flush.await.unwrap(); + } + }) + .await + .unwrap(); + assert_eq!(lock(&seen).len(), 200); + assert_eq!(processed.load(Ordering::Relaxed), 200); + assert!(service.is_empty()); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn racing_producers_cannot_overbook_capacity() { + let (service, _entered, gate) = gated(8); + let completions = std::thread::scope(|scope| { + let threads: Vec<_> = (0..32) + .map(|number| { + let service = &service; + scope.spawn(move || service.enqueue(number)) + }) + .collect(); + threads + .into_iter() + .map(|thread| thread.join().unwrap()) + .collect::>() + }); + assert_eq!( + completions + .iter() + .filter(|completion| completion.is_ok()) + .count(), + 8 + ); + assert_eq!(service.outstanding(), 8); + gate.add_permits(8); + let mut delivered = 0; + for completion in completions { + if completion.await.is_ok() { + delivered += 1; + } + } + assert_eq!(delivered, 8); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + assert!(service.is_empty()); + } +} diff --git a/src/internal/delivery/tests.rs b/src/internal/delivery/tests.rs deleted file mode 100644 index 197515f..0000000 --- a/src/internal/delivery/tests.rs +++ /dev/null @@ -1,630 +0,0 @@ -use super::*; -use std::{ - future::IntoFuture, - sync::atomic::{AtomicBool, Ordering}, -}; -use tokio::sync::mpsc; - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn admission_does_not_wait_for_coordinator_lock() { - let (service, mut entered, gate) = gated(4); - let coordinator = lock(&service.producer.inner.state); - let producer = service.clone(); - let (sent, received) = std::sync::mpsc::channel(); - let thread = std::thread::spawn(move || { - let completion = producer.enqueue(7); - sent.send(completion.submission().is_ok()).unwrap(); - }); - let admitted = received.recv_timeout(Duration::from_secs(2)); - drop(coordinator); - thread.join().unwrap(); - assert!(admitted.unwrap()); - assert_eq!(entered.recv().await.unwrap(), vec![7]); - gate.add_permits(1); - service.shutdown(Duration::from_secs(2)).await.unwrap(); -} - -#[tokio::test] -async fn shutdown_includes_a_publisher_already_in_progress() { - let (service, mut entered, gate) = gated(4); - let first = service.enqueue(1); - assert_eq!(entered.recv().await.unwrap(), vec![1]); - - service.producer.inner.admit().unwrap(); - let (sender, second) = Completion::pair(); - let shutdown = service.shutdown(Duration::from_secs(2)); - let mut shutdown = Box::pin(shutdown.into_future()); - assert!( - tokio::time::timeout(Duration::from_millis(20), &mut shutdown) - .await - .is_err() - ); - service - .producer - .inner - .publish(Accepted { command: 2, sender }); - gate.add_permits(1); - assert!(first.await.is_ok()); - assert_eq!(entered.recv().await.unwrap(), vec![2]); - gate.add_permits(1); - assert!(second.await.is_ok()); - assert_eq!(shutdown.await.unwrap().delivered(), 2); -} - -#[tokio::test] -async fn flush_swaps_generation_and_waits_for_its_late_publisher() { - let (service, mut entered, gate) = gated(8); - let first = service.enqueue(1); - assert_eq!(entered.recv().await.unwrap(), vec![1]); - - service.producer.inner.admit().unwrap(); - let old_epoch = service.producer.inner.ingress.load_full(); - assert!(old_epoch.enter()); - let (sender, before) = Completion::pair(); - let mut flush = Box::pin(service.flush().into_future()); - assert!(!Arc::ptr_eq( - &old_epoch, - &service.producer.inner.ingress.load_full() - )); - let after = service.enqueue(3); - service - .producer - .inner - .publish_to_epoch(&old_epoch, Accepted { command: 2, sender }); - - gate.add_permits(1); - first.await.unwrap(); - assert_eq!(entered.recv().await.unwrap(), vec![2]); - assert!(futures::poll!(flush.as_mut()).is_pending()); - gate.add_permits(1); - before.await.unwrap(); - assert_eq!(flush.await.unwrap().delivered(), 2); - assert_eq!(entered.recv().await.unwrap(), vec![3]); - gate.add_permits(1); - after.await.unwrap(); - service.shutdown(Duration::from_secs(2)).await.unwrap(); -} - -#[tokio::test] -async fn flush_detaches_without_draining_the_queue() { - let (service, mut entered, gate) = gated(8); - let writes: Vec<_> = (0..4).map(|item| service.enqueue(item)).collect(); - let old_epoch = service.producer.inner.ingress.load_full(); - let flush = service.flush(); - assert!(!Arc::ptr_eq( - &old_epoch, - &service.producer.inner.ingress.load_full() - )); - // On the current-thread runtime, the worker has not run yet. Flush - // must not have taken the first item out of the detached queue. - assert!(old_epoch.queue.peek(&sdd::Guard::new()).is_some()); - assert_eq!(entered.recv().await.unwrap(), vec![0, 1, 2, 3]); - gate.add_permits(1); - assert_eq!(flush.await.unwrap().delivered(), 4); - for write in writes { - write.await.unwrap(); - } - service.shutdown(Duration::from_secs(2)).await.unwrap(); -} - -#[tokio::test] -async fn worker_failure_completes_a_late_publisher() { - let (service, mut entered, _gate) = gated(4); - let first = service.enqueue(1); - assert_eq!(entered.recv().await.unwrap(), vec![1]); - service.producer.inner.admit().unwrap(); - let old_epoch = service.producer.inner.ingress.load_full(); - assert!(old_epoch.enter()); - let (sender, late) = Completion::pair(); - service.producer.abort.abort(); - assert!(matches!( - first.await, - Err(WriteError::Execution(ExecutionError::WorkerLost)) - )); - service - .producer - .inner - .publish_to_epoch(&old_epoch, Accepted { command: 2, sender }); - assert!(matches!( - late.await, - Err(WriteError::Execution(ExecutionError::WorkerLost)) - )); - assert!(service.is_empty()); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn timeout_fails_items_taken_and_not_yet_taken_by_worker() { - let (service, _entered, _gate) = gated(4); - let (taken, wait_taken) = std::sync::mpsc::channel(); - let (resume, wait_resume) = std::sync::mpsc::channel(); - *lock(&service.producer.inner.after_take) = Some((taken, wait_resume)); - let coordinator = lock(&service.producer.inner.state); - let first = service.enqueue(1); - let second = service.enqueue(2); - drop(coordinator); - wait_taken.recv_timeout(Duration::from_secs(2)).unwrap(); - let shutdown = service.shutdown(Duration::from_millis(50)); - assert!(matches!( - shutdown.await, - Err(ShutdownError::Execution(ExecutionError::TimedOut)) - )); - resume.send(()).unwrap(); - for completion in [first, second] { - assert!(matches!( - completion.await, - Err(WriteError::Execution(ExecutionError::TimedOut)) - )); - } - let _ = service.producer.exit.clone().await; - assert!(service.is_empty()); -} - -#[tokio::test] -async fn terminal_stop_is_preserved_when_the_processor_returns_success() { - let (service, mut entered, gate) = gated(4); - let active = service.enqueue(1); - assert_eq!(entered.recv().await.unwrap(), vec![1]); - let pending = service.enqueue(2); - let flush = service.flush(); - - service.producer.inner.stop(ExecutionError::TimedOut); - gate.add_permits(1); // The processor returns Ok after its batch was failed. - for write in [active, pending] { - assert!(matches!( - write.await, - Err(WriteError::Execution(ExecutionError::TimedOut)) - )); - } - assert!(matches!( - flush.await, - Err(FlushError::Execution(ExecutionError::TimedOut)) - )); - assert_eq!( - service.producer.exit.clone().await, - Err(ExecutionError::TimedOut) - ); - assert!(service.is_empty()); - assert!(matches!( - service.flush().submission(), - Err(FlushError::Execution(ExecutionError::TimedOut)) - )); - assert!(matches!( - service.shutdown(Duration::from_secs(1)).await, - Err(ShutdownError::Execution(ExecutionError::TimedOut)) - )); -} - -fn gated( - capacity: usize, -) -> ( - WriteService, - mpsc::UnboundedReceiver>, - Arc, -) { - let (entered, receiver) = mpsc::unbounded_channel(); - let gate = Arc::new(Semaphore::new(0)); - let process_gate = gate.clone(); - let service = WriteService::new(capacity, Handle::current(), move |commands| { - let gate = process_gate.clone(); - let entered = entered.clone(); - Box::pin(async move { - entered.send(commands).unwrap(); - gate.acquire().await.unwrap().forget(); - Ok(()) - }) - }); - (service, receiver, gate) -} - -#[tokio::test] -async fn concurrent_flushes_observe_fixed_cutoffs_including_active_work() { - let (service, mut entered, gate) = gated(8); - let one = service.enqueue(1); - let first = service.flush(); - assert_eq!(entered.recv().await.unwrap(), vec![1]); - let concurrent = service.flush(); - let two = service.enqueue(2); - let mut second = Box::pin(service.flush().into_future()); - let three = service.enqueue(3); // beyond second's cutoff - gate.add_permits(1); - assert_eq!(first.await.unwrap().delivered, 1); - assert_eq!(concurrent.await.unwrap().delivered, 1); - one.await.unwrap(); - assert_eq!(entered.recv().await.unwrap(), vec![2]); - assert!(futures::poll!(second.as_mut()).is_pending()); - gate.add_permits(1); - assert_eq!(second.await.unwrap().delivered, 2); - two.await.unwrap(); - assert_eq!(entered.recv().await.unwrap(), vec![3]); - // The second flush finished without waiting for the later arrival. - gate.add_permits(1); - three.await.unwrap(); - assert_eq!(service.flush().await.unwrap().delivered, 0); - service.shutdown(Duration::from_secs(1)).await.unwrap(); -} - -#[tokio::test] -async fn a_new_flush_does_not_inherit_a_failed_prefix_from_an_older_active_flush() { - let (entered, mut receiver) = mpsc::unbounded_channel(); - let gate = Arc::new(Semaphore::new(0)); - let process_gate = gate.clone(); - let service = WriteService::new(4, Handle::current(), move |commands: Vec| { - let entered = entered.clone(); - let gate = process_gate.clone(); - Box::pin(async move { - entered.send(commands.clone()).unwrap(); - gate.acquire().await.unwrap().forget(); - if commands == [1] { - anyhow::bail!("first batch failed"); - } - Ok(()) - }) - }); - let first = service.enqueue(1); - assert_eq!(receiver.recv().await.unwrap(), vec![1]); - let second = service.enqueue(2); - let covering = service.flush(); - gate.add_permits(1); - assert!(matches!(first.await, Err(WriteError::Delivery(_)))); - assert_eq!(receiver.recv().await.unwrap(), vec![2]); - let fresh = service.flush(); - gate.add_permits(1); - second.await.unwrap(); - let Err(FlushError::Delivery(failure)) = covering.await else { - panic!("expected captured failure"); - }; - assert_eq!(failure.failed, 1); - assert_eq!(fresh.await.unwrap().delivered, 1); - service.shutdown(Duration::from_secs(1)).await.unwrap(); -} - -#[tokio::test] -async fn graceful_shutdown_reports_its_failed_batch_after_notifying_all_writes() { - let service = WriteService::new(4, Handle::current(), |_: Vec| { - Box::pin(async { - anyhow::bail!("terminal failure"); - }) - }); - let first = service.enqueue(1); - let second = service.enqueue(2); - let shutdown = service.shutdown(Duration::from_secs(1)); - let Err(ShutdownError::Delivery(failure)) = shutdown.await else { - panic!("expected delivery failure"); - }; - assert_eq!(failure.failed, 2); - for completion in [first, second] { - assert!(matches!(completion.await, Err(WriteError::Delivery(_)))); - } - assert!(service.is_empty()); -} - -#[tokio::test] -async fn capacity_includes_in_flight_work_and_rejection_is_observable_twice() { - let (service, mut entered, gate) = gated(1); - let first = service.enqueue(1); - entered.recv().await.unwrap(); - let rejected = service.enqueue(2); - assert!(matches!( - &*rejected, - Err(WriteError::Admission(EnqueueError::Full)) - )); - assert!(matches!( - rejected.await, - Err(WriteError::Admission(EnqueueError::Full)) - )); - assert_eq!(service.outstanding(), 1); - gate.add_permits(1); - first.await.unwrap(); - assert!(service.is_empty()); - gate.add_permits(1); - service.enqueue(3).await.unwrap(); - service.shutdown(Duration::from_secs(1)).await.unwrap(); -} - -#[tokio::test] -async fn dropped_observer_does_not_cancel_automatic_delivery_or_fill_the_queue() { - let (service, mut entered, gate) = gated(2); - drop(service.enqueue(1)); - assert_eq!(entered.recv().await.unwrap(), vec![1]); - let flush = service.flush(); - drop(flush); // lifecycle observers also do not cancel accepted work - gate.add_permits(1); - service.shutdown(Duration::from_secs(1)).await.unwrap(); - assert!(service.is_empty()); -} - -#[tokio::test] -async fn every_failed_batch_notifies_all_members_releases_errors_and_allows_fresh_work() { - let fail = Arc::new(AtomicBool::new(true)); - let process_fail = fail.clone(); - let service = WriteService::new(4, Handle::current(), move |_: Vec| { - let fail = process_fail.load(Ordering::Relaxed); - Box::pin(async move { - if fail { - anyhow::bail!("terminal failure"); - } - Ok(()) - }) - }); - for _ in 0..100 { - let first = service.enqueue(1); - let second = service.enqueue(2); - let flush = service.flush(); - let Err(FlushError::Delivery(failure)) = flush.await else { - panic!("expected failure"); - }; - assert_eq!(failure.failed, 2); - let weak = Arc::downgrade(&failure.cause); - for completion in [first, second] { - let Err(WriteError::Delivery(error)) = completion.await else { - panic!("expected failure"); - }; - assert!(Arc::ptr_eq(&error, &failure.cause)); - } - drop(failure); - assert!(weak.upgrade().is_none(), "SDK must release terminal errors"); - assert_eq!(service.outstanding(), 0); - assert_eq!(service.flush().await.unwrap().delivered, 0); - } - fail.store(false, Ordering::Relaxed); - let fresh = service.enqueue(3); - assert_eq!(service.flush().await.unwrap().delivered, 1); - fresh.await.unwrap(); - service.shutdown(Duration::from_secs(1)).await.unwrap(); -} - -#[tokio::test] -async fn shutdown_stops_admission_immediately_drains_and_is_shared_by_clones() { - let (service, mut entered, gate) = gated(4); - let write = service.enqueue(1); - let shutdown = service.shutdown(Duration::from_secs(1)); - let concurrent = service.clone().shutdown(Duration::from_secs(1)); - assert!(shutdown.is_ok()); - assert!(matches!( - service.enqueue(2).await, - Err(WriteError::Admission(EnqueueError::Closed)) - )); - entered.recv().await.unwrap(); - gate.add_permits(1); - write.await.unwrap(); - assert_eq!(shutdown.await.unwrap().delivered, 1); - assert_eq!(concurrent.await.unwrap().delivered, 1); - assert!(service.is_empty()); -} - -#[tokio::test] -async fn shutdown_timeout_fails_active_work_even_when_flush_slots_are_full() { - let (service, mut entered, _) = gated(1); - let write = service.enqueue(1); - entered.recv().await.unwrap(); - let flush = service.flush(); - let rejected = service.flush(); - assert!(matches!( - &*rejected, - Err(FlushError::Admission(EnqueueError::Full)) - )); - let shutdown = service.shutdown(Duration::from_millis(10)); - assert!( - shutdown.is_ok(), - "shutdown must not compete with flush slots" - ); - assert!(matches!( - shutdown.await, - Err(ShutdownError::Execution(ExecutionError::TimedOut)) - )); - assert!(matches!( - write.await, - Err(WriteError::Execution(ExecutionError::TimedOut)) - )); - assert!(matches!( - flush.await, - Err(FlushError::Execution(ExecutionError::TimedOut)) - )); - assert!(service.is_empty()); - assert!(matches!( - service.enqueue(2).await, - Err(WriteError::Admission(EnqueueError::Closed)) - )); -} - -#[tokio::test] -async fn shutdown_timeout_also_retires_batches_queued_behind_active_work() { - let (service, mut entered, _) = gated(4); - let first = service.enqueue(1); - entered.recv().await.unwrap(); - let second = service.enqueue(2); - let third = service.enqueue(3); - let flush = service.flush(); - let shutdown = service.shutdown(Duration::from_millis(10)); - assert!(matches!( - shutdown.await, - Err(ShutdownError::Execution(ExecutionError::TimedOut)) - )); - for completion in [first, second, third] { - assert!(matches!( - completion.await, - Err(WriteError::Execution(ExecutionError::TimedOut)) - )); - } - assert!(matches!( - flush.await, - Err(FlushError::Execution(ExecutionError::TimedOut)) - )); - assert_eq!(service.outstanding(), 0); - assert!(lock(&service.producer.inner.state).barriers.is_empty()); - assert!(lock(&service.producer.inner.state).ready.is_empty()); -} - -#[tokio::test] -async fn very_large_write_capacity_does_not_overflow_lifecycle_control_limits() { - let (service, _, _) = gated(usize::MAX); - assert_eq!(service.producer.control_limit, 1024); - service.shutdown(Duration::from_secs(1)).await.unwrap(); -} - -#[tokio::test] -async fn worker_cancellation_before_first_poll_completes_every_accepted_sender() { - let (service, _, _) = gated(4); - let first = service.enqueue(1); - let second = service.enqueue(2); - let flush = service.flush(); - service.producer.abort.abort(); - for completion in [first, second] { - assert!(matches!( - completion.await, - Err(WriteError::Execution(ExecutionError::WorkerLost)) - )); - } - assert!(matches!( - flush.await, - Err(FlushError::Execution(ExecutionError::WorkerLost)) - )); - assert!(service.is_empty()); - assert!(matches!( - service.enqueue(3).await, - Err(WriteError::Admission(EnqueueError::Closed)) - )); -} - -#[tokio::test] -async fn worker_panic_notifies_active_and_pending_members() { - let (entered, mut receiver) = mpsc::unbounded_channel(); - let gate = Arc::new(Semaphore::new(0)); - let process_gate = gate.clone(); - let service = WriteService::new(4, Handle::current(), move |_: Vec| { - let entered = entered.clone(); - let gate = process_gate.clone(); - Box::pin(async move { - entered.send(()).unwrap(); - gate.acquire().await.unwrap().forget(); - panic!("test processor failed"); - }) - }); - let active = service.enqueue(1); - receiver.recv().await.unwrap(); - let pending = service.enqueue(2); - let flush = service.flush(); - gate.add_permits(1); - for completion in [active, pending] { - assert!(matches!( - completion.await, - Err(WriteError::Execution(ExecutionError::WorkerLost)) - )); - } - assert!(matches!( - flush.await, - Err(FlushError::Execution(ExecutionError::WorkerLost)) - )); - assert!(service.is_empty()); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn dropping_the_last_producers_drains_work_and_stops_the_worker() { - for _ in 0..64 { - let (service, mut entered, gate) = gated(4); - let write = service.enqueue(1); - let clone = service.clone(); - let exit = service.producer.exit.clone(); - std::thread::scope(|scope| { - scope.spawn(move || drop(service)); - scope.spawn(move || drop(clone)); - }); - entered.recv().await.unwrap(); - gate.add_permits(1); - write.await.unwrap(); - tokio::time::timeout(Duration::from_secs(1), exit) - .await - .unwrap() - .unwrap(); - } -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn fast_workers_and_racing_flushes_cannot_lose_registered_completions() { - let seen = Arc::new(Mutex::new(std::collections::HashSet::new())); - let processed = Arc::new(std::sync::atomic::AtomicUsize::new(0)); - let process_seen = seen.clone(); - let process_count = processed.clone(); - let service = WriteService::new(1024, Handle::current(), move |commands: Vec| { - lock(&process_seen).extend(commands.iter().copied()); - process_count.fetch_add(commands.len(), Ordering::Relaxed); - Box::pin(async { Ok(()) }) - }); - let (writes, flushes) = std::thread::scope(|scope| { - let producers: Vec<_> = (0..4) - .map(|producer| { - let service = &service; - scope.spawn(move || { - (0..50) - .map(|item| service.enqueue(producer * 50 + item)) - .collect::>() - }) - }) - .collect(); - let flushers: Vec<_> = (0..2) - .map(|_| { - let service = &service; - scope.spawn(move || (0..50).map(|_| service.flush()).collect::>()) - }) - .collect(); - ( - producers - .into_iter() - .flat_map(|producer| producer.join().unwrap()) - .collect::>(), - flushers - .into_iter() - .flat_map(|flusher| flusher.join().unwrap()) - .collect::>(), - ) - }); - tokio::time::timeout(Duration::from_secs(5), async { - for write in writes { - write.await.unwrap(); - } - for flush in flushes { - flush.await.unwrap(); - } - }) - .await - .unwrap(); - assert_eq!(lock(&seen).len(), 200); - assert_eq!(processed.load(Ordering::Relaxed), 200); - assert!(service.is_empty()); - service.shutdown(Duration::from_secs(1)).await.unwrap(); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn racing_producers_cannot_overbook_capacity() { - let (service, _entered, gate) = gated(8); - let completions = std::thread::scope(|scope| { - let threads: Vec<_> = (0..32) - .map(|number| { - let service = &service; - scope.spawn(move || service.enqueue(number)) - }) - .collect(); - threads - .into_iter() - .map(|thread| thread.join().unwrap()) - .collect::>() - }); - assert_eq!( - completions - .iter() - .filter(|completion| completion.is_ok()) - .count(), - 8 - ); - assert_eq!(service.outstanding(), 8); - gate.add_permits(8); - let mut delivered = 0; - for completion in completions { - if completion.await.is_ok() { - delivered += 1; - } - } - assert_eq!(delivered, 8); - service.shutdown(Duration::from_secs(1)).await.unwrap(); - assert!(service.is_empty()); -}