diff --git a/Cargo.toml b/Cargo.toml index db3ccc5..b31c851 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,13 +14,14 @@ categories = ["api-bindings", "asynchronous", "development-tools"] anyhow = "1" base64 = "0.23" bon = "3" -arc-swap = "1" 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 = "0.8" +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/completion.rs b/src/completion.rs new file mode 100644 index 0000000..b9e6e81 --- /dev/null +++ b/src/completion.rs @@ -0,0 +1,141 @@ +//! 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::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. +/// +/// ```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.submission() { +/// 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 { + /// 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(); + ( + 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.submission(), &Err(Error::Admission)); + assert_eq!(&*token, token.submission()); + 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.submission().is_ok()); + 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..2a67f04 --- /dev/null +++ b/src/delivery.rs @@ -0,0 +1,120 @@ +//! Outcomes of queued ingestion and its lifecycle operations. +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, + #[error("delivery queue is closed")] + Closed, +} +#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] +#[non_exhaustive] +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)] +#[non_exhaustive] +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 { + 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(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(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), + #[error(transparent)] + Execution(#[from] ExecutionError), + #[error(transparent)] + Delivery(FlushFailure), +} +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct ShutdownReport { + 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), + #[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..3d346bb --- /dev/null +++ b/src/internal/delivery.rs @@ -0,0 +1,1242 @@ +//! 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, +}; +use arc_swap::{ArcSwap, Guard}; +use crossbeam_utils::atomic::AtomicCell; +use futures::{ + future::{BoxFuture, Shared}, + task::AtomicWaker, + FutureExt, +}; +use std::{ + collections::VecDeque, + sync::{ + atomic::{AtomicBool, AtomicU8, AtomicUsize, Ordering}, + Arc, Mutex, MutexGuard, + }, + time::Duration, +}; +use tokio::{ + runtime::Handle, + sync::{oneshot, Semaphore}, + task::AbortHandle, +}; +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 { + producer: Arc>, +} +impl Clone for WriteService { + fn clone(&self) -> Self { + Self { + producer: self.producer.clone(), + } + } +} +/// 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 Producer { + fn drop(&mut self) { + self.inner.close(&mut lock(&self.inner.state)); + } +} +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(), + }), + ingress: ArcSwap::from_pointee(QueueEpoch::new(false)), + outstanding: AtomicUsize::new(0), + terminal: AtomicU8::new(0), + depth_warned: AtomicBool::new(false), + wake: Wake::new(), + capacity, + #[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 }); + Self { + producer: Arc::new(Producer { + inner, + abort: task.abort_handle(), + runtime, + exit, + control_limit, + shutdown_controls: Arc::new(Semaphore::new(control_limit)), + }), + } + } + + #[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 producer = &self.producer; + let (sender, mut completion) = Completion::pair(); + match producer.inner.admit() { + Err(error) => { + completion.record_issuance(Err(error.into())); + rejected(command, error); + } + 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) + { + warn!( + outstanding = depth, + capacity = producer.inner.capacity, + "delivery queue has reached 80% of capacity" + ); + } + } + } + completion + } + + pub fn flush(&self) -> Completion { + let producer = &self.producer; + let (sender, mut completion) = Completion::pair(); + 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.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, false); + 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) = 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(&producer.inner.state); + if let Some(error) = producer.inner.terminal_error() { + completion.record_issuance(Err(error.into())); + return completion; + } + 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(); + let exit = producer.exit.clone(); + producer.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.producer.inner.outstanding.load(Ordering::Acquire) == 0 + } + #[cfg(test)] + fn outstanding(&self) -> usize { + self.producer.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, + 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 + } +} +struct Accepted { + command: C, + sender: WriteSender, +} +type IngressItem = AtomicCell>>>; +type EpochGuard = Guard>>; +struct QueueEpoch { + queue: sdd::Queue>, + publishers: AtomicUsize, +} +impl QueueEpoch { + fn new(closed: bool) -> Self { + Self { + queue: sdd::Queue::new(), + publishers: AtomicUsize::new(if closed { CLOSED } else { 0 }), + } + } + fn enter(&self) -> bool { + enter(&self.publishers) + } + 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 + } +} +struct Batch { + epoch: Arc>, + senders: Vec, +} +fn lock(value: &Mutex) -> MutexGuard<'_, T> { + value + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} +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>, + sender: FlushSender, + delivered: u64, + failure: Option, +} +impl Barrier { + fn record(&mut self, count: u64, result: &Result<(), Arc>) { + match result { + Ok(()) => self.delivered += count, + Err(error) => match &mut self.failure { + Some(failure) => failure.failed += count, + None => { + self.failure = Some(FlushFailure { + failed: count, + cause: error.clone(), + }) + } + }, + } + } + 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: VecDeque>, +} +impl State { + fn barrier(&mut self, sender: FlushSender, through: Arc>) { + self.barriers.push_back(Barrier { + through, + sender, + delivered: 0, + failure: None, + }); + } + 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. + for barrier in &mut self.barriers { + barrier.record(count as u64, result); + } + // 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 + } +} +struct Inner { + state: Mutex>, + ingress: ArcSwap>, + outstanding: AtomicUsize, + terminal: AtomicU8, + depth_warned: AtomicBool, + wake: Wake, + capacity: usize, + #[cfg(test)] + 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), + } + } + + #[allow(deprecated, reason = "try_update requires Rust 1.95")] + 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); + } + } + } + } + + fn end_publication(&self, epoch: &QueueEpoch) { + epoch.publishers.fetch_sub(1, Ordering::SeqCst); + self.wake.notify_one(); + } + + 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, closing: bool) -> Arc> { + let epoch = self.ingress.swap(Arc::new(QueueEpoch::new(closing))); + epoch.close(); + state.ready.push_back(epoch.clone()); + epoch + } + + 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, epoch: &QueueEpoch, accepted: Accepted) + where + C: 'static, + { + epoch.queue.push(AtomicCell::new(Some(Box::new(accepted)))); + // 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); + } + } + + fn stop(&self, error: ExecutionError) -> ExecutionError { + let code = match error { + ExecutionError::WorkerLost => 1, + ExecutionError::TimedOut => 2, + }; + let mut state = lock(&self.state); + if let Some(error) = self.terminal_error() { + return 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.close(&mut state); + if let Some(active) = state.active.take() { + let count = active.senders.len(); + finish_writes(active.senders, Err(error.into())); + self.release(count); + self.fail_epoch(&active.epoch, error); + } + for epoch in state.ready.drain(..) { + self.fail_epoch(&epoch, error); + } + for barrier in state.barriers.drain(..) { + let _ = barrier.sender.send(Err(error.into())); + } + error + } + fn close(&self, state: &mut State) { + if !self.ingress.load().closed() { + self.detach(state, true); + } + self.wake.notify_one(); + } +} + +struct WorkerGuard { + inner: Arc>, + exit: Option>>, +} +impl Drop for WorkerGuard { + fn drop(&mut self) { + if let Some(exit) = self.exit.take() { + let error = self.inner.stop(ExecutionError::WorkerLost); + 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); + if guard.inner.terminal_error().is_some() { + break; + } + 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 { + None + } + }; + 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"); + } + 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); + let _ = exit.send(result); + } +} +#[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.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/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..de69ff0 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; @@ -27,86 +27,25 @@ 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, 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, - dropped_count: AtomicUsize, - last_drop_log_time: AtomicU64, - queue_depth_warning_emitted: AtomicBool, +/// Preparation and protocol orchestration, independent of admission/lifecycle. +struct SpanBatchProcessor { 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 { +impl SpanBatchProcessor { fn new( config: LogQueueConfig, login_state: LoginState, @@ -114,131 +53,17 @@ 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 +80,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 +107,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 +143,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 +242,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. @@ -725,40 +540,16 @@ 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) - } +/// 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, @@ -790,23 +581,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 +716,232 @@ 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>, + rejections: 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 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(), + 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 { rejections, 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.rejections + .dropped_count + .fetch_add(1, Ordering::Relaxed); + if matches!(error, crate::EnqueueError::Full) { + self.rejections.log_drop_warning(); + } else { + warn!(%error, "span submission rejected"); + } + self.rejections.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.submission() { + 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..493e23e 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 @@ -612,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. @@ -630,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( @@ -639,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()) @@ -732,8 +746,11 @@ impl BraintrustClient { extra: event_data.extra, }; - self.submit_payload(token, payload, None); - Ok(()) + let completion = self.submit_payload(token, payload, None); + if let Err(error) = completion.submission() { + return Err(BraintrustError::Background(error.to_string())); + } + Ok(completion) } /// Perform login synchronously. @@ -909,24 +926,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 +958,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 +1518,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..7fad7f9 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::new())); + } + 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..ec1dcd5 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::new())); + 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..82457dd 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,91 @@ 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 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; + 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); }