From 927e318a0c04253acf669e68bb632ba5ed711b3f Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 21 Oct 2026 00:03:20 +0800 Subject: [PATCH] Schedule traced callbacks eagerly with scoped shutdown writes --- src/experiments/experiment.rs | 57 ++++++-- src/internal/delivery.rs | 231 +++++++++++++++++++++++++++--- src/internal/ports.rs | 25 ++++ src/log_queue/queue.rs | 58 ++++++-- src/span.rs | 28 ++++ src/test_utils.rs | 8 ++ tests/eager_traced.rs | 260 ++++++++++++++++++++++++++++++++++ tests/experiments.rs | 7 +- 8 files changed, 626 insertions(+), 48 deletions(-) create mode 100644 tests/eager_traced.rs diff --git a/src/experiments/experiment.rs b/src/experiments/experiment.rs index de704a8..726a32e 100644 --- a/src/experiments/experiment.rs +++ b/src/experiments/experiment.rs @@ -67,6 +67,10 @@ struct RecordingMode { selected: bool, } impl ExperimentState { + fn pending_start(self: &Arc) -> PendingSpanStart { + self.recording().pending_starts += 1; + PendingSpanStart(self.clone()) + } fn recording(&self) -> std::sync::MutexGuard<'_, RecordingMode> { self.recording .lock() @@ -157,8 +161,7 @@ impl Experiment { /// Rejected or failed starts release their pending mode reservation. pub fn start_span(&self) -> crate::Completion { self.execution.response_try(|| { - self.state.recording().pending_starts += 1; - let pending = PendingSpanStart(self.state.clone()); + let pending = self.state.pending_start(); let experiment = self.clone(); Ok(async move { // Capture the guard before spawning, including runtime loss @@ -365,6 +368,10 @@ impl Experiment { /// Finalizes the span and flushes active work when the callback completes. /// Setup, final-write, and final-flush errors are returned. Individual writes expose their /// delivery results through the completion tokens returned inside the callback. + /// Submission is immediate. The owned callback runs on the selected runtime + /// even if its token is dropped. Its span and child builders can finish + /// dependent writes while shutdown drains accepted work; that permission + /// ends with the operation. Ordinary flush cutoffs remain fixed. /// /// # Example /// @@ -374,20 +381,42 @@ impl Experiment { /// compute_result() /// }).await; /// ``` - pub async fn traced(&self, callback: F) -> Result + pub fn traced(&self, callback: F) -> crate::Completion where - F: FnOnce(SpanHandle) -> Fut, - Fut: std::future::Future, + F: FnOnce(SpanHandle) -> Fut + Send + 'static, + Fut: std::future::Future + Send + 'static, + R: Send + 'static, { - let span = self.start_span().await?.name("root").build(); - let result = callback(span.clone()).await; - let end = span.end(); - let flush = span.flush(); - let ended = end.await; - let flushed = flush.await; - ended?; - flushed?; - Ok(result) + self.execution.response_try(|| { + let operation = self.submitter.operation()?; + let pending = self.state.pending_start(); + let experiment = self.clone(); + Ok(async move { + // Both guards are captured before first poll. Dropped observers + // do not affect them; cancellation/loss revokes write permission. + let operation = operation; + let pending = pending; + let metadata = experiment.register_metadata().await?; + experiment.state.recording().selected = true; + drop(pending); + let span = SpanBuilder::new(operation.submitter.clone(), experiment.scope.clone()) + .parent_info(ParentSpanInfo::Experiment { + object_id: metadata.experiment_id.clone(), + }) + .span_type(SpanType::Eval) + .name("root") + .build(); + let result = callback(span.clone()).await; + let end = span.end(); + let flush = span.flush(); + let ended = end.await; + let flushed = flush.await; + ended?; + flushed?; + drop(operation); + Ok(result) + }) + }) } /// Update an existing span by ID, merging new data with existing data. diff --git a/src/internal/delivery.rs b/src/internal/delivery.rs index 9ef7241..b378650 100644 --- a/src/internal/delivery.rs +++ b/src/internal/delivery.rs @@ -8,7 +8,7 @@ use futures::{ FutureExt, }; use std::{ - collections::VecDeque, + collections::{HashSet, VecDeque}, sync::{Arc, Mutex, MutexGuard}, time::Duration, }; @@ -59,13 +59,13 @@ impl BatchCompletion { } } /// One bounded summary per active flush, rather than cloning every covered -/// batch into each observer. Its cutoff is the last batch present at admission. +/// batch into each observer. Flushes capture a cutoff; shutdown also covers +/// dependent writes admitted through live operation leases. struct Barrier { - through: Arc, + through: Option>, sender: Option, delivered: u64, failure: Option, - shutdown: bool, } impl Barrier { fn record(&mut self, count: u64, result: &Result<(), BatchError>) { @@ -108,6 +108,8 @@ struct State { accepting: bool, terminal: Option, depth_warned: bool, + leases: HashSet, + next_lease: u64, } impl State { fn seal(&mut self) { @@ -129,23 +131,37 @@ impl State { } fn barrier(&mut self, sender: FlushSender, shutdown: bool) { self.seal(); - if let Some(through) = self - .ready - .back() - .map(|batch| batch.completion.clone()) - .or_else(|| self.active.clone()) - { + let through = if shutdown { + None + } else { + self.ready + .back() + .map(|batch| batch.completion.clone()) + .or_else(|| self.active.clone()) + }; + if through.is_some() || (shutdown && (self.outstanding > 0 || !self.leases.is_empty())) { self.barriers.push(Barrier { through, sender: Some(sender), delivered: 0, failure: None, - shutdown, }); } else { let _ = sender.send(Ok(FlushReport::default())); } } + fn finish_shutdowns(&mut self) { + if self.outstanding == 0 && self.leases.is_empty() { + self.barriers.retain_mut(|barrier| { + if barrier.through.is_none() { + barrier.finish(); + false + } else { + true + } + }); + } + } fn finish(&mut self, batch: &Arc, result: Result<(), BatchError>) { if self .active @@ -162,13 +178,18 @@ impl State { // are removed here, before admitting any subsequent flush. self.barriers.retain_mut(|barrier| { barrier.record(batch.count as u64, &result); - if Arc::ptr_eq(&barrier.through, batch) { + if barrier + .through + .as_ref() + .is_some_and(|through| Arc::ptr_eq(through, batch)) + { barrier.finish(); false } else { true } }); + self.finish_shutdowns(); } } } @@ -190,6 +211,7 @@ impl Inner { } state.terminal = Some(error); state.outstanding = 0; + state.leases.clear(); ( std::mem::take(&mut state.pending), state.active.take(), @@ -238,6 +260,39 @@ impl Clone for WriteService { } } } +/// Permission for dependent writes from one already admitted operation. +/// Explicit revocation is independent of handles retaining cloned permissions. +pub(crate) struct WriteLease(Arc>); +struct Lease { + inner: Arc>, + id: u64, +} +impl Clone for WriteLease { + fn clone(&self) -> Self { + Self(self.0.clone()) + } +} +impl WriteLease { + pub fn revoke(&self) { + self.0.revoke(); + } +} +impl Lease { + fn revoke(&self) { + let mut state = lock(&self.inner.state); + if state.leases.remove(&self.id) { + state.finish_shutdowns(); + } + drop(state); + self.inner.wake.notify_one(); + } +} +impl Drop for Lease { + fn drop(&mut self) { + self.revoke(); + } +} + struct ProducerOwner { inner: std::sync::Weak>, } @@ -264,6 +319,8 @@ impl WriteService { accepting: true, terminal: None, depth_warned: false, + leases: HashSet::new(), + next_lease: 0, }), wake: Notify::new(), capacity, @@ -294,6 +351,31 @@ impl WriteService { } } + pub fn lease(&self) -> Result, EnqueueError> { + let mut state = lock(&self.inner.state); + if !state.accepting || state.terminal.is_some() { + return Err(EnqueueError::Closed); + } + if state.leases.len() >= self.inner.control_limit { + return Err(EnqueueError::Full); + } + let id = state.next_lease; + state.next_lease = id.checked_add(1).ok_or(EnqueueError::Full)?; + state.leases.insert(id); + Ok(WriteLease(Arc::new(Lease { + inner: self.inner.clone(), + id, + }))) + } + + fn accepts(&self, state: &State, lease: Option<&WriteLease>) -> bool { + state.terminal.is_none() + && (state.accepting + || lease.is_some_and(|lease| { + Arc::ptr_eq(&lease.0.inner, &self.inner) && state.leases.contains(&lease.0.id) + })) + } + #[cfg(test)] pub fn enqueue(&self, command: C) -> Completion { self.enqueue_with_rejection(command, None, |_, _| {}) @@ -304,11 +386,21 @@ impl WriteService { command: C, outcome: Option, rejected: impl FnOnce(C, EnqueueError), + ) -> Completion { + self.enqueue_with_lease(command, outcome, None, rejected) + } + + pub fn enqueue_with_lease( + &self, + command: C, + outcome: Option, + lease: Option<&WriteLease>, + rejected: impl FnOnce(C, EnqueueError), ) -> Completion { let (sender, mut completion) = Completion::pair(); let sender = WriteResponder::new(sender, outcome); let mut state = lock(&self.inner.state); - let error = if !state.accepting { + let error = if !self.accepts(&state, lease) { Some(EnqueueError::Closed) } else if state.outstanding >= self.inner.capacity { Some(EnqueueError::Full) @@ -340,14 +432,21 @@ impl WriteService { } pub fn flush(&self) -> Completion { + self.flush_with_lease(None) + } + + pub fn flush_with_lease( + &self, + lease: Option<&WriteLease>, + ) -> Completion { let (sender, mut completion) = Completion::pair(); let mut state = lock(&self.inner.state); - if !state.accepting { + if !self.accepts(&state, lease) { completion.record_issuance(Err(EnqueueError::Closed.into())); } else if state .barriers .iter() - .filter(|barrier| !barrier.shutdown) + .filter(|barrier| barrier.through.is_some()) .count() >= self.inner.control_limit { @@ -450,7 +549,7 @@ where state.active = Some(batch.completion.clone()); Some(batch) } - None if !state.accepting => break, + None if !state.accepting && state.leases.is_empty() => break, None => None, } }; @@ -505,6 +604,106 @@ mod tests { (service, receiver, gate) } + #[tokio::test] + async fn shutdown_waits_for_leases_and_their_later_writes_but_flush_does_not() { + let (service, mut entered, gate) = gated(4); + let lease = service.lease().unwrap(); + let escaped = lease.clone(); + assert_eq!(service.flush().await.unwrap().delivered, 0); + let mut shutdown = Box::pin(service.shutdown(Duration::from_secs(5)).into_future()); + assert!(matches!( + service.enqueue(0).await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); + assert!(futures::poll!(shutdown.as_mut()).is_pending()); + let write = service.enqueue_with_lease(1, None, Some(&lease), |_, _| {}); + let flush = service.flush_with_lease(Some(&lease)); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + gate.add_permits(1); + write.await.unwrap(); + assert_eq!(flush.await.unwrap().delivered, 1); + assert!(futures::poll!(shutdown.as_mut()).is_pending()); + lease.revoke(); + lease.revoke(); + assert_eq!(shutdown.await.unwrap().delivered, 1); + assert!(matches!( + service + .enqueue_with_lease(2, None, Some(&escaped), |_, _| {}) + .await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); + assert!(matches!( + service.flush_with_lease(Some(&escaped)).await, + Err(FlushError::Admission(EnqueueError::Closed)) + )); + } + + #[tokio::test] + async fn leases_are_bounded_service_specific_and_revoked_on_cancellation() { + let (service, _, _) = gated(1); + let lease = service.lease().unwrap(); + assert!(matches!(service.lease(), Err(EnqueueError::Full))); + lease.revoke(); + let replacement = service.lease().unwrap(); + let (other, _, _) = gated(1); + let foreign = other.lease().unwrap(); + service.inner.close(); + assert!(matches!( + service + .enqueue_with_lease(1, None, Some(&foreign), |_, _| {}) + .await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); + service.cancel(ExecutionError::TimedOut); + assert!(lock(&service.inner.state).leases.is_empty()); + assert!(matches!( + service + .enqueue_with_lease(2, None, Some(&replacement), |_, _| {}) + .await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); + } + + #[tokio::test] + async fn shutdown_counts_terminal_failures_from_later_leased_writes_without_replay() { + let service = WriteService::new( + 2, + ExecutionBinding::external(tokio::runtime::Handle::current()), + |commands: Vec| { + Box::pin(async move { + if commands == [1] { + anyhow::bail!("terminal batch"); + } + Ok(()) + }) + }, + ); + let lease = service.lease().unwrap(); + let shutdown = service.shutdown(Duration::from_secs(5)); + let failed = service.enqueue_with_lease(1, None, Some(&lease), |_, _| {}); + assert!(matches!(failed.await, Err(WriteError::Delivery(_)))); + assert!(service.is_empty()); + assert_eq!( + service + .flush_with_lease(Some(&lease)) + .await + .unwrap() + .delivered, + 0 + ); + service + .enqueue_with_lease(2, None, Some(&lease), |_, _| {}) + .await + .unwrap(); + lease.revoke(); + let Err(ShutdownError::Delivery(failure)) = shutdown.await else { + panic!("shutdown must report its covered failed batch"); + }; + assert_eq!(failure.failed, 1); + assert!(service.is_empty()); + assert!(lock(&service.inner.state).barriers.is_empty()); + } + #[test] fn cancel_drops_pending_and_sealed_commands_outside_the_state_lock() { struct Reenter(Box); diff --git a/src/internal/ports.rs b/src/internal/ports.rs index 6122c8f..ba9da91 100644 --- a/src/internal/ports.rs +++ b/src/internal/ports.rs @@ -24,6 +24,29 @@ pub(crate) struct PreparedSpan { pub parent_info: Option, } +pub(crate) struct SpanOperation { + pub submitter: std::sync::Arc, + release: Option>, +} +impl SpanOperation { + pub fn new( + submitter: std::sync::Arc, + release: impl FnOnce() + Send + Sync + 'static, + ) -> Self { + Self { + submitter, + release: Some(Box::new(release)), + } + } +} +impl Drop for SpanOperation { + fn drop(&mut self) { + if let Some(release) = self.release.take() { + release(); + } + } +} + /// Shared delivery lifecycle, independent of resource HTTP operations. #[async_trait] pub(crate) trait DeliveryControl: Send + Sync { @@ -33,6 +56,8 @@ pub(crate) trait DeliveryControl: Send + Sync { #[async_trait] pub(crate) trait SpanSubmitter: DeliveryControl { + fn operation(&self) -> std::result::Result; + /// Queue the write now and resolve its immutable scope/destination during /// preparation. Rejected work drops the future without polling it. fn submit_deferred( diff --git a/src/log_queue/queue.rs b/src/log_queue/queue.rs index 44ba6b7..b0d2226 100644 --- a/src/log_queue/queue.rs +++ b/src/log_queue/queue.rs @@ -6,9 +6,9 @@ use super::http::{fetch_version_info, save_payload_debug, send_batch_with_retry} use super::merge::merge_row_into; use super::row_key::RowKey; use crate::error::{BraintrustError, Result}; -use crate::internal::delivery::WriteService; +use crate::internal::delivery::{WriteLease, WriteService}; use crate::internal::execution::ExecutionBinding; -use crate::internal::ports::{DeliveryControl, PreparedSpan, SpanSubmitter}; +use crate::internal::ports::{DeliveryControl, PreparedSpan, SpanOperation, SpanSubmitter}; use crate::internal::session::CredentialScope; use crate::internal::shared_outcome::WriteOutcome; use crate::span_components::{project_logs_identifier, ProjectLogsIdentifier}; @@ -644,6 +644,7 @@ impl LogQueueCore { pub struct LogQueue { core: Arc, writes: Arc>, + lease: Option>, } enum QueueState { Dormant, @@ -660,6 +661,7 @@ impl LogQueue { Self { core: LogQueueCore::new(config, client, execution), writes: Arc::new(Mutex::new(QueueState::Dormant)), + lease: None, } } fn state(&self) -> MutexGuard<'_, QueueState> { @@ -685,15 +687,10 @@ impl LogQueue { ) } - fn submit_command( - &self, - command: SubmitCommand, - outcome: Option, - ) -> Completion { - let mut state = self.state(); - if matches!(*state, QueueState::Dormant) { + fn ensure_running(&self, state: &mut MutexGuard<'_, QueueState>) { + if matches!(**state, QueueState::Dormant) { let processor = self.core.clone(); - *state = QueueState::Running(WriteService::new( + **state = QueueState::Running(WriteService::new( self.core.config.queue_max_size(), self.core.execution.clone(), move |commands| { @@ -702,12 +699,26 @@ impl LogQueue { }, )); } + } + + fn submit_command( + &self, + command: SubmitCommand, + outcome: Option, + ) -> Completion { + let mut state = self.state(); + self.ensure_running(&mut state); match &*state { - QueueState::Running(writes) => { - writes.enqueue_with_rejection(command, outcome, |command, error| { + QueueState::Running(writes) => match &self.lease { + Some(lease) => { + writes.enqueue_with_lease(command, outcome, Some(lease), |command, error| { + self.reject(command, error) + }) + } + None => 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())); @@ -731,7 +742,10 @@ impl LogQueue { pub fn flush(&self) -> Completion { match &*self.state() { - QueueState::Running(writes) => writes.flush(), + QueueState::Running(writes) => match &self.lease { + Some(lease) => writes.flush_with_lease(Some(lease)), + None => writes.flush(), + }, state => { let (sender, mut completion) = Completion::pair(); if matches!(state, QueueState::Closed) { @@ -802,6 +816,20 @@ impl DeliveryControl for LogQueue { } #[async_trait::async_trait] impl SpanSubmitter for LogQueue { + fn operation(&self) -> std::result::Result { + let lease = { + let mut state = self.state(); + self.ensure_running(&mut state); + match &*state { + QueueState::Running(writes) => writes.lease()?, + _ => return Err(crate::EnqueueError::Closed), + } + }; + let mut scoped = self.clone(); + scoped.lease = Some(lease.clone()); + Ok(SpanOperation::new(Arc::new(scoped), move || lease.revoke())) + } + fn submit_deferred( &self, row_id: String, diff --git a/src/span.rs b/src/span.rs index e347103..7d6158f 100644 --- a/src/span.rs +++ b/src/span.rs @@ -571,6 +571,22 @@ impl Clone for SpanHandle { } impl SpanHandle { + /// Create a child builder with this span's context and execution scope. + /// The child has its own finish state. Retaining it does not extend a + /// callback's permission to submit dependent writes during shutdown. + pub fn child_builder(&self) -> Result { + let parent = self.export()?.to_parent_span_info()?; + let inner = self + .inner + .lock() + .expect("span mutex should not be poisoned"); + let mut builder = + SpanBuilder::new(self.submitter.clone(), self.scope.clone()).parent_info(parent); + builder.environment = inner.environment.clone(); + builder.span_origin = inner.span_origin.clone(); + Ok(builder) + } + /// Return the stable row ID for this span, usable with `log_feedback()` and `update_span()`. pub fn row_id(&self) -> &str { &self.row_id @@ -1104,6 +1120,12 @@ mod tests { } #[async_trait] impl SpanSubmitter for Submitter { + fn operation( + &self, + ) -> std::result::Result + { + unreachable!("this test submitter only accepts ready span snapshots") + } fn submit_deferred( &self, _: String, @@ -1190,6 +1212,12 @@ mod tests { } #[async_trait] impl SpanSubmitter for Submitter { + fn operation( + &self, + ) -> std::result::Result + { + unreachable!("this test submitter only accepts ready span snapshots") + } fn submit_deferred( &self, _: String, diff --git a/src/test_utils.rs b/src/test_utils.rs index b936cae..a6007c0 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -60,6 +60,14 @@ impl MockBraintrustClient { #[async_trait] impl SpanSubmitter for MockBraintrustClient { + fn operation( + &self, + ) -> std::result::Result { + Ok(crate::internal::ports::SpanOperation::new( + Arc::new(self.clone()), + || {}, + )) + } fn submit_deferred( &self, _: String, diff --git a/tests/eager_traced.rs b/tests/eager_traced.rs new file mode 100644 index 0000000..fb08434 --- /dev/null +++ b/tests/eager_traced.rs @@ -0,0 +1,260 @@ +use braintrust_sdk_rust::{ + BraintrustClient, BraintrustError, EnqueueError, ExecutionError, Experiment, ShutdownError, + SpanLog, WriteError, +}; +use serde_json::json; +use std::{future::IntoFuture, sync::Arc, time::Duration}; +use tokio::sync::{oneshot, Notify}; +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() +} +async fn endpoints(server: &MockServer) { + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "project":{"id":"project","name":"project"}, + "experiment":{"id":"experiment","name":"experiment"} + }))) + .expect(1) + .mount(server) + .await; + 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, capacity: 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(capacity) + .build() + .await + .unwrap() +} + +#[tokio::test] +async fn shutdown_before_first_poll_includes_callback_writes_and_revokes_escaped_children() { + let server = MockServer::start().await; + endpoints(&server).await; + let client = client(&server, 8).await; + let experiment = experiment(&client); + let traced = experiment.traced(|root| async move { + root.log(SpanLog::builder().input(json!(1)).build().unwrap()) + .await + .unwrap(); + let child = root.child_builder().unwrap().name("child").build(); + child + .log(SpanLog::builder().output(json!(2)).build().unwrap()) + .await + .unwrap(); + child + }); + assert!(traced.is_ok()); + let cutoff = client.flush(); + let shutdown = client.shutdown(Duration::from_secs(5)); + assert_eq!(cutoff.await.unwrap().delivered, 0); + assert_eq!(shutdown.await.unwrap().delivered, 3); + let child = traced.await.unwrap(); + assert!(child.end_time().is_none()); + assert!(matches!( + child.end().await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); + 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(); + let child_row = rows + .iter() + .find(|r| r["span_attributes"]["name"] == "child") + .unwrap(); + let root_row = rows + .iter() + .find(|r| r["span_attributes"]["name"] == "root") + .unwrap(); + assert_eq!(child_row["root_span_id"], root_row["root_span_id"]); + assert_eq!(child_row["span_parents"][0], root_row["span_id"]); + assert!(rows + .iter() + .any(|r| r["span_attributes"]["name"] == "root" && r["metrics"]["end"].is_number())); +} + +#[tokio::test] +async fn owned_runtime_runs_dropped_callback_observers_from_a_plain_thread() { + let server = MockServer::start().await; + endpoints(&server).await; + let uri = server.uri(); + let (done, observed) = oneshot::channel(); + let 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); + drop(experiment.traced(move |span| async move { + span.log(SpanLog::builder().input(json!("owned")).build().unwrap()) + .await + .unwrap(); + done.send(()).unwrap(); + })); + client.shutdown(Duration::from_secs(5)) + }) + .join() + .unwrap(); + assert_eq!(shutdown.await.unwrap().delivered, 2); + observed.await.unwrap(); +} + +#[tokio::test] +async fn timeout_cancels_callback_and_revokes_its_escaped_write_permission() { + let server = MockServer::start().await; + endpoints(&server).await; + let client = client(&server, 4).await; + let experiment = experiment(&client); + let (sent, escaped) = oneshot::channel(); + let traced = experiment.traced(move |root| async move { + assert!(sent.send(root.child_builder().unwrap().build()).is_ok()); + std::future::pending::<()>().await; + }); + let child = escaped.await.unwrap(); + assert!(matches!( + client.shutdown(Duration::from_millis(20)).await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + traced.await, + Err(BraintrustError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + child.end().await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); + assert!(server + .received_requests() + .await + .unwrap() + .iter() + .all(|r| r.url.path() != "/logs3")); +} + +#[tokio::test] +async fn lease_overload_rejects_before_callback_and_releases_capacity_after_completion() { + let server = MockServer::start().await; + endpoints(&server).await; + let client = client(&server, 1).await; + let experiment = experiment(&client); + let gate = Arc::new(Notify::new()); + let waiting = gate.clone(); + let first = experiment.traced(move |_| async move { + waiting.notified().await; + }); + let rejected = experiment.traced(|_| async { panic!("rejected callback ran") }); + assert!(matches!( + rejected.await, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + gate.notify_one(); + first.await.unwrap(); + assert_eq!(experiment.traced(|_| async { 42 }).await.unwrap(), 42); + client.shutdown(Duration::from_secs(5)).await.unwrap(); +} + +#[tokio::test] +async fn panic_releases_the_operation_and_a_later_callback_can_finish() { + let server = MockServer::start().await; + endpoints(&server).await; + let client = client(&server, 1).await; + let experiment = experiment(&client); + assert!(matches!( + experiment + .traced(|_| async { panic!("callback panic") }) + .await, + Err(BraintrustError::Execution(ExecutionError::WorkerLost)) + )); + assert_eq!( + experiment + .traced(|_| async { Err::<(), _>("user error") }) + .await + .unwrap(), + Err("user error") + ); + client.shutdown(Duration::from_secs(5)).await.unwrap(); +} + +#[test] +fn runtime_loss_before_first_poll_retires_callback_captures_without_deadlocking() { + struct Capture(std::sync::mpsc::Sender<()>); + impl Drop for Capture { + fn drop(&mut self) { + let _ = self.0.send(()); + } + } + 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() + .runtime(Some(runtime.handle().clone())) + .skip_login(true) + .api_url("http://127.0.0.1:9") + .build() + .into_future(), + ) + .unwrap(); + let experiment = experiment(&client); + let (dropped, observed) = std::sync::mpsc::channel(); + let capture = Capture(dropped); + let traced = experiment.traced(move |_| async move { + drop(capture); + panic!("callback must not run on a stopped runtime"); + }); + assert!(traced.is_ok()); + drop(experiment); + drop(client); + drop(runtime); + observed.recv_timeout(Duration::from_secs(5)).unwrap(); + send.send(traced).unwrap(); + }); + let traced = receive + .recv_timeout(Duration::from_secs(5)) + .expect("dropping an unpolled callback must release its lease without deadlocking"); + assert!(matches!( + futures::executor::block_on(traced.into_future()), + Err(BraintrustError::Execution(ExecutionError::WorkerLost)) + )); + thread.join().unwrap(); +} diff --git a/tests/experiments.rs b/tests/experiments.rs index 513fcc8..1212758 100644 --- a/tests/experiments.rs +++ b/tests/experiments.rs @@ -361,11 +361,12 @@ async fn every_registration_dependent_operation_returns_the_real_failure() { ); assert_api_error(experiment.fetch_base_experiment().await, 400); assert_api_error(experiment.summarize().await, 400); - let called = std::sync::atomic::AtomicBool::new(false); + let called = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + let observed = called.clone(); assert_api_error( experiment - .traced(|_| async { - called.store(true, std::sync::atomic::Ordering::Relaxed); + .traced(move |_| async move { + observed.store(true, std::sync::atomic::Ordering::Relaxed); 42 }) .await,