From 09a79df78d46a12db1af2102ef7b665abdaf59df Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Tue, 20 Oct 2026 23:16:58 +0800 Subject: [PATCH] Admit experiment events before resolving registration --- src/error.rs | 5 + src/experiments/builder.rs | 6 + src/experiments/experiment.rs | 130 ++++++------ src/internal/delivery.rs | 79 ++++++- src/internal/execution.rs | 20 ++ src/internal/ports.rs | 9 + src/internal/resources.rs | 14 +- src/log_queue/queue.rs | 73 +++++-- src/logger.rs | 1 + src/span.rs | 27 +++ src/test_utils.rs | 23 +++ tests/eager_experiment_events.rs | 340 +++++++++++++++++++++++++++++++ 12 files changed, 632 insertions(+), 95 deletions(-) create mode 100644 tests/eager_experiment_events.rs diff --git a/src/error.rs b/src/error.rs index 205da7a..fa2bab7 100644 --- a/src/error.rs +++ b/src/error.rs @@ -44,6 +44,11 @@ impl From for BraintrustError { impl From for BraintrustError { fn from(error: crate::WriteError) -> Self { + match &error { + crate::WriteError::Admission(error) => return Self::Admission(*error), + crate::WriteError::Execution(error) => return Self::Execution(*error), + crate::WriteError::Delivery(_) => {} + } if let crate::WriteError::Delivery(cause) = &error { if let Some(Self::Api { status, message }) = cause.cause().downcast_ref::() { return Self::Api { diff --git a/src/experiments/builder.rs b/src/experiments/builder.rs index fe97b25..c4d0d49 100644 --- a/src/experiments/builder.rs +++ b/src/experiments/builder.rs @@ -7,6 +7,7 @@ use serde_json::{Map, Value}; use crate::experiments::experiment::{Experiment, ExperimentState}; use crate::experiments::metadata::RepoInfo; +use crate::internal::execution::ExecutionBinding; use crate::internal::ports::{ExperimentApi, SpanSubmitter}; use crate::internal::session::CredentialScope; @@ -36,6 +37,7 @@ impl std::error::Error for ExperimentBuilderError {} /// credentials belong to the resource and do not mutate the client's session. pub struct ExperimentBuilder { api: Arc, + execution: ExecutionBinding, submitter: Arc, scope: CredentialScope, project_name: Option, @@ -52,6 +54,7 @@ impl Clone for ExperimentBuilder { fn clone(&self) -> Self { Self { api: self.api.clone(), + execution: self.execution.clone(), submitter: self.submitter.clone(), scope: self.scope.clone(), project_name: self.project_name.clone(), @@ -70,11 +73,13 @@ impl ExperimentBuilder { /// Create a new ExperimentBuilder (internal use). pub(crate) fn new( api: Arc, + execution: ExecutionBinding, submitter: Arc, scope: CredentialScope, ) -> Self { Self { api, + execution, submitter, scope, project_name: None, @@ -164,6 +169,7 @@ impl ExperimentBuilder { Ok(Experiment { api: self.api, + execution: self.execution, submitter: self.submitter, scope: self.scope, project_name, diff --git a/src/experiments/experiment.rs b/src/experiments/experiment.rs index 7bdcf2d..71e52f0 100644 --- a/src/experiments/experiment.rs +++ b/src/experiments/experiment.rs @@ -17,6 +17,7 @@ use crate::experiments::metadata::{ }; use crate::experiments::span_builder::ExperimentSpanBuilder; use crate::experiments::summary::{ExperimentSummary, MetricSummary, ScoreSummary}; +use crate::internal::execution::ExecutionBinding; use crate::internal::ports::{ExperimentApi, SpanSubmitter}; use crate::internal::session::CredentialScope; use crate::span::{SpanBuilder, SpanHandle, SpanLog}; @@ -35,6 +36,7 @@ pub(crate) struct ExperimentMetadata { /// Use this to log evaluation events, create spans, log feedback, and get summaries. pub struct Experiment { pub(crate) api: Arc, + pub(crate) execution: ExecutionBinding, pub(crate) submitter: Arc, pub(crate) scope: CredentialScope, pub(crate) project_name: String, @@ -61,6 +63,7 @@ impl Clone for Experiment { fn clone(&self) -> Self { Self { api: self.api.clone(), + execution: self.execution.clone(), submitter: self.submitter.clone(), scope: self.scope.clone(), project_name: self.project_name.clone(), @@ -83,50 +86,55 @@ impl Experiment { /// /// Returns an error if registration/delivery fails, or if start_span() has /// already selected the span-recording interface for this experiment. - pub async fn log(&self, event: ExperimentLog) -> Result { - if self.state.called_start_span.load(Ordering::Acquire) { - return Err(crate::BraintrustError::InvalidConfig( - "cannot use log() after start_span(); use the span to log events".into(), - )); - } - - // Get or create experiment registration - let metadata = self.ensure_registered().await?; - - // Create a span for this log event - let start_time = f64::from_bits(self.state.last_start_time.load(Ordering::Acquire)); - let start_time = if start_time == 0.0 { - SystemTime::now() + /// Admission happens immediately, before registration. Dropping the token + /// does not cancel the write; awaiting includes full batch delivery. + pub fn log(&self, event: ExperimentLog) -> crate::Completion { + self.execution.response_try(|| { + if self.state.called_start_span.load(Ordering::Acquire) { + return Err(crate::BraintrustError::InvalidConfig( + "cannot use log() after start_span(); use the span to log events".into(), + )); + } + let start_time = f64::from_bits(self.state.last_start_time.load(Ordering::Acquire)); + let now = SystemTime::now() .duration_since(UNIX_EPOCH) - .map(|d| d.as_secs_f64()) - .unwrap_or(0.0) - } else { - start_time - }; - - let span = self.create_internal_span(metadata, start_time); - - // Get the row_id before logging - it's stable from span creation - let row_id = span.row_id().to_string(); - - let end_time = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|duration| duration.as_secs_f64()) - .unwrap_or(0.0); - let mut span_log = event.into_span_log(); - span_log - .metrics - .get_or_insert_with(HashMap::new) - .entry("end".into()) - .or_insert(end_time); - // A single complete event snapshot owns the delivery obligation. - span.log(span_log).await?; - - self.state - .last_start_time - .store(end_time.to_bits(), Ordering::Release); - - Ok(row_id) + .map(|duration| duration.as_secs_f64()) + .unwrap_or(0.0); + let start_time = if start_time == 0.0 { now } else { start_time }; + let mut event = event.into_span_log(); + event + .metrics + .get_or_insert_with(HashMap::new) + .entry("end".into()) + .or_insert(now); + let payload = SpanBuilder::new(self.submitter.clone(), self.scope.clone()) + .span_type(SpanType::Eval) + .start_time(start_time) + .into_payload(event); + let row_id = payload.row_id.clone(); + let experiment = self.clone(); + let delivery = self.submitter.submit_resolving_parent( + self.scope.clone(), + payload, + Box::pin(async move { + let metadata = experiment.register_metadata().await?; + Ok(ParentSpanInfo::Experiment { + object_id: metadata.experiment_id.clone(), + }) + }), + ); + if let Err(error) = &*delivery { + return Err(error.clone().into()); + } + let state = self.state.clone(); + Ok(async move { + delivery.await?; + state + .last_start_time + .store(now.to_bits(), Ordering::Release); + Ok(row_id) + }) + }) } /// Start a new root span for this experiment. @@ -278,10 +286,9 @@ impl Experiment { Ok(summary) } - /// Flush all pending experiment data to Braintrust. - pub async fn flush(&self) -> Result<()> { - self.submitter.flush().await?; - Ok(()) + /// Capture pending writes through this call, including unresolved registration. + pub fn flush(&self) -> crate::Completion { + self.submitter.flush() } /// Get the experiment ID (after registration). @@ -407,7 +414,18 @@ impl Experiment { } /// Ensure the experiment is registered, returning cached metadata. - pub(crate) async fn ensure_registered(&self) -> Result<&ExperimentMetadata> { + pub(crate) async fn ensure_registered(&self) -> Result { + if let Some(metadata) = self.state.metadata.get() { + return Ok(metadata.clone()); + } + let experiment = self.clone(); + self.execution + .response(async move { experiment.register_metadata().await.cloned() }) + .await + } + + /// Polled only on the selected context, by preparation or a bound operation. + async fn register_metadata(&self) -> Result<&ExperimentMetadata> { self.state .metadata .get_or_try_init(|| async { @@ -431,20 +449,4 @@ impl Experiment { }) .await } - - /// Create an internal span for log() calls. - fn create_internal_span(&self, metadata: &ExperimentMetadata, start_time: f64) -> SpanHandle { - let parent_info = ParentSpanInfo::Experiment { - object_id: metadata.experiment_id.clone(), - }; - - // Note: When using ParentSpanInfo::Experiment, we don't need project_name - // because the destination is determined by the experiment_id. - let builder = SpanBuilder::new(Arc::clone(&self.submitter), self.scope.clone()) - .parent_info(parent_info) - .span_type(SpanType::Eval) - .start_time(start_time); - - builder.build() - } } diff --git a/src/internal/delivery.rs b/src/internal/delivery.rs index 08351f8..9ef7241 100644 --- a/src/internal/delivery.rs +++ b/src/internal/delivery.rs @@ -182,26 +182,36 @@ struct Inner { } impl Inner { fn stop(&self, error: ExecutionError) { - let mut state = lock(&self.state); - state.accepting = false; - if state.terminal.is_some() { - return; - } - state.terminal = Some(error); - for accepted in state.pending.drain(..) { + let (pending, active, ready, barriers) = { + let mut state = lock(&self.state); + state.accepting = false; + if state.terminal.is_some() { + return; + } + state.terminal = Some(error); + state.outstanding = 0; + ( + std::mem::take(&mut state.pending), + state.active.take(), + std::mem::take(&mut state.ready), + std::mem::take(&mut state.barriers), + ) + }; + // Deferred commands can own producers of this same service. Destroying + // them, or waking completion observers, must happen outside its lock. + for accepted in pending { let _ = accepted.sender.send(Err(error.into())); } let result = Err(BatchError::Execution(error)); - if let Some(active) = state.active.take() { + if let Some(active) = active { active.finish(&result); } - for batch in state.ready.drain(..) { + for batch in ready { batch.completion.finish(&result); } - for mut barrier in state.barriers.drain(..) { + for mut barrier in barriers { barrier.record(0, &result); } - state.outstanding = 0; self.wake.notify_one(); } fn close(&self) { @@ -495,6 +505,53 @@ mod tests { (service, receiver, gate) } + #[test] + fn cancel_drops_pending_and_sealed_commands_outside_the_state_lock() { + struct Reenter(Box); + impl Drop for Reenter { + fn drop(&mut self) { + (self.0)(); + } + } + let (done, completed) = std::sync::mpsc::channel(); + let thread = std::thread::spawn(move || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let service = WriteService::new( + 4, + ExecutionBinding::external(runtime.handle().clone()), + |_| Box::pin(async { Ok(()) }), + ); + let mut receipts = Vec::new(); + for _ in 0..2 { + let inner = Arc::downgrade(&service.inner); + let done = done.clone(); + receipts.push(service.enqueue(Reenter(Box::new(move || { + inner.upgrade().unwrap().close(); + done.send(()).unwrap(); + })))); + if receipts.len() == 1 { + drop(service.flush()); + } + } + service.cancel(ExecutionError::TimedOut); + for receipt in receipts { + assert!(matches!( + futures::executor::block_on(receipt.into_future()), + Err(WriteError::Execution(ExecutionError::TimedOut)) + )); + } + }); + for _ in 0..2 { + completed + .recv_timeout(Duration::from_secs(5)) + .expect("command destruction must not deadlock while reentering the service"); + } + thread.join().unwrap(); + } + #[tokio::test] async fn concurrent_flushes_observe_fixed_cutoffs_including_active_work() { let (service, mut entered, gate) = gated(8); diff --git a/src/internal/execution.rs b/src/internal/execution.rs index 7e6f41e..eebfde4 100644 --- a/src/internal/execution.rs +++ b/src/internal/execution.rs @@ -146,6 +146,19 @@ impl ExecutionBinding { T: Send + 'static, E: From + From + Send + 'static, F: Future> + Send + 'static, + { + self.response_try(|| Ok(work)) + } + + /// Reserve response capacity before synchronously admitting dependent work. + /// A rejected reservation never invokes the factory; a rejected dependency + /// retires the reservation and reports that same issuance failure. + pub fn response_try(&self, make: M) -> Completion + where + T: Send + 'static, + E: From + From + Send + 'static, + F: Future> + Send + 'static, + M: FnOnce() -> Result, { let (sender, mut completion) = Completion::pair(); let (guard, response) = match self.responses.admit(sender) { @@ -155,6 +168,13 @@ impl ExecutionBinding { return completion; } }; + let work = match make() { + Ok(work) => work, + Err(error) => { + completion.record_issuance(Err(error)); + return completion; // The guard retires this reserved slot. + } + }; let id = guard.id; let task = self.spawn(async move { guard.complete(&response, work.await); diff --git a/src/internal/ports.rs b/src/internal/ports.rs index 72a42bf..2d4702a 100644 --- a/src/internal/ports.rs +++ b/src/internal/ports.rs @@ -25,6 +25,15 @@ pub(crate) trait DeliveryControl: Send + Sync { #[async_trait] pub(crate) trait SpanSubmitter: DeliveryControl { + /// Queue the write now, resolving its destination during preparation. + /// Rejected work must drop the resolver without polling it. + fn submit_resolving_parent( + &self, + scope: CredentialScope, + payload: SpanPayload, + resolve: futures::future::BoxFuture<'static, Result>, + ) -> crate::Completion; + fn submit( &self, scope: CredentialScope, diff --git a/src/internal/resources.rs b/src/internal/resources.rs index 7efca52..9b2fa89 100644 --- a/src/internal/resources.rs +++ b/src/internal/resources.rs @@ -39,15 +39,13 @@ impl ExperimentApi for ResourceTransport { scope: &CredentialScope, request: ExperimentRegisterRequest, ) -> Result { - let api = self.resource_api(scope)?; - self.execution - .response(async move { - api.experiments() - .register(request) - .await - .map_err(BraintrustError::from) - }) + // Registration is part of admitted event preparation or an explicitly + // bound metadata operation; it must remain valid during queue drain. + self.resource_api(scope)? + .experiments() + .register(request) .await + .map_err(BraintrustError::from) } async fn fetch_base_experiment( diff --git a/src/log_queue/queue.rs b/src/log_queue/queue.rs index 3938100..3c1d566 100644 --- a/src/log_queue/queue.rs +++ b/src/log_queue/queue.rs @@ -16,7 +16,7 @@ use crate::types::{LogDestination, Logs3Row, ParentSpanInfo, SpanObjectType, Spa use crate::{ Completion, DeliveryReceipt, FlushError, FlushReport, ShutdownError, ShutdownReport, WriteError, }; -use futures::future::join_all; +use futures::future::{join_all, BoxFuture}; use indexmap::IndexMap; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, MutexGuard}; @@ -43,7 +43,11 @@ fn cache_insert(cache: &mut IndexMap, key: K, struct SubmitCommand { scope: CredentialScope, payload: SpanPayload, - parent_info: Option, + parent_info: ParentResolution, +} +enum ParentResolution { + Ready(Option), + Deferred(BoxFuture<'static, Result>), } /// Preparation and protocol orchestration, independent of admission/lifecycle. @@ -113,7 +117,15 @@ impl LogQueueCore { payload, parent_info, } = cmd; - match self.prepare_row(&scope, payload, parent_info).await { + let prepared = async { + let parent = match parent_info { + ParentResolution::Ready(parent) => parent, + ParentResolution::Deferred(resolve) => Some(resolve.await?), + }; + self.prepare_row(&scope, payload, parent).await + } + .await; + match prepared { Ok(row) => { groups.entry(scope).or_default().push(row); } @@ -506,6 +518,11 @@ impl LogQueueCore { parent_info, } = cmd; + let ParentResolution::Ready(parent_info) = parent_info else { + // Rejection must not trigger registration just for a debug dump. + warn!(row_id = %payload.row_id, "cannot dump rejected row with unresolved destination"); + return; + }; let span_components = payload.span_components.clone(); let span_id = payload.span_id.clone(); @@ -658,11 +675,21 @@ impl LogQueue { parent_info: Option, outcome: Option, ) -> Completion { - let command = SubmitCommand { - scope, - payload, - parent_info, - }; + self.submit_command( + SubmitCommand { + scope, + payload, + parent_info: ParentResolution::Ready(parent_info), + }, + outcome, + ) + } + + fn submit_command( + &self, + command: SubmitCommand, + outcome: Option, + ) -> Completion { let mut state = self.state(); if matches!(*state, QueueState::Dormant) { let processor = self.core.clone(); @@ -730,11 +757,17 @@ impl LogQueue { completion } pub fn cancel(&self, error: crate::ExecutionError) { - let mut state = self.state(); - if let QueueState::Running(writes) = &*state { + let writes = { + let mut state = self.state(); + if let QueueState::Running(writes) = &*state { + Some(writes.clone()) + } else { + *state = QueueState::Closed; + None + } + }; + if let Some(writes) = writes { writes.cancel(error); - } else { - *state = QueueState::Closed; } } @@ -769,6 +802,22 @@ impl DeliveryControl for LogQueue { } #[async_trait::async_trait] impl SpanSubmitter for LogQueue { + fn submit_resolving_parent( + &self, + scope: CredentialScope, + payload: SpanPayload, + resolve: BoxFuture<'static, Result>, + ) -> Completion { + self.submit_command( + SubmitCommand { + scope, + payload, + parent_info: ParentResolution::Deferred(resolve), + }, + None, + ) + } + fn submit( &self, scope: CredentialScope, diff --git a/src/logger.rs b/src/logger.rs index 4abf02a..192fff8 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -992,6 +992,7 @@ impl BraintrustClient { let submitter = self.span_submitter(); ExperimentBuilder::new( self.resource_transport(), + self.inner.execution.clone(), submitter, self.explicit_scope(token, org_id, None), ) diff --git a/src/span.rs b/src/span.rs index 18c993d..d54f8b5 100644 --- a/src/span.rs +++ b/src/span.rs @@ -481,6 +481,17 @@ impl SpanBuilder { self } + /// Normalize a complete event whose destination will resolve later. + pub(crate) fn into_payload(self, event: SpanLog) -> SpanPayload { + let span = self.build(); + let mut inner = span + .inner + .lock() + .expect("span mutex should not be poisoned"); + apply_span_log_to_data(&mut inner, event); + current_span_payload(&mut inner) + } + pub fn build(self) -> SpanHandle { use std::time::{SystemTime, UNIX_EPOCH}; @@ -1093,6 +1104,14 @@ mod tests { } #[async_trait] impl SpanSubmitter for Submitter { + fn submit_resolving_parent( + &self, + _: CredentialScope, + _: SpanPayload, + _: futures::future::BoxFuture<'static, Result>, + ) -> crate::Completion { + unreachable!("this test submitter only accepts ready span snapshots") + } fn submit( &self, _: CredentialScope, @@ -1169,6 +1188,14 @@ mod tests { } #[async_trait] impl SpanSubmitter for Submitter { + fn submit_resolving_parent( + &self, + _: CredentialScope, + _: SpanPayload, + _: futures::future::BoxFuture<'static, Result>, + ) -> crate::Completion { + unreachable!("this test submitter only accepts ready span snapshots") + } fn submit( &self, _: CredentialScope, diff --git a/src/test_utils.rs b/src/test_utils.rs index 9c91732..21909b9 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -60,6 +60,29 @@ impl MockBraintrustClient { #[async_trait] impl SpanSubmitter for MockBraintrustClient { + fn submit_resolving_parent( + &self, + scope: CredentialScope, + payload: SpanPayload, + resolve: futures::future::BoxFuture<'static, Result>, + ) -> crate::Completion { + let collector = self.collector.clone(); + crate::internal::execution::ExecutionBinding::external(tokio::runtime::Handle::current()) + .response(async move { + let parent = resolve.await.map_err(|source| { + crate::WriteError::Delivery(Arc::new(crate::DeliveryError { + source: source.into(), + })) + })?; + collector.push(CapturedSpan { + token: scope.token().to_owned(), + payload, + parent: Some(parent), + }); + Ok(crate::DeliveryReceipt) + }) + } + fn submit( &self, scope: CredentialScope, diff --git a/tests/eager_experiment_events.rs b/tests/eager_experiment_events.rs new file mode 100644 index 0000000..601f6a2 --- /dev/null +++ b/tests/eager_experiment_events.rs @@ -0,0 +1,340 @@ +use braintrust_sdk_rust::{ + BraintrustClient, BraintrustError, EnqueueError, ExecutionError, Experiment, ExperimentLog, + FlushError, ParentSpanInfo, ShutdownError, WriteError, +}; +use serde_json::json; +use std::{future::IntoFuture, sync::Arc, time::Duration}; +use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, +}; + +fn experiment(client: &BraintrustClient) -> Experiment { + client + .experiment_builder_with_credentials("token", "org") + .project_name("project") + .build() + .unwrap() +} +fn registration() -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(json!({ + "project":{"id":"project","name":"project"}, + "experiment":{"id":"experiment","name":"experiment"} + })) +} +async fn ingestion(server: &MockServer) { + Mock::given(method("GET")) + .and(path("/version")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({}))) + .mount(server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(204)) + .mount(server) + .await; +} +async fn client(server: &MockServer, queue_size: usize) -> BraintrustClient { + BraintrustClient::builder() + .runtime(braintrust_sdk_rust::current_runtime()) + .skip_login(true) + .app_url(server.uri()) + .api_url(server.uri()) + .with_queue_max_size(queue_size) + .build() + .await + .unwrap() +} + +#[tokio::test] +async fn plain_thread_issuance_drains_unobserved_registration_and_delivery_after_drop() { + let server = MockServer::start().await; + ingestion(&server).await; + let arrived = Arc::new(tokio::sync::Notify::new()); + let observed = arrived.clone(); + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(move |_: &wiremock::Request| { + observed.notify_one(); + registration() + }) + .expect(1) + .mount(&server) + .await; + let uri = server.uri(); + let (flush, shutdown) = std::thread::spawn(move || { + let client = futures::executor::block_on( + BraintrustClient::builder() + .skip_login(true) + .app_url(&uri) + .api_url(&uri) + .build() + .into_future(), + ) + .unwrap(); + let experiment = experiment(&client); + let event = experiment.log(ExperimentLog::builder().input(json!("accepted")).build()); + assert!(event.is_ok()); + drop(event); + let flush = experiment.flush(); + let shutdown = client.shutdown(Duration::from_secs(5)); + drop(experiment); + drop(client); + (flush, shutdown) + }) + .join() + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), arrived.notified()) + .await + .unwrap(); + assert_eq!(flush.await.unwrap().delivered, 1); + shutdown.await.unwrap(); + let requests = server.received_requests().await.unwrap(); + let rows: Vec = requests + .iter() + .filter(|r| r.url.path() == "/logs3") + .flat_map(|r| { + serde_json::from_slice::(&r.body).unwrap()["rows"] + .as_array() + .unwrap() + .clone() + }) + .collect(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0]["experiment_id"], "experiment"); + assert_eq!(rows[0]["input"], "accepted"); + assert!(rows[0]["metrics"]["start"].is_number()); + assert!(rows[0]["metrics"]["end"].is_number()); + assert_eq!(rows[0]["span_attributes"]["type"], "eval"); +} + +#[tokio::test] +async fn shutdown_before_first_poll_covers_registration_and_the_final_write() { + let server = MockServer::start().await; + ingestion(&server).await; + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(registration()) + .expect(1) + .mount(&server) + .await; + let client = client(&server, 10).await; + let experiment = experiment(&client); + let event = experiment.log(ExperimentLog::default()); + let shutdown = client.shutdown(Duration::from_secs(5)); + assert!(event.is_ok()); + let rejected = experiment.log(ExperimentLog::default()); + assert!(matches!( + &*rejected, + Err(BraintrustError::Admission(EnqueueError::Closed)) + )); + assert_eq!(shutdown.await.unwrap().delivered, 1); + let id = event.await.unwrap(); + let requests = server.received_requests().await.unwrap(); + let body: serde_json::Value = serde_json::from_slice( + &requests + .iter() + .find(|r| r.url.path() == "/logs3") + .unwrap() + .body, + ) + .unwrap(); + assert_eq!(body["rows"][0]["id"], id); +} + +#[tokio::test] +async fn registration_failure_fails_the_batch_then_retires_before_a_fresh_event() { + let server = MockServer::start().await; + ingestion(&server).await; + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(ResponseTemplate::new(503)) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(registration()) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + let client = client(&server, 10).await; + let experiment = experiment(&client); + let event = experiment.log(ExperimentLog::default()); + let neighbor = client + .span_builder_with_credentials("token", "org") + .parent_info(ParentSpanInfo::ProjectLogs { + object_id: "project".into(), + }) + .build() + .end(); + let flush = experiment.flush(); + assert!(matches!( + event.await, + Err(BraintrustError::Api { status: 503, .. }) + )); + assert!(matches!(neighbor.await, Err(WriteError::Delivery(_)))); + assert!(matches!(flush.await, Err(FlushError::Delivery(failure)) if failure.failed==2)); + assert_eq!(experiment.flush().await.unwrap().delivered, 0); + let next = experiment.clone().log(ExperimentLog::default()); + let fresh = experiment.flush(); + assert!(next.is_ok()); + assert!(!next.await.unwrap().is_empty()); + assert_eq!(fresh.await.unwrap().delivered, 1); + client.shutdown(Duration::from_secs(5)).await.unwrap(); +} + +#[tokio::test] +async fn queue_rejection_never_registers_and_the_reserved_response_slot_is_retired() { + let server = MockServer::start().await; + ingestion(&server).await; + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(registration()) + .expect(1) + .mount(&server) + .await; + let client = client(&server, 1).await; + let experiment = experiment(&client); + let existing = client + .span_builder_with_credentials("token", "org") + .parent_info(ParentSpanInfo::ProjectLogs { + object_id: "project".into(), + }) + .build() + .end(); + let rejected = experiment.log(ExperimentLog::default()); + assert!(matches!( + &*rejected, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + assert!(server.received_requests().await.unwrap().is_empty()); + assert!(matches!( + rejected.await, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + client.flush().await.unwrap(); + existing.await.unwrap(); + let retry = experiment.log(ExperimentLog::default()); + let flush = experiment.flush(); + assert!(retry.is_ok()); + retry.await.unwrap(); + assert_eq!(flush.await.unwrap().delivered, 1); + client.shutdown(Duration::from_secs(5)).await.unwrap(); +} + +#[tokio::test] +async fn full_response_admission_never_creates_an_extra_write_obligation() { + let server = MockServer::start().await; + ingestion(&server).await; + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(registration()) + .expect(1) + .mount(&server) + .await; + let client = client(&server, 2048).await; + let experiment = experiment(&client); + for _ in 0..1024 { + let event = experiment.log(ExperimentLog::default()); + assert!(event.is_ok()); + drop(event); + } + let rejected = experiment.log(ExperimentLog::default()); + assert!(matches!( + &*rejected, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + let shutdown = client.shutdown(Duration::from_secs(10)); + assert_eq!(shutdown.await.unwrap().delivered, 1024); + let count: usize = server + .received_requests() + .await + .unwrap() + .iter() + .filter(|r| r.url.path() == "/logs3") + .map(|r| { + serde_json::from_slice::(&r.body).unwrap()["rows"] + .as_array() + .unwrap() + .len() + }) + .sum(); + assert_eq!(count, 1024); +} + +#[tokio::test] +async fn deadline_cancels_deferred_registration_and_never_submits_its_row() { + let server = MockServer::start().await; + let arrived = Arc::new(tokio::sync::Notify::new()); + let observed = arrived.clone(); + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(move |_: &wiremock::Request| { + observed.notify_one(); + registration().set_delay(Duration::from_secs(60)) + }) + .expect(1) + .mount(&server) + .await; + let client = client(&server, 10).await; + let experiment = experiment(&client); + let event = experiment.log(ExperimentLog::default()); + tokio::time::timeout(Duration::from_secs(5), arrived.notified()) + .await + .unwrap(); + let shutdown = client.shutdown(Duration::from_millis(20)); + assert!(matches!( + shutdown.await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + event.await, + Err(BraintrustError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + &*experiment.log(ExperimentLog::default()), + Err(BraintrustError::Admission(EnqueueError::Closed)) + )); + assert_eq!(server.received_requests().await.unwrap().len(), 1); +} + +#[test] +fn stopped_runtime_retires_an_unpolled_event_that_owns_the_last_queue_producer() { + let (send, receive) = std::sync::mpsc::channel(); + let thread = std::thread::spawn(move || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let client = runtime + .block_on( + BraintrustClient::builder() + .skip_login(true) + .runtime(Some(runtime.handle().clone())) + .build() + .into_future(), + ) + .unwrap(); + let experiment = experiment(&client); + let event = experiment.log(ExperimentLog::default()); + assert!(event.is_ok()); + drop(experiment); + drop(client); + // Deferred preparation owns the only remaining queue producer. Runtime + // loss drops it while canceling the worker, before its first poll. + drop(runtime); + send.send(event).unwrap(); + }); + let event = receive + .recv_timeout(Duration::from_secs(5)) + .expect("runtime destruction must not deadlock on deferred producer cleanup"); + assert!(matches!( + futures::executor::block_on(event.into_future()), + Err(BraintrustError::Execution(ExecutionError::WorkerLost)) + )); + thread.join().unwrap(); +}