From 6d18a20b3ccfc02a35bda5ff9dd6f5981a4dcb9d Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Tue, 20 Oct 2026 22:24:53 +0800 Subject: [PATCH 1/2] Bind SDK work to an explicit or owned runtime --- src/error.rs | 4 + src/internal/delivery.rs | 128 ++++++++++++-------- src/internal/execution.rs | 244 ++++++++++++++++++++++++++++++++++++++ src/internal/mod.rs | 1 + src/internal/resources.rs | 112 ++++++++++------- src/lib.rs | 2 + src/log_queue/queue.rs | 212 +++++++++++++++++++++++++-------- src/logger.rs | 103 ++++++++++++++-- src/runtime.rs | 10 ++ tests/dataset_test.rs | 1 + tests/finished_spans.rs | 1 + tests/http_flow.rs | 4 + tests/runtime_binding.rs | 173 +++++++++++++++++++++++++++ 13 files changed, 843 insertions(+), 152 deletions(-) create mode 100644 src/internal/execution.rs create mode 100644 src/runtime.rs create mode 100644 tests/runtime_binding.rs diff --git a/src/error.rs b/src/error.rs index 7d2c366..205da7a 100644 --- a/src/error.rs +++ b/src/error.rs @@ -2,6 +2,10 @@ use thiserror::Error; #[derive(Debug, Error)] pub enum BraintrustError { + #[error(transparent)] + Execution(#[from] crate::ExecutionError), + #[error(transparent)] + Admission(#[from] crate::EnqueueError), #[error("invalid configuration: {0}")] InvalidConfig(String), #[error("http error: {0}")] diff --git a/src/internal/delivery.rs b/src/internal/delivery.rs index 87cb815..56489e2 100644 --- a/src/internal/delivery.rs +++ b/src/internal/delivery.rs @@ -13,12 +13,12 @@ use std::{ time::Duration, }; use tokio::{ - runtime::Handle, sync::{oneshot, Notify, Semaphore}, task::AbortHandle, }; use tracing::warn; +use crate::internal::execution::ExecutionBinding; use crate::internal::write_outcome::{WriteOutcome, WriteResponder}; type WriteSender = WriteResponder; type FlushSender = oneshot::Sender>; @@ -215,7 +215,7 @@ impl Inner { pub(crate) struct WriteService { inner: Arc>, owners: Arc>, - runtime: Handle, + runtime: ExecutionBinding, exit: Shared>>, } impl Clone for WriteService { @@ -239,7 +239,7 @@ impl Drop for ProducerOwner { } } impl WriteService { - pub fn new(capacity: usize, runtime: Handle, process: F) -> Self + pub fn new(capacity: usize, runtime: ExecutionBinding, process: F) -> Self where F: Fn(Vec) -> BoxFuture<'static, anyhow::Result<()>> + Send + 'static, { @@ -472,15 +472,19 @@ mod tests { 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(()) - }) - }); + let service = WriteService::new( + capacity, + ExecutionBinding::external(tokio::runtime::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) } @@ -516,18 +520,22 @@ mod tests { 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 service = WriteService::new( + 4, + ExecutionBinding::external(tokio::runtime::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); @@ -548,11 +556,15 @@ mod tests { #[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 service = WriteService::new( + 4, + ExecutionBinding::external(tokio::runtime::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)); @@ -605,15 +617,19 @@ mod tests { 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(()) - }) - }); + let service = WriteService::new( + 4, + ExecutionBinding::external(tokio::runtime::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); @@ -765,15 +781,19 @@ mod tests { 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 service = WriteService::new( + 4, + ExecutionBinding::external(tokio::runtime::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); @@ -817,11 +837,15 @@ mod tests { 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 service = WriteService::new( + 1024, + ExecutionBinding::external(tokio::runtime::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| { diff --git a/src/internal/execution.rs b/src/internal/execution.rs new file mode 100644 index 0000000..9ad8a35 --- /dev/null +++ b/src/internal/execution.rs @@ -0,0 +1,244 @@ +//! Fixed execution binding and ownership of a continuously driven SDK runtime. +use crate::{BraintrustError, Completion, EnqueueError, ExecutionError}; +use std::{future::Future, sync::Arc, time::Duration}; +use tokio::{ + runtime::{Handle, Runtime}, + sync::{oneshot, Semaphore}, + task::JoinHandle, +}; + +const MAX_RESPONSE_TASKS: usize = 1024; + +#[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, +} +struct RuntimeOwner { + stop: Option>, + #[cfg(test)] + exit: Arc, +} +#[cfg(test)] +#[derive(Default)] +struct ExitProbe { + done: std::sync::atomic::AtomicBool, + wake: tokio::sync::Notify, +} +#[cfg(test)] +impl ExitProbe { + async fn wait(&self) { + let wake = self.wake.notified(); + if !self.done.load(std::sync::atomic::Ordering::Acquire) { + wake.await; + } + } + fn finish(&self) { + self.done.store(true, std::sync::atomic::Ordering::Release); + self.wake.notify_one(); + } +} +impl Drop for RuntimeOwner { + fn drop(&mut self) { + if let Some(stop) = self.stop.take() { + let _ = stop.send(()); + } + } +} +impl ExecutionBinding { + pub fn new(runtime: Option) -> Result { + if let Some(handle) = runtime { + return Ok(Self::external(handle)); + } + let (initialized, receiver) = std::sync::mpsc::sync_channel(1); + let (stop, stopped) = oneshot::channel(); + #[cfg(test)] + let exit = Arc::new(ExitProbe::default()); + #[cfg(test)] + let host_exit = exit.clone(); + std::thread::Builder::new() + .name("braintrust-runtime".into()) + .spawn(move || { + let runtime = match tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + { + Ok(runtime) => runtime, + Err(error) => { + let _ = initialized.send(Err(error)); + return; + } + }; + if initialized.send(Ok(runtime.handle().clone())).is_err() { + return; + } + drive(runtime, stopped); + #[cfg(test)] + host_exit.finish(); + }) + .map_err(|error| { + BraintrustError::InvalidConfig(format!("failed to start SDK runtime: {error}")) + })?; + let handle = receiver + .recv() + .map_err(|_| BraintrustError::Execution(ExecutionError::WorkerLost))? + .map_err(|error| { + BraintrustError::InvalidConfig(format!("failed to initialize SDK runtime: {error}")) + })?; + Ok(Self { + handle, + _owner: Some(Arc::new(RuntimeOwner { + stop: Some(stop), + #[cfg(test)] + exit, + })), + responses: Arc::new(Semaphore::new(MAX_RESPONSE_TASKS)), + }) + } + pub fn external(handle: Handle) -> Self { + Self { + handle, + _owner: None, + responses: Arc::new(Semaphore::new(MAX_RESPONSE_TASKS)), + } + } + pub fn spawn(&self, future: F) -> JoinHandle + where + F: Future + Send + 'static, + F::Output: Send + 'static, + { + let lease = self.clone(); + self.handle.spawn(async move { + let _lease = lease; + future.await + }) + } + pub fn spawn_blocking(&self, work: F) -> JoinHandle + where + F: FnOnce() -> T + Send + 'static, + T: Send + 'static, + { + let lease = self.clone(); + self.handle.spawn_blocking(move || { + let _lease = lease; + work() + }) + } + pub fn response(&self, work: F) -> Completion + where + T: Send + 'static, + E: From + From + Send + 'static, + 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; + }; + self.spawn(async move { + let _permit = permit; + let _ = sender.send(work.await); + }); + completion + } +} +// Runtime destruction occurs on its host thread after leaving the runtime +// context. No join or Runtime::drop runs on an SDK worker or caller's executor. +fn drive(runtime: Runtime, stopped: oneshot::Receiver<()>) { + runtime.block_on(async { + let _ = stopped.await; + }); + runtime.shutdown_timeout(Duration::from_millis(100)); +} + +#[cfg(test)] +mod tests { + use super::*; + use tokio::sync::mpsc; + + #[tokio::test] + async fn default_binding_ignores_ambient_runtime_and_explicit_binding_honors_it() { + let ambient = Handle::current().id(); + let owned = ExecutionBinding::new(None).unwrap(); + let selected = owned.spawn(async { Handle::current().id() }).await.unwrap(); + assert_ne!(selected, ambient); + assert_eq!( + owned + .spawn_blocking(|| Handle::current().id()) + .await + .unwrap(), + selected + ); + let explicit = ExecutionBinding::new(crate::current_runtime()).unwrap(); + assert_eq!( + explicit + .spawn(async { Handle::current().id() }) + .await + .unwrap(), + ambient + ); + assert!(owned._owner.is_some()); + assert!(explicit._owner.is_none()); + } + + #[tokio::test] + async fn accepted_response_tasks_retain_the_owner_and_bound_unobserved_work() { + let binding = ExecutionBinding::new(None).unwrap(); + let owner = Arc::downgrade(binding._owner.as_ref().unwrap()); + let exit = binding._owner.as_ref().unwrap().exit.clone(); + let gate = Arc::new(Semaphore::new(0)); + let (finished, mut done) = mpsc::unbounded_channel(); + let mut tokens = Vec::new(); + for _ in 0..MAX_RESPONSE_TASKS { + let gate = gate.clone(); + let finished = finished.clone(); + let token = binding.response::<_, BraintrustError, _>(async move { + gate.acquire().await.unwrap().forget(); + finished.send(()).unwrap(); + Ok(()) + }); + assert!(token.is_ok()); + tokens.push(token); + } + drop(tokens.pop()); + assert!(matches!( + binding + .response::<(), BraintrustError, _>(async { Ok(()) }) + .await, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + drop(binding); + assert!( + owner.upgrade().is_some(), + "accepted tasks must retain the runtime owner" + ); + gate.add_permits(MAX_RESPONSE_TASKS); + for token in tokens { + token.await.unwrap(); + } + for _ in 0..MAX_RESPONSE_TASKS { + done.recv().await.unwrap(); + } + tokio::time::timeout(Duration::from_secs(5), exit.wait()) + .await + .unwrap(); + assert!( + owner.upgrade().is_none(), + "runtime host must stop after accepted work and owners retire" + ); + } + + #[test] + fn current_runtime_outside_tokio_selects_the_internal_default() { + assert!(crate::current_runtime().is_none()); + let binding = ExecutionBinding::new(crate::current_runtime()).unwrap(); + futures::executor::block_on(async { + binding + .response::<_, BraintrustError, _>(async { Ok(()) }) + .await + .unwrap(); + }); + } +} diff --git a/src/internal/mod.rs b/src/internal/mod.rs index a55e1dc..838118c 100644 --- a/src/internal/mod.rs +++ b/src/internal/mod.rs @@ -1,5 +1,6 @@ pub(crate) mod batching; pub(crate) mod delivery; +pub(crate) mod execution; pub(crate) mod ingest; pub(crate) mod object_store; pub(crate) mod ports; diff --git a/src/internal/resources.rs b/src/internal/resources.rs index 557aec6..7591a65 100644 --- a/src/internal/resources.rs +++ b/src/internal/resources.rs @@ -1,4 +1,5 @@ //! Resource ports adapt SDK policy to the shared one-attempt API contracts. +use crate::internal::execution::ExecutionBinding; use crate::internal::session::CredentialScope; use async_trait::async_trait; @@ -16,14 +17,15 @@ use crate::error::{BraintrustError, Result}; use crate::experiments::BaseExperimentInfo; use crate::internal::ports::{DatasetApi, ExperimentApi}; -/// Shared HTTP pool without retaining client/session/delivery state. +/// Shared HTTP pool and execution lease without client/session/delivery state. #[derive(Clone)] pub(crate) struct ResourceTransport { http: reqwest::Client, + execution: ExecutionBinding, } impl ResourceTransport { - pub fn new(http: reqwest::Client) -> Self { - Self { http } + pub fn new(http: reqwest::Client, execution: ExecutionBinding) -> Self { + Self { http, execution } } fn resource_api(&self, scope: &CredentialScope) -> Result { scope.api(&self.http) @@ -37,11 +39,15 @@ impl ExperimentApi for ResourceTransport { scope: &CredentialScope, request: ExperimentRegisterRequest, ) -> Result { - self.resource_api(scope)? - .experiments() - .register(request) + let api = self.resource_api(scope)?; + self.execution + .response(async move { + api.experiments() + .register(request) + .await + .map_err(BraintrustError::from) + }) .await - .map_err(Into::into) } async fn fetch_base_experiment( @@ -49,22 +55,23 @@ impl ExperimentApi for ResourceTransport { scope: &CredentialScope, experiment_id: &str, ) -> Result> { - let response = match self - .resource_api(scope)? - .experiments() - .base(BaseExperimentRequest { - id: experiment_id.to_owned(), + let api = self.resource_api(scope)?; + let request = BaseExperimentRequest { + id: experiment_id.to_owned(), + }; + self.execution + .response(async move { + let response = match api.experiments().base(request).await { + // The SDK interprets 400 as no base; the API layer preserves the status. + Err(crate::api::Error::Http { status: 400, .. }) => return Ok(None), + result => result.map_err(BraintrustError::from)?, + }; + Ok(match (response.base_exp_id, response.base_exp_name) { + (Some(id), Some(name)) => Some(BaseExperimentInfo { id, name }), + _ => None, + }) }) .await - { - // The SDK interprets 400 as no base; the API layer preserves the status. - Err(crate::api::Error::Http { status: 400, .. }) => return Ok(None), - result => result.map_err(BraintrustError::from)?, - }; - Ok(match (response.base_exp_id, response.base_exp_name) { - (Some(id), Some(name)) => Some(BaseExperimentInfo { id, name }), - _ => None, - }) } async fn fetch_experiment_comparison( @@ -73,11 +80,17 @@ impl ExperimentApi for ResourceTransport { experiment_id: &str, base_experiment_id: Option<&str>, ) -> Result { - self.resource_api(scope)? - .experiments() - .compare(experiment_id, base_experiment_id) + let api = self.resource_api(scope)?; + let experiment_id = experiment_id.to_owned(); + let base_experiment_id = base_experiment_id.map(str::to_owned); + self.execution + .response(async move { + api.experiments() + .compare(&experiment_id, base_experiment_id.as_deref()) + .await + .map_err(BraintrustError::from) + }) .await - .map_err(Into::into) } } @@ -88,11 +101,15 @@ impl DatasetApi for ResourceTransport { scope: &CredentialScope, request: DatasetRegisterRequest, ) -> Result { - self.resource_api(scope)? - .datasets() - .register(request) + let api = self.resource_api(scope)?; + self.execution + .response(async move { + api.datasets() + .register(request) + .await + .map_err(BraintrustError::from) + }) .await - .map_err(Into::into) } async fn fetch_dataset_records( @@ -100,11 +117,15 @@ impl DatasetApi for ResourceTransport { scope: &CredentialScope, query: BTQLQuery, ) -> Result { - self.resource_api(scope)? - .btql() - .query_dataset_records(query) + let api = self.resource_api(scope)?; + self.execution + .response(async move { + api.btql() + .query_dataset_records(query) + .await + .map_err(BraintrustError::from) + }) .await - .map_err(Into::into) } async fn fetch_dataset_summary( @@ -112,18 +133,27 @@ impl DatasetApi for ResourceTransport { scope: &CredentialScope, dataset_id: &str, ) -> Result { - self.resource_api(scope)? - .datasets() - .summarize(dataset_id) + let api = self.resource_api(scope)?; + let dataset_id = dataset_id.to_owned(); + self.execution + .response(async move { + api.datasets() + .summarize(&dataset_id) + .await + .map_err(BraintrustError::from) + }) .await - .map_err(Into::into) } async fn write_dataset_record(&self, scope: &CredentialScope, row: DatasetWrite) -> Result<()> { - self.resource_api(scope)? - .datasets() - .write(&[row]) + let api = self.resource_api(scope)?; + self.execution + .response(async move { + api.datasets() + .write(&[row]) + .await + .map_err(BraintrustError::from) + }) .await - .map_err(Into::into) } } diff --git a/src/lib.rs b/src/lib.rs index 15c7583..504c5b0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -14,6 +14,7 @@ mod json_merge; mod log_queue; mod logger; mod logs3; +mod runtime; mod span; mod span_components; mod stream; @@ -44,6 +45,7 @@ pub use logger::{ DEFAULT_APP_URL, }; pub use logs3::{Logs3BatchUploader, Logs3UploadResult}; +pub use runtime::current_runtime; pub use span::{ SpanBuilder, SpanHandle, SpanLog, SpanLogBuilder, SpanLogBuilderError, SpanOrigin, SpanOriginEnvironment, diff --git a/src/log_queue/queue.rs b/src/log_queue/queue.rs index 24ee715..c84bff1 100644 --- a/src/log_queue/queue.rs +++ b/src/log_queue/queue.rs @@ -7,6 +7,7 @@ use super::merge::merge_row_into; use super::row_key::RowKey; use crate::error::{BraintrustError, Result}; 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; @@ -18,7 +19,7 @@ use crate::{ use futures::future::join_all; use indexmap::IndexMap; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; -use std::sync::Arc; +use std::sync::{Arc, Mutex, MutexGuard}; use std::time::{SystemTime, UNIX_EPOCH}; use tokio::sync::Mutex as TokioMutex; use tracing::warn; @@ -51,16 +52,22 @@ struct LogQueueCore { last_drop_log_time: AtomicU64, config: LogQueueConfig, client: reqwest::Client, + execution: ExecutionBinding, project_cache: TokioMutex>, version_info: TokioMutex>, } impl LogQueueCore { - fn new(config: LogQueueConfig, client: reqwest::Client) -> Arc { + fn new( + config: LogQueueConfig, + client: reqwest::Client, + execution: ExecutionBinding, + ) -> Arc { Arc::new(Self { dropped_count: AtomicUsize::new(0), last_drop_log_time: AtomicU64::new(0), config, client, + execution, project_cache: TokioMutex::new(IndexMap::new()), version_info: TokioMutex::new(IndexMap::new()), }) @@ -141,25 +148,27 @@ impl LogQueueCore { for chunk in chunks { let config = self.config.clone(); - let batches = tokio::task::spawn_blocking(move || { - let mut merged: IndexMap = IndexMap::new(); - for row in chunk { - let key = RowKey::from_row(&row); - if let Some(existing) = merged.get_mut(&key) { - if row.is_merge.unwrap_or(false) { - merge_row_into(existing, row); + let batches = self + .execution + .spawn_blocking(move || { + let mut merged: IndexMap = IndexMap::new(); + for row in chunk { + let key = RowKey::from_row(&row); + if let Some(existing) = merged.get_mut(&key) { + if row.is_merge.unwrap_or(false) { + merge_row_into(existing, row); + } else { + *existing = row; + } } else { - *existing = row; + merged.insert(key, row); } - } else { - merged.insert(key, row); } - } - let rows: Vec = merged.into_values().collect(); - batch_and_serialize_rows(rows, &config, effective_batch_bytes) - }) - .await; + let rows: Vec = merged.into_values().collect(); + batch_and_serialize_rows(rows, &config, effective_batch_bytes) + }) + .await; let batches = match batches { Ok(Ok(batches)) => batches, Ok(Err(error)) => { @@ -490,10 +499,7 @@ impl LogQueueCore { if dirs.is_empty() { return; } - let Ok(handle) = tokio::runtime::Handle::try_current() else { - return; - }; - handle.spawn(async move { + self.execution.spawn(async move { let SubmitCommand { scope: _, payload, @@ -620,22 +626,29 @@ impl LogQueueCore { #[derive(Clone)] pub struct LogQueue { core: Arc, - writes: WriteService, + writes: Arc>, +} +enum QueueState { + Dormant, + Running(WriteService), + Closed, } impl LogQueue { - pub fn new(config: LogQueueConfig, client: reqwest::Client, _worker_queue_size: usize) -> Self { - let capacity = config.queue_max_size(); - let core = LogQueueCore::new(config, client); - let processor = core.clone(); - let writes = WriteService::new( - capacity, - tokio::runtime::Handle::current(), - move |commands| { - let processor = processor.clone(); - Box::pin(async move { processor.process(commands).await }) - }, - ); - Self { core, writes } + pub fn new( + config: LogQueueConfig, + client: reqwest::Client, + _worker_queue_size: usize, + execution: ExecutionBinding, + ) -> Self { + Self { + core: LogQueueCore::new(config, client, execution), + writes: Arc::new(Mutex::new(QueueState::Dormant)), + } + } + fn state(&self) -> MutexGuard<'_, QueueState> { + self.writes + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) } pub fn submit( @@ -650,30 +663,82 @@ impl LogQueue { payload, parent_info, }; - self.writes - .enqueue_with_rejection(command, outcome, |command, error| { - self.core.dropped_count.fetch_add(1, Ordering::Relaxed); - if matches!(error, crate::EnqueueError::Full) { - self.core.log_drop_warning(); - } else { - warn!(%error, "span submission rejected"); - } - self.core.dump_dropped_row_if_configured(command); - }) + let mut state = self.state(); + if matches!(*state, QueueState::Dormant) { + let processor = self.core.clone(); + *state = QueueState::Running(WriteService::new( + self.core.config.queue_max_size(), + self.core.execution.clone(), + move |commands| { + let processor = processor.clone(); + Box::pin(async move { processor.process(commands).await }) + }, + )); + } + match &*state { + QueueState::Running(writes) => { + writes.enqueue_with_rejection(command, outcome, |command, error| { + self.reject(command, error) + }) + } + 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) + .send(Err(crate::EnqueueError::Closed.into())); + self.reject(command, crate::EnqueueError::Closed); + completion + } + QueueState::Dormant => unreachable!("worker installed before admission"), + } + } + fn reject(&self, command: SubmitCommand, error: crate::EnqueueError) { + self.core.dropped_count.fetch_add(1, Ordering::Relaxed); + if matches!(error, crate::EnqueueError::Full) { + self.core.log_drop_warning(); + } else { + warn!(%error, "span submission rejected"); + } + self.core.dump_dropped_row_if_configured(command); } pub fn flush(&self) -> Completion { - self.writes.flush() + match &*self.state() { + QueueState::Running(writes) => writes.flush(), + state => { + let (sender, mut completion) = Completion::pair(); + if matches!(state, QueueState::Closed) { + completion.record_issuance(Err(crate::EnqueueError::Closed.into())); + } else { + let _ = sender.send(Ok(FlushReport::default())); + } + completion + } + } } pub fn shutdown( &self, timeout: std::time::Duration, ) -> Completion { - self.writes.shutdown(timeout) + let mut state = self.state(); + if let QueueState::Running(writes) = &*state { + return writes.shutdown(timeout); + } + *state = QueueState::Closed; + let (sender, completion) = Completion::pair(); + let _ = sender.send(Ok(ShutdownReport::default())); + completion } #[cfg(test)] pub(crate) fn is_empty(&self) -> bool { - self.writes.is_empty() + match &*self.state() { + QueueState::Running(writes) => writes.is_empty(), + _ => true, + } + } + #[cfg(test)] + pub(crate) fn is_started(&self) -> bool { + matches!(*self.state(), QueueState::Running(_)) } pub async fn trigger_flush_command(&self) -> Result<()> { let token = self.flush(); @@ -743,6 +808,7 @@ mod tests { LogQueueConfig::builder().num_retries(0).build(), reqwest::Client::new(), 32, + ExecutionBinding::external(tokio::runtime::Handle::current()), ) } fn scope(server: &MockServer, token: &str) -> CredentialScope { @@ -760,6 +826,56 @@ mod tests { }) } + #[tokio::test] + async fn concurrent_first_submissions_share_one_bounded_lazy_worker() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&server) + .await; + let queue = LogQueue::new( + LogQueueConfig::builder() + .queue_max_size(8) + .num_retries(0) + .build(), + reqwest::Client::new(), + 8, + ExecutionBinding::external(tokio::runtime::Handle::current()), + ); + assert!(!queue.is_started()); + let tokens = std::thread::scope(|threads| { + (0..32) + .map(|number| { + let queue = &queue; + let server = &server; + threads.spawn(move || { + queue.submit( + scope(server, "token"), + payload(&format!("row-{number}"), false), + parent(), + None, + ) + }) + }) + .collect::>() + .into_iter() + .map(|thread| thread.join().unwrap()) + .collect::>() + }); + assert_eq!(tokens.iter().filter(|token| token.is_ok()).count(), 8); + assert!(queue.is_started()); + assert_eq!(queue.flush().await.unwrap().delivered, 8); + for token in tokens { + let _ = token.await; + } + queue + .shutdown(std::time::Duration::from_secs(1)) + .await + .unwrap(); + } + #[tokio::test] async fn late_route_failure_fails_merged_members_and_successful_siblings_then_retires() { let server = MockServer::start().await; diff --git a/src/logger.rs b/src/logger.rs index 3699e2f..0a13895 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -11,6 +11,7 @@ use crate::dataset::DatasetBuilder; use crate::error::{BraintrustError, Result}; use crate::experiments::ExperimentBuilder; use crate::http::{build_http_client, CaBundleConfig}; +use crate::internal::execution::ExecutionBinding; use crate::internal::ports::SpanSubmitter; use crate::internal::resources::ResourceTransport; use crate::internal::retry::{exponential, RetryExecutor}; @@ -188,6 +189,7 @@ pub struct BraintrustClientBuilder { environment: Option, span_origin: Option, user_agent: Option, + runtime: Option, } impl BraintrustClientBuilder { @@ -219,6 +221,7 @@ impl BraintrustClientBuilder { environment: None, span_origin: None, user_agent: None, + runtime: None, } } @@ -318,10 +321,10 @@ impl BraintrustClientBuilder { self } - /// Block until login completes (default: false, login happens in background). + /// Wait for login before initialization completes (default: false). /// - /// When `false` (default), login happens asynchronously in a background task. - /// When `true`, the `build()` method waits for login to complete before returning. + /// When `false`, the build completion returns a client while login continues + /// in the background. When `true`, that completion also observes login. pub fn blocking_login(mut self, blocking: bool) -> Self { self.blocking_login = blocking; self @@ -356,11 +359,44 @@ impl BraintrustClientBuilder { self } - /// Build the client, performing login. - /// - /// If `blocking_login` is true, waits for login to complete. - /// Otherwise, login happens in the background with retry logic. - pub async fn build(self) -> Result { + /// Select one runtime for initialization, resource I/O and delivery. + /// `None` (the default) owns and drives an internal runtime. An explicit + /// handle requires its caller-owned runtime to stay alive and progressing. + pub fn runtime(mut self, runtime: Option) -> Self { + self.runtime = runtime; + self + } + + /// Schedule client initialization immediately and observe its result. + /// Dropping the completion does not cancel initialization or login. + pub fn build(self) -> crate::Completion { + let validate = || -> Result<()> { + Url::parse(self.app_url.as_deref().unwrap_or(DEFAULT_APP_URL)) + .map_err(|_| BraintrustError::InvalidConfig("invalid app_url".into()))?; + Url::parse(self.api_url.as_deref().unwrap_or(DEFAULT_API_URL)) + .map_err(|_| BraintrustError::InvalidConfig("invalid api_url".into()))?; + if !self.skip_login && self.api_key.is_none() { + return Err(BraintrustError::InvalidConfig( + "API key required: set BRAINTRUST_API_KEY or call .api_key()".into(), + )); + } + Ok(()) + }; + let binding = validate().and_then(|()| ExecutionBinding::new(self.runtime.clone())); + match binding { + Ok(binding) => { + let context = binding.clone(); + binding.response(async move { self.initialize(context).await }) + } + Err(error) => { + let (_, mut completion) = crate::Completion::pair(); + completion.record_issuance(Err(error)); + completion + } + } + } + + async fn initialize(self, execution: ExecutionBinding) -> Result { let app_url_str = self.app_url.unwrap_or_else(|| DEFAULT_APP_URL.into()); let api_url_str = self.api_url.unwrap_or_else(|| DEFAULT_API_URL.into()); @@ -388,12 +424,18 @@ impl BraintrustClientBuilder { .build(); // Commands carry their own immutable scope; the worker has no session fallback. - let queue = LogQueue::new(log_config, http_client.clone(), self.queue_size); + let queue = LogQueue::new( + log_config, + http_client.clone(), + self.queue_size, + execution.clone(), + ); let client = BraintrustClient { inner: Arc::new(ClientInner { api_url, app_url, + execution, queue, login_state, login_notify: Notify::new(), @@ -438,6 +480,7 @@ pub struct BraintrustClient { } struct ClientInner { + execution: ExecutionBinding, #[allow(dead_code)] api_url: Url, app_url: Url, @@ -463,7 +506,10 @@ impl std::fmt::Debug for ClientInner { impl BraintrustClient { fn resource_transport(&self) -> Arc { - Arc::new(ResourceTransport::new(self.inner.http_client.clone())) + Arc::new(ResourceTransport::new( + self.inner.http_client.clone(), + self.inner.execution.clone(), + )) } fn span_submitter(&self) -> Arc { @@ -832,7 +878,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(); - tokio::spawn(async move { + self.inner.execution.spawn(async move { const MAX_ATTEMPTS: usize = 10; let outcome = RetryExecutor::new( MAX_ATTEMPTS - 1, @@ -1146,6 +1192,41 @@ mod tests { ResponseTemplate::new(200).set_body_json(serde_json::json!({ "org_info": org_info })) } + #[tokio::test] + async fn reads_and_empty_lifecycle_calls_do_not_start_a_log_worker() { + let server = MockServer::start().await; + Mock::given(method("POST")).and(path("/api/dataset/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({"project":{"id":"project","name":"p"},"dataset":{"id":"dataset","name":"d"}}))) + .expect(1).mount(&server).await; + let client = BraintrustClient::builder() + .skip_login(true) + .app_url(server.uri()) + .api_url(server.uri()) + .build() + .await + .unwrap(); + assert!(!client.inner.queue.is_started()); + client.flush().await.unwrap(); + let dataset = client + .dataset_builder_with_credentials("token", "tenant") + .project_name("p") + .dataset_name("d") + .build() + .unwrap(); + assert_eq!(dataset.dataset_id().await.unwrap(), "dataset"); + assert!(!client.inner.queue.is_started()); + client.shutdown(Duration::from_secs(1)).await.unwrap(); + assert!(!client.inner.queue.is_started()); + let span = client + .span_builder_with_credentials("token", "org") + .parent_info(ParentSpanInfo::ProjectLogs { + object_id: "project".into(), + }) + .build(); + assert!(span.log(SpanLog::default()).is_err()); + assert!(!client.inner.queue.is_started()); + } + #[tokio::test] async fn resource_ports_do_not_retain_the_whole_client_or_login_state() { let client = BraintrustClient::builder() diff --git a/src/runtime.rs b/src/runtime.rs new file mode 100644 index 0000000..90adc9e --- /dev/null +++ b/src/runtime.rs @@ -0,0 +1,10 @@ +//! Explicit selection of a caller-owned Tokio runtime. +/// Capture the current Tokio runtime, or `None` outside Tokio. +/// +/// Pass this to [`crate::BraintrustClientBuilder::runtime`] to explicitly opt +/// into the caller's runtime. The default `None` always uses an SDK-owned runtime, +/// including when construction happens inside Tokio. With an explicit handle, +/// callers must keep that runtime alive and driving accepted work. +pub fn current_runtime() -> Option { + tokio::runtime::Handle::try_current().ok() +} diff --git a/tests/dataset_test.rs b/tests/dataset_test.rs index 6de0cd4..330f57f 100644 --- a/tests/dataset_test.rs +++ b/tests/dataset_test.rs @@ -288,6 +288,7 @@ async fn dataset_summarize_returns_summary() { async fn isolated_client(server: &MockServer) -> BraintrustClient { BraintrustClient::builder() + .runtime(braintrust_sdk_rust::current_runtime()) .skip_login(true) .app_url(server.uri()) .api_url(server.uri()) diff --git a/tests/finished_spans.rs b/tests/finished_spans.rs index a2fa11f..90221ec 100644 --- a/tests/finished_spans.rs +++ b/tests/finished_spans.rs @@ -12,6 +12,7 @@ use wiremock::{ async fn client(server: &MockServer) -> BraintrustClient { BraintrustClient::builder() + .runtime(braintrust_sdk_rust::current_runtime()) .skip_login(true) .api_url(server.uri()) .app_url(server.uri()) diff --git a/tests/http_flow.rs b/tests/http_flow.rs index 1721184..115d081 100644 --- a/tests/http_flow.rs +++ b/tests/http_flow.rs @@ -23,6 +23,7 @@ async fn submits_with_bearer_token() { .await; let client = BraintrustClient::builder() + .runtime(braintrust_sdk_rust::current_runtime()) .api_key("test-key") .user_agent("bt/0.20.0") .api_url(server.uri()) @@ -75,6 +76,7 @@ async fn preparation_failure_reaches_write_and_covering_flush() { .await; let client = BraintrustClient::builder() + .runtime(braintrust_sdk_rust::current_runtime()) .api_key("test-key") .api_url(server.uri()) .app_url(server.uri()) @@ -122,6 +124,7 @@ async fn automatic_delivery_reports_backend_rejection_through_completion() { .await; let client = BraintrustClient::builder() + .runtime(braintrust_sdk_rust::current_runtime()) .skip_login(true) .api_url(server.uri()) .app_url(server.uri()) @@ -159,6 +162,7 @@ async fn submission_and_flush_can_be_issued_without_a_call_site_runtime() { .mount(&server) .await; let client = BraintrustClient::builder() + .runtime(braintrust_sdk_rust::current_runtime()) .skip_login(true) .api_url(server.uri()) .app_url(server.uri()) diff --git a/tests/runtime_binding.rs b/tests/runtime_binding.rs new file mode 100644 index 0000000..3855762 --- /dev/null +++ b/tests/runtime_binding.rs @@ -0,0 +1,173 @@ +use braintrust_sdk_rust::{ + current_runtime, BraintrustClient, BraintrustError, EnqueueError, ExecutionError, + ParentSpanInfo, WriteError, +}; +use serde_json::json; +use std::future::IntoFuture; +use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, +}; + +#[tokio::test] +async fn synchronous_callers_deliver_without_a_tokio_runtime_and_observe_on_another_executor() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&server) + .await; + let uri = server.uri(); + let (sender, result) = tokio::sync::oneshot::channel(); + let thread = std::thread::spawn(move || { + assert!(current_runtime().is_none()); + let initialization = BraintrustClient::builder() + .skip_login(true) + .api_url(&uri) + .app_url(&uri) + .build(); + assert!(initialization.is_ok()); + let client = futures::executor::block_on(initialization.into_future()).unwrap(); + let span = client + .span_builder_with_credentials("token", "org") + .parent_info(ParentSpanInfo::ProjectLogs { + object_id: "project".into(), + }) + .build(); + let end = span.end(); + let shutdown = client.shutdown(std::time::Duration::from_secs(5)); + drop(span); + drop(client); + sender.send((end, shutdown)).unwrap(); + }); + let (end, shutdown) = result.await.unwrap(); + end.await.unwrap(); + shutdown.await.unwrap(); + thread.join().unwrap(); +} + +#[tokio::test] +async fn unobserved_construction_starts_login_without_polling_its_token() { + let server = MockServer::start().await; + let arrived = std::sync::Arc::new(tokio::sync::Notify::new()); + let observed = arrived.clone(); + Mock::given(method("POST")) + .and(path("/api/apikey/login")) + .respond_with(move |_: &wiremock::Request| { + observed.notify_one(); + ResponseTemplate::new(200) + .set_body_json(json!({"org_info":[{"id":"org","name":"tenant"}]})) + }) + .expect(1) + .mount(&server) + .await; + drop( + BraintrustClient::builder() + .api_key("token") + .blocking_login(true) + .api_url(server.uri()) + .app_url(server.uri()) + .build(), + ); + tokio::time::timeout(std::time::Duration::from_secs(5), arrived.notified()) + .await + .unwrap(); +} + +#[test] +fn explicit_current_thread_runtime_must_be_driven_and_observation_does_not_move_work() { + let selected = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let observer = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let initialization = BraintrustClient::builder() + .skip_login(true) + .runtime(Some(selected.handle().clone())) + .build(); + assert!(initialization.is_ok()); + let mut initialization = Box::pin(initialization.into_future()); + observer.block_on(async { + assert!(futures::poll!(initialization.as_mut()).is_pending()); + }); + selected.block_on(async { + tokio::task::yield_now().await; + }); + let client = observer.block_on(initialization).unwrap(); + let shutdown = client.shutdown(std::time::Duration::from_secs(1)); + observer.block_on(async { + shutdown.await.unwrap(); + }); + drop(client); + // SDK shutdown did not stop the caller-owned runtime. + assert_eq!(selected.block_on(async { 42 }), 42); +} + +#[tokio::test] +async fn stopped_explicit_runtime_fails_delivery_and_resource_requests_without_fallback() { + let server = MockServer::start().await; + let uri = server.uri(); + let client = std::thread::spawn(move || { + let selected = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + selected + .block_on( + BraintrustClient::builder() + .skip_login(true) + .runtime(Some(selected.handle().clone())) + .api_url(&uri) + .app_url(&uri) + .build() + .into_future(), + ) + .unwrap() + }) + .join() + .unwrap(); + let span = client + .span_builder_with_credentials("token", "org") + .parent_info(ParentSpanInfo::ProjectLogs { + object_id: "project".into(), + }) + .build(); + // A stopped runtime may cancel the worker before or after queue admission. + let end = span.end().await; + assert!( + matches!( + end, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + | Err(WriteError::Admission(EnqueueError::Closed)) + ), + "unexpected end result: {end:?}" + ); + let dataset = client + .dataset_builder_with_credentials("token", "tenant") + .project_name("p") + .dataset_name("d") + .build() + .unwrap(); + assert!(matches!( + dataset.dataset_id().await, + Err(BraintrustError::Execution(ExecutionError::WorkerLost)) + )); + assert!(server.received_requests().await.unwrap().is_empty()); +} + +#[test] +fn invalid_construction_reports_issuance_failure_synchronously() { + let build = BraintrustClient::builder() + .skip_login(true) + .api_url("invalid URL") + .build(); + assert!(matches!(&*build, Err(BraintrustError::InvalidConfig(_)))); + assert!(matches!( + futures::executor::block_on(build.into_future()), + Err(BraintrustError::InvalidConfig(_)) + )); +} From 3b254dbc457356d7bd1875b26495fee4767149f7 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Tue, 20 Oct 2026 22:44:22 +0800 Subject: [PATCH 2/2] Coordinate response tasks and delivery during shutdown --- src/internal/delivery.rs | 9 +- src/internal/execution.rs | 193 +++++++++++++++++++++++++++++++-- src/internal/lifecycle.rs | 65 +++++++++++ src/internal/mod.rs | 4 +- src/internal/ports.rs | 2 +- src/internal/response_tasks.rs | 190 ++++++++++++++++++++++++++++++++ src/internal/shared_outcome.rs | 183 +++++++++++++++++++++++++++++++ src/internal/write_outcome.rs | 133 ----------------------- src/log_queue/queue.rs | 13 ++- src/logger.rs | 193 ++++++++++++++++++++++++++++++++- src/span.rs | 6 +- src/test_utils.rs | 2 +- tests/response_lifecycle.rs | 131 ++++++++++++++++++++++ tests/runtime_binding.rs | 12 +- 14 files changed, 976 insertions(+), 160 deletions(-) create mode 100644 src/internal/lifecycle.rs create mode 100644 src/internal/response_tasks.rs create mode 100644 src/internal/shared_outcome.rs delete mode 100644 src/internal/write_outcome.rs create mode 100644 tests/response_lifecycle.rs 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);