Coordinate response tasks and delivery during shutdown

src/internal/delivery.rs

@@ -19,7 +19,7 @@ use tokio::{
 use tracing::warn;
 
 use crate::internal::execution::ExecutionBinding;
-use crate::internal::write_outcome::{WriteOutcome, WriteResponder};
+use crate::internal::shared_outcome::{WriteOutcome, WriteResponder};
 type WriteSender = WriteResponder;
 type FlushSender = oneshot::Sender<Result<FlushReport, FlushError>>;
 #[derive(Clone)]
@@ -394,6 +394,13 @@ impl<C: Send + 'static> WriteService<C> {
         completion
     }
 
+    pub fn cancel(&self, error: ExecutionError) {
+        self.inner.stop(error);
+        if let Some(abort) = lock(&self.inner.abort).as_ref() {
+            abort.abort();
+        }
+    }
+
     #[cfg(test)]
     pub fn is_empty(&self) -> bool {
         lock(&self.inner.state).outstanding == 0

src/internal/execution.rs

@@ -3,18 +3,20 @@ use crate::{BraintrustError, Completion, EnqueueError, ExecutionError};
 use std::{future::Future, sync::Arc, time::Duration};
 use tokio::{
     runtime::{Handle, Runtime},
-    sync::{oneshot, Semaphore},
+    sync::oneshot,
     task::JoinHandle,
 };
 
-const MAX_RESPONSE_TASKS: usize = 1024;
+use crate::internal::response_tasks::ResponseTasks;
+#[cfg(test)]
+use crate::internal::response_tasks::MAX_RESPONSE_TASKS;
 
 #[derive(Clone)]
 pub(crate) struct ExecutionBinding {
     handle: Handle,
     // Handles alone do not own a runtime. Every accepted task retains this lease.
     _owner: Option<Arc<RuntimeOwner>>,
-    responses: Arc<Semaphore>,
+    responses: Arc<ResponseTasks>,
 }
 struct RuntimeOwner {
     stop: Option<oneshot::Sender<()>>,
@@ -94,14 +96,14 @@ impl ExecutionBinding {
                 #[cfg(test)]
                 exit,
             })),
-            responses: Arc::new(Semaphore::new(MAX_RESPONSE_TASKS)),
+            responses: Arc::new(ResponseTasks::default()),
         })
     }
     pub fn external(handle: Handle) -> Self {
         Self {
             handle,
             _owner: None,
-            responses: Arc::new(Semaphore::new(MAX_RESPONSE_TASKS)),
+            responses: Arc::new(ResponseTasks::default()),
         }
     }
     pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
@@ -126,21 +128,38 @@ impl ExecutionBinding {
             work()
         })
     }
+    pub fn close_responses(&self) {
+        self.responses.close();
+    }
+    pub async fn responses_drained(&self) -> Result<(), ExecutionError> {
+        self.responses.drained().await
+    }
+    pub fn cancel_responses(&self, error: ExecutionError) {
+        self.responses.cancel(error);
+    }
+    #[cfg(test)]
+    pub fn outstanding_responses(&self) -> usize {
+        self.responses.outstanding()
+    }
     pub fn response<T, E, F>(&self, work: F) -> Completion<T, E>
     where
         T: Send + 'static,
         E: From<ExecutionError> + From<EnqueueError> + Send + 'static,
         F: Future<Output = Result<T, E>> + Send + 'static,
     {
         let (sender, mut completion) = Completion::pair();
-        let Ok(permit) = self.responses.clone().try_acquire_owned() else {
-            completion.record_issuance(Err(EnqueueError::Full.into()));
-            return completion;
+        let (guard, response) = match self.responses.admit(sender) {
+            Ok(admitted) => admitted,
+            Err(error) => {
+                completion.record_issuance(Err(error.into()));
+                return completion;
+            }
         };
-        self.spawn(async move {
-            let _permit = permit;
-            let _ = sender.send(work.await);
+        let id = guard.id;
+        let task = self.spawn(async move {
+            guard.complete(&response, work.await);
         });
+        self.responses.install_abort(id, task.abort_handle());
         completion
     }
 }
@@ -156,7 +175,7 @@ fn drive(runtime: Runtime, stopped: oneshot::Receiver<()>) {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use tokio::sync::mpsc;
+    use tokio::sync::{mpsc, Semaphore};
 
     #[tokio::test]
     async fn default_binding_ignores_ambient_runtime_and_explicit_binding_honors_it() {
@@ -230,6 +249,156 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn close_rejects_new_responses_but_drains_unobserved_accepted_work() {
+        use futures::FutureExt;
+        let binding = ExecutionBinding::external(Handle::current());
+        let gate = Arc::new(Semaphore::new(0));
+        let work_gate = gate.clone();
+        drop(binding.response::<(), BraintrustError, _>(async move {
+            work_gate.acquire().await.unwrap().forget();
+            Ok(())
+        }));
+        binding.close_responses();
+        assert_eq!(binding.outstanding_responses(), 1);
+        assert!(matches!(
+            binding
+                .response::<(), BraintrustError, _>(async { Ok(()) })
+                .await,
+            Err(BraintrustError::Admission(EnqueueError::Closed))
+        ));
+        assert!(binding.responses_drained().now_or_never().is_none());
+        gate.add_permits(1);
+        binding.responses_drained().await.unwrap();
+        assert_eq!(binding.outstanding_responses(), 0);
+    }
+
+    #[tokio::test]
+    async fn timeout_before_first_poll_notifies_tokens_and_drops_their_work() {
+        struct DropSignal(Option<oneshot::Sender<()>>);
+        impl Drop for DropSignal {
+            fn drop(&mut self) {
+                let _ = self.0.take().unwrap().send(());
+            }
+        }
+        let binding = ExecutionBinding::external(Handle::current());
+        let (sender, dropped) = oneshot::channel();
+        let signal = DropSignal(Some(sender));
+        let response = binding.response::<(), BraintrustError, _>(async move {
+            let _signal = signal;
+            std::future::pending().await
+        });
+        binding.close_responses();
+        binding.cancel_responses(ExecutionError::TimedOut);
+        assert!(matches!(
+            response.await,
+            Err(BraintrustError::Execution(ExecutionError::TimedOut))
+        ));
+        assert_eq!(binding.outstanding_responses(), 0);
+        assert_eq!(
+            binding.responses_drained().await,
+            Err(ExecutionError::TimedOut)
+        );
+        tokio::time::timeout(Duration::from_secs(1), dropped)
+            .await
+            .unwrap()
+            .unwrap();
+    }
+
+    #[tokio::test]
+    async fn task_loss_retires_its_entry_and_does_not_poison_fresh_work() {
+        let binding = ExecutionBinding::external(Handle::current());
+        let response = binding
+            .response::<(), BraintrustError, _>(async { panic!("response processor failed") });
+        assert!(matches!(
+            response.await,
+            Err(BraintrustError::Execution(ExecutionError::WorkerLost))
+        ));
+        assert_eq!(binding.outstanding_responses(), 0);
+        binding
+            .response::<(), BraintrustError, _>(async { Ok(()) })
+            .await
+            .unwrap();
+        binding.close_responses();
+        binding.responses_drained().await.unwrap();
+    }
+
+    #[tokio::test]
+    async fn racing_admission_and_close_cannot_lose_accepted_response_tasks() {
+        use std::sync::atomic::{AtomicUsize, Ordering};
+        let binding = ExecutionBinding::external(Handle::current());
+        let processed = Arc::new(AtomicUsize::new(0));
+        let first = binding.response::<(), BraintrustError, _>(async { Ok(()) });
+        let tokens = std::thread::scope(|threads| {
+            let mut producers = Vec::new();
+            for _ in 0..4 {
+                let binding = binding.clone();
+                let processed = processed.clone();
+                producers.push(threads.spawn(move || {
+                    (0..100)
+                        .map(|_| {
+                            let processed = processed.clone();
+                            binding.response::<(), BraintrustError, _>(async move {
+                                processed.fetch_add(1, Ordering::Relaxed);
+                                Ok(())
+                            })
+                        })
+                        .collect::<Vec<_>>()
+                }));
+            }
+            let close = threads.spawn(|| binding.close_responses());
+            let tokens = producers
+                .into_iter()
+                .flat_map(|thread| thread.join().unwrap())
+                .collect::<Vec<_>>();
+            close.join().unwrap();
+            tokens
+        });
+        let admitted = tokens.iter().filter(|token| token.is_ok()).count();
+        first.await.unwrap();
+        for token in tokens {
+            let _ = token.await;
+        }
+        binding.responses_drained().await.unwrap();
+        assert_eq!(processed.load(Ordering::Relaxed), admitted);
+        assert_eq!(binding.outstanding_responses(), 0);
+    }
+
+    #[tokio::test]
+    async fn unobserved_results_can_reenter_scheduling_from_their_destructor() {
+        struct Reenter {
+            binding: ExecutionBinding,
+            done: Option<oneshot::Sender<()>>,
+        }
+        impl Drop for Reenter {
+            fn drop(&mut self) {
+                assert_eq!(
+                    self.binding.outstanding_responses(),
+                    1,
+                    "drain must include result cleanup"
+                );
+                assert!(self
+                    .binding
+                    .response::<(), BraintrustError, _>(async { Ok(()) })
+                    .is_ok());
+                let _ = self.done.take().unwrap().send(());
+            }
+        }
+        let binding = ExecutionBinding::external(Handle::current());
+        let (done, finished) = oneshot::channel();
+        let value = Reenter {
+            binding: binding.clone(),
+            done: Some(done),
+        };
+        drop(binding.response::<_, BraintrustError, _>(async move { Ok(value) }));
+        tokio::time::timeout(Duration::from_secs(1), finished)
+            .await
+            .unwrap()
+            .unwrap();
+        binding.close_responses();
+        binding.responses_drained().await.unwrap();
+    }
+
     #[test]
     fn current_runtime_outside_tokio_selects_the_internal_default() {
         assert!(crate::current_runtime().is_none());

src/internal/lifecycle.rs

@@ -0,0 +1,65 @@
+//! One client shutdown coordinates response tasks and ingestion under one deadline.
+use crate::internal::execution::ExecutionBinding;
+use crate::internal::shared_outcome::{SharedOutcome, SharedResponder};
+use crate::log_queue::LogQueue;
+use crate::{Completion, ExecutionError, ShutdownError, ShutdownReport};
+use std::sync::Mutex;
+use std::time::Duration;
+
+type Outcome = SharedOutcome<ShutdownReport, ShutdownError>;
+pub(crate) struct Lifecycle {
+    execution: ExecutionBinding,
+    writes: LogQueue,
+    shutdown: Mutex<Option<Outcome>>,
+}
+impl Lifecycle {
+    pub fn new(execution: ExecutionBinding, writes: LogQueue) -> Self {
+        Self {
+            execution,
+            writes,
+            shutdown: Mutex::new(None),
+        }
+    }
+    pub fn shutdown(&self, timeout: Duration) -> Completion<ShutdownReport, ShutdownError> {
+        let mut state = self
+            .shutdown
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        if let Some(outcome) = &*state {
+            return outcome.observe();
+        }
+        let outcome = Outcome::default();
+        let (sender, completion) = Completion::pair();
+        let responder = SharedResponder::new(sender, Some(outcome.clone()));
+        *state = Some(outcome);
+        // Close both admission domains before returning any shutdown observer.
+        self.execution.close_responses();
+        let writes = self.writes.shutdown(timeout);
+        let execution = self.execution.clone();
+        let queue = self.writes.clone();
+        let deadline = tokio::time::Instant::now().checked_add(timeout);
+        self.execution.spawn(async move {
+            let work = async {
+                // A write failure must not stop waiting for accepted HTTP work.
+                let (writes, responses) =
+                    futures::join!(async { writes.await }, execution.responses_drained());
+                responses.map_err(ShutdownError::Execution)?;
+                writes
+            };
+            let result = match deadline {
+                Some(deadline) => tokio::time::timeout_at(deadline, work).await,
+                None => Ok(work.await),
+            };
+            let result = match result {
+                Ok(result) => result,
+                Err(_) => {
+                    execution.cancel_responses(ExecutionError::TimedOut);
+                    queue.cancel(ExecutionError::TimedOut);
+                    Err(ShutdownError::Execution(ExecutionError::TimedOut))
+                }
+            };
+            let _ = responder.send(result);
+        });
+        completion
+    }
+}

src/internal/mod.rs

@@ -2,10 +2,12 @@ pub(crate) mod batching;
 pub(crate) mod delivery;
 pub(crate) mod execution;
 pub(crate) mod ingest;
+pub(crate) mod lifecycle;
 pub(crate) mod object_store;
 pub(crate) mod ports;
 pub(crate) mod resources;
+pub(crate) mod response_tasks;
 pub(crate) mod retry;
 pub(crate) mod session;
+pub(crate) mod shared_outcome;
 pub(crate) mod upload;
-pub(crate) mod write_outcome;

src/internal/ports.rs

@@ -13,7 +13,7 @@ use crate::api::{
 use crate::error::Result;
 use crate::experiments::BaseExperimentInfo;
 use crate::internal::session::CredentialScope;
-use crate::internal::write_outcome::WriteOutcome;
+use crate::internal::shared_outcome::WriteOutcome;
 use crate::types::{ParentSpanInfo, SpanPayload};
 
 /// Shared delivery lifecycle, independent of resource HTTP operations.