Admit experiment events before resolving registration

src/error.rs

@@ -44,6 +44,11 @@ impl From<crate::FlushError> for BraintrustError {
 
 impl From<crate::WriteError> 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::<Self>() {
                 return Self::Api {

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<dyn ExperimentApi>,
+    execution: ExecutionBinding,
     submitter: Arc<dyn SpanSubmitter>,
     scope: CredentialScope,
     project_name: Option<String>,
@@ -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<dyn ExperimentApi>,
+        execution: ExecutionBinding,
         submitter: Arc<dyn SpanSubmitter>,
         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,

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<dyn ExperimentApi>,
+    pub(crate) execution: ExecutionBinding,
     pub(crate) submitter: Arc<dyn SpanSubmitter>,
     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<String> {
-        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<String, crate::BraintrustError> {
+        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<crate::FlushReport, crate::FlushError> {
+        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<ExperimentMetadata> {
+        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()
-    }
 }

src/internal/delivery.rs

@@ -182,26 +182,36 @@ struct Inner<C> {
 }
 impl<C> Inner<C> {
     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<dyn Fn() + Send>);
+        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);

src/internal/execution.rs

@@ -146,6 +146,19 @@ impl ExecutionBinding {
         T: Send + 'static,
         E: From<ExecutionError> + From<EnqueueError> + Send + 'static,
         F: Future<Output = Result<T, E>> + 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<T, E, F, M>(&self, make: M) -> Completion<T, E>
+    where
+        T: Send + 'static,
+        E: From<ExecutionError> + From<EnqueueError> + Send + 'static,
+        F: Future<Output = Result<T, E>> + Send + 'static,
+        M: FnOnce() -> Result<F, E>,
     {
         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);

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<ParentSpanInfo>>,
+    ) -> crate::Completion<crate::DeliveryReceipt, crate::WriteError>;
+
     fn submit(
         &self,
         scope: CredentialScope,