Add bounded delivery admission, completions, and lifecycle barriers

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"] }

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<BraintrustClient>)
+/// # -> Result<(), Box<dyn std::error::Error>> {
+/// 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<T, E> {
+    issued: Result<(), E>,
+    receiver: oneshot::Receiver<Result<T, E>>,
+}
+impl<T, E> Completion<T, E> {
+    /// 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<Result<T, E>>, 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<T, E> Deref for Completion<T, E> {
+    type Target = Result<(), E>;
+    fn deref(&self) -> &Self::Target {
+        &self.issued
+    }
+}
+impl<T, E> std::fmt::Debug for Completion<T, E>
+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<T, E> {
+    issued: Option<Result<(), E>>,
+    receiver: oneshot::Receiver<Result<T, E>>,
+}
+// Neither field is structurally pinned; the channel receiver is Unpin.
+impl<T, E> Unpin for CompletionFuture<T, E> {}
+impl<T, E: From<ExecutionError>> IntoFuture for Completion<T, E> {
+    type Output = Result<T, E>;
+    type IntoFuture = CompletionFuture<T, E>;
+    fn into_future(self) -> Self::IntoFuture {
+        CompletionFuture {
+            issued: Some(self.issued),
+            receiver: self.receiver,
+        }
+    }
+}
+impl<T, E: From<ExecutionError>> Future for CompletionFuture<T, E> {
+    type Output = Result<T, E>;
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        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<ExecutionError> 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));
+    }
+}

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<DeliveryError>),
+}
+/// 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<DeliveryError>,
+}
+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<DeliveryError> {
+        &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<FlushError> 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),
+        }
+    }
+}

src/error.rs

@@ -19,3 +19,21 @@ pub enum BraintrustError {
 }
 
 pub type Result<T> = std::result::Result<T, BraintrustError>;
+
+// Compatibility for resource APIs that still return BraintrustError while the
+// public operation surfaces migrate to typed completion errors.
+impl From<crate::FlushError> 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::<Self>()
+            {
+                return Self::Api {
+                    status: *status,
+                    message: message.clone(),
+                };
+            }
+        }
+        Self::Background(error.to_string())
+    }
+}