diff --git a/src/internal/delivery.rs b/src/internal/delivery.rs index 56489e2..08351f8 100644 --- a/src/internal/delivery.rs +++ b/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>; #[derive(Clone)] @@ -394,6 +394,13 @@ impl WriteService { 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 diff --git a/src/internal/execution.rs b/src/internal/execution.rs index 9ad8a35..7e6f41e 100644 --- a/src/internal/execution.rs +++ b/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>, - responses: Arc, + responses: Arc, } struct RuntimeOwner { stop: Option>, @@ -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(&self, future: F) -> JoinHandle @@ -126,6 +128,19 @@ 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(&self, work: F) -> Completion where T: Send + 'static, @@ -133,14 +148,18 @@ impl ExecutionBinding { F: Future> + 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>); + 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::>() + })); + } + let close = threads.spawn(|| binding.close_responses()); + let tokens = producers + .into_iter() + .flat_map(|thread| thread.join().unwrap()) + .collect::>(); + 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>, + } + 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()); diff --git a/src/internal/lifecycle.rs b/src/internal/lifecycle.rs new file mode 100644 index 0000000..77d0db3 --- /dev/null +++ b/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; +pub(crate) struct Lifecycle { + execution: ExecutionBinding, + writes: LogQueue, + shutdown: Mutex>, +} +impl Lifecycle { + pub fn new(execution: ExecutionBinding, writes: LogQueue) -> Self { + Self { + execution, + writes, + shutdown: Mutex::new(None), + } + } + pub fn shutdown(&self, timeout: Duration) -> Completion { + 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 + } +} diff --git a/src/internal/mod.rs b/src/internal/mod.rs index 838118c..864bdec 100644 --- a/src/internal/mod.rs +++ b/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; diff --git a/src/internal/ports.rs b/src/internal/ports.rs index b6a6a34..7626cd8 100644 --- a/src/internal/ports.rs +++ b/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. diff --git a/src/internal/response_tasks.rs b/src/internal/response_tasks.rs new file mode 100644 index 0000000..0d45a31 --- /dev/null +++ b/src/internal/response_tasks.rs @@ -0,0 +1,190 @@ +//! Bounded response-task admission, completion publication and cancellation. +use crate::{EnqueueError, ExecutionError}; +use std::{ + collections::HashMap, + sync::{Arc, Mutex, MutexGuard}, +}; +use tokio::{ + sync::{oneshot, Notify}, + task::AbortHandle, +}; + +pub(crate) const MAX_RESPONSE_TASKS: usize = 1024; +trait CancelResponse: Send + Sync { + fn cancel(&self, error: ExecutionError); +} +struct Entry { + response: Arc, + abort: Option, +} +pub(crate) struct Response(Mutex>>>); +impl Response { + fn send(&self, result: Result) -> Option> { + let sender = lock(&self.0).take(); + match sender { + Some(sender) => sender.send(result).err(), + None => Some(result), + } + } +} +impl + Send> CancelResponse for Response { + fn cancel(&self, error: ExecutionError) { + drop(self.send(Err(error.into()))); + } +} +struct State { + accepting: bool, + terminal: Option, + next_id: u64, + active: HashMap, + publishing: usize, +} +pub(crate) struct ResponseTasks { + state: Mutex, + changed: Notify, +} +fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} +impl Default for ResponseTasks { + fn default() -> Self { + Self { + state: Mutex::new(State { + accepting: true, + terminal: None, + next_id: 0, + active: HashMap::new(), + publishing: 0, + }), + changed: Notify::new(), + } + } +} +impl ResponseTasks { + pub fn admit + Send + 'static>( + self: &Arc, + sender: oneshot::Sender>, + ) -> Result<(TaskGuard, Arc>), EnqueueError> { + let mut state = lock(&self.state); + if !state.accepting { + return Err(EnqueueError::Closed); + } + if state.active.len() + state.publishing >= MAX_RESPONSE_TASKS { + return Err(EnqueueError::Full); + } + let id = state.next_id; + state.next_id = id.checked_add(1).ok_or(EnqueueError::Full)?; + let response = Arc::new(Response(Mutex::new(Some(sender)))); + state.active.insert( + id, + Entry { + response: response.clone(), + abort: None, + }, + ); + Ok(( + TaskGuard { + tasks: self.clone(), + id, + finished: false, + }, + response, + )) + } + pub fn install_abort(&self, id: u64, abort: AbortHandle) { + let mut state = lock(&self.state); + if let Some(entry) = state.active.get_mut(&id) { + entry.abort = Some(abort); + } else { + abort.abort(); + } // Completion/cancellation can win the spawn race. + } + pub fn close(&self) { + lock(&self.state).accepting = false; + } + pub async fn drained(&self) -> Result<(), ExecutionError> { + loop { + let changed = self.changed.notified(); + { + let state = lock(&self.state); + if let Some(error) = state.terminal { + return Err(error); + } + if state.active.is_empty() && state.publishing == 0 { + return Ok(()); + } + } + changed.await; + } + } + pub fn cancel(&self, error: ExecutionError) { + let entries = { + let mut state = lock(&self.state); + state.accepting = false; + state.terminal.get_or_insert(error); + std::mem::take(&mut state.active) + }; + for entry in entries.into_values() { + entry.response.cancel(error); + if let Some(abort) = entry.abort { + abort.abort(); + } + } + self.changed.notify_waiters(); + } + // Reserve completion atomically against cancellation, then publish outside + // the lock. Drain still counts the reservation until publication is done. + fn completing(self: &Arc, id: u64) -> Option<(Entry, Publication)> { + let mut state = lock(&self.state); + let entry = state.active.remove(&id)?; + state.publishing += 1; + Some((entry, Publication(self.clone()))) + } + fn lost(self: &Arc, id: u64) { + if let Some((entry, _publication)) = self.completing(id) { + entry.response.cancel(ExecutionError::WorkerLost); + } + } + #[cfg(test)] + pub fn outstanding(&self) -> usize { + let state = lock(&self.state); + state.active.len() + state.publishing + } +} +struct Publication(Arc); +impl Drop for Publication { + fn drop(&mut self) { + lock(&self.0.state).publishing -= 1; + self.0.changed.notify_waiters(); + } +} +/// Constructed before spawning so runtime loss before first poll is observable. +pub(crate) struct TaskGuard { + tasks: Arc, + pub id: u64, + finished: bool, +} +impl TaskGuard { + pub fn complete(mut self, response: &Response, result: Result) { + self.finished = true; + if let Some((entry, publication)) = self.tasks.completing(self.id) { + let unused = response.send(result); + drop(entry); + // Neither waking an observer nor dropping its unobserved value runs + // under the registry lock; both can reenter SDK scheduling. + drop(unused); + drop(publication); + } else { + drop(result); // Cancellation already published this task's outcome. + } + } +} +impl Drop for TaskGuard { + fn drop(&mut self) { + if !self.finished { + self.tasks.lost(self.id); + } + } +} diff --git a/src/internal/shared_outcome.rs b/src/internal/shared_outcome.rs new file mode 100644 index 0000000..9b73d10 --- /dev/null +++ b/src/internal/shared_outcome.rs @@ -0,0 +1,183 @@ +//! A bounded final outcome for repeated observers of one logical operation. +use std::sync::{Arc, Mutex, MutexGuard}; +use tokio::sync::oneshot; + +use crate::{Completion, EnqueueError, ExecutionError}; + +type Sender = oneshot::Sender>; +const MAX_PENDING_OBSERVERS: usize = 1024; + +/// Lives with its operation owner, not with an active-work registry. +pub(crate) struct SharedOutcome(Arc>>); +struct State { + result: Option>, + observers: Vec>, +} +impl Clone for SharedOutcome { + fn clone(&self) -> Self { + Self(self.0.clone()) + } +} +impl Default for SharedOutcome { + fn default() -> Self { + Self(Arc::new(Mutex::new(State { + result: None, + observers: Vec::new(), + }))) + } +} +fn lock(value: &Mutex) -> MutexGuard<'_, T> { + value + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} +impl> SharedOutcome { + pub fn observe(&self) -> Completion { + let (sender, mut completion) = Completion::pair(); + let mut state = lock(&self.0); + if let Some(result) = &state.result { + let result = result.clone(); + drop(state); + let _ = sender.send(result); + } else { + state.observers.retain(|observer| !observer.is_closed()); + if state.observers.len() >= MAX_PENDING_OBSERVERS { + completion.record_issuance(Err(EnqueueError::Full.into())); + } else { + state.observers.push(sender); + } + } + completion + } + + pub fn finish(&self, result: Result) { + let mut state = lock(&self.0); + if state.result.is_some() { + return; + } + let observers = std::mem::take(&mut state.observers); + state.result = Some(result.clone()); + drop(state); + for sender in observers { + let _ = sender.send(result.clone()); + } + } +} + +/// Publish to both a one-shot observer and an optional shared final outcome. +/// No relay task or polling of a completion token is needed. +pub(crate) struct SharedResponder + From> { + sender: Option>, + outcome: Option>, +} +impl + From> SharedResponder { + pub fn new(sender: Sender, outcome: Option>) -> Self { + Self { + sender: Some(sender), + outcome, + } + } + pub fn send(mut self, result: Result) -> Result<(), Result> { + if let Some(outcome) = self.outcome.take() { + outcome.finish(result.clone()); + } + match self.sender.take() { + Some(sender) => sender.send(result), + None => Err(result), + } + } +} +impl + From> Drop + for SharedResponder +{ + fn drop(&mut self) { + if let Some(outcome) = self.outcome.take() { + outcome.finish(Err(ExecutionError::WorkerLost.into())); + } + } +} + +pub(crate) type WriteOutcome = SharedOutcome; +pub(crate) type WriteResponder = SharedResponder; + +#[cfg(test)] +mod tests { + use super::*; + use crate::{DeliveryReceipt, WriteError}; + + #[test] + fn publication_wakes_observers_outside_the_outcome_lock() { + use std::future::{Future, IntoFuture}; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::task::Context; + + struct Reenter { + outcome: WriteOutcome, + wakes: AtomicUsize, + } + impl futures::task::ArcWake for Reenter { + fn wake_by_ref(this: &Arc) { + assert!(this.outcome.observe().is_ok()); + this.wakes.fetch_add(1, Ordering::SeqCst); + } + } + let outcome = WriteOutcome::default(); + let observer = Arc::new(Reenter { + outcome: outcome.clone(), + wakes: AtomicUsize::new(0), + }); + let waker = futures::task::waker(observer.clone()); + let mut context = Context::from_waker(&waker); + let mut completion = Box::pin(outcome.observe().into_future()); + assert!(completion.as_mut().poll(&mut context).is_pending()); + outcome.finish(Ok(DeliveryReceipt)); + assert_eq!(observer.wakes.load(Ordering::SeqCst), 1); + assert!(completion.as_mut().poll(&mut context).is_ready()); + } + + #[tokio::test] + async fn pending_observers_are_bounded_and_dropped_tokens_release_slots() { + let outcome = WriteOutcome::default(); + let mut observers: Vec<_> = (0..MAX_PENDING_OBSERVERS) + .map(|_| outcome.observe()) + .collect(); + let full = outcome.observe(); + assert!(matches!( + &*full, + Err(WriteError::Admission(EnqueueError::Full)) + )); + assert!(matches!( + full.await, + Err(WriteError::Admission(EnqueueError::Full)) + )); + drop(observers.pop()); + let replacement = outcome.observe(); + assert!(replacement.is_ok()); + outcome.finish(Ok(DeliveryReceipt)); + for observer in observers { + observer.await.unwrap(); + } + replacement.await.unwrap(); + outcome.observe().await.unwrap(); + } + + #[tokio::test] + async fn producer_loss_completes_present_and_future_observers_without_a_relay_task() { + let outcome = WriteOutcome::default(); + let before = outcome.observe(); + let (sender, initial) = Completion::pair(); + drop(WriteResponder::new(sender, Some(outcome.clone()))); + assert!(matches!( + initial.await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + assert!(matches!( + before.await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + assert!(matches!( + outcome.observe().await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + } +} diff --git a/src/internal/write_outcome.rs b/src/internal/write_outcome.rs deleted file mode 100644 index 12415f1..0000000 --- a/src/internal/write_outcome.rs +++ /dev/null @@ -1,133 +0,0 @@ -//! One final write outcome shared by repeated span-end observers. -use std::sync::{Arc, Mutex, MutexGuard}; -use tokio::sync::oneshot; - -use crate::{Completion, DeliveryReceipt, EnqueueError, ExecutionError, WriteError}; - -type Outcome = Result; -type Sender = oneshot::Sender; -const MAX_PENDING_OBSERVERS: usize = 1024; - -/// Lives with a finished span, not with the queue. Holds no command or payload. -#[derive(Clone, Default)] -pub(crate) struct WriteOutcome(Arc>); -#[derive(Default)] -struct State { - result: Option, - observers: Vec, -} -fn lock(value: &Mutex) -> MutexGuard<'_, T> { - value - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) -} -impl WriteOutcome { - pub fn observe(&self) -> Completion { - let (sender, mut completion) = Completion::pair(); - let mut state = lock(&self.0); - if let Some(result) = &state.result { - let _ = sender.send(result.clone()); - } else { - state.observers.retain(|observer| !observer.is_closed()); - if state.observers.len() >= MAX_PENDING_OBSERVERS { - completion.record_issuance(Err(EnqueueError::Full.into())); - } else { - state.observers.push(sender); - } - } - completion - } - - fn finish(&self, result: Outcome) { - let mut state = lock(&self.0); - if state.result.is_some() { - return; - } - for sender in state.observers.drain(..) { - let _ = sender.send(result.clone()); - } - state.result = Some(result); - } -} - -/// Connect the existing batch sender to a span's optional shared final outcome. -/// No relay task or polling of a completion token is needed. -pub(crate) struct WriteResponder { - sender: Option, - outcome: Option, -} -impl WriteResponder { - pub fn new(sender: Sender, outcome: Option) -> Self { - Self { - sender: Some(sender), - outcome, - } - } - pub fn send(mut self, result: Outcome) -> Result<(), Outcome> { - if let Some(outcome) = self.outcome.take() { - outcome.finish(result.clone()); - } - match self.sender.take() { - Some(sender) => sender.send(result), - None => Err(result), - } - } -} -impl Drop for WriteResponder { - fn drop(&mut self) { - if let Some(outcome) = self.outcome.take() { - outcome.finish(Err(ExecutionError::WorkerLost.into())); - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[tokio::test] - async fn pending_observers_are_bounded_and_dropped_tokens_release_slots() { - let outcome = WriteOutcome::default(); - let mut observers: Vec<_> = (0..MAX_PENDING_OBSERVERS) - .map(|_| outcome.observe()) - .collect(); - let full = outcome.observe(); - assert!(matches!( - &*full, - Err(WriteError::Admission(EnqueueError::Full)) - )); - assert!(matches!( - full.await, - Err(WriteError::Admission(EnqueueError::Full)) - )); - drop(observers.pop()); - let replacement = outcome.observe(); - assert!(replacement.is_ok()); - outcome.finish(Ok(DeliveryReceipt)); - for observer in observers { - observer.await.unwrap(); - } - replacement.await.unwrap(); - outcome.observe().await.unwrap(); - } - - #[tokio::test] - async fn producer_loss_completes_present_and_future_observers_without_a_relay_task() { - let outcome = WriteOutcome::default(); - let before = outcome.observe(); - let (sender, initial) = Completion::pair(); - drop(WriteResponder::new(sender, Some(outcome.clone()))); - assert!(matches!( - initial.await, - Err(WriteError::Execution(ExecutionError::WorkerLost)) - )); - assert!(matches!( - before.await, - Err(WriteError::Execution(ExecutionError::WorkerLost)) - )); - assert!(matches!( - outcome.observe().await, - Err(WriteError::Execution(ExecutionError::WorkerLost)) - )); - } -} diff --git a/src/log_queue/queue.rs b/src/log_queue/queue.rs index c84bff1..3938100 100644 --- a/src/log_queue/queue.rs +++ b/src/log_queue/queue.rs @@ -10,7 +10,7 @@ use crate::internal::delivery::WriteService; use crate::internal::execution::ExecutionBinding; use crate::internal::ports::{DeliveryControl, SpanSubmitter}; use crate::internal::session::CredentialScope; -use crate::internal::write_outcome::WriteOutcome; +use crate::internal::shared_outcome::WriteOutcome; use crate::span_components::{project_logs_identifier, ProjectLogsIdentifier}; use crate::types::{LogDestination, Logs3Row, ParentSpanInfo, SpanObjectType, SpanPayload}; use crate::{ @@ -684,7 +684,7 @@ impl LogQueue { QueueState::Closed => { let (sender, mut completion) = Completion::pair(); completion.record_issuance(Err(crate::EnqueueError::Closed.into())); - let _ = crate::internal::write_outcome::WriteResponder::new(sender, outcome) + let _ = crate::internal::shared_outcome::WriteResponder::new(sender, outcome) .send(Err(crate::EnqueueError::Closed.into())); self.reject(command, crate::EnqueueError::Closed); completion @@ -729,6 +729,15 @@ impl LogQueue { let _ = sender.send(Ok(ShutdownReport::default())); completion } + pub fn cancel(&self, error: crate::ExecutionError) { + let mut state = self.state(); + if let QueueState::Running(writes) = &*state { + writes.cancel(error); + } else { + *state = QueueState::Closed; + } + } + #[cfg(test)] pub(crate) fn is_empty(&self) -> bool { match &*self.state() { diff --git a/src/logger.rs b/src/logger.rs index 0a13895..6ce2f14 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -12,6 +12,7 @@ use crate::error::{BraintrustError, Result}; use crate::experiments::ExperimentBuilder; use crate::http::{build_http_client, CaBundleConfig}; use crate::internal::execution::ExecutionBinding; +use crate::internal::lifecycle::Lifecycle; use crate::internal::ports::SpanSubmitter; use crate::internal::resources::ResourceTransport; use crate::internal::retry::{exponential, RetryExecutor}; @@ -435,6 +436,7 @@ impl BraintrustClientBuilder { inner: Arc::new(ClientInner { api_url, app_url, + lifecycle: Lifecycle::new(execution.clone(), queue.clone()), execution, queue, login_state, @@ -480,6 +482,7 @@ pub struct BraintrustClient { } struct ClientInner { + lifecycle: Lifecycle, execution: ExecutionBinding, #[allow(dead_code)] api_url: Url, @@ -878,7 +881,7 @@ impl BraintrustClient { /// Start login in a background task with retry logic. fn start_background_login(&self, api_key: String, org_name: Option) { let client = self.clone(); - self.inner.execution.spawn(async move { + self.inner.execution.response(async move { const MAX_ATTEMPTS: usize = 10; let outcome = RetryExecutor::new( MAX_ATTEMPTS - 1, @@ -898,7 +901,7 @@ impl BraintrustClient { }, ) .await; - match outcome.result { + match &outcome.result { Ok(()) => tracing::debug!("Background login completed successfully"), Err(error) => tracing::error!( "Background login failed after {} attempts: {}. \ @@ -907,6 +910,7 @@ impl BraintrustClient { error ), } + outcome.result }); } @@ -984,13 +988,15 @@ impl BraintrustClient { self.inner.queue.submit(scope, payload, parent_info, None) } - /// Stop accepting writes, drain accepted work and stop the delivery worker. + /// Stop response-task and write admission, drain accepted work, and stop delivery. + /// Concurrent calls share the first call's deadline and outcome. Individual + /// HTTP operation results stay on their own tokens; `delivered` counts writes. /// 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) + self.inner.lifecycle.shutdown(timeout) } /// Capture a cutoff immediately and observe delivery of the covered batch. @@ -1192,6 +1198,185 @@ mod tests { ResponseTemplate::new(200).set_body_json(serde_json::json!({ "org_info": org_info })) } + #[tokio::test] + async fn shutdown_covers_accepted_response_tasks_even_without_a_log_worker() { + let client = BraintrustClient::builder() + .skip_login(true) + .build() + .await + .unwrap(); + let response = client + .inner + .execution + .response::<(), BraintrustError, _>(std::future::pending()); + assert!(response.is_ok()); + let shutdown = client.shutdown(Duration::from_millis(10)); + assert!(matches!( + shutdown.await, + Err(crate::ShutdownError::Execution( + crate::ExecutionError::TimedOut + )) + )); + assert!(matches!( + response.await, + Err(BraintrustError::Execution(crate::ExecutionError::TimedOut)) + )); + assert_eq!(client.inner.execution.outstanding_responses(), 0); + } + + #[tokio::test] + async fn concurrent_shutdowns_share_the_first_deadline_and_reject_both_admission_domains() { + let client = BraintrustClient::builder() + .skip_login(true) + .build() + .await + .unwrap(); + let response = client + .inner + .execution + .response::<(), BraintrustError, _>(std::future::pending()); + let first = client.shutdown(Duration::from_millis(10)); + let concurrent = client.clone().shutdown(Duration::from_secs(60)); + assert!(matches!( + client + .inner + .execution + .response::<(), BraintrustError, _>(async { Ok(()) }) + .await, + Err(BraintrustError::Admission(crate::EnqueueError::Closed)) + )); + let span = client + .span_builder_with_credentials("token", "org") + .parent_info(ParentSpanInfo::ProjectLogs { + object_id: "project".into(), + }) + .build(); + assert!(matches!( + span.end().await, + Err(crate::WriteError::Admission(crate::EnqueueError::Closed)) + )); + for token in [first, concurrent, client.shutdown(Duration::ZERO)] { + assert!(matches!( + token.await, + Err(crate::ShutdownError::Execution( + crate::ExecutionError::TimedOut + )) + )); + } + assert!(matches!( + response.await, + Err(BraintrustError::Execution(crate::ExecutionError::TimedOut)) + )); + assert_eq!(client.inner.execution.outstanding_responses(), 0); + } + + #[tokio::test] + async fn a_write_failure_does_not_short_circuit_draining_response_tasks() { + use futures::FutureExt; + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(401)) + .expect(1) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .skip_login(true) + .runtime(crate::current_runtime()) + .api_url(server.uri()) + .app_url(server.uri()) + .build() + .await + .unwrap(); + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + let work_gate = gate.clone(); + let response = client + .inner + .execution + .response::<(), BraintrustError, _>(async move { + work_gate.acquire().await.unwrap().forget(); + Ok(()) + }); + let span = client + .span_builder_with_credentials("token", "org") + .parent_info(ParentSpanInfo::ProjectLogs { + object_id: "project".into(), + }) + .build(); + let write = span.end(); + let shutdown = client.shutdown(Duration::from_secs(5)); + assert!(write.await.is_err()); + let mut shutdown = Box::pin(std::future::IntoFuture::into_future(shutdown)); + assert!(shutdown.as_mut().now_or_never().is_none()); + gate.add_permits(1); + response.await.unwrap(); + assert!( + matches!(shutdown.await, Err(crate::ShutdownError::Delivery(failure)) if failure.failed == 1) + ); + assert_eq!(client.inner.execution.outstanding_responses(), 0); + } + + #[tokio::test] + async fn shutdown_has_reserved_control_capacity_when_response_admission_is_full() { + let client = BraintrustClient::builder() + .skip_login(true) + .runtime(crate::current_runtime()) + .build() + .await + .unwrap(); + let responses = (0..crate::internal::response_tasks::MAX_RESPONSE_TASKS) + .map(|_| { + let response = client + .inner + .execution + .response::<(), BraintrustError, _>(std::future::pending()); + assert!(response.is_ok()); + response + }) + .collect::>(); + let shutdown = client.shutdown(Duration::from_millis(10)); + assert!(shutdown.is_ok()); + assert!(matches!( + shutdown.await, + Err(crate::ShutdownError::Execution( + crate::ExecutionError::TimedOut + )) + )); + for response in responses { + assert!(matches!( + response.await, + Err(BraintrustError::Execution(crate::ExecutionError::TimedOut)) + )); + } + assert_eq!(client.inner.execution.outstanding_responses(), 0); + } + + #[tokio::test] + async fn response_operation_errors_stay_on_their_tokens_instead_of_a_shutdown_failure_ledger() { + let client = BraintrustClient::builder() + .skip_login(true) + .runtime(crate::current_runtime()) + .build() + .await + .unwrap(); + let response = client + .inner + .execution + .response::<(), BraintrustError, _>(async { + Err(BraintrustError::Api { + status: 503, + message: "operation failed".into(), + }) + }); + let shutdown = client.shutdown(Duration::from_secs(1)); + assert!(matches!( + response.await, + Err(BraintrustError::Api { status: 503, .. }) + )); + assert_eq!(shutdown.await.unwrap().delivered, 0); + assert_eq!(client.inner.execution.outstanding_responses(), 0); + } + #[tokio::test] async fn reads_and_empty_lifecycle_calls_do_not_start_a_log_worker() { let server = MockServer::start().await; diff --git a/src/span.rs b/src/span.rs index 7d8e7e4..18c993d 100644 --- a/src/span.rs +++ b/src/span.rs @@ -13,7 +13,7 @@ use uuid::Uuid; use crate::error::Result; use crate::internal::ports::SpanSubmitter; use crate::internal::session::CredentialScope; -use crate::internal::write_outcome::WriteOutcome; +use crate::internal::shared_outcome::WriteOutcome; use crate::span_components::SpanComponents; use crate::types::{ ParentSpanInfo, SpanAttributes, SpanEventData, SpanObjectType, SpanPayload, SpanType, @@ -1101,7 +1101,7 @@ mod tests { outcome: Option, ) -> crate::Completion { let (sender, mut completion) = crate::Completion::pair(); - let sender = crate::internal::write_outcome::WriteResponder::new(sender, outcome); + let sender = crate::internal::shared_outcome::WriteResponder::new(sender, outcome); if self.reject.load(Ordering::Relaxed) { completion.record_issuance(Err(crate::EnqueueError::Full.into())); } else { @@ -1154,7 +1154,7 @@ mod tests { #[tokio::test] async fn concurrent_ends_share_one_snapshot_and_one_delivery_result() { - use crate::internal::write_outcome::WriteResponder; + use crate::internal::shared_outcome::WriteResponder; struct Submitter(Mutex>); #[async_trait] impl crate::internal::ports::DeliveryControl for Submitter { diff --git a/src/test_utils.rs b/src/test_utils.rs index dc527db..9c91732 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -8,7 +8,7 @@ use crate::span::{SpanBuilder, SpanHandle}; use crate::types::{ParentSpanInfo, SpanPayload}; use crate::internal::ports::{DeliveryControl, SpanSubmitter}; -use crate::internal::write_outcome::{WriteOutcome, WriteResponder}; +use crate::internal::shared_outcome::{WriteOutcome, WriteResponder}; #[derive(Clone, Default)] pub(crate) struct TestSpanCollector { diff --git a/tests/response_lifecycle.rs b/tests/response_lifecycle.rs new file mode 100644 index 0000000..0c86885 --- /dev/null +++ b/tests/response_lifecycle.rs @@ -0,0 +1,131 @@ +use braintrust_sdk_rust::{ + BraintrustClient, BraintrustError, EnqueueError, ExecutionError, ParentSpanInfo, ShutdownError, + WriteError, +}; +use std::{sync::Arc, time::Duration}; +use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, +}; + +fn held_response( + arrived: Arc, + status: u16, +) -> impl Fn(&wiremock::Request) -> ResponseTemplate { + move |_| { + arrived.notify_one(); + ResponseTemplate::new(status).set_delay(Duration::from_secs(60)) + } +} + +#[tokio::test] +async fn one_deadline_fails_active_http_and_ingestion_then_rejects_new_work() { + let server = MockServer::start().await; + let http_arrived = Arc::new(tokio::sync::Notify::new()); + let write_arrived = Arc::new(tokio::sync::Notify::new()); + Mock::given(method("POST")) + .and(path("/api/dataset/register")) + .respond_with(held_response(http_arrived.clone(), 200)) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/version")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({}))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(held_response(write_arrived.clone(), 204)) + .expect(1) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .skip_login(true) + .app_url(server.uri()) + .api_url(server.uri()) + .build() + .await + .unwrap(); + let dataset = client + .dataset_builder_with_credentials("token", "tenant") + .project_name("p") + .dataset_name("d") + .build() + .unwrap(); + let pending = dataset.clone(); + let http = tokio::spawn(async move { pending.dataset_id().await }); + let span = client + .span_builder_with_credentials("token", "org") + .parent_info(ParentSpanInfo::ProjectLogs { + object_id: "project".into(), + }) + .build(); + let write = span.end(); + tokio::time::timeout(Duration::from_secs(5), async { + http_arrived.notified().await; + write_arrived.notified().await; + }) + .await + .unwrap(); + let first = client.shutdown(Duration::from_millis(20)); + let concurrent = client.shutdown(Duration::from_secs(60)); + for shutdown in [first, concurrent] { + assert!(matches!( + shutdown.await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); + } + assert!(matches!( + write.await, + Err(WriteError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + http.await.unwrap(), + Err(BraintrustError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + dataset.dataset_id().await, + Err(BraintrustError::Admission(EnqueueError::Closed)) + )); + let fresh = client + .span_builder_with_credentials("token", "org") + .parent_info(ParentSpanInfo::ProjectLogs { + object_id: "project".into(), + }) + .build(); + assert!(matches!( + fresh.end().await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); + assert_eq!(server.received_requests().await.unwrap().len(), 3); +} + +#[tokio::test] +async fn background_login_is_accepted_work_covered_by_shutdown() { + let server = MockServer::start().await; + let arrived = Arc::new(tokio::sync::Notify::new()); + Mock::given(method("POST")) + .and(path("/api/apikey/login")) + .respond_with(held_response(arrived.clone(), 200)) + .expect(1) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .api_key("token") + .blocking_login(false) + .app_url(server.uri()) + .api_url(server.uri()) + .build() + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), arrived.notified()) + .await + .unwrap(); + assert!(matches!( + client.shutdown(Duration::from_millis(20)).await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); + assert!(!client.is_logged_in().await); + assert_eq!(server.received_requests().await.unwrap().len(), 1); +} diff --git a/tests/runtime_binding.rs b/tests/runtime_binding.rs index 3855762..0224c0d 100644 --- a/tests/runtime_binding.rs +++ b/tests/runtime_binding.rs @@ -98,10 +98,18 @@ fn explicit_current_thread_runtime_must_be_driven_and_observation_does_not_move_ tokio::task::yield_now().await; }); let client = observer.block_on(initialization).unwrap(); - let shutdown = client.shutdown(std::time::Duration::from_secs(1)); + let mut shutdown = Box::pin( + client + .shutdown(std::time::Duration::from_secs(1)) + .into_future(), + ); observer.block_on(async { - shutdown.await.unwrap(); + assert!(futures::poll!(shutdown.as_mut()).is_pending()); + }); + selected.block_on(async { + tokio::task::yield_now().await; }); + observer.block_on(shutdown).unwrap(); drop(client); // SDK shutdown did not stop the caller-owned runtime. assert_eq!(selected.block_on(async { 42 }), 42);