diff --git a/src/completion.rs b/src/completion.rs index 9a35ef8..f51bb83 100644 --- a/src/completion.rs +++ b/src/completion.rs @@ -35,6 +35,11 @@ pub struct Completion { receiver: oneshot::Receiver>, } impl Completion { + pub(crate) fn rejected(error: E) -> Self { + let (_, mut completion) = Self::pair(); + completion.record_issuance(Err(error)); + completion + } pub(crate) fn pair() -> (oneshot::Sender>, Self) { let (sender, receiver) = oneshot::channel(); ( diff --git a/src/dataset.rs b/src/dataset.rs index 5bcacf7..dc4a6a4 100644 --- a/src/dataset.rs +++ b/src/dataset.rs @@ -27,15 +27,18 @@ use std::collections::VecDeque; use std::fmt; use std::sync::Arc; +use futures::{future::BoxFuture, FutureExt}; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use tokio::sync::OnceCell; use uuid::Uuid; use crate::api::datasets::{DatasetFields, DatasetWrite}; -use crate::error::Result; +use crate::error::{BraintrustError, Result}; +use crate::internal::execution::ExecutionBinding; use crate::internal::ports::{DatasetApi, DeliveryControl}; use crate::internal::session::CredentialScope; +use crate::Completion; // ============================================================================ // Dataset Insert Types @@ -275,6 +278,7 @@ impl std::error::Error for DatasetBuilderError {} /// Builder for creating datasets with lazy registration. pub struct DatasetBuilder { api: Arc, + execution: ExecutionBinding, lifecycle: Arc, scope: CredentialScope, project_name: Option, @@ -287,6 +291,7 @@ impl Clone for DatasetBuilder { fn clone(&self) -> Self { Self { api: self.api.clone(), + execution: self.execution.clone(), lifecycle: self.lifecycle.clone(), scope: self.scope.clone(), project_name: self.project_name.clone(), @@ -301,11 +306,13 @@ impl DatasetBuilder { /// Create a new DatasetBuilder (internal use). pub(crate) fn new( api: Arc, + execution: ExecutionBinding, lifecycle: Arc, scope: CredentialScope, ) -> Self { Self { api, + execution, lifecycle, scope, project_name: None, @@ -358,6 +365,7 @@ impl DatasetBuilder { Ok(Dataset { api: self.api, + execution: self.execution, lifecycle: self.lifecycle, scope: self.scope, project_name, @@ -375,9 +383,12 @@ impl DatasetBuilder { /// A dataset for storing test/evaluation data. /// -/// Datasets are lazily registered with the API on first use. +/// Datasets are lazily registered with the API on first use. Operations submit +/// immediately to the client’s selected runtime and return completion tokens. +/// Awaiting observes accepted work; dropping a token does not cancel it. pub struct Dataset { api: Arc, + execution: ExecutionBinding, lifecycle: Arc, scope: CredentialScope, project_name: String, @@ -391,6 +402,7 @@ impl Clone for Dataset { fn clone(&self) -> Self { Self { api: self.api.clone(), + execution: self.execution.clone(), lifecycle: self.lifecycle.clone(), scope: self.scope.clone(), project_name: self.project_name.clone(), @@ -434,7 +446,13 @@ impl Dataset { } /// Get the dataset ID (triggers registration if needed). - pub async fn dataset_id(&self) -> Result { + pub fn dataset_id(&self) -> Completion { + let dataset = self.clone(); + self.execution + .response(async move { dataset.dataset_id_inner().await }) + } + + async fn dataset_id_inner(&self) -> Result { Ok(self.ensure_registered().await?.dataset_id.clone()) } @@ -449,7 +467,13 @@ impl Dataset { } /// Insert a record and return its generated ID after the API acknowledges it. - pub async fn insert(&self, record: DatasetInsert) -> Result { + pub fn insert(&self, record: DatasetInsert) -> Completion { + let dataset = self.clone(); + self.execution + .response(async move { dataset.insert_inner(record).await }) + } + + async fn insert_inner(&self, record: DatasetInsert) -> Result { let dataset_id = self.ensure_registered().await?.dataset_id.clone(); let id = Uuid::new_v4().to_string(); self.api @@ -466,7 +490,14 @@ impl Dataset { } /// Merge supplied fields into a record, returning after API acknowledgement. - pub async fn update(&self, id: &str, record: DatasetInsert) -> Result<()> { + pub fn update(&self, id: &str, record: DatasetInsert) -> Completion<(), BraintrustError> { + let dataset = self.clone(); + let id = id.to_owned(); + self.execution + .response(async move { dataset.update_inner(&id, record).await }) + } + + async fn update_inner(&self, id: &str, record: DatasetInsert) -> Result<()> { let dataset_id = self.ensure_registered().await?.dataset_id.clone(); self.api .write_dataset_record( @@ -481,7 +512,14 @@ impl Dataset { } /// Delete a record, returning after API acknowledgement. - pub async fn delete(&self, id: &str) -> Result<()> { + pub fn delete(&self, id: &str) -> Completion<(), BraintrustError> { + let dataset = self.clone(); + let id = id.to_owned(); + self.execution + .response(async move { dataset.delete_inner(&id).await }) + } + + async fn delete_inner(&self, id: &str) -> Result<()> { let dataset_id = self.ensure_registered().await?.dataset_id.clone(); self.api .write_dataset_record( @@ -497,10 +535,17 @@ impl Dataset { /// Fetch records with pagination. /// /// Returns an iterator that fetches records in batches. - pub async fn fetch(&self, batch_size: Option) -> Result { + pub fn fetch(&self, batch_size: Option) -> Completion { + let dataset = self.clone(); + self.execution + .response(async move { dataset.fetch_inner(batch_size).await }) + } + + async fn fetch_inner(&self, batch_size: Option) -> Result { let metadata = self.ensure_registered().await?; Ok(DatasetIterator::new( self.api.clone(), + self.execution.clone(), self.scope.clone(), metadata.dataset_id.clone(), batch_size.unwrap_or(1000).max(1), @@ -508,7 +553,13 @@ impl Dataset { } /// Get dataset summary statistics. - pub async fn summarize(&self) -> Result { + pub fn summarize(&self) -> Completion { + let dataset = self.clone(); + self.execution + .response(async move { dataset.summarize_inner().await }) + } + + async fn summarize_inner(&self) -> Result { let metadata = self.ensure_registered().await?; // Fetch summary from API @@ -539,6 +590,16 @@ impl Dataset { // DatasetIterator // ============================================================================ +struct DatasetCursor { + fetcher: Arc, + scope: CredentialScope, + dataset_id: String, + batch_size: usize, + cursor: Option, + buffer: VecDeque, + done: bool, +} + /// Iterator for fetching dataset records with pagination. /// /// ```no_run @@ -552,35 +613,86 @@ impl Dataset { /// # } /// ``` pub struct DatasetIterator { - fetcher: Arc, - scope: CredentialScope, - dataset_id: String, - batch_size: usize, - cursor: Option, - buffer: VecDeque, - done: bool, + execution: ExecutionBinding, + cursor: Arc>, + tail: futures::future::Shared>, } impl DatasetIterator { fn new( fetcher: Arc, + execution: ExecutionBinding, scope: CredentialScope, dataset_id: String, batch_size: usize, ) -> Self { Self { - fetcher, - scope, - dataset_id, - batch_size, - cursor: None, - buffer: VecDeque::new(), - done: false, + execution, + tail: futures::future::ready(()).boxed().shared(), + cursor: Arc::new(tokio::sync::Mutex::new(DatasetCursor { + fetcher, + scope, + dataset_id, + batch_size, + cursor: None, + buffer: VecDeque::new(), + done: false, + })), } } + /// Submit a request for the next record, fetching a page if needed. + /// Accepted cursor operations run in issuance order; dropping a token still consumes + /// its record. An API error preserves the cursor for a later retry. + #[allow(clippy::should_implement_trait)] // Observation is async, unlike Iterator::next. + pub fn next(&mut self) -> Completion, BraintrustError> { + self.submit(|mut cursor| async move { cursor.next().await }) + } + + /// Submit collection of all remaining records. Dropping the returned token + /// does not cancel collection; page failures remain operation errors. + pub fn collect(mut self) -> Completion, BraintrustError> { + self.submit(|mut cursor| async move { + let mut records = Vec::new(); + while let Some(record) = cursor.next().await? { + records.push(record); + } + Ok(records) + }) + } + + fn submit(&mut self, work: F) -> Completion + where + T: Send + 'static, + F: FnOnce(tokio::sync::OwnedMutexGuard) -> Fut + Send + 'static, + Fut: std::future::Future> + Send + 'static, + { + let cursor = self.cursor.clone(); + let previous = self.tail.clone(); + let (release, released) = tokio::sync::oneshot::channel::<()>(); + let completion = self.execution.response(async move { + // Holding the sender across the whole operation preserves issuance + // order even if runtime task polling is reordered. Drop releases + // the next operation on success, error, panic or cancellation. + let _release = release; + previous.await; + work(cursor.lock_owned().await).await + }); + if completion.is_ok() { + self.tail = async move { + let _ = released.await; + } + .boxed() + .shared(); + } + // Rejected work leaves the existing chain and cursor untouched. + completion + } +} + +impl DatasetCursor { /// Get the next record, fetching a new batch if needed. - pub async fn next(&mut self) -> Result> { + async fn next(&mut self) -> Result> { loop { if let Some(record) = self.buffer.pop_front() { return Ok(Some(record)); @@ -612,15 +724,6 @@ impl DatasetIterator { // Empty pages with a continuation are not end-of-data. } } - - /// Collect all remaining records, preserving any page failure as an error. - pub async fn collect(mut self) -> Result> { - let mut records = Vec::new(); - while let Some(record) = self.next().await? { - records.push(record); - } - Ok(records) - } } #[cfg(test)] diff --git a/src/error.rs b/src/error.rs index 205da7a..c6f1be9 100644 --- a/src/error.rs +++ b/src/error.rs @@ -28,6 +28,11 @@ pub type Result = std::result::Result; // public operation surfaces migrate to typed completion errors. impl From for BraintrustError { fn from(error: crate::FlushError) -> Self { + match &error { + crate::FlushError::Admission(error) => return Self::Admission(*error), + crate::FlushError::Execution(error) => return Self::Execution(*error), + crate::FlushError::Delivery(_) => {} + } if let crate::FlushError::Delivery(failure) = &error { if let Some(Self::Api { status, message }) = failure.cause.cause().downcast_ref::() @@ -44,6 +49,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..726a32e 100644 --- a/src/experiments/experiment.rs +++ b/src/experiments/experiment.rs @@ -1,7 +1,7 @@ //! Main Experiment type and implementation. use std::collections::HashMap; -use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; @@ -17,7 +17,8 @@ use crate::experiments::metadata::{ }; use crate::experiments::span_builder::ExperimentSpanBuilder; use crate::experiments::summary::{ExperimentSummary, MetricSummary, ScoreSummary}; -use crate::internal::ports::{ExperimentApi, SpanSubmitter}; +use crate::internal::execution::ExecutionBinding; +use crate::internal::ports::{ExperimentApi, PreparedSpan, SpanSubmitter}; use crate::internal::session::CredentialScope; use crate::span::{SpanBuilder, SpanHandle, SpanLog}; use crate::types::{ParentSpanInfo, SpanAttributes, SpanPayload, SpanType}; @@ -33,8 +34,11 @@ pub(crate) struct ExperimentMetadata { /// Handle to a Braintrust experiment. /// /// Use this to log evaluation events, create spans, log feedback, and get summaries. +/// Scheduled methods submit immediately on the selected runtime. Awaiting only +/// observes their completion; dropping a token does not cancel accepted work. pub struct Experiment { pub(crate) api: Arc, + pub(crate) execution: ExecutionBinding, pub(crate) submitter: Arc, pub(crate) scope: CredentialScope, pub(crate) project_name: String, @@ -54,13 +58,38 @@ pub struct Experiment { pub(crate) struct ExperimentState { metadata: OnceCell, last_start_time: AtomicU64, - called_start_span: AtomicBool, + recording: std::sync::Mutex, +} + +#[derive(Default)] +struct RecordingMode { + pending_starts: usize, + 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() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } +} +/// A rejected/failed/lost start must not permanently select span recording. +struct PendingSpanStart(Arc); +impl Drop for PendingSpanStart { + fn drop(&mut self) { + self.0.recording().pending_starts -= 1; + } } 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,68 +112,78 @@ 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(|| { + let span_recording = { + let mode = self.state.recording(); + mode.selected || mode.pending_starts > 0 + }; + if span_recording { + 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 delivery = self.admit_payload(payload)?; + 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. /// - /// Use this for more control over span lifecycle. Once called, - /// you cannot use `log()` on this experiment. - pub async fn start_span(&self) -> Result { - let metadata = self.ensure_registered().await?; - self.state.called_start_span.store(true, Ordering::Release); - Ok(ExperimentSpanBuilder::new(self, metadata.clone())) + /// Use this for more control over span lifecycle. Pending and successful + /// starts prevent `log()` on this experiment. + /// Rejected or failed starts release their pending mode reservation. + pub fn start_span(&self) -> crate::Completion { + self.execution.response_try(|| { + let pending = self.state.pending_start(); + let experiment = self.clone(); + Ok(async move { + // Capture the guard before spawning, including runtime loss + // before first poll. Only successful registration commits mode. + let _pending = pending; + let metadata = experiment.register_metadata().await?; + experiment.state.recording().selected = true; + Ok(ExperimentSpanBuilder::new(&experiment, metadata.clone())) + }) + }) } /// Log feedback on an existing span/record. /// /// Use this to add scores, comments, or other feedback after the initial log. - pub async fn log_feedback(&self, id: &str, feedback: Feedback) -> Result<()> { - let metadata = self.ensure_registered().await?; - + /// Admission happens before registration. Dropping the token does not cancel + /// the patch; awaiting reports full batch delivery. + pub fn log_feedback( + &self, + id: &str, + feedback: Feedback, + ) -> crate::Completion<(), crate::BraintrustError> { // Create a merge payload for the existing span let mut scores = HashMap::new(); if let Some(fb_scores) = feedback.scores { @@ -187,28 +226,36 @@ impl Experiment { extra: HashMap::new(), }; - let parent_info = ParentSpanInfo::Experiment { - object_id: metadata.experiment_id.clone(), - }; - - self.submitter - .submit(self.scope.clone(), payload, Some(parent_info), None) - .await?; - Ok(()) + self.submit_patch(payload) } /// Get a summary of the experiment's performance. /// /// This flushes pending data and fetches aggregated scores from the API. /// If a base experiment is configured, comparison data will be included. - pub async fn summarize(&self) -> Result { - // Flush all pending data first - self.flush().await?; + /// The flush cutoff is captured at issuance, before the returned token is polled. + pub fn summarize(&self) -> crate::Completion { + self.execution.response_try(|| { + let flush = self.flush(); + if let Err(error) = &*flush { + return Err(error.clone().into()); + } + let experiment = self.clone(); + Ok(async move { + flush.await?; + experiment.summarize_inner().await + }) + }) + } - let metadata = self.ensure_registered().await?; + async fn summarize_inner(&self) -> Result { + let metadata = self.register_metadata().await?; // Try to get base experiment for comparison - let base_exp_id = self.fetch_base_experiment().await?.map(|base| base.id); + let base_exp_id = self + .fetch_base_experiment_inner() + .await? + .map(|base| base.id); // Fetch comparison data from API let comparison = self @@ -278,10 +325,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). @@ -303,8 +349,14 @@ impl Experiment { } /// Get the project containing this experiment. - pub async fn project(&self) -> Result { - let metadata = self.ensure_registered().await?; + pub fn project(&self) -> crate::Completion { + let experiment = self.clone(); + self.execution + .response(async move { experiment.project_inner().await }) + } + + async fn project_inner(&self) -> Result { + let metadata = self.register_metadata().await?; Ok(ProjectMetadata { id: metadata.project_id.clone(), name: self.project_name.clone(), @@ -316,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 /// @@ -325,29 +381,55 @@ 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. /// /// This is useful for adding information to a span after it has been created, /// such as adding scores or metadata computed later. - pub async fn update_span(&self, id: &str, event: SpanLog) -> Result<()> { - let metadata = self.ensure_registered().await?; - + /// Admission happens before registration. Dropping the token does not cancel + /// the patch; awaiting reports full batch delivery. + pub fn update_span( + &self, + id: &str, + event: SpanLog, + ) -> crate::Completion<(), crate::BraintrustError> { let payload = SpanPayload { row_id: id.to_string(), span_id: id.to_string(), @@ -374,21 +456,71 @@ impl Experiment { extra: HashMap::new(), }; - let parent_info = ParentSpanInfo::Experiment { - object_id: metadata.experiment_id.clone(), - }; + self.submit_patch(payload) + } - self.submitter - .submit(self.scope.clone(), payload, Some(parent_info), None) - .await?; - Ok(()) + fn submit_patch(&self, payload: SpanPayload) -> crate::Completion<(), crate::BraintrustError> { + self.execution.response_try(|| { + let delivery = self.admit_payload(payload)?; + Ok(async move { + delivery.await?; + Ok(()) + }) + }) + } + + /// Called only after reserving response capacity, so a rejected response + /// never leaves a write behind. The write owns all remaining preparation. + fn admit_payload( + &self, + payload: SpanPayload, + ) -> Result> { + let delivery = if let Some(metadata) = self.state.metadata.get() { + self.submitter.submit( + self.scope.clone(), + payload, + Some(ParentSpanInfo::Experiment { + object_id: metadata.experiment_id.clone(), + }), + None, + ) + } else { + let experiment = self.clone(); + let scope = self.scope.clone(); + self.submitter.submit_deferred( + payload.row_id.clone(), + Box::pin(async move { + let metadata = experiment.register_metadata().await?; + let parent_info = Some(ParentSpanInfo::Experiment { + object_id: metadata.experiment_id.clone(), + }); + Ok(PreparedSpan { + scope, + payload, + parent_info, + }) + }), + ) + }; + if let Err(error) = &*delivery { + return Err(error.clone().into()); + } + Ok(delivery) } /// Fetch the base experiment used for comparison. /// /// Returns `None` if no base experiment has been configured or detected. - pub async fn fetch_base_experiment(&self) -> Result> { - let metadata = self.ensure_registered().await?; + pub fn fetch_base_experiment( + &self, + ) -> crate::Completion, crate::BraintrustError> { + let experiment = self.clone(); + self.execution + .response(async move { experiment.fetch_base_experiment_inner().await }) + } + + async fn fetch_base_experiment_inner(&self) -> Result> { + let metadata = self.register_metadata().await?; self.api .fetch_base_experiment(&self.scope, &metadata.experiment_id) .await @@ -398,16 +530,22 @@ impl Experiment { /// /// Use this to pass experiment context to another process or service, /// allowing spans created there to be linked to this experiment. - pub async fn export(&self) -> Result { - let metadata = self.ensure_registered().await?; + pub fn export(&self) -> crate::Completion { + let experiment = self.clone(); + self.execution + .response(async move { experiment.export_inner().await }) + } + + async fn export_inner(&self) -> Result { + let metadata = self.register_metadata().await?; Ok(ExportedExperiment { object_type: "experiment".to_string(), object_id: metadata.experiment_id.clone(), }) } - /// Ensure the experiment is registered, returning cached metadata. - pub(crate) async fn ensure_registered(&self) -> Result<&ExperimentMetadata> { + /// 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 +569,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..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(); } } } @@ -182,26 +203,37 @@ 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; + state.leases.clear(); + ( + 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) { @@ -228,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>, } @@ -254,6 +319,8 @@ impl WriteService { accepting: true, terminal: None, depth_warned: false, + leases: HashSet::new(), + next_lease: 0, }), wake: Notify::new(), capacity, @@ -284,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, |_, _| {}) @@ -294,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) @@ -330,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 { @@ -440,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, } }; @@ -495,6 +604,153 @@ 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); + 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 7626cd8..ba9da91 100644 --- a/src/internal/ports.rs +++ b/src/internal/ports.rs @@ -16,6 +16,37 @@ use crate::internal::session::CredentialScope; use crate::internal::shared_outcome::WriteOutcome; use crate::types::{ParentSpanInfo, SpanPayload}; +/// Fully scoped input to row preparation. Unresolved identity stays in the +/// preparation future rather than being represented by placeholder credentials. +pub(crate) struct PreparedSpan { + pub scope: CredentialScope, + pub payload: SpanPayload, + 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 { @@ -25,6 +56,16 @@ 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( + &self, + row_id: String, + prepare: futures::future::BoxFuture<'static, Result>, + ) -> crate::Completion; + fn submit( &self, scope: CredentialScope, @@ -34,6 +75,8 @@ pub(crate) trait SpanSubmitter: DeliveryControl { ) -> crate::Completion; } +/// One-attempt dataset I/O, polled inside an already admitted resource operation +/// on its execution binding. Subrequests do not reserve additional task slots. #[async_trait] pub(crate) trait DatasetApi: Send + Sync { async fn register_dataset( @@ -54,6 +97,8 @@ pub(crate) trait DatasetApi: Send + Sync { async fn write_dataset_record(&self, scope: &CredentialScope, row: DatasetWrite) -> Result<()>; } +/// One-attempt experiment I/O, polled inside admitted preparation or a bound +/// resource operation. HTTP subrequests do not reserve additional task slots. #[async_trait] pub(crate) trait ExperimentApi: Send + Sync { async fn register_experiment( diff --git a/src/internal/resources.rs b/src/internal/resources.rs index 7591a65..423ec18 100644 --- a/src/internal/resources.rs +++ b/src/internal/resources.rs @@ -1,5 +1,4 @@ //! 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; @@ -17,15 +16,14 @@ use crate::error::{BraintrustError, Result}; use crate::experiments::BaseExperimentInfo; use crate::internal::ports::{DatasetApi, ExperimentApi}; -/// Shared HTTP pool and execution lease without client/session/delivery state. +/// Shared HTTP pool without execution, client, session or delivery state. #[derive(Clone)] pub(crate) struct ResourceTransport { http: reqwest::Client, - execution: ExecutionBinding, } impl ResourceTransport { - pub fn new(http: reqwest::Client, execution: ExecutionBinding) -> Self { - Self { http, execution } + pub fn new(http: reqwest::Client) -> Self { + Self { http } } fn resource_api(&self, scope: &CredentialScope) -> Result { scope.api(&self.http) @@ -39,15 +37,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( @@ -59,19 +55,15 @@ impl ExperimentApi for ResourceTransport { 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 + 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, + }) } async fn fetch_experiment_comparison( @@ -80,20 +72,17 @@ impl ExperimentApi for ResourceTransport { experiment_id: &str, base_experiment_id: Option<&str>, ) -> Result { - 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) - }) + self.resource_api(scope)? + .experiments() + .compare(experiment_id, base_experiment_id) .await + .map_err(BraintrustError::from) } } +// Dataset operations already own one response-task admission for their entire +// registration/request sequence. Subrequests must not reenter admission: they +// remain valid when shutdown closes the public boundary while draining them. #[async_trait] impl DatasetApi for ResourceTransport { async fn register_dataset( @@ -101,15 +90,11 @@ impl DatasetApi for ResourceTransport { scope: &CredentialScope, request: DatasetRegisterRequest, ) -> Result { - let api = self.resource_api(scope)?; - self.execution - .response(async move { - api.datasets() - .register(request) - .await - .map_err(BraintrustError::from) - }) + self.resource_api(scope)? + .datasets() + .register(request) .await + .map_err(BraintrustError::from) } async fn fetch_dataset_records( @@ -117,15 +102,11 @@ impl DatasetApi for ResourceTransport { scope: &CredentialScope, query: BTQLQuery, ) -> Result { - let api = self.resource_api(scope)?; - self.execution - .response(async move { - api.btql() - .query_dataset_records(query) - .await - .map_err(BraintrustError::from) - }) + self.resource_api(scope)? + .btql() + .query_dataset_records(query) .await + .map_err(BraintrustError::from) } async fn fetch_dataset_summary( @@ -133,27 +114,18 @@ impl DatasetApi for ResourceTransport { scope: &CredentialScope, dataset_id: &str, ) -> Result { - 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) - }) + self.resource_api(scope)? + .datasets() + .summarize(dataset_id) .await + .map_err(BraintrustError::from) } async fn write_dataset_record(&self, scope: &CredentialScope, row: DatasetWrite) -> Result<()> { - let api = self.resource_api(scope)?; - self.execution - .response(async move { - api.datasets() - .write(&[row]) - .await - .map_err(BraintrustError::from) - }) + self.resource_api(scope)? + .datasets() + .write(&[row]) .await + .map_err(BraintrustError::from) } } diff --git a/src/log_queue/queue.rs b/src/log_queue/queue.rs index 3938100..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, 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}; @@ -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}; @@ -40,10 +40,12 @@ fn cache_insert(cache: &mut IndexMap, key: K, /// /// Fields are private — construction and destructuring happen only within this /// module. `worker.rs` never touches the fields directly. -struct SubmitCommand { - scope: CredentialScope, - payload: SpanPayload, - parent_info: Option, +enum SubmitCommand { + Ready(Box), + Deferred { + row_id: String, + prepare: BoxFuture<'static, Result>, + }, } /// Preparation and protocol orchestration, independent of admission/lifecycle. @@ -108,13 +110,21 @@ impl LogQueueCore { // exact scope captured by each producer when its handle was created. let mut groups: IndexMap> = IndexMap::new(); for cmd in cmds { - let SubmitCommand { - scope, - payload, - parent_info, - } = cmd; - match self.prepare_row(&scope, payload, parent_info).await { - Ok(row) => { + let prepared = async { + let PreparedSpan { + scope, + payload, + parent_info, + } = match cmd { + SubmitCommand::Ready(span) => *span, + SubmitCommand::Deferred { prepare, .. } => prepare.await?, + }; + let row = self.prepare_row(&scope, payload, parent_info).await?; + Ok::<_, anyhow::Error>((scope, row)) + } + .await; + match prepared { + Ok((scope, row)) => { groups.entry(scope).or_default().push(row); } Err(e) => { @@ -500,12 +510,19 @@ impl LogQueueCore { return; } self.execution.spawn(async move { - let SubmitCommand { + let PreparedSpan { scope: _, payload, parent_info, - } = cmd; - + } = match cmd { + SubmitCommand::Ready(span) => *span, + SubmitCommand::Deferred { row_id, .. } => { + // Rejection must not initiate login or registration merely + // to serialize a concrete wire row for a debug dump. + warn!(%row_id, "cannot dump rejected row with unresolved scope or destination"); + return; + } + }; let span_components = payload.span_components.clone(); let span_id = payload.span_id.clone(); @@ -627,6 +644,7 @@ impl LogQueueCore { pub struct LogQueue { core: Arc, writes: Arc>, + lease: Option>, } enum QueueState { Dormant, @@ -643,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> { @@ -658,15 +677,20 @@ impl LogQueue { parent_info: Option, outcome: Option, ) -> Completion { - let command = SubmitCommand { - scope, - payload, - parent_info, - }; - let mut state = self.state(); - if matches!(*state, QueueState::Dormant) { + self.submit_command( + SubmitCommand::Ready(Box::new(PreparedSpan { + scope, + payload, + parent_info, + })), + outcome, + ) + } + + 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| { @@ -675,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())); @@ -704,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) { @@ -730,11 +771,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 +816,28 @@ 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, + prepare: BoxFuture<'static, Result>, + ) -> Completion { + self.submit_command(SubmitCommand::Deferred { row_id, prepare }, None) + } + fn submit( &self, scope: CredentialScope, diff --git a/src/logger.rs b/src/logger.rs index 6ce2f14..3b22eb2 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -13,7 +13,7 @@ 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::ports::{PreparedSpan, SpanSubmitter}; use crate::internal::resources::ResourceTransport; use crate::internal::retry::{exponential, RetryExecutor}; use crate::internal::session::CredentialScope; @@ -389,11 +389,7 @@ impl BraintrustClientBuilder { 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 - } + Err(error) => crate::Completion::rejected(error), } } @@ -481,6 +477,50 @@ pub struct BraintrustClient { inner: Arc, } +/// Validated, owned patch data without unresolved or placeholder credentials. +struct SpanUpdate { + row_id: String, + span_id: String, + components: SpanComponents, + event_data: SpanEventData, +} +impl SpanUpdate { + fn into_prepared(self, scope: CredentialScope) -> PreparedSpan { + let Self { + row_id, + span_id, + components, + event_data, + } = self; + let payload = SpanPayload { + row_id, + span_id, + is_merge: true, + span_components: Some(components), + org_id: scope.org_id().to_owned(), + org_name: scope.org_name().map(str::to_owned), + project_name: None, + input: event_data.input, + output: event_data.output, + expected: event_data.expected, + error: event_data.error, + scores: event_data.scores, + metadata: event_data.metadata, + metrics: event_data.metrics, + tags: event_data.tags, + context: event_data.context, + span_attributes: event_data.span_attributes, + extra: event_data.extra, + }; + + PreparedSpan { + scope, + payload, + parent_info: None, + } + } +} + struct ClientInner { lifecycle: Lifecycle, execution: ExecutionBinding, @@ -509,10 +549,7 @@ impl std::fmt::Debug for ClientInner { impl BraintrustClient { fn resource_transport(&self) -> Arc { - Arc::new(ResourceTransport::new( - self.inner.http_client.clone(), - self.inner.execution.clone(), - )) + Arc::new(ResourceTransport::new(self.inner.http_client.clone())) } fn span_submitter(&self) -> Arc { @@ -553,31 +590,37 @@ impl BraintrustClient { BraintrustClientBuilder::new() } - /// Check if the client is logged in. - pub async fn is_logged_in(&self) -> bool { + /// Check local login state immediately, without scheduling work. + pub fn is_logged_in(&self) -> bool { self.inner.login_state.is_logged_in() } /// Wait for login to complete (useful if using background login). /// - /// Returns the login state once available, or an error if login times out. - pub async fn wait_for_login(&self) -> Result { + /// Submits the wait immediately on the selected runtime. Awaiting observes + /// the login state or a timeout; dropping the token does not cancel the wait. + /// Login-dependent methods reject issuance when login was skipped. + pub fn wait_for_login(&self) -> crate::Completion { if self.inner.login_skipped { - return Err(BraintrustError::InvalidConfig( + return crate::Completion::rejected(BraintrustError::InvalidConfig( "wait_for_login() not available when skip_login is enabled".into(), )); } - self.wait_for_login_state().await + let client = self.clone(); + self.inner + .execution + .response(async move { client.wait_for_login_state().await }) } - /// Get the current login state. - pub async fn login_state(&self) -> LoginState { + /// Get a shared view of local login state without scheduling work. + pub fn login_state(&self) -> LoginState { self.inner.login_state.clone() } /// Create a span builder using the logged-in state and default project. /// - /// This waits for login to complete if it hasn't already. + /// Submits immediately, waiting for login on the selected runtime if needed. + /// Dropping the returned completion does not cancel accepted work. /// /// # Example /// @@ -595,12 +638,19 @@ impl BraintrustClient { /// # Ok(()) /// # } /// ``` - pub async fn span_builder(&self) -> Result { + pub fn span_builder(&self) -> crate::Completion { if self.inner.login_skipped { - return Err(BraintrustError::InvalidConfig( + return crate::Completion::rejected(BraintrustError::InvalidConfig( "span_builder() requires login; use span_builder_with_credentials() when skip_login is enabled".into(), )); } + let client = self.clone(); + self.inner + .execution + .response(async move { client.span_builder_inner().await }) + } + + async fn span_builder_inner(&self) -> Result { let state = self.wait_for_login_state().await?; let mut builder = crate::span::SpanBuilder::new(self.span_submitter(), state.scope()?); if let Some(environment) = &self.inner.environment { @@ -617,14 +667,23 @@ impl BraintrustClient { /// Create a dataset builder using the logged-in state. /// - /// This waits for login to complete if it hasn't already. - pub async fn dataset_builder(&self) -> Result { + /// Submits immediately, waiting for login on the selected runtime if needed. + /// Dropping the returned completion does not cancel accepted work. + pub fn dataset_builder(&self) -> crate::Completion { if self.inner.login_skipped { - return Err(BraintrustError::InvalidConfig("dataset_builder() requires login; use dataset_builder_with_credentials() when skip_login is enabled".into())); + return crate::Completion::rejected(BraintrustError::InvalidConfig("dataset_builder() requires login; use dataset_builder_with_credentials() when skip_login is enabled".into())); } + let client = self.clone(); + self.inner + .execution + .response(async move { client.dataset_builder_inner().await }) + } + + async fn dataset_builder_inner(&self) -> Result { let state = self.wait_for_login_state().await?; Ok(DatasetBuilder::new( self.resource_transport(), + self.inner.execution.clone(), Arc::new(self.inner.queue.clone()), state.scope()?, )) @@ -639,6 +698,7 @@ impl BraintrustClient { ) -> DatasetBuilder { DatasetBuilder::new( self.resource_transport(), + self.inner.execution.clone(), Arc::new(self.inner.queue.clone()), self.explicit_scope(token.into(), String::new(), Some(org_name.into())), ) @@ -670,47 +730,88 @@ impl BraintrustClient { builder } - /// Update an existing span using the output of `SpanHandle::export()`. - /// - /// Validation happens before queuing. The actual upload remains asynchronous. - pub async fn update_span(&self, exported: &str, event: crate::span::SpanLog) -> Result<()> { - let login_state = if self.inner.login_state.is_logged_in() { - self.inner.login_state.clone() + /// Submit an update using exported span context and the shared login state. + /// Validation and admission happen immediately. The completion includes + /// login (if pending), preparation and full batch delivery. + pub fn update_span( + &self, + exported: &str, + event: crate::span::SpanLog, + ) -> crate::Completion<(), BraintrustError> { + let scope = if self.inner.login_state.is_logged_in() { + match self.inner.login_state.scope() { + Ok(scope) => Some(scope), + Err(error) => return crate::Completion::rejected(error), + } } else if self.inner.login_skipped { - return Err(BraintrustError::InvalidConfig( + return crate::Completion::rejected(BraintrustError::InvalidConfig( "update_span() requires a logged-in client; use update_span_with_credentials() when skip_login is enabled".into(), )); } else { - self.wait_for_login_state().await? + None }; - - self.update_span_internal(login_state.scope()?, exported, event) + self.schedule_span_update(scope, exported, event) } - /// Update an existing span using explicit credentials instead of shared login state. - /// - /// This is the safe entrypoint for multi-tenant `skip_login` clients: the - /// supplied credentials travel with the queued row and are NOT written into - /// the client's shared login state. + /// Submit an update with immutable explicit credentials and configured routes. + /// Credentials travel with this write and never populate shared login state. + /// Dropping the token does not cancel an admitted update. pub fn update_span_with_credentials( &self, token: impl Into, org_id: impl Into, exported: &str, event: crate::span::SpanLog, - ) -> Result<()> { - let token = token.into(); - let org_id = org_id.into(); + ) -> crate::Completion<(), BraintrustError> { + let scope = self.explicit_scope(token.into(), org_id.into(), None); + self.schedule_span_update(Some(scope), exported, event) + } - self.update_span_internal(self.explicit_scope(token, org_id, None), exported, event) + // None means the immutable session scope is still waiting for login. The + // queue owns that resolution before any flush or shutdown can capture it. + fn schedule_span_update( + &self, + scope: Option, + exported: &str, + event: crate::span::SpanLog, + ) -> crate::Completion<(), BraintrustError> { + self.inner.execution.response_try(|| { + let update = self.prepare_span_update(exported, event)?; + let delivery = match scope { + Some(scope) => { + let PreparedSpan { + scope, + payload, + parent_info, + } = update.into_prepared(scope); + self.submit_payload(scope, payload, parent_info) + } + None => { + let client = self.clone(); + self.inner.queue.submit_deferred( + update.row_id.clone(), + Box::pin(async move { + let state = client.wait_for_login_state().await?; + Ok(update.into_prepared(state.scope()?)) + }), + ) + } + }; + if let Err(error) = &*delivery { + return Err(error.clone().into()); + } + Ok(async move { + delivery.await?; + Ok(()) + }) + }) } - fn update_span_internal( + fn prepare_span_update( &self, - scope: CredentialScope, exported: &str, event: crate::span::SpanLog, - ) -> Result<()> { + ) -> Result { let components = SpanComponents::parse(exported)?; let row_id = components.row_id.clone().ok_or_else(|| { BraintrustError::InvalidConfig("Exported span must have a row_id".into()) @@ -782,32 +883,12 @@ impl BraintrustClient { event_data.apply_propagated_event(propagated_event); } - let payload = SpanPayload { + Ok(SpanUpdate { row_id, span_id, - is_merge: true, - span_components: Some(components), - org_id: scope.org_id().to_owned(), - org_name: scope.org_name().map(str::to_owned), - project_name: None, - input: event_data.input, - output: event_data.output, - expected: event_data.expected, - error: event_data.error, - scores: event_data.scores, - metadata: event_data.metadata, - metrics: event_data.metrics, - tags: event_data.tags, - context: event_data.context, - span_attributes: event_data.span_attributes, - extra: event_data.extra, - }; - - let completion = self.submit_payload(scope, payload, None); - if let Err(error) = &*completion { - return Err(BraintrustError::Background(error.to_string())); - } - Ok(()) + components, + event_data, + }) } /// Perform login synchronously. @@ -973,6 +1054,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), ) @@ -1699,7 +1781,7 @@ mod tests { .await .expect("client"); - let login_state = client.login_state().await; + let login_state = client.login_state(); assert!(login_state.is_logged_in(), "should be logged in"); assert_eq!(login_state.org_id().as_deref(), Some("org-1")); assert_eq!(login_state.org_name().as_deref(), Some("First Org")); @@ -1729,7 +1811,7 @@ mod tests { .await .expect("client"); - let login_state = client.login_state().await; + let login_state = client.login_state(); assert!(login_state.is_logged_in(), "should be logged in"); assert_eq!(login_state.org_id().as_deref(), Some("org-2")); assert_eq!(login_state.org_name().as_deref(), Some("Second Org")); @@ -1822,7 +1904,7 @@ mod tests { .expect("client"); // Initially should not be logged in (background login hasn't completed) - let is_logged_in = client.is_logged_in().await; + let is_logged_in = client.is_logged_in(); assert!(!is_logged_in); } @@ -1847,7 +1929,7 @@ mod tests { // Wait for background login to complete let login_state = client.wait_for_login().await.expect("login should succeed"); assert_eq!(login_state.org_id().as_deref(), Some("org-1")); - assert!(client.is_logged_in().await); + assert!(client.is_logged_in()); } #[tokio::test] @@ -2016,7 +2098,7 @@ mod tests { // No credential leak into shared global state. assert!( - !client.is_logged_in().await, + !client.is_logged_in(), "per-request credentials must not populate the shared login state" ); diff --git a/src/span.rs b/src/span.rs index 18c993d..7d6158f 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}; @@ -560,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 @@ -1093,6 +1120,22 @@ 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, + _: futures::future::BoxFuture< + 'static, + Result, + >, + ) -> crate::Completion { + unreachable!("this test submitter only accepts ready span snapshots") + } fn submit( &self, _: CredentialScope, @@ -1169,6 +1212,22 @@ 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, + _: 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..a6007c0 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -7,7 +7,7 @@ use crate::internal::session::CredentialScope; use crate::span::{SpanBuilder, SpanHandle}; use crate::types::{ParentSpanInfo, SpanPayload}; -use crate::internal::ports::{DeliveryControl, SpanSubmitter}; +use crate::internal::ports::{DeliveryControl, PreparedSpan, SpanSubmitter}; use crate::internal::shared_outcome::{WriteOutcome, WriteResponder}; #[derive(Clone, Default)] @@ -60,6 +60,40 @@ 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, + prepare: 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 PreparedSpan { + scope, + payload, + parent_info, + } = prepare.await.map_err(|source| { + crate::WriteError::Delivery(Arc::new(crate::DeliveryError { + source: source.into(), + })) + })?; + collector.push(CapturedSpan { + token: scope.token().to_owned(), + payload, + parent: parent_info, + }); + Ok(crate::DeliveryReceipt) + }) + } + fn submit( &self, scope: CredentialScope, diff --git a/tests/credential_scopes.rs b/tests/credential_scopes.rs index b9a4c65..d2e748f 100644 --- a/tests/credential_scopes.rs +++ b/tests/credential_scopes.rs @@ -43,7 +43,7 @@ async fn explicit_handles_keep_their_routes_after_shared_login_state_is_populate let builder = client .span_builder_with_credentials("explicit", "org") .project_name("p"); - assert!(client.login_state().await.set( + assert!(client.login_state().set( "session-secret".into(), "other-org".into(), "other".into(), @@ -335,7 +335,7 @@ async fn experiment_clones_and_child_spans_preserve_the_explicit_scope() { let span = experiment.start_span().await.unwrap().build(); span.log(SpanLog::default()).await.unwrap(); experiment.clone().summarize().await.unwrap(); - assert!(!client.is_logged_in().await); + assert!(!client.is_logged_in()); } #[tokio::test] diff --git a/tests/dataset_test.rs b/tests/dataset_test.rs index 330f57f..6142031 100644 --- a/tests/dataset_test.rs +++ b/tests/dataset_test.rs @@ -337,7 +337,7 @@ async fn failed_dataset_registration_is_retryable_and_clone_registration_is_shar assert_eq!(first.unwrap(), "dataset"); assert_eq!(second.unwrap(), "dataset"); assert_eq!(dataset.clone().dataset_id().await.unwrap(), "dataset"); - assert!(!client.is_logged_in().await); + assert!(!client.is_logged_in()); } #[tokio::test] diff --git a/tests/eager_client.rs b/tests/eager_client.rs new file mode 100644 index 0000000..21451fc --- /dev/null +++ b/tests/eager_client.rs @@ -0,0 +1,248 @@ +//! Login-dependent client operations are accepted before their tokens are polled. +use braintrust_sdk_rust::{ + BraintrustClient, BraintrustError, EnqueueError, ExecutionError, ShutdownError, +}; +use serde_json::json; +use std::{future::IntoFuture, time::Duration}; +use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, +}; + +fn is_closed(completion: &braintrust_sdk_rust::Completion) -> bool { + matches!( + &**completion, + Err(BraintrustError::Admission(EnqueueError::Closed)) + ) +} + +#[test] +fn skipped_login_rejects_issuance_without_polling_or_a_call_site_runtime() { + let client = futures::executor::block_on( + BraintrustClient::builder() + .skip_login(true) + .build() + .into_future(), + ) + .unwrap(); + assert!(!client.is_logged_in()); + assert!(!client.login_state().is_logged_in()); + let wait = client.wait_for_login(); + let span = client.span_builder(); + let dataset = client.dataset_builder(); + assert!(matches!(&*wait, Err(BraintrustError::InvalidConfig(_)))); + assert!(matches!(&*span, Err(BraintrustError::InvalidConfig(_)))); + assert!(matches!(&*dataset, Err(BraintrustError::InvalidConfig(_)))); + assert!(matches!( + futures::executor::block_on(wait.into_future()), + Err(BraintrustError::InvalidConfig(_)) + )); + assert!(matches!( + futures::executor::block_on(span.into_future()), + Err(BraintrustError::InvalidConfig(_)) + )); + assert!(matches!( + futures::executor::block_on(dataset.into_future()), + Err(BraintrustError::InvalidConfig(_)) + )); +} + +#[tokio::test] +async fn synchronous_issuance_and_cross_executor_observation_share_login_and_shutdown() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/apikey/login")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "org_info":[{"id":"org","name":"tenant"}] + }))) + .expect(1) + .mount(&server) + .await; + let uri = server.uri(); + let (wait, span, dataset, shutdown) = std::thread::spawn(move || { + assert!(braintrust_sdk_rust::current_runtime().is_none()); + let client = futures::executor::block_on( + BraintrustClient::builder() + .api_key("token") + .blocking_login(false) + .app_url(&uri) + .api_url(&uri) + .build() + .into_future(), + ) + .unwrap(); + let wait = client.wait_for_login(); + let span = client.span_builder(); + let dataset = client.dataset_builder(); + assert!(wait.is_ok()); + assert!(span.is_ok()); + assert!(dataset.is_ok()); + // No observation starts these operations. Shutdown must drain the + // background login and each accepted builder/wait before completing. + let shutdown = client.shutdown(Duration::from_secs(5)); + drop(client); + (wait, span, dataset, shutdown) + }) + .join() + .unwrap(); + shutdown.await.unwrap(); + assert_eq!(wait.await.unwrap().org_id().as_deref(), Some("org")); + let _span = span.await.unwrap().name("accepted").build(); + let dataset = dataset + .await + .unwrap() + .project_name("p") + .dataset_name("d") + .build() + .unwrap(); + assert_eq!(dataset.dataset_name(), "d"); + // Obtaining/building the handle didn't register the dataset or write a span. + assert_eq!(server.received_requests().await.unwrap().len(), 1); +} + +#[tokio::test] +async fn shutdown_deadline_resolves_unobserved_login_waits_and_resource_builder_requests() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/apikey/login")) + .respond_with(ResponseTemplate::new(200).set_delay(Duration::from_secs(60))) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .api_key("token") + .blocking_login(false) + .app_url(server.uri()) + .api_url(server.uri()) + .build() + .await + .unwrap(); + let wait = client.wait_for_login(); + let span = client.span_builder(); + let dataset = client.dataset_builder(); + assert!(wait.is_ok()); + assert!(span.is_ok()); + assert!(dataset.is_ok()); + assert!(matches!( + client.shutdown(Duration::from_millis(20)).await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + wait.await, + Err(BraintrustError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + span.await, + Err(BraintrustError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + dataset.await, + Err(BraintrustError::Execution(ExecutionError::TimedOut)) + )); + assert!(is_closed(&client.wait_for_login())); + assert!(is_closed(&client.span_builder())); + assert!(is_closed(&client.dataset_builder())); + // Local metadata remains available after work admission closes. + assert!(!client.is_logged_in()); + assert!(!client.login_state().is_logged_in()); +} + +#[tokio::test] +async fn unobserved_client_requests_share_the_bounded_task_budget() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/apikey/login")) + .respond_with(ResponseTemplate::new(200).set_delay(Duration::from_secs(60))) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .runtime(braintrust_sdk_rust::current_runtime()) + .api_key("token") + .blocking_login(false) + .app_url(server.uri()) + .api_url(server.uri()) + .build() + .await + .unwrap(); + // Background login already occupies one slot. Dropped observations keep + // their accepted waits active, so all 1,023 remaining slots are consumed. + for _ in 0..1023 { + let wait = client.wait_for_login(); + assert!(wait.is_ok()); + drop(wait); + } + let span = client.span_builder(); + let dataset = client.dataset_builder(); + assert!(matches!( + &*span, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + assert!(matches!( + &*dataset, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + assert!(matches!( + span.await, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + assert!(matches!( + dataset.await, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + assert!(matches!( + client.shutdown(Duration::from_millis(20)).await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); +} + +#[tokio::test] +async fn cached_login_does_not_move_scheduled_work_off_a_stopped_explicit_runtime() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/apikey/login")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "org_info":[{"id":"org","name":"tenant"}] + }))) + .expect(1) + .mount(&server) + .await; + let uri = server.uri(); + let client = tokio::task::spawn_blocking(move || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + runtime + .block_on( + BraintrustClient::builder() + .runtime(Some(runtime.handle().clone())) + .api_key("token") + .blocking_login(true) + .app_url(&uri) + .api_url(&uri) + .build() + .into_future(), + ) + .unwrap() + // The selected runtime is dropped before returning the client. + }) + .await + .unwrap(); + assert!(client.is_logged_in()); + assert_eq!(client.login_state().org_id().as_deref(), Some("org")); + let wait = client.wait_for_login(); + let span = client.span_builder(); + let dataset = client.dataset_builder(); + assert!(matches!( + wait.await, + Err(BraintrustError::Execution(ExecutionError::WorkerLost)) + )); + assert!(matches!( + span.await, + Err(BraintrustError::Execution(ExecutionError::WorkerLost)) + )); + assert!(matches!( + dataset.await, + Err(BraintrustError::Execution(ExecutionError::WorkerLost)) + )); + assert_eq!(server.received_requests().await.unwrap().len(), 1); +} diff --git a/tests/eager_client_updates.rs b/tests/eager_client_updates.rs new file mode 100644 index 0000000..6bcdd40 --- /dev/null +++ b/tests/eager_client_updates.rs @@ -0,0 +1,420 @@ +use braintrust_sdk_rust::{ + BraintrustClient, BraintrustError, EnqueueError, ExecutionError, FlushError, ParentSpanInfo, + ShutdownError, SpanComponents, SpanLog, SpanObjectType, WriteError, +}; +use serde_json::{json, Value}; +use std::{future::IntoFuture, time::Duration}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use wiremock::{ + matchers::{header, method, path}, + Match, Mock, MockServer, ResponseTemplate, +}; + +fn exported(id: &str) -> String { + SpanComponents { + object_type: SpanObjectType::ProjectLogs, + object_id: Some("project".into()), + compute_object_metadata_args: None, + row_id: Some(id.into()), + span_id: Some(format!("span-{id}")), + root_span_id: Some(format!("root-{id}")), + span_parents: None, + propagated_event: None, + } + .to_str() +} +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, 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() +} + +// Keep login unresolved until the test has issued update, flush and shutdown. +// This async HTTP gate avoids relying on response latency or blocking a runtime. +async fn gated_login( + body: Value, +) -> ( + String, + tokio::sync::oneshot::Receiver<()>, + tokio::sync::oneshot::Sender<()>, + tokio::task::JoinHandle<()>, +) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let uri = format!("http://{}", listener.local_addr().unwrap()); + let (arrived, arrival) = tokio::sync::oneshot::channel(); + let (release, released) = tokio::sync::oneshot::channel(); + let task = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = Vec::new(); + while !request.windows(4).any(|part| part == b"\r\n\r\n") { + assert!(stream.read_buf(&mut request).await.unwrap() > 0); + } + assert!(request.starts_with(b"POST /api/apikey/login ")); + arrived.send(()).unwrap(); + released.await.unwrap(); + let body = body.to_string(); + let response = format!("HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", body.len(), body); + stream.write_all(response.as_bytes()).await.unwrap(); + stream.shutdown().await.unwrap(); + }); + (uri, arrival, release, task) +} + +#[tokio::test] +async fn plain_thread_explicit_updates_progress_unobserved_and_own_exported_context() { + let server = MockServer::start().await; + ingestion(&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 update = client.update_span_with_credentials( + "explicit", + "org", + &exported("row"), + SpanLog::builder() + .output(json!("accepted")) + .build() + .unwrap(), + ); + assert!(update.is_ok()); + assert!(!client.is_logged_in()); + drop(update); + let flush = client.flush(); + let shutdown = client.shutdown(Duration::from_secs(5)); + drop(client); + (flush, shutdown) + }) + .join() + .unwrap(); + assert_eq!(flush.await.unwrap().delivered, 1); + shutdown.await.unwrap(); + let requests = server.received_requests().await.unwrap(); + let request = requests.iter().find(|r| r.url.path() == "/logs3").unwrap(); + assert!(header("authorization", "Bearer explicit").matches(request)); + let body: Value = serde_json::from_slice(&request.body).unwrap(); + let row = &body["rows"][0]; + assert_eq!(row["id"], "row"); + assert_eq!(row["span_id"], "span-row"); + assert_eq!(row["root_span_id"], "root-row"); + assert_eq!(row["org_id"], "org"); + assert_eq!(row["project_id"], "project"); + assert_eq!(row["output"], "accepted"); + assert_eq!(row["_is_merge"], true); + assert!(requests.iter().all(|r| r.url.path() != "/api/apikey/login")); +} + +#[tokio::test] +async fn pending_login_update_captures_shutdown_and_uses_the_selected_org_route() { + let configured = MockServer::start().await; + let selected = MockServer::start().await; + ingestion(&selected).await; + Mock::given(method("POST")) + .and(path("/logs3")) + .and(header("authorization", "Bearer session")) + .respond_with(ResponseTemplate::new(204)) + .with_priority(1) + .expect(1) + .mount(&selected) + .await; + let (control, arrived, release, login) = gated_login(json!({ + "org_info":[{"id":"selected-org","name":"tenant","api_url":selected.uri()}] + })) + .await; + let client = BraintrustClient::builder() + .runtime(braintrust_sdk_rust::current_runtime()) + .api_key("session") + .blocking_login(false) + .app_url(control) + .api_url(configured.uri()) + .build() + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), arrived) + .await + .unwrap() + .unwrap(); + assert!(!client.is_logged_in()); + let update = client.update_span(&exported("pending"), SpanLog::default()); + assert!(update.is_ok()); + let flush = client.flush(); + let shutdown = client.shutdown(Duration::from_secs(5)); + release.send(()).unwrap(); + assert_eq!(shutdown.await.unwrap().delivered, 1); + assert_eq!(flush.await.unwrap().delivered, 1); + update.await.unwrap(); + login.await.unwrap(); + assert!(configured.received_requests().await.unwrap().is_empty()); + let requests = selected.received_requests().await.unwrap(); + let body: Value = serde_json::from_slice( + &requests + .iter() + .find(|r| r.url.path() == "/logs3") + .unwrap() + .body, + ) + .unwrap(); + assert_eq!(body["rows"][0]["org_id"], "selected-org"); + assert_eq!(body["rows"][0]["org_name"], "tenant"); + assert!(matches!( + &*client.update_span(&exported("closed"), SpanLog::default()), + Err(BraintrustError::Admission(EnqueueError::Closed)) + )); +} + +#[tokio::test] +async fn delivery_failure_resolves_update_and_covering_flush_then_fresh_work_succeeds() { + let server = MockServer::start().await; + ingestion(&server).await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(400)) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + let client = client(&server, 10).await; + let update = client.update_span_with_credentials( + "token", + "org", + &exported("failed"), + SpanLog::default(), + ); + let neighbor = client + .span_builder_with_credentials("token", "org") + .parent_info(ParentSpanInfo::ProjectLogs { + object_id: "project".into(), + }) + .build() + .end(); + let flush = client.flush(); + assert!(update.is_ok()); + assert!(matches!( + update.await, + Err(BraintrustError::Api { status: 400, .. }) + )); + assert!(matches!(neighbor.await, Err(WriteError::Delivery(_)))); + assert!(matches!(flush.await, Err(FlushError::Delivery(failure)) if failure.failed==2)); + assert_eq!(client.flush().await.unwrap().delivered, 0); + let fresh = + client.update_span_with_credentials("token", "org", &exported("fresh"), SpanLog::default()); + let flush = client.flush(); + fresh.await.unwrap(); + assert_eq!(flush.await.unwrap().delivered, 1); + client.shutdown(Duration::from_secs(5)).await.unwrap(); +} + +#[tokio::test] +async fn validation_and_queue_overload_report_issuance_failure_without_network_work() { + let server = MockServer::start().await; + ingestion(&server).await; + let client = client(&server, 1).await; + let invalid = + client.update_span_with_credentials("token", "org", "invalid", SpanLog::default()); + assert!(matches!(&*invalid, Err(BraintrustError::InvalidConfig(_)))); + let missing_login = client.update_span(&exported("row"), SpanLog::default()); + assert!(matches!( + &*missing_login, + Err(BraintrustError::InvalidConfig(_)) + )); + let accepted = client.update_span_with_credentials( + "token", + "org", + &exported("accepted"), + SpanLog::default(), + ); + let full = client.update_span_with_credentials( + "token", + "org", + &exported("rejected"), + SpanLog::default(), + ); + assert!(accepted.is_ok()); + assert!(matches!( + &*full, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + assert!(server.received_requests().await.unwrap().is_empty()); + client.flush().await.unwrap(); + accepted.await.unwrap(); + assert!(matches!( + full.await, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + client.shutdown(Duration::from_secs(5)).await.unwrap(); + let closed = client.update_span_with_credentials( + "token", + "org", + &exported("closed"), + SpanLog::default(), + ); + assert!(matches!( + &*closed, + Err(BraintrustError::Admission(EnqueueError::Closed)) + )); + let requests = server.received_requests().await.unwrap(); + let body: Value = serde_json::from_slice( + &requests + .iter() + .find(|r| r.url.path() == "/logs3") + .unwrap() + .body, + ) + .unwrap(); + assert_eq!(body["rows"].as_array().unwrap().len(), 1); + assert_eq!(body["rows"][0]["id"], "accepted"); +} + +#[tokio::test] +async fn full_response_budget_never_admits_an_extra_update() { + let server = MockServer::start().await; + ingestion(&server).await; + let client = client(&server, 2048).await; + for index in 0..1024 { + let update = client.update_span_with_credentials( + "token", + "org", + &exported(&format!("row-{index}")), + SpanLog::default(), + ); + assert!(update.is_ok()); + drop(update); + } + let rejected = client.update_span_with_credentials( + "token", + "org", + &exported("rejected"), + SpanLog::default(), + ); + assert!(matches!( + &*rejected, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + assert_eq!( + client + .shutdown(Duration::from_secs(10)) + .await + .unwrap() + .delivered, + 1024 + ); + let rows: 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!(rows, 1024); +} + +#[tokio::test] +async fn shutdown_timeout_fails_updates_still_waiting_for_login_without_ingestion() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/apikey/login")) + .respond_with(ResponseTemplate::new(200).set_delay(Duration::from_secs(60))) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .runtime(braintrust_sdk_rust::current_runtime()) + .api_key("session") + .blocking_login(false) + .app_url(server.uri()) + .api_url(server.uri()) + .build() + .await + .unwrap(); + let update = client.update_span(&exported("pending"), SpanLog::default()); + assert!(update.is_ok()); + let flush = client.flush(); + assert!(matches!( + client.shutdown(Duration::from_millis(20)).await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + update.await, + Err(BraintrustError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + flush.await, + Err(FlushError::Execution(ExecutionError::TimedOut)) + )); + assert!(server + .received_requests() + .await + .unwrap() + .iter() + .all(|r| r.url.path() != "/logs3")); +} + +#[test] +fn stopped_runtime_retires_a_pending_login_update_that_owns_the_client() { + 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())) + .api_key("session") + .blocking_login(false) + .app_url("http://127.0.0.1:9") + .api_url("http://127.0.0.1:9") + .build() + .into_future(), + ) + .unwrap(); + let update = client.update_span(&exported("pending"), SpanLog::default()); + assert!(update.is_ok()); + drop(client); + // Pending scope resolution retains the whole client, including its + // lifecycle coordinator. Cancellation must release it without locking + // the delivery state again from inside command destruction. + drop(runtime); + send.send(update).unwrap(); + }); + let update = receive + .recv_timeout(Duration::from_secs(5)) + .expect("runtime loss must not deadlock while dropping an unresolved update"); + assert!(matches!( + futures::executor::block_on(update.into_future()), + Err(BraintrustError::Execution(ExecutionError::WorkerLost)) + )); + thread.join().unwrap(); +} diff --git a/tests/eager_datasets.rs b/tests/eager_datasets.rs new file mode 100644 index 0000000..8d8712b --- /dev/null +++ b/tests/eager_datasets.rs @@ -0,0 +1,241 @@ +//! Eager dataset operations own admission through all of their subrequests. +use braintrust_sdk_rust::{ + BraintrustClient, BraintrustError, Dataset, DatasetInsert, EnqueueError, +}; +use serde_json::json; +use std::{future::IntoFuture, sync::Arc, time::Duration}; +use wiremock::{ + matchers::{body_partial_json, method, path}, + Mock, MockServer, ResponseTemplate, +}; + +fn dataset(client: &BraintrustClient) -> Dataset { + client + .dataset_builder_with_credentials("token", "tenant") + .project_name("project") + .dataset_name("dataset") + .build() + .unwrap() +} +fn record() -> DatasetInsert { + DatasetInsert::builder().input(json!(1)).build().unwrap() +} +async fn registration(server: &MockServer) { + Mock::given(method("POST")) + .and(path("/api/dataset/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "project":{"id":"project","name":"project"}, + "dataset":{"id":"dataset","name":"dataset"} + }))) + .expect(1) + .mount(server) + .await; +} + +#[tokio::test] +async fn synchronous_callers_issue_unobserved_operations_and_shutdown_drains_them() { + let server = MockServer::start().await; + registration(&server).await; + let arrived = Arc::new(tokio::sync::Notify::new()); + let observed = arrived.clone(); + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(move |_: &wiremock::Request| { + observed.notify_one(); + ResponseTemplate::new(204) + }) + .expect(3) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/dataset-summary")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "project_name":"project", "dataset_name":"dataset", + "project_id":"project", "dataset_id":"dataset" + }))) + .expect(1) + .mount(&server) + .await; + let uri = server.uri(); + 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 dataset = dataset(&client); + for operation in [dataset.update("old", record()), dataset.delete("old")] { + assert!(operation.is_ok()); + drop(operation); + } + let insert = dataset.insert(record()); + assert!(insert.is_ok()); + drop(insert); + drop(dataset.summarize()); + let shutdown = client.shutdown(Duration::from_secs(5)); + drop(dataset); + drop(client); + shutdown + }) + .join() + .unwrap(); + // Network progress occurs before observing any accepted operation token. + tokio::time::timeout(Duration::from_secs(5), arrived.notified()) + .await + .unwrap(); + assert_eq!(shutdown.await.unwrap().delivered, 0); +} + +#[tokio::test] +async fn saturated_operations_drain_after_shutdown_without_nested_http_admission() { + let server = MockServer::start().await; + registration(&server).await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .runtime(braintrust_sdk_rust::current_runtime()) + .skip_login(true) + .app_url(server.uri()) + .api_url(server.uri()) + .build() + .await + .unwrap(); + let dataset = dataset(&client); + // No selected-runtime polling between these issuance calls and shutdown. + let ids: Vec<_> = (0..1023).map(|_| dataset.dataset_id()).collect(); + assert!(ids.iter().all(|id| id.is_ok())); + let insert = dataset.insert(record()); + assert!(insert.is_ok()); + let full = dataset.fetch(None); + assert!(matches!( + &*full, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + let shutdown = client.shutdown(Duration::from_secs(5)); + let closed = dataset.dataset_id(); + assert!(matches!( + &*closed, + Err(BraintrustError::Admission(EnqueueError::Closed)) + )); + shutdown.await.unwrap(); + for id in ids { + assert_eq!(id.await.unwrap(), "dataset"); + } + assert!(!insert.await.unwrap().is_empty()); + assert!(matches!( + full.await, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + assert!(matches!( + closed.await, + Err(BraintrustError::Admission(EnqueueError::Closed)) + )); + // One shared registration followed by the accepted write, even though + // public admission was closed before either request started. + assert_eq!(server.received_requests().await.unwrap().len(), 2); +} + +#[tokio::test] +async fn unobserved_cursor_commands_consume_records_in_order_and_collect_across_shutdown() { + let server = MockServer::start().await; + registration(&server).await; + Mock::given(method("POST")) + .and(path("/btql")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "data":[{"id":"first"}], "cursor":"next" + }))) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/btql")) + .and(body_partial_json(json!({"query":{"cursor":"next"}}))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "data":[{"id":"second"},{"id":"third"}] + }))) + .with_priority(1) + .expect(1) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .runtime(braintrust_sdk_rust::current_runtime()) + .skip_login(true) + .app_url(server.uri()) + .api_url(server.uri()) + .build() + .await + .unwrap(); + let dataset = dataset(&client); + let mut cursor = dataset.fetch(Some(1)).await.unwrap(); + drop(cursor.next()); + let next = cursor.next(); + let rest = cursor.collect(); + assert!(next.is_ok()); + assert!(rest.is_ok()); + let shutdown = client.shutdown(Duration::from_secs(5)); + drop(dataset); + drop(client); + shutdown.await.unwrap(); + assert_eq!(next.await.unwrap().unwrap().id, "second"); + assert_eq!( + rest.await + .unwrap() + .iter() + .map(|r| r.id.as_str()) + .collect::>(), + ["third"] + ); +} + +#[tokio::test] +async fn rejected_cursor_admission_preserves_its_position_and_allows_retry() { + let server = MockServer::start().await; + registration(&server).await; + Mock::given(method("POST")) + .and(path("/btql")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "data":[{"id":"first"},{"id":"second"}] + }))) + .expect(1) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .runtime(braintrust_sdk_rust::current_runtime()) + .skip_login(true) + .app_url(server.uri()) + .api_url(server.uri()) + .build() + .await + .unwrap(); + let dataset = dataset(&client); + let mut cursor = dataset.fetch(None).await.unwrap(); + let first = cursor.next(); + let ids: Vec<_> = (0..1023).map(|_| dataset.dataset_id()).collect(); + let rejected = cursor.next(); + assert!(matches!( + &*rejected, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + for id in ids { + id.await.unwrap(); + } + assert_eq!(first.await.unwrap().unwrap().id, "first"); + assert!(matches!( + rejected.await, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + let retry = cursor.next(); + assert!(retry.is_ok()); + assert_eq!(retry.await.unwrap().unwrap().id, "second"); + assert!(cursor.next().await.unwrap().is_none()); + client.shutdown(Duration::from_secs(5)).await.unwrap(); +} diff --git a/tests/eager_experiment_events.rs b/tests/eager_experiment_events.rs new file mode 100644 index 0000000..921c206 --- /dev/null +++ b/tests/eager_experiment_events.rs @@ -0,0 +1,579 @@ +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)) + )); + let feedback = experiment.log_feedback( + "row", + braintrust_sdk_rust::Feedback::builder().build().unwrap(), + ); + let update = experiment.update_span("row", braintrust_sdk_rust::SpanLog::default()); + for patch in [feedback, update] { + assert!(matches!( + &*patch, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + assert!(matches!( + patch.await, + 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 feedback = experiment.log_feedback( + "row", + braintrust_sdk_rust::Feedback::builder().build().unwrap(), + ); + let update = experiment.update_span("row", braintrust_sdk_rust::SpanLog::default()); + for patch in [feedback, update] { + assert!(matches!( + &*patch, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + assert!(matches!( + patch.await, + 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()); + let feedback = experiment.log_feedback( + "feedback", + braintrust_sdk_rust::Feedback::builder().build().unwrap(), + ); + let update = experiment.update_span("update", braintrust_sdk_rust::SpanLog::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)) + )); + for patch in [feedback, update] { + assert!(matches!( + patch.await, + Err(BraintrustError::Execution(ExecutionError::TimedOut)) + )); + } + 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(); +} + +#[tokio::test] +async fn unobserved_patches_own_their_ids_and_share_registration_across_shutdown() { + use braintrust_sdk_rust::{Feedback, SpanLog}; + 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 feedback = experiment.log_feedback( + &String::from("feedback"), + Feedback::builder() + .expected(json!("correct")) + .score("quality", 1.0) + .metadata_entry("source_detail", "test") + .comment("reviewed") + .source("external") + .tag("checked") + .build() + .unwrap(), + ); + let update = experiment.clone().update_span( + &String::from("update"), + SpanLog::builder() + .input(json!("input")) + .output(json!("output")) + .expected(json!("expected")) + .error(json!("error")) + .score("quality", 0.5) + .metadata_entry("extra", "value") + .metric("tokens", 12.0) + .tag("patched") + .name("updated span") + .build() + .unwrap(), + ); + assert!(feedback.is_ok()); + assert!(update.is_ok()); + drop(feedback); + drop(update); + let flush = experiment.flush(); + let shutdown = client.shutdown(Duration::from_secs(5)); + let closed = experiment.log_feedback("after", Feedback::builder().build().unwrap()); + assert!(matches!( + &*closed, + Err(BraintrustError::Admission(EnqueueError::Closed)) + )); + assert!(matches!( + &*experiment.update_span("after", SpanLog::default()), + Err(BraintrustError::Admission(EnqueueError::Closed)) + )); + drop(experiment); + drop(client); + assert_eq!(flush.await.unwrap().delivered, 2); + assert_eq!(shutdown.await.unwrap().delivered, 2); + 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(), 2); + for row in &rows { + assert_eq!(row["experiment_id"], "experiment"); + assert_eq!(row["_is_merge"], true); + } + let feedback = rows.iter().find(|row| row["id"] == "feedback").unwrap(); + assert_eq!(feedback["expected"], "correct"); + assert_eq!(feedback["scores"]["quality"], 1.0); + assert_eq!( + feedback["metadata"], + json!({"source_detail":"test","comment":"reviewed","source":"external"}) + ); + assert_eq!(feedback["tags"], json!(["checked"])); + let update = rows.iter().find(|row| row["id"] == "update").unwrap(); + assert_eq!(update["input"], "input"); + assert_eq!(update["output"], "output"); + assert_eq!(update["expected"], "expected"); + assert_eq!(update["error"], "error"); + assert_eq!(update["scores"]["quality"], 0.5); + assert_eq!(update["metadata"]["extra"], "value"); + assert_eq!(update["metrics"]["tokens"], 12.0); + assert_eq!(update["tags"], json!(["patched"])); + assert_eq!(update["span_attributes"]["name"], "updated span"); +} + +#[tokio::test] +async fn patch_registration_failure_retires_the_batch_and_allows_fresh_retry() { + use braintrust_sdk_rust::{Feedback, SpanLog}; + let server = MockServer::start().await; + ingestion(&server).await; + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(ResponseTemplate::new(400)) + .with_priority(1) + .up_to_n_times(2) + .expect(2) + .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 feedback = experiment.log_feedback("row", Feedback::builder().build().unwrap()); + let update = experiment.update_span("row", SpanLog::default()); + let flush = experiment.flush(); + for patch in [feedback, update] { + assert!(patch.is_ok()); + assert!(matches!( + patch.await, + Err(BraintrustError::Api { status: 400, .. }) + )); + } + assert!(matches!(flush.await, Err(FlushError::Delivery(failure)) if failure.failed==2)); + assert_eq!(experiment.flush().await.unwrap().delivered, 0); + assert_eq!(experiment.experiment_id(), None); + let retry = experiment.update_span( + "row", + SpanLog::builder().output(json!("retry")).build().unwrap(), + ); + let fresh = experiment.flush(); + retry.await.unwrap(); + assert_eq!(fresh.await.unwrap().delivered, 1); + let requests = server.received_requests().await.unwrap(); + assert_eq!( + requests.iter().filter(|r| r.url.path() == "/logs3").count(), + 1 + ); + client.shutdown(Duration::from_secs(5)).await.unwrap(); +} + +#[tokio::test] +async fn explicit_patches_remain_independent_of_a_finished_live_span() { + use braintrust_sdk_rust::{Feedback, SpanLog}; + 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 span = experiment.start_span().await.unwrap().build(); + let end = span.end(); + let update = experiment.update_span( + span.row_id(), + SpanLog::builder() + .output(json!("after finish")) + .build() + .unwrap(), + ); + let feedback = experiment.log_feedback( + span.row_id(), + Feedback::builder().score("review", 1.0).build().unwrap(), + ); + let flush = experiment.flush(); + assert!(update.is_ok()); + assert!(feedback.is_ok()); + end.await.unwrap(); + update.await.unwrap(); + feedback.await.unwrap(); + assert_eq!(flush.await.unwrap().delivered, 3); + 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"].as_array().unwrap().len(), 1); + let row = &body["rows"][0]; + assert_eq!(row["id"], span.row_id()); + assert_eq!(row["output"], "after finish"); + assert_eq!(row["scores"]["review"], 1.0); + assert!(row["metrics"]["end"].is_number()); + client.shutdown(Duration::from_secs(5)).await.unwrap(); +} diff --git a/tests/eager_experiment_reads.rs b/tests/eager_experiment_reads.rs new file mode 100644 index 0000000..e361a23 --- /dev/null +++ b/tests/eager_experiment_reads.rs @@ -0,0 +1,339 @@ +use braintrust_sdk_rust::{ + BraintrustClient, BraintrustError, EnqueueError, ExecutionError, Experiment, ExperimentLog, + 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 endpoints(server: &MockServer) { + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(registration()) + .expect(1) + .mount(server) + .await; + Mock::given(method("POST")) + .and(path("/api/base_experiment/get_id")) + .respond_with(ResponseTemplate::new(400)) + .mount(server) + .await; + Mock::given(method("GET")) + .and(path("/experiment-comparison2")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({}))) + .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 plain_thread_reads_complete_unobserved_through_shutdown_after_owner_drop() { + let server = MockServer::start().await; + endpoints(&server).await; + let uri = server.uri(); + let (project, exported, base, summary, span, 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 project = experiment.project(); + let exported = experiment.export(); + let base = experiment.fetch_base_experiment(); + let summary = experiment.summarize(); + let span = experiment.start_span(); + assert!(project.is_ok()); + assert!(exported.is_ok()); + assert!(base.is_ok()); + assert!(summary.is_ok()); + assert!(span.is_ok()); + let shutdown = client.shutdown(Duration::from_secs(5)); + drop(experiment); + drop(client); + (project, exported, base, summary, span, shutdown) + }) + .join() + .unwrap(); + shutdown.await.unwrap(); + assert_eq!(project.await.unwrap().id, "project"); + assert_eq!(exported.await.unwrap().object_id, "experiment"); + assert!(base.await.unwrap().is_none()); + assert_eq!(summary.await.unwrap().experiment_id(), Some("experiment")); + // An accepted builder request may finish during shutdown, but its returned + // live handle does not acquire permission to submit new writes afterward. + assert!(matches!( + &*span.await.unwrap().build().end(), + Err(WriteError::Admission(EnqueueError::Closed)) + )); + let requests = server.received_requests().await.unwrap(); + assert_eq!( + requests + .iter() + .filter(|r| r.url.path() == "/api/base_experiment/get_id") + .count(), + 2 + ); + assert_eq!( + requests + .iter() + .filter(|r| r.url.path() == "/experiment-comparison2") + .count(), + 1 + ); + assert!(requests.iter().all(|r| r.url.path() != "/logs3")); +} + +#[tokio::test] +async fn full_task_budget_allows_summary_subrequests_and_rejected_starts_do_not_select_mode() { + let server = MockServer::start().await; + endpoints(&server).await; + let client = client(&server, 10).await; + let experiment = experiment(&client); + let projects: Vec<_> = (0..1023).map(|_| experiment.project()).collect(); + assert!(projects.iter().all(|project| project.is_ok())); + let summary = experiment.summarize(); + assert!(summary.is_ok()); + let rejected = experiment.start_span(); + assert!(matches!( + &*rejected, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + assert!(matches!( + &*experiment.fetch_base_experiment(), + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + summary.await.unwrap(); + for project in projects { + assert_eq!(project.await.unwrap().id, "project"); + } + assert!(matches!( + rejected.await, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + let event = experiment.log(ExperimentLog::default()); + assert!(event.is_ok()); // The rejected start left event mode available. + let flush = experiment.flush(); + event.await.unwrap(); + assert_eq!(flush.await.unwrap().delivered, 1); + client.shutdown(Duration::from_secs(5)).await.unwrap(); +} + +#[tokio::test] +async fn summary_captures_its_cutoff_before_a_later_write_fails() { + let server = MockServer::start().await; + endpoints(&server).await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(204)) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(400)) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + let client = client(&server, 10).await; + let experiment = experiment(&client); + let prior = experiment.log(ExperimentLog::default()); + let summary = experiment.summarize(); + let later = experiment.log(ExperimentLog::default()); + assert!(summary.is_ok()); + prior.await.unwrap(); + assert!(matches!( + later.await, + Err(BraintrustError::Api { status: 400, .. }) + )); + assert_eq!(summary.await.unwrap().experiment_id(), Some("experiment")); + assert_eq!(experiment.flush().await.unwrap().delivered, 0); + client.shutdown(Duration::from_secs(5)).await.unwrap(); +} + +#[tokio::test] +async fn pending_span_starts_block_events_and_failed_registration_releases_mode() { + let server = MockServer::start().await; + endpoints(&server).await; + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(ResponseTemplate::new(400)) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + let client = client(&server, 10).await; + let experiment = experiment(&client); + let start = experiment.start_span(); + assert!(start.is_ok()); + let blocked = experiment.clone().log(ExperimentLog::default()); + assert!(matches!(&*blocked, Err(BraintrustError::InvalidConfig(_)))); + assert!(matches!( + start.await, + Err(BraintrustError::Api { status: 400, .. }) + )); + let event = experiment.log(ExperimentLog::default()); + assert!(event.is_ok()); + event.await.unwrap(); + let accepted = experiment.start_span(); + assert!(matches!( + &*experiment.log(ExperimentLog::default()), + Err(BraintrustError::InvalidConfig(_)) + )); + drop(accepted.await.unwrap()); + assert!(matches!( + &*experiment.clone().log(ExperimentLog::default()), + Err(BraintrustError::InvalidConfig(_)) + )); + client.shutdown(Duration::from_secs(5)).await.unwrap(); +} + +#[tokio::test] +async fn deadline_resolves_all_accepted_read_and_builder_tokens_without_nested_admission() { + 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 project = experiment.project(); + let exported = experiment.export(); + let base = experiment.fetch_base_experiment(); + let summary = experiment.summarize(); + let span = experiment.start_span(); + 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!(matches!( + project.await, + Err(BraintrustError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + exported.await, + Err(BraintrustError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + base.await, + Err(BraintrustError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + summary.await, + Err(BraintrustError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + span.await, + Err(BraintrustError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + &*experiment.project(), + Err(BraintrustError::Admission(EnqueueError::Closed)) + )); + assert!(matches!( + &*experiment.export(), + Err(BraintrustError::Admission(EnqueueError::Closed)) + )); + assert!(matches!( + &*experiment.fetch_base_experiment(), + Err(BraintrustError::Admission(EnqueueError::Closed)) + )); + assert!(matches!( + &*experiment.summarize(), + Err(BraintrustError::Admission(EnqueueError::Closed)) + )); + assert!(matches!( + &*experiment.start_span(), + Err(BraintrustError::Admission(EnqueueError::Closed)) + )); + assert_eq!(server.received_requests().await.unwrap().len(), 1); +} + +#[tokio::test] +async fn summary_flush_rejection_reports_issuance_failure_without_starting_http_work() { + let server = MockServer::start().await; + endpoints(&server).await; + let client = client(&server, 1).await; + let experiment = experiment(&client); + let write = client + .span_builder_with_credentials("token", "org") + .parent_info(ParentSpanInfo::ProjectLogs { + object_id: "project".into(), + }) + .build() + .end(); + let occupied = client.flush(); + let rejected = experiment.summarize(); + assert!(matches!( + &*rejected, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + occupied.await.unwrap(); + write.await.unwrap(); + assert!(matches!( + rejected.await, + Err(BraintrustError::Admission(EnqueueError::Full)) + )); + assert!(server + .received_requests() + .await + .unwrap() + .iter() + .all(|r| r.url.path() != "/api/experiment/register")); + experiment.summarize().await.unwrap(); + client.shutdown(Duration::from_secs(5)).await.unwrap(); +} 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 9de77b4..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, @@ -456,7 +457,7 @@ async fn explicit_experiment_credentials_never_prime_the_shared_session() { .mount(&server) .await; let client = isolated_client(&server).await; - let state = client.login_state().await; + let state = client.login_state(); let experiment = client .experiment_builder_with_credentials("private-token", "tenant") .project_name("project") @@ -465,7 +466,7 @@ async fn explicit_experiment_credentials_never_prime_the_shared_session() { assert!(!state.is_logged_in()); experiment.project().await.unwrap(); assert!(!state.is_logged_in()); - assert!(!client.is_logged_in().await); + assert!(!client.is_logged_in()); assert!(client.dataset_builder().await.is_err()); let exported = experiment.export().await.unwrap(); assert_eq!(exported.object_id, "experiment"); diff --git a/tests/finished_spans.rs b/tests/finished_spans.rs index 90221ec..4c9fbb1 100644 --- a/tests/finished_spans.rs +++ b/tests/finished_spans.rs @@ -134,6 +134,7 @@ async fn finishing_a_parent_does_not_finish_its_child_and_explicit_updates_remai .build() .unwrap(), ) + .await .unwrap(); client.flush().await.unwrap(); let bodies: Vec = server diff --git a/tests/response_lifecycle.rs b/tests/response_lifecycle.rs index 0c86885..338dc21 100644 --- a/tests/response_lifecycle.rs +++ b/tests/response_lifecycle.rs @@ -126,6 +126,6 @@ async fn background_login_is_accepted_work_covered_by_shutdown() { client.shutdown(Duration::from_millis(20)).await, Err(ShutdownError::Execution(ExecutionError::TimedOut)) )); - assert!(!client.is_logged_in().await); + assert!(!client.is_logged_in()); assert_eq!(server.received_requests().await.unwrap().len(), 1); } diff --git a/tests/span_lifecycle.rs b/tests/span_lifecycle.rs index b7e876e..4dc3f94 100644 --- a/tests/span_lifecycle.rs +++ b/tests/span_lifecycle.rs @@ -109,6 +109,7 @@ async fn client_update_span_uses_exported_ids_for_project_logs() { .build() .expect("build"), ) + .await .expect("update"); client.flush().await.expect("flush"); @@ -183,6 +184,7 @@ async fn client_update_span_with_credentials_works_without_priming_login_state() .build() .expect("build"), ) + .await .expect("update"); client.flush().await.expect("flush"); @@ -268,6 +270,7 @@ async fn client_update_span_skip_span_origin_preserves_only_explicit_context() { &exported, event.build().expect("build"), ) + .await .expect("update"); client.flush().await.expect("flush"); } @@ -406,6 +409,7 @@ async fn client_update_span_includes_exported_span_parents() { .build() .expect("build"), ) + .await .expect("update"); client.flush().await.expect("flush"); @@ -466,6 +470,7 @@ async fn client_update_span_with_credentials_includes_exported_span_parents() { .build() .expect("build"), ) + .await .expect("update"); client.flush().await.expect("flush"); @@ -559,6 +564,7 @@ async fn client_update_span_with_credentials_inherits_exported_propagated_event( .build() .expect("build"), ) + .await .expect("update"); client.flush().await.expect("flush"); } @@ -773,6 +779,7 @@ async fn client_update_span_resolves_project_name_from_exported_compute_metadata .build() .expect("build"), ) + .await .expect("update"); client.flush().await.expect("flush");