diff --git a/src/experiments/experiment.rs b/src/experiments/experiment.rs index 08906cd..de704a8 100644 --- a/src/experiments/experiment.rs +++ b/src/experiments/experiment.rs @@ -18,7 +18,7 @@ use crate::experiments::metadata::{ use crate::experiments::span_builder::ExperimentSpanBuilder; use crate::experiments::summary::{ExperimentSummary, MetricSummary, ScoreSummary}; use crate::internal::execution::ExecutionBinding; -use crate::internal::ports::{ExperimentApi, SpanSubmitter}; +use crate::internal::ports::{ExperimentApi, PreparedSpan, SpanSubmitter}; use crate::internal::session::CredentialScope; use crate::span::{SpanBuilder, SpanHandle, SpanLog}; use crate::types::{ParentSpanInfo, SpanAttributes, SpanPayload, SpanType}; @@ -457,13 +457,18 @@ impl Experiment { ) } else { let experiment = self.clone(); - self.submitter.submit_resolving_parent( - self.scope.clone(), - payload, + let scope = self.scope.clone(); + self.submitter.submit_deferred( + payload.row_id.clone(), Box::pin(async move { let metadata = experiment.register_metadata().await?; - Ok(ParentSpanInfo::Experiment { + let parent_info = Some(ParentSpanInfo::Experiment { object_id: metadata.experiment_id.clone(), + }); + Ok(PreparedSpan { + scope, + payload, + parent_info, }) }), ) diff --git a/src/internal/ports.rs b/src/internal/ports.rs index e952669..6122c8f 100644 --- a/src/internal/ports.rs +++ b/src/internal/ports.rs @@ -16,6 +16,14 @@ 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, +} + /// Shared delivery lifecycle, independent of resource HTTP operations. #[async_trait] pub(crate) trait DeliveryControl: Send + Sync { @@ -25,13 +33,12 @@ pub(crate) trait DeliveryControl: Send + Sync { #[async_trait] pub(crate) trait SpanSubmitter: DeliveryControl { - /// Queue the write now, resolving its destination during preparation. - /// Rejected work must drop the resolver without polling it. - fn submit_resolving_parent( + /// Queue the write now and resolve its immutable scope/destination during + /// preparation. Rejected work drops the future without polling it. + fn submit_deferred( &self, - scope: CredentialScope, - payload: SpanPayload, - resolve: futures::future::BoxFuture<'static, Result>, + row_id: String, + prepare: futures::future::BoxFuture<'static, Result>, ) -> crate::Completion; fn submit( diff --git a/src/log_queue/queue.rs b/src/log_queue/queue.rs index 3c1d566..44ba6b7 100644 --- a/src/log_queue/queue.rs +++ b/src/log_queue/queue.rs @@ -8,7 +8,7 @@ use super::row_key::RowKey; use crate::error::{BraintrustError, Result}; use crate::internal::delivery::WriteService; use crate::internal::execution::ExecutionBinding; -use crate::internal::ports::{DeliveryControl, SpanSubmitter}; +use crate::internal::ports::{DeliveryControl, PreparedSpan, SpanSubmitter}; use crate::internal::session::CredentialScope; use crate::internal::shared_outcome::WriteOutcome; use crate::span_components::{project_logs_identifier, ProjectLogsIdentifier}; @@ -40,14 +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: ParentResolution, -} -enum ParentResolution { - Ready(Option), - Deferred(BoxFuture<'static, Result>), +enum SubmitCommand { + Ready(Box), + Deferred { + row_id: String, + prepare: BoxFuture<'static, Result>, + }, } /// Preparation and protocol orchestration, independent of admission/lifecycle. @@ -112,21 +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; let prepared = async { - let parent = match parent_info { - ParentResolution::Ready(parent) => parent, - ParentResolution::Deferred(resolve) => Some(resolve.await?), + let PreparedSpan { + scope, + payload, + parent_info, + } = match cmd { + SubmitCommand::Ready(span) => *span, + SubmitCommand::Deferred { prepare, .. } => prepare.await?, }; - self.prepare_row(&scope, payload, parent).await + let row = self.prepare_row(&scope, payload, parent_info).await?; + Ok::<_, anyhow::Error>((scope, row)) } .await; match prepared { - Ok(row) => { + Ok((scope, row)) => { groups.entry(scope).or_default().push(row); } Err(e) => { @@ -512,16 +510,18 @@ impl LogQueueCore { return; } self.execution.spawn(async move { - let SubmitCommand { + let PreparedSpan { scope: _, payload, parent_info, - } = cmd; - - let ParentResolution::Ready(parent_info) = parent_info else { - // Rejection must not trigger registration just for a debug dump. - warn!(row_id = %payload.row_id, "cannot dump rejected row with unresolved destination"); - return; + } = 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(); @@ -676,11 +676,11 @@ impl LogQueue { outcome: Option, ) -> Completion { self.submit_command( - SubmitCommand { + SubmitCommand::Ready(Box::new(PreparedSpan { scope, payload, - parent_info: ParentResolution::Ready(parent_info), - }, + parent_info, + })), outcome, ) } @@ -802,20 +802,12 @@ impl DeliveryControl for LogQueue { } #[async_trait::async_trait] impl SpanSubmitter for LogQueue { - fn submit_resolving_parent( + fn submit_deferred( &self, - scope: CredentialScope, - payload: SpanPayload, - resolve: BoxFuture<'static, Result>, + row_id: String, + prepare: BoxFuture<'static, Result>, ) -> Completion { - self.submit_command( - SubmitCommand { - scope, - payload, - parent_info: ParentResolution::Deferred(resolve), - }, - None, - ) + self.submit_command(SubmitCommand::Deferred { row_id, prepare }, None) } fn submit( diff --git a/src/logger.rs b/src/logger.rs index d9eef1f..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; @@ -477,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, @@ -686,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()) @@ -798,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. diff --git a/src/span.rs b/src/span.rs index d54f8b5..e347103 100644 --- a/src/span.rs +++ b/src/span.rs @@ -1104,11 +1104,13 @@ mod tests { } #[async_trait] impl SpanSubmitter for Submitter { - fn submit_resolving_parent( + fn submit_deferred( &self, - _: CredentialScope, - _: SpanPayload, - _: futures::future::BoxFuture<'static, Result>, + _: String, + _: futures::future::BoxFuture< + 'static, + Result, + >, ) -> crate::Completion { unreachable!("this test submitter only accepts ready span snapshots") } @@ -1188,11 +1190,13 @@ mod tests { } #[async_trait] impl SpanSubmitter for Submitter { - fn submit_resolving_parent( + fn submit_deferred( &self, - _: CredentialScope, - _: SpanPayload, - _: futures::future::BoxFuture<'static, Result>, + _: String, + _: futures::future::BoxFuture< + 'static, + Result, + >, ) -> crate::Completion { unreachable!("this test submitter only accepts ready span snapshots") } diff --git a/src/test_utils.rs b/src/test_utils.rs index 21909b9..b936cae 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,16 +60,19 @@ impl MockBraintrustClient { #[async_trait] impl SpanSubmitter for MockBraintrustClient { - fn submit_resolving_parent( + fn submit_deferred( &self, - scope: CredentialScope, - payload: SpanPayload, - resolve: futures::future::BoxFuture<'static, Result>, + _: 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 parent = resolve.await.map_err(|source| { + let PreparedSpan { + scope, + payload, + parent_info, + } = prepare.await.map_err(|source| { crate::WriteError::Delivery(Arc::new(crate::DeliveryError { source: source.into(), })) @@ -77,7 +80,7 @@ impl SpanSubmitter for MockBraintrustClient { collector.push(CapturedSpan { token: scope.token().to_owned(), payload, - parent: Some(parent), + parent: parent_info, }); Ok(crate::DeliveryReceipt) }) 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/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/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");