diff --git a/src/error.rs b/src/error.rs index fa2bab7..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::() diff --git a/src/experiments/experiment.rs b/src/experiments/experiment.rs index 93ef4e4..08906cd 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}; @@ -34,6 +34,8 @@ 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, @@ -56,7 +58,27 @@ 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 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 { @@ -90,7 +112,11 @@ impl Experiment { /// does not cancel the write; awaiting includes full batch delivery. pub fn log(&self, event: ExperimentLog) -> crate::Completion { self.execution.response_try(|| { - if self.state.called_start_span.load(Ordering::Acquire) { + 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(), )); @@ -126,12 +152,23 @@ impl Experiment { /// 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(|| { + self.state.recording().pending_starts += 1; + let pending = PendingSpanStart(self.state.clone()); + 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. @@ -193,14 +230,29 @@ impl Experiment { /// /// 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 @@ -294,8 +346,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(), @@ -419,8 +477,16 @@ impl Experiment { /// 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 @@ -430,25 +496,20 @@ 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 { - if let Some(metadata) = self.state.metadata.get() { - return Ok(metadata.clone()); - } - let experiment = self.clone(); - self.execution - .response(async move { experiment.register_metadata().await.cloned() }) - .await - } - /// Polled only on the selected context, by preparation or a bound operation. async fn register_metadata(&self) -> Result<&ExperimentMetadata> { self.state diff --git a/src/internal/ports.rs b/src/internal/ports.rs index 2d4702a..e952669 100644 --- a/src/internal/ports.rs +++ b/src/internal/ports.rs @@ -65,6 +65,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 9b2fa89..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) @@ -57,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( @@ -78,17 +72,11 @@ 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) } } diff --git a/src/logger.rs b/src/logger.rs index 192fff8..d9eef1f 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -505,10 +505,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 { 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(); +}