From c04c7b5ae279beb202f0cb7657797bdefac7c12c Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Tue, 20 Oct 2026 20:45:17 +0800 Subject: [PATCH] Unify SDK ingestion upload orchestration --- src/api/client.rs | 49 +++- src/http.rs | 2 + src/internal/mod.rs | 2 + src/internal/object_store.rs | 84 ++++++ src/internal/retry.rs | 77 +++++- src/internal/upload.rs | 514 +++++++++++++++++++++++++++++++++++ src/log_queue/http.rs | 425 +++-------------------------- src/logs3.rs | 362 ++---------------------- src/types.rs | 5 +- 9 files changed, 776 insertions(+), 744 deletions(-) create mode 100644 src/internal/object_store.rs create mode 100644 src/internal/upload.rs diff --git a/src/api/client.rs b/src/api/client.rs index 2b59a28..1d088d8 100644 --- a/src/api/client.rs +++ b/src/api/client.rs @@ -34,18 +34,7 @@ impl Client { if config.api_key.trim().is_empty() { return Err(Error::Config("api_key must not be empty".into())); } - let mut headers = HeaderMap::new(); - let mut token = HeaderValue::from_str(&format!("Bearer {}", config.api_key)) - .map_err(|_| Error::Config("invalid api_key header".into()))?; - token.set_sensitive(true); - headers.insert(AUTHORIZATION, token); - if let Some(org) = config.org_name { - headers.insert( - "x-bt-org-name", - HeaderValue::from_str(&org) - .map_err(|_| Error::Config("invalid organization header".into()))?, - ); - } + let headers = scoped_headers(&config.api_key, config.org_name.as_deref())?; let ca = match (config.custom_ca_bundle_path, config.custom_ca_bundle_pem) { (Some(_), Some(_)) => { return Err(Error::Config( @@ -72,6 +61,26 @@ impl Client { }), }) } + /// Reuse the SDK's configured transport without copying session state into + /// the API layer. This scope is immutable and used only for ingestion. + pub(crate) fn for_ingestion( + http: reqwest::Client, + api_url: &Url, + token: &str, + org_name: Option<&str>, + ) -> Result { + let api_url = base_url(api_url.as_str())?; + let headers = scoped_headers(token, org_name)?; + Ok(Self { + inner: Arc::new(Inner { + http, + app_url: api_url.clone(), + api_url, + headers, + }), + }) + } + pub fn logs3(&self) -> super::logs3::Logs3 { super::logs3::Logs3 { client: self.clone(), @@ -199,6 +208,22 @@ impl Client { Ok(()) } } +fn scoped_headers(token: &str, org_name: Option<&str>) -> Result { + let mut headers = HeaderMap::new(); + let mut token = HeaderValue::from_str(&format!("Bearer {token}")) + .map_err(|_| Error::Config("invalid api_key header".into()))?; + token.set_sensitive(true); + headers.insert(AUTHORIZATION, token); + if let Some(org) = org_name { + headers.insert( + "x-bt-org-name", + HeaderValue::from_str(org) + .map_err(|_| Error::Config("invalid organization header".into()))?, + ); + } + Ok(headers) +} + fn base_url(value: &str) -> Result { let mut url = Url::parse(value).map_err(|_| Error::Config("invalid endpoint URL".into()))?; if !matches!(url.scheme(), "http" | "https") diff --git a/src/http.rs b/src/http.rs index 9af9589..ec76f9a 100644 --- a/src/http.rs +++ b/src/http.rs @@ -69,6 +69,8 @@ pub(crate) fn build_http_client_builder( user_agent: Option<&str>, ) -> Result { let mut builder = Client::builder() + // SDK retries own request repetition; redirects must not add attempts. + .redirect(reqwest::redirect::Policy::none()) .timeout(timeout) .user_agent(user_agent.unwrap_or(USER_AGENT)); diff --git a/src/internal/mod.rs b/src/internal/mod.rs index f24492a..0b285bc 100644 --- a/src/internal/mod.rs +++ b/src/internal/mod.rs @@ -1,3 +1,5 @@ pub(crate) mod batching; pub(crate) mod ingest; +pub(crate) mod object_store; pub(crate) mod retry; +pub(crate) mod upload; diff --git a/src/internal/object_store.rs b/src/internal/object_store.rs new file mode 100644 index 0000000..51ff4af --- /dev/null +++ b/src/internal/object_store.rs @@ -0,0 +1,84 @@ +//! One-attempt signed object uploads, independent of Braintrust credentials. +use crate::api::logs3::Logs3OverflowUpload; +use crate::error::{BraintrustError, Result}; +use bytes::Bytes; +use url::Url; + +/// Only signed grant information reaches object storage. Braintrust credentials +/// are scoped to the API client, never added by this transport. +pub(crate) struct ObjectStoreTransport { + pub(crate) http: reqwest::Client, +} +impl ObjectStoreTransport { + pub(crate) async fn upload(&self, grant: &Logs3OverflowUpload, payload: Bytes) -> Result<()> { + let url = Url::parse(&grant.signed_url) + .map_err(|_| BraintrustError::InvalidConfig("invalid overflow signed URL".into()))?; + let mut headers = reqwest::header::HeaderMap::new(); + if let Some(pairs) = &grant.headers { + for (key, value) in pairs { + let key = + reqwest::header::HeaderName::from_bytes(key.as_bytes()).map_err(|_| { + BraintrustError::InvalidConfig("invalid overflow header name".into()) + })?; + let value = reqwest::header::HeaderValue::from_str(value).map_err(|_| { + BraintrustError::InvalidConfig("invalid overflow header value".into()) + })?; + headers.insert(key, value); + } + } + let request = if grant.method.eq_ignore_ascii_case("POST") { + let fields = grant.fields.as_ref().ok_or_else(|| { + BraintrustError::InvalidConfig("overflow POST upload missing fields".into()) + })?; + let content_type = fields + .get("Content-Type") + .map(String::as_str) + .unwrap_or("application/json"); + let mut form = reqwest::multipart::Form::new(); + for (key, value) in fields { + form = form.text(key.clone(), value.clone()); + } + form = form.part( + "file", + reqwest::multipart::Part::bytes(payload.to_vec()) + .mime_str(content_type) + .map_err(|_| { + BraintrustError::InvalidConfig("invalid overflow content type".into()) + })?, + ); + headers.remove(reqwest::header::CONTENT_TYPE); + self.http.post(url).headers(headers).multipart(form) + } else if grant.method.eq_ignore_ascii_case("PUT") { + if url + .host_str() + .is_some_and(|host| host.ends_with(".blob.core.windows.net")) + && !headers.contains_key("x-ms-blob-type") + { + headers.insert( + "x-ms-blob-type", + reqwest::header::HeaderValue::from_static("BlockBlob"), + ); + } + self.http + .put(url) + .header("content-type", "application/json") + .headers(headers) + .body(payload) + } else { + return Err(BraintrustError::InvalidConfig( + "unsupported overflow upload method".into(), + )); + }; + let response = request + .send() + .await + .map_err(|error| BraintrustError::Http(error.without_url()))?; + if !response.status().is_success() { + return Err(BraintrustError::Api { + status: response.status().as_u16(), + message: response.text().await.unwrap_or_default(), + }); + } + Ok(()) + } +} diff --git a/src/internal/retry.rs b/src/internal/retry.rs index 08d4324..6c6653f 100644 --- a/src/internal/retry.rs +++ b/src/internal/retry.rs @@ -65,6 +65,23 @@ impl, S: Sleeper> RetryExecutor { pub(crate) async fn run( mut self, + operation: F, + retryable: C, + on_retry: N, + ) -> RetryOutcome + where + F: FnMut(Attempt) -> Fut, + Fut: Future>, + C: FnMut(&E) -> bool, + N: FnMut(Attempt, &E, Duration), + { + self.run_shared(operation, retryable, on_retry).await + } + + /// Execute one stage, sharing the remaining retry allowance with later + /// stages. Successful stages do not consume retries or reset the schedule. + pub(crate) async fn run_shared( + &mut self, mut operation: F, mut retryable: C, mut on_retry: N, @@ -75,10 +92,13 @@ impl, S: Sleeper> RetryExecutor { C: FnMut(&E) -> bool, N: FnMut(Attempt, &E, Duration), { - // Construction always provides the initial attempt, even for an empty - // delay schedule. A classified error is returned without cloning it. - while let Some(attempt) = self.next().await { - let attempts = attempt.index.saturating_add(1); + if self.next_index == Some(0) { + self.next().await; + } + let mut attempts: usize = 0; + loop { + let attempt = Attempt { index: attempts }; + attempts = attempts.saturating_add(1); match operation(attempt).await { Ok(value) => { return RetryOutcome { @@ -87,7 +107,7 @@ impl, S: Sleeper> RetryExecutor { } } Err(error) => { - if attempt.index == self.max_retries || !retryable(&error) { + if !retryable(&error) || self.next_index.is_none() { return RetryOutcome { result: Err(error), attempts, @@ -100,10 +120,15 @@ impl, S: Sleeper> RetryExecutor { }; }; on_retry(attempt, &error, delay); + if self.next().await.is_none() { + return RetryOutcome { + result: Err(error), + attempts, + }; + } } } } - unreachable!("initial attempt exists and a terminal error returns immediately") } } @@ -146,6 +171,46 @@ mod tests { } } + #[tokio::test] + async fn stages_share_retries_without_resetting_the_delay_schedule() { + let sleeper = FakeSleeper::default(); + let mut executor = RetryExecutor::with_sleeper( + 2, + exponential(Duration::from_secs(1), Duration::from_secs(8)), + sleeper.clone(), + ); + for _ in 0..2 { + let outcome = executor + .run_shared( + |attempt| async move { + if attempt.index == 0 { + Err("transient") + } else { + Ok(()) + } + }, + |_| true, + |_, _, _| {}, + ) + .await; + assert_eq!(outcome.attempts, 2); + assert_eq!(outcome.result, Ok(())); + } + let outcome = executor + .run_shared( + |_| async { Err::<(), _>("exhausted") }, + |_| true, + |_, _, _| {}, + ) + .await; + assert_eq!(outcome.attempts, 1); + assert_eq!(outcome.result, Err("exhausted")); + assert_eq!( + *sleeper.delays.lock().unwrap(), + vec![Duration::from_secs(1), Duration::from_secs(2)] + ); + } + #[tokio::test] async fn zero_retries_still_attempts_once_without_sleeping() { let sleeper = FakeSleeper::default(); diff --git a/src/internal/upload.rs b/src/internal/upload.rs new file mode 100644 index 0000000..3028a2f --- /dev/null +++ b/src/internal/upload.rs @@ -0,0 +1,514 @@ +//! SDK ingestion orchestration, built on the one-attempt API. +use super::ingest::{EncodedBatch, UploadPlan, UploadPlanner}; +use super::object_store::ObjectStoreTransport; +use super::retry::{exponential, jittered_exponential, RetryExecutor}; +use crate::api::{ + self, + logs3::{Logs3OverflowReference, Logs3OverflowRequest, Logs3OverflowUploadRequest}, +}; +use crate::error::{BraintrustError, Result}; +use crate::types::LOGS_API_VERSION; +use std::{future::Future, time::Duration}; +use tracing::warn; +use url::Url; + +#[derive(Clone, Copy, Debug)] +pub(crate) struct Capabilities { + pub max_request_size: usize, + pub overflow: bool, +} + +/// Existing schedules remain caller policy; orchestration and classification +/// are shared. Retries are shared across all stages and splits of one batch. +#[derive(Clone, Copy)] +pub(crate) enum UploadPolicy { + Queue { retries: usize }, + Standalone, +} +impl UploadPolicy { + fn budget(self) -> RetryBudget { + let (retries, delays): (_, Box + Send>) = match self { + Self::Queue { retries } => ( + retries, + Box::new(exponential(Duration::from_secs(1), Duration::from_secs(30))), + ), + Self::Standalone => ( + 4, + Box::new(jittered_exponential( + Duration::from_millis(300), + Duration::from_secs(8), + 0.2, + )), + ), + }; + RetryBudget { + executor: RetryExecutor::new(retries, delays), + requests: 0, + } + } +} +struct RetryBudget { + executor: RetryExecutor + Send>>, + requests: usize, +} +impl RetryBudget { + async fn run(&mut self, mut operation: F) -> Result + where + F: FnMut() -> Fut, + Fut: Future>, + { + let outcome = self.executor.run_shared( + |_| operation(), retryable, + |attempt, error, _| warn!(%error, retry = attempt.index + 1, "ingestion request failed, retrying"), + ).await; + self.requests += outcome.attempts; + outcome.result + } +} +fn retryable(error: &BraintrustError) -> bool { + match error { + BraintrustError::Http(error) => !error.is_builder(), + BraintrustError::Api { status, .. } => *status == 429 || (500..600).contains(status), + _ => false, + } +} +pub(crate) fn api_error(error: api::Error) -> BraintrustError { + match error { + api::Error::Config(message) => BraintrustError::InvalidConfig(message), + api::Error::Transport { source, .. } => BraintrustError::Http(source), + api::Error::Http { status, body, .. } => BraintrustError::Api { + status, + message: body, + }, + api::Error::Decode { + operation, + status, + source, + } => BraintrustError::Api { + status, + message: format!("{operation}: invalid response: {source}"), + }, + } +} + +#[derive(Default)] +pub(crate) struct UploadResult { + pub rows_uploaded: usize, + pub bytes_processed: usize, + pub requests_sent: usize, +} + +pub(crate) struct UploadCoordinator { + logs: api::logs3::Logs3, + objects: ObjectStoreTransport, + policy: UploadPolicy, +} +impl UploadCoordinator { + pub fn new( + http: reqwest::Client, + api_url: &Url, + token: &str, + org_name: Option<&str>, + policy: UploadPolicy, + ) -> Result { + let logs = api::Client::for_ingestion(http.clone(), api_url, token, org_name) + .map_err(api_error)? + .logs3(); + Ok(Self { + logs, + objects: ObjectStoreTransport { http }, + policy, + }) + } + + pub async fn capabilities( + &self, + fallback: usize, + override_bytes: Option, + ) -> Capabilities { + let server_limit = match self.logs.version().await { + Ok(version) => version.logs3_payload_max_bytes.filter(|limit| *limit > 0), + Err(error) => { + warn!(%error, "failed to fetch ingestion capabilities"); + None + } + }; + let max_request_size = match (server_limit, override_bytes) { + (Some(server), Some(local)) => server.min(local), + (Some(server), None) => server, + (None, Some(local)) => local, + (None, None) => fallback, + } + .max(1); + Capabilities { + max_request_size, + overflow: server_limit.is_some(), + } + } + + pub async fn upload( + &self, + batch: EncodedBatch, + capabilities: Capabilities, + ) -> Result { + let planner = UploadPlanner { + max_request_bytes: capabilities.max_request_size, + overflow: capabilities.overflow, + }; + let mut pending = vec![planner.plan(batch)?]; + let mut budget = self.policy.budget(); + let mut result = UploadResult::default(); + while let Some(plan) = pending.pop() { + match plan { + UploadPlan::Split(left, right) => { + pending.push(planner.plan(right)?); + pending.push(planner.plan(left)?); + } + UploadPlan::Direct(batch) => { + match budget + .run(|| async { + self.logs + .post_raw(batch.data.to_vec()) + .await + .map_err(api_error) + }) + .await + { + Ok(()) => { + result.rows_uploaded += batch.len(); + result.bytes_processed += batch.data.len(); + } + Err(BraintrustError::Api { status: 413, .. }) => { + pending.push(planner.rejected(batch)?) + } + Err(error) => return Err(error), + } + } + UploadPlan::Overflow(batch) => { + self.overflow(&batch, &mut budget).await?; + result.rows_uploaded += batch.len(); + result.bytes_processed += batch.data.len(); + } + } + } + result.requests_sent = budget.requests; + Ok(result) + } + + async fn overflow(&self, batch: &EncodedBatch, budget: &mut RetryBudget) -> Result<()> { + let grant = budget + .run(|| async { + self.logs + .request_overflow_upload(Logs3OverflowUploadRequest { + content_type: "application/json".into(), + size_bytes: batch.data.len(), + rows: batch.overflow_rows.clone(), + }) + .await + .map_err(api_error) + }) + .await?; + budget + .run(|| self.objects.upload(&grant, batch.data.clone())) + .await?; + budget + .run(|| async { + self.logs + .commit_overflow(Logs3OverflowRequest { + rows: Logs3OverflowReference { + reference_type: "logs3_overflow", + key: grant.key.clone(), + }, + api_version: LOGS_API_VERSION, + }) + .await + .map_err(api_error) + }) + .await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::api::logs3::Logs3OverflowUpload; + use crate::internal::ingest::EncodedRow; + use serde_json::{json, Value}; + use wiremock::{ + matchers::{header, method, path}, + Mock, MockServer, ResponseTemplate, + }; + + fn batch() -> EncodedBatch { + EncodedBatch::partition( + vec![EncodedRow::from_map( + json!({"id":"row", "project_id":"project", "log_id":"g", "input":"é"}) + .as_object() + .unwrap(), + ) + .unwrap()], + 100, + usize::MAX, + ) + .pop() + .unwrap() + } + fn coordinator(server: &MockServer, policy: UploadPolicy) -> UploadCoordinator { + UploadCoordinator::new( + crate::http::build_http_client(Duration::from_secs(3), None, None).unwrap(), + &Url::parse(&server.uri()).unwrap(), + "secret", + Some("tenant"), + policy, + ) + .unwrap() + } + fn capabilities() -> Capabilities { + Capabilities { + max_request_size: 1, + overflow: true, + } + } + + #[tokio::test] + async fn both_callers_reuse_successful_overflow_stages_and_preserve_scopes() { + for policy in [UploadPolicy::Queue { retries: 2 }, UploadPolicy::Standalone] { + let server = MockServer::start().await; + Mock::given(method("POST")).and(path("/logs3/overflow")) + .and(header("authorization", "Bearer secret")).and(header("x-bt-org-name", "tenant")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"method":"PUT", "signedUrl":format!("{}/object", server.uri()), "key":"key"}))) + .expect(1).mount(&server).await; + Mock::given(method("PUT")) + .and(path("/object")) + .respond_with(ResponseTemplate::new(503)) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("PUT")) + .and(path("/object")) + .respond_with(ResponseTemplate::new(200)) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(503)) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(200)) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + let batch = batch(); + let data = batch.data.clone(); + let result = coordinator(&server, policy) + .upload(batch, capabilities()) + .await + .unwrap(); + assert_eq!(result.requests_sent, 5); + assert_eq!(result.rows_uploaded, 1); + assert_eq!(result.bytes_processed, data.len()); + let requests = server.received_requests().await.unwrap(); + for request in requests + .iter() + .filter(|request| request.url.path() == "/object") + { + assert!(!request + .headers + .keys() + .any(|key| key.as_str() == "authorization")); + assert!(!request + .headers + .keys() + .any(|key| key.as_str() == "x-bt-org-name")); + assert_eq!(request.body, data); + } + let grant: Value = serde_json::from_slice(&requests[0].body).unwrap(); + assert_eq!(grant["size_bytes"], data.len()); + assert_eq!(grant["rows"][0]["object_ids"]["project_id"], "project"); + } + } + + #[tokio::test] + async fn retry_budget_is_shared_across_grant_object_and_reference() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3/overflow")) + .respond_with(ResponseTemplate::new(503)) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3/overflow")) + .respond_with(ResponseTemplate::new(200).set_body_json( + json!({"method":"PUT", "signedUrl":format!("{}/object",server.uri()), "key":"key"}), + )) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + Mock::given(method("PUT")) + .and(path("/object")) + .respond_with(ResponseTemplate::new(503)) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("PUT")) + .and(path("/object")) + .respond_with(ResponseTemplate::new(200)) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(503)) + .expect(1) + .mount(&server) + .await; + let error = coordinator(&server, UploadPolicy::Queue { retries: 2 }) + .upload(batch(), capabilities()) + .await + .err() + .unwrap(); + assert!(matches!(error, BraintrustError::Api { status: 503, .. })); + assert_eq!(server.received_requests().await.unwrap().len(), 5); + } + + #[tokio::test] + async fn permanent_errors_stop_at_each_overflow_stage() { + for stage in ["grant", "object", "reference"] { + let server = MockServer::start().await; + let grant = if stage == "grant" { + ResponseTemplate::new(400) + } else { + ResponseTemplate::new(200).set_body_json(json!({"method":"PUT","signedUrl":format!("{}/object", server.uri()),"key":"key"})) + }; + Mock::given(method("POST")) + .and(path("/logs3/overflow")) + .respond_with(grant) + .expect(1) + .mount(&server) + .await; + Mock::given(method("PUT")) + .and(path("/object")) + .respond_with(ResponseTemplate::new(if stage == "object" { + 400 + } else { + 200 + })) + .expect(if stage == "grant" { 0 } else { 1 }) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(400)) + .expect(if stage == "reference" { 1 } else { 0 }) + .mount(&server) + .await; + assert!(matches!( + coordinator(&server, UploadPolicy::Queue { retries: 3 }) + .upload(batch(), capabilities()) + .await, + Err(BraintrustError::Api { status: 400, .. }) + )); + } + } + + #[tokio::test] + async fn signed_post_uses_multipart_boundary_and_preserves_grant_headers() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/object")) + .and(header("x-signed", "value")) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&server) + .await; + let payload = batch().data; + ObjectStoreTransport { + http: reqwest::Client::new(), + } + .upload( + &Logs3OverflowUpload { + method: "POST".into(), + signed_url: format!("{}/object", server.uri()), + headers: Some( + [ + ("content-type".into(), "must-be-replaced".into()), + ("x-signed".into(), "value".into()), + ] + .into(), + ), + fields: Some( + [ + ("Content-Type".into(), "application/json".into()), + ("key".into(), "object-key".into()), + ] + .into(), + ), + key: "key".into(), + }, + payload.clone(), + ) + .await + .unwrap(); + let requests = server.received_requests().await.unwrap(); + let content_type = requests[0] + .headers + .iter() + .find(|(key, _)| key.as_str() == "content-type") + .unwrap() + .1 + .get(0) + .unwrap() + .as_str(); + assert!(content_type.starts_with("multipart/form-data; boundary=")); + let body = String::from_utf8(requests[0].body.clone()).unwrap(); + assert!(body.contains("name=\"file\"")); + assert!(body.contains(std::str::from_utf8(&payload).unwrap())); + assert!(body.contains("object-key")); + } + + #[tokio::test] + async fn redirects_do_not_create_hidden_attempts_or_forward_credentials() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(307).insert_header("location", "/redirected")) + .expect(1) + .mount(&server) + .await; + Mock::given(path("/redirected")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&server) + .await; + let result = coordinator(&server, UploadPolicy::Queue { retries: 2 }) + .upload( + batch(), + Capabilities { + max_request_size: usize::MAX, + overflow: false, + }, + ) + .await; + assert!(matches!( + result, + Err(BraintrustError::Api { status: 307, .. }) + )); + } +} diff --git a/src/log_queue/http.rs b/src/log_queue/http.rs index 782664a..806355b 100644 --- a/src/log_queue/http.rs +++ b/src/log_queue/http.rs @@ -1,10 +1,8 @@ use super::batching::SerializedBatch; use super::config::LogQueueConfig; -use crate::internal::ingest::{UploadPlan, UploadPlanner}; -use crate::internal::retry::{exponential, RetryExecutor}; -use crate::types::{Logs3OverflowReference, Logs3OverflowRequest, Logs3OverflowUpload}; +use crate::internal::upload::{Capabilities, UploadCoordinator, UploadPolicy}; use std::path::PathBuf; -use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; +use std::time::{SystemTime, UNIX_EPOCH}; use tracing::warn; use url::Url; @@ -28,73 +26,32 @@ pub(crate) fn save_payload_debug( Ok(()) } -/// Fetch version info from the server and return (max_request_size, can_use_overflow). -/// -/// `can_use_overflow` is true when the server returns a positive `logs3_payload_max_bytes`, -/// indicating that the overflow upload path is supported. +/// Compatibility adapter for the queue; protocol behavior lives in one coordinator. pub(crate) async fn fetch_version_info( client: &reqwest::Client, api_url: &Url, token: &str, org_name: Option<&str>, ) -> (usize, bool) { - let version_url = match api_url.join("version") { - Ok(url) => url, - Err(e) => { - warn!(error = %e, "Invalid API URL for version fetch"); - return (crate::log_queue::config::DEFAULT_BATCH_MAX_BYTES * 2, false); + let fallback = crate::log_queue::config::DEFAULT_BATCH_MAX_BYTES * 2; + match UploadCoordinator::new( + client.clone(), + api_url, + token, + org_name, + UploadPolicy::Queue { retries: 0 }, + ) { + Ok(coordinator) => { + let capabilities = coordinator.capabilities(fallback, None).await; + (capabilities.max_request_size, capabilities.overflow) } - }; - - let mut request = client.get(version_url).bearer_auth(token); - if let Some(org) = org_name { - request = request.header("x-bt-org-name", org); - } - let response = request.send().await; - - match response { - Ok(resp) if resp.status().is_success() => match resp.json::().await { - Ok(json) => { - let server_limit = json - .get("logs3_payload_max_bytes") - .and_then(|v| v.as_u64()) - .filter(|&v| v > 0) - .map(|v| v as usize); - - match server_limit { - Some(limit) => (limit, true), - None => (crate::log_queue::config::DEFAULT_BATCH_MAX_BYTES * 2, false), - } - } - Err(e) => { - warn!(error = %e, "Failed to parse version response"); - (crate::log_queue::config::DEFAULT_BATCH_MAX_BYTES * 2, false) - } - }, - Ok(resp) => { - warn!(status = %resp.status(), "Version fetch returned non-success status"); - (crate::log_queue::config::DEFAULT_BATCH_MAX_BYTES * 2, false) - } - Err(e) => { - warn!(error = %e, "Failed to fetch version info"); - (crate::log_queue::config::DEFAULT_BATCH_MAX_BYTES * 2, false) + Err(error) => { + warn!(%error, "invalid ingestion scope"); + (fallback, false) } } } -/// Send a batch with retry logic (matching TypeScript SDK behavior). -/// -/// If the batch exceeds `max_request_size` and `can_use_overflow` is true, -/// uses the overflow upload path (POST logs3/overflow → signed URL → POST logs3). -/// -/// Retry semantics match the TypeScript SDK: -/// - `config.num_retries()` is the number of *retries* (not total attempts) -/// - Total attempts = num_retries + 1 -/// - Retry delay starts at 1 second, doubles each attempt: `1s * 2^attempt` -/// -/// Note: the TypeScript SDK retries on all errors including 4xx. This implementation -/// intentionally diverges: only 429 and 5xx are retried since retrying on 400 Bad -/// Request (malformed payload) cannot succeed and wastes time. #[allow(clippy::too_many_arguments)] pub(crate) async fn send_batch_with_retry( client: &reqwest::Client, @@ -106,332 +63,30 @@ pub(crate) async fn send_batch_with_retry( max_request_size: usize, can_use_overflow: bool, ) -> Result<(), anyhow::Error> { - let planner = UploadPlanner { - max_request_bytes: max_request_size, - overflow: can_use_overflow, - }; - let mut pending = vec![planner.plan(batch)?]; - while let Some(plan) = pending.pop() { - match plan { - UploadPlan::Split(left, right) => { - pending.push(planner.plan(right)?); - pending.push(planner.plan(left)?); - } - UploadPlan::Direct(batch) => { - let result = - send_planned_batch(client, api_url, token, org_name, &batch, config, false) - .await; - match result { - Err(error) - if matches!( - error.downcast_ref::(), - Some(crate::error::BraintrustError::Api { status: 413, .. }) - ) => - { - pending.push(planner.rejected(batch)?); - } - result => result?, - } - } - UploadPlan::Overflow(batch) => { - send_planned_batch(client, api_url, token, org_name, &batch, config, true).await? - } - } - } - Ok(()) -} - -#[allow(clippy::too_many_arguments)] -async fn send_planned_batch( - client: &reqwest::Client, - api_url: &Url, - token: &str, - org_name: Option<&str>, - batch: &SerializedBatch, - config: &LogQueueConfig, - use_overflow: bool, -) -> Result<(), anyhow::Error> { - let json_bytes = &batch.data; - let overflow_rows = &batch.overflow_rows; - - // Save to debug directories if configured - save_payload_debug(json_bytes, config.all_publish_payloads_dir(), "all").ok(); - - let logs_url = api_url - .join("logs3") - .map_err(|e| anyhow::anyhow!("invalid logs url: {e}"))?; - - // If overflow is needed, the signed URL is requested once and reused across retries - // (matches TypeScript SDK's `overflowUpload` caching pattern). - let mut overflow_upload: Option = None; - - let mut retries = RetryExecutor::new( - config.num_retries(), - exponential(Duration::from_secs(1), Duration::from_secs(30)), - ); - while let Some(attempt) = retries.next().await { - let attempt = attempt.index; - - if use_overflow { - // Get or create the overflow upload (cached across retries) - if overflow_upload.is_none() { - let upload_result = - request_overflow_upload(client, api_url, token, json_bytes, overflow_rows) - .await; - - match upload_result { - Ok(upload) => { - match upload_overflow_payload(client, &upload, json_bytes).await { - Ok(()) => overflow_upload = Some(upload), - Err(e) => { - if attempt < config.num_retries() { - warn!(error = %e, attempt, "overflow upload failed, retrying"); - continue; - } - save_payload_debug( - json_bytes, - config.failed_publish_payloads_dir(), - "failed", - ) - .ok(); - return Err(e); - } - } - } - Err(e) => { - if attempt < config.num_retries() { - warn!(error = %e, attempt, "overflow URL request failed, retrying"); - continue; - } - save_payload_debug( - json_bytes, - config.failed_publish_payloads_dir(), - "failed", - ) - .ok(); - return Err(e); - } - } - } - - // Send the overflow reference to logs3 - let key = overflow_upload.as_ref().unwrap().key.clone(); - let overflow_ref = Logs3OverflowRequest { - rows: Logs3OverflowReference { - reference_type: "logs3_overflow", - key, - }, - api_version: crate::types::LOGS_API_VERSION, - }; - let overflow_bytes = serde_json::to_vec(&overflow_ref) - .map_err(|e| anyhow::anyhow!("failed to serialize overflow ref: {e}"))?; - - let mut overflow_req = client - .post(logs_url.clone()) - .bearer_auth(token) - .header("content-type", "application/json") - .body(overflow_bytes); - if let Some(org) = org_name { - overflow_req = overflow_req.header("x-bt-org-name", org); - } - let req_start = Instant::now(); - match overflow_req.send().await { - Ok(resp) if resp.status().is_success() => { - tracing::debug!( - bytes = json_bytes.len(), - elapsed_ms = req_start.elapsed().as_millis(), - "overflow batch sent" - ); - return Ok(()); - } - Ok(resp) => { - let status = resp.status(); - let body = resp.text().await.unwrap_or_default(); - let is_retryable = status == reqwest::StatusCode::TOO_MANY_REQUESTS - || status.is_server_error(); - if is_retryable && attempt < config.num_retries() { - warn!(status = %status, attempt, "overflow ref send failed, retrying"); - continue; - } - save_payload_debug(json_bytes, config.failed_publish_payloads_dir(), "failed") - .ok(); - return Err(crate::error::BraintrustError::Api { - status: status.as_u16(), - message: body, - } - .into()); - } - Err(e) => { - if attempt < config.num_retries() { - warn!(error = %e, attempt, "overflow ref send network error, retrying"); - continue; - } - save_payload_debug(json_bytes, config.failed_publish_payloads_dir(), "failed") - .ok(); - return Err(e.into()); - } - } - } else { - // Normal path: POST directly to logs3 - let mut normal_req = client - .post(logs_url.clone()) - .bearer_auth(token) - .header("content-type", "application/json") - .body(json_bytes.clone()); - if let Some(org) = org_name { - normal_req = normal_req.header("x-bt-org-name", org); - } - let req_start = Instant::now(); - match normal_req.send().await { - Ok(resp) if resp.status().is_success() => { - tracing::debug!( - bytes = json_bytes.len(), - elapsed_ms = req_start.elapsed().as_millis(), - "batch sent" - ); - return Ok(()); - } - Ok(resp) => { - let status = resp.status(); - let body = resp.text().await.unwrap_or_default(); - let is_retryable = status == reqwest::StatusCode::TOO_MANY_REQUESTS - || status.is_server_error(); - if is_retryable && attempt < config.num_retries() { - warn!(status = %status, attempt, "batch send failed with retryable error, retrying"); - continue; - } - save_payload_debug(json_bytes, config.failed_publish_payloads_dir(), "failed") - .ok(); - return Err(crate::error::BraintrustError::Api { - status: status.as_u16(), - message: body, - } - .into()); - } - Err(e) => { - if attempt < config.num_retries() { - warn!(error = %e, attempt, "batch send network error, retrying"); - continue; - } - save_payload_debug(json_bytes, config.failed_publish_payloads_dir(), "failed") - .ok(); - return Err(e.into()); - } - } - } - } - - unreachable!("retry loop should always return or continue") -} - -/// Request an overflow upload URL from the server. -pub(crate) async fn request_overflow_upload( - client: &reqwest::Client, - api_url: &Url, - token: &str, - json_bytes: &[u8], - overflow_rows: &[crate::types::Logs3OverflowInputRow], -) -> Result { - let overflow_url = api_url - .join("logs3/overflow") - .map_err(|e| anyhow::anyhow!("invalid overflow url: {e}"))?; - - let request_body = serde_json::json!({ - "content_type": "application/json", - "size_bytes": json_bytes.len(), - "rows": overflow_rows, - }); - - let resp = client - .post(overflow_url) - .bearer_auth(token) - .json(&request_body) - .send() - .await - .map_err(|e| anyhow::anyhow!("overflow URL request failed: {e}"))?; - - if !resp.status().is_success() { - let status = resp.status(); - let body = resp.text().await.unwrap_or_default(); - return Err(crate::error::BraintrustError::Api { - status: status.as_u16(), - message: body, - } - .into()); - } - - let upload: Logs3OverflowUpload = resp - .json() - .await - .map_err(|e| anyhow::anyhow!("failed to parse overflow response: {e}"))?; - - Ok(upload) -} - -/// Upload the payload to the signed URL provided by the overflow response. -pub(crate) async fn upload_overflow_payload( - client: &reqwest::Client, - upload: &Logs3OverflowUpload, - payload: &[u8], -) -> Result<(), anyhow::Error> { - let resp = if upload.method.eq_ignore_ascii_case("POST") { - // Multipart POST (e.g. S3) - let fields = upload - .fields - .as_ref() - .ok_or_else(|| anyhow::anyhow!("overflow POST upload missing fields"))?; - - let mut form = reqwest::multipart::Form::new(); - for (key, value) in fields { - form = form.text(key.clone(), value.clone()); - } - let content_type = fields - .get("Content-Type") - .map(|s| s.as_str()) - .unwrap_or("application/json"); - let part = reqwest::multipart::Part::bytes(payload.to_vec()) - .mime_str(content_type) - .map_err(|e| anyhow::anyhow!("invalid content-type: {e}"))?; - form = form.part("file", part); - - let mut req = client.post(&upload.signed_url).multipart(form); - if let Some(headers) = &upload.headers { - for (key, value) in headers { - if key.to_lowercase() != "content-type" { - req = req.header(key, value); - } - } - } - req.send().await - } else { - // PUT (e.g. Azure, GCS) - let mut req = client - .put(&upload.signed_url) - .header("content-type", "application/json") - .body(payload.to_vec()); - if let Some(headers) = &upload.headers { - for (key, value) in headers { - req = req.header(key, value); - } - } - // Azure Blob Storage requires `x-ms-blob-type: BlockBlob` on SAS-URL PUT - // requests; without it the server returns 400. The TypeScript SDK uses the - // same hostname-based check (`addAzureBlobHeaders`). A more robust signal - // would be an `upload_type` field in the server's overflow response, but - // that is a server-side change outside this SDK's scope for now. - if upload.signed_url.contains("blob.core.windows.net") { - req = req.header("x-ms-blob-type", "BlockBlob"); - } - req.send().await - }; - - let resp = resp.map_err(|e| anyhow::anyhow!("overflow upload failed: {e}"))?; - if !resp.status().is_success() { - let status = resp.status(); - let body = resp.text().await.unwrap_or_default(); - anyhow::bail!("overflow upload failed: [{status}] {body}"); + save_payload_debug(&batch.data, config.all_publish_payloads_dir(), "all").ok(); + let data = batch.data.clone(); + let coordinator = UploadCoordinator::new( + client.clone(), + api_url, + token, + org_name, + UploadPolicy::Queue { + retries: config.num_retries(), + }, + )?; + let result = coordinator + .upload( + batch, + Capabilities { + max_request_size, + overflow: can_use_overflow, + }, + ) + .await; + if result.is_err() { + save_payload_debug(&data, config.failed_publish_payloads_dir(), "failed").ok(); } - Ok(()) + result.map(|_| ()).map_err(Into::into) } #[cfg(test)] diff --git a/src/logs3.rs b/src/logs3.rs index f2c6669..1582f31 100644 --- a/src/logs3.rs +++ b/src/logs3.rs @@ -1,23 +1,13 @@ -use std::collections::HashMap; -use std::time::Duration; - -use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; -use reqwest::{multipart, Client, Url}; -use serde::Deserialize; -use serde_json::{json, Map, Value}; - use crate::error::{BraintrustError, Result}; use crate::http::{build_http_client, CaBundleConfig}; -use crate::internal::ingest::{EncodedBatch, EncodedRow, UploadPlan, UploadPlanner}; -use crate::internal::retry::{jittered_exponential, RetryExecutor}; -use crate::types::Logs3OverflowInputRow; +use crate::internal::ingest::{EncodedBatch, EncodedRow}; +use crate::internal::upload::{Capabilities, UploadCoordinator, UploadPolicy}; +use reqwest::{Client, Url}; +use serde_json::{Map, Value}; +use std::time::Duration; const DEFAULT_MAX_REQUEST_SIZE: usize = 6 * 1024 * 1024; -const LOGS_API_VERSION: u8 = 2; -const LOGS3_OVERFLOW_REFERENCE_TYPE: &str = "logs3_overflow"; -const MAX_ATTEMPTS: usize = 5; -const BASE_BACKOFF_MS: u64 = 300; -const MAX_BACKOFF_MS: u64 = 8_000; + #[derive(Debug, Clone, Default)] pub struct Logs3UploadResult { pub rows_uploaded: usize, @@ -32,36 +22,7 @@ pub struct Logs3BatchUploader { api_key: String, org_name: Option, max_request_size_override: Option, - payload_limits: Option, -} - -#[derive(Debug, Clone)] -struct PayloadLimits { - max_request_size: usize, - can_use_overflow: bool, -} - -#[derive(Debug, Deserialize)] -struct VersionInfo { - #[serde(default)] - logs3_payload_max_bytes: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "UPPERCASE")] -enum OverflowMethod { - Put, - Post, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -struct Logs3OverflowUpload { - method: OverflowMethod, - signed_url: String, - headers: Option>, - fields: Option>, - key: String, + payload_limits: Option, } impl Logs3BatchUploader { @@ -128,7 +89,23 @@ impl Logs3BatchUploader { return Ok(Logs3UploadResult::default()); } - let limits = self.get_payload_limits().await; + let coordinator = UploadCoordinator::new( + self.client.clone(), + &self.api_url, + &self.api_key, + self.org_name.as_deref(), + UploadPolicy::Standalone, + )?; + let limits = match self.payload_limits { + Some(limits) => limits, + None => { + let limits = coordinator + .capabilities(DEFAULT_MAX_REQUEST_SIZE, self.max_request_size_override) + .await; + self.payload_limits = Some(limits); + limits + } + }; let row_payloads = rows .iter() .map(EncodedRow::from_map) @@ -144,302 +121,13 @@ impl Logs3BatchUploader { let mut result = Logs3UploadResult::default(); for batch in batches { - let batch_result = self.submit_row_batch(batch, &limits).await?; + let batch_result = coordinator.upload(batch, limits).await?; result.bytes_processed += batch_result.bytes_processed; result.requests_sent += batch_result.requests_sent; result.rows_uploaded += batch_result.rows_uploaded; } Ok(result) } - - async fn get_payload_limits(&mut self) -> PayloadLimits { - if let Some(existing) = &self.payload_limits { - return existing.clone(); - } - let fetched = self.fetch_payload_limits().await; - self.payload_limits = Some(fetched.clone()); - fetched - } - - async fn fetch_payload_limits(&self) -> PayloadLimits { - let mut server_limit: Option = None; - if let Ok(url) = self.api_url.join("version") { - let mut request = self.client.get(url).bearer_auth(&self.api_key); - if let Some(org_name) = &self.org_name { - request = request.header("x-bt-org-name", org_name); - } - if let Ok(response) = request.send().await { - if response.status().is_success() { - if let Ok(version) = response.json::().await { - server_limit = version.logs3_payload_max_bytes.filter(|v| *v > 0); - } - } - } - } - - let can_use_overflow = server_limit.is_some(); - let mut max_request_size = DEFAULT_MAX_REQUEST_SIZE; - if let Some(override_bytes) = self.max_request_size_override { - max_request_size = if let Some(server_limit) = server_limit { - override_bytes.min(server_limit) - } else { - override_bytes - }; - } else if let Some(server_limit) = server_limit { - max_request_size = server_limit; - } - PayloadLimits { - max_request_size: max_request_size.max(1), - can_use_overflow, - } - } - - async fn submit_row_batch( - &self, - batch: EncodedBatch, - limits: &PayloadLimits, - ) -> Result { - let planner = UploadPlanner { - max_request_bytes: limits.max_request_size, - overflow: limits.can_use_overflow, - }; - let mut pending = vec![planner.plan(batch)?]; - let mut result = Logs3UploadResult::default(); - while let Some(plan) = pending.pop() { - match plan { - UploadPlan::Split(left, right) => { - pending.push(planner.plan(right)?); - pending.push(planner.plan(left)?); - } - UploadPlan::Overflow(batch) => { - let uploaded = self.submit_overflow(&batch).await?; - result.rows_uploaded += uploaded.rows_uploaded; - result.bytes_processed += uploaded.bytes_processed; - result.requests_sent += uploaded.requests_sent; - } - UploadPlan::Direct(batch) => { - let (posted, attempts) = - run_with_backoff(|| self.post_logs3_raw(&batch.data)).await; - result.requests_sent += attempts; - match posted { - Ok(()) => { - result.rows_uploaded += batch.len(); - result.bytes_processed += batch.data.len(); - } - Err(BraintrustError::Api { status: 413, .. }) => { - pending.push(planner.rejected(batch)?) - } - Err(error) => return Err(error), - } - } - } - } - Ok(result) - } - - async fn submit_overflow(&self, batch: &EncodedBatch) -> Result { - let mut requests_sent = 0; - let (upload_result, attempts) = run_with_backoff(|| { - self.request_overflow_upload(&batch.overflow_rows, batch.data.len()) - }) - .await; - requests_sent += attempts; - let upload = upload_result?; - let (upload_result, attempts) = - run_with_backoff(|| self.upload_overflow_payload(&upload, &batch.data)).await; - requests_sent += attempts; - upload_result?; - let reference = serde_json::to_vec(&json!({ - "rows": {"type": LOGS3_OVERFLOW_REFERENCE_TYPE, "key": upload.key}, - "api_version": LOGS_API_VERSION, - })) - .map_err(|error| BraintrustError::InvalidConfig(error.to_string()))?; - let (posted, attempts) = run_with_backoff(|| self.post_logs3_raw(&reference)).await; - requests_sent += attempts; - posted?; - Ok(Logs3UploadResult { - rows_uploaded: batch.len(), - bytes_processed: batch.data.len(), - requests_sent, - }) - } - - async fn request_overflow_upload( - &self, - rows: &[Logs3OverflowInputRow], - payload_bytes: usize, - ) -> Result { - let url = self - .api_url - .join("logs3/overflow") - .map_err(|e| BraintrustError::InvalidConfig(format!("invalid overflow URL: {e}")))?; - - let request_body = json!({ - "content_type": "application/json", - "size_bytes": payload_bytes, - "rows": rows, - }); - let mut request = self - .client - .post(url) - .bearer_auth(&self.api_key) - .header("content-type", "application/json") - .json(&request_body); - if let Some(org_name) = &self.org_name { - request = request.header("x-bt-org-name", org_name); - } - let response = request.send().await?; - if !response.status().is_success() { - let status = response.status().as_u16(); - let message = response.text().await.unwrap_or_default(); - return Err(BraintrustError::Api { status, message }); - } - response - .json::() - .await - .map_err(|e| BraintrustError::Api { - status: 200, - message: format!("failed to parse logs3 overflow response: {e}"), - }) - } - - async fn upload_overflow_payload( - &self, - upload: &Logs3OverflowUpload, - payload: &[u8], - ) -> Result<()> { - let signed_url = Url::parse(&upload.signed_url).map_err(|e| { - BraintrustError::InvalidConfig(format!("invalid logs3 overflow signed URL: {e}")) - })?; - match upload.method { - OverflowMethod::Post => { - let fields = upload.fields.clone().ok_or_else(|| { - BraintrustError::InvalidConfig( - "logs3 overflow POST upload missing form fields".to_string(), - ) - })?; - let content_type = fields - .get("Content-Type") - .cloned() - .unwrap_or_else(|| "application/json".to_string()); - let mut form = multipart::Form::new(); - for (key, value) in &fields { - form = form.text(key.clone(), value.clone()); - } - let file_part = multipart::Part::bytes(payload.to_vec()) - .mime_str(&content_type) - .map_err(|e| { - BraintrustError::InvalidConfig(format!( - "invalid overflow content-type '{content_type}': {e}" - )) - })?; - form = form.part("file", file_part); - - let mut request = self.client.post(signed_url).multipart(form); - let mut headers = header_map_from_pairs(upload.headers.as_ref())?; - headers.remove("content-type"); - if !headers.is_empty() { - request = request.headers(headers); - } - let response = request.send().await?; - if !response.status().is_success() { - let status = response.status().as_u16(); - let message = response.text().await.unwrap_or_default(); - return Err(BraintrustError::Api { status, message }); - } - } - OverflowMethod::Put => { - let mut headers = header_map_from_pairs(upload.headers.as_ref())?; - if upload.signed_url.contains("blob.core.windows.net") - && !headers.contains_key("x-ms-blob-type") - { - headers.insert("x-ms-blob-type", HeaderValue::from_static("BlockBlob")); - } - - let mut request = self.client.put(signed_url).body(payload.to_vec()); - if !headers.is_empty() { - request = request.headers(headers); - } - let response = request.send().await?; - if !response.status().is_success() { - let status = response.status().as_u16(); - let message = response.text().await.unwrap_or_default(); - return Err(BraintrustError::Api { status, message }); - } - } - } - Ok(()) - } - - async fn post_logs3_raw(&self, payload: &[u8]) -> Result<()> { - let url = self - .api_url - .join("logs3") - .map_err(|e| BraintrustError::InvalidConfig(format!("invalid logs3 URL: {e}")))?; - let mut request = self - .client - .post(url) - .bearer_auth(&self.api_key) - .header("content-type", "application/json") - .body(payload.to_vec()); - if let Some(org_name) = &self.org_name { - request = request.header("x-bt-org-name", org_name); - } - let response = request.send().await?; - if !response.status().is_success() { - let status = response.status().as_u16(); - let message = response.text().await.unwrap_or_default(); - return Err(BraintrustError::Api { status, message }); - } - Ok(()) - } -} - -fn header_map_from_pairs(headers: Option<&HashMap>) -> Result { - let mut out = HeaderMap::new(); - if let Some(headers) = headers { - for (key, value) in headers { - let name = HeaderName::from_bytes(key.as_bytes()).map_err(|e| { - BraintrustError::InvalidConfig(format!("invalid HTTP header name '{key}': {e}")) - })?; - let header_value = HeaderValue::from_str(value).map_err(|e| { - BraintrustError::InvalidConfig(format!( - "invalid HTTP header value for '{key}': {e}" - )) - })?; - out.insert(name, header_value); - } - } - Ok(out) -} - -fn should_retry(err: &BraintrustError) -> bool { - match err { - BraintrustError::Http(http_err) => http_err.is_timeout() || http_err.is_connect(), - BraintrustError::Network(_) - | BraintrustError::ChannelClosed - | BraintrustError::Background(_) => true, - BraintrustError::Api { status, .. } => *status == 429 || (*status >= 500 && *status <= 599), - BraintrustError::InvalidConfig(_) | BraintrustError::StreamAggregation(_) => false, - } -} - -async fn run_with_backoff(mut operation: F) -> (Result, usize) -where - F: FnMut() -> Fut, - Fut: std::future::Future>, -{ - let outcome = RetryExecutor::new( - MAX_ATTEMPTS - 1, - jittered_exponential( - Duration::from_millis(BASE_BACKOFF_MS), - Duration::from_millis(MAX_BACKOFF_MS), - 0.2, - ), - ) - .run(|_| operation(), should_retry, |_, _, _| {}) - .await; - (outcome.result, outcome.attempts) } #[cfg(test)] diff --git a/src/types.rs b/src/types.rs index 8a510fe..d7d3b23 100644 --- a/src/types.rs +++ b/src/types.rs @@ -340,10 +340,7 @@ impl LogDestination { } } -pub(crate) use crate::api::logs3::{ - Logs3OverflowInputRow, Logs3OverflowInputRowMeta, Logs3OverflowReference, Logs3OverflowRequest, - Logs3OverflowUpload, -}; +pub(crate) use crate::api::logs3::{Logs3OverflowInputRow, Logs3OverflowInputRowMeta}; /// The object ID keys that identify a row's destination, used for overflow metadata. /// Matches the TypeScript SDK's OBJECT_ID_KEYS.