diff --git a/src/internal/batching.rs b/src/internal/batching.rs new file mode 100644 index 0000000..9991d14 --- /dev/null +++ b/src/internal/batching.rs @@ -0,0 +1,110 @@ +/// Limits include the container and separators supplied by the protocol encoder. +#[derive(Clone, Copy)] +pub(crate) struct BatchLimits { + pub max_items: usize, + pub max_bytes: usize, + pub container_bytes: usize, + pub separator_bytes: usize, +} + +/// Ordered groups: oversized items remain in their original position. +pub(crate) enum Partition { + Batch(Vec), + Oversized(T), +} + +pub(crate) struct Batcher; + +impl Batcher { + pub fn partition( + items: Vec, + limits: BatchLimits, + weight: impl Fn(&T) -> usize, + ) -> Vec> { + let mut output = Vec::new(); + let mut batch = Vec::new(); + let mut bytes = limits.container_bytes; + for item in items { + let size = weight(&item); + let singleton = limits.container_bytes.checked_add(size); + let next = bytes + .checked_add(if batch.is_empty() { + 0 + } else { + limits.separator_bytes + }) + .and_then(|bytes| bytes.checked_add(size)); + let fits = next.is_some_and(|bytes| bytes <= limits.max_bytes) + && batch.len() < limits.max_items.max(1); + if !batch.is_empty() && !fits { + output.push(Partition::Batch(std::mem::take(&mut batch))); + bytes = limits.container_bytes; + } + if singleton.is_none_or(|bytes| bytes > limits.max_bytes) { + output.push(Partition::Oversized(item)); + } else { + bytes += if batch.is_empty() { + 0 + } else { + limits.separator_bytes + }; + bytes += size; + batch.push(item); + } + } + if !batch.is_empty() { + output.push(Partition::Batch(batch)); + } + output + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn partitions_in_order_with_explicit_oversized_items() { + let parts = Batcher::partition( + vec![2, 3, 50, 1, 1, 1], + BatchLimits { + max_items: 2, + max_bytes: 8, + container_bytes: 2, + separator_bytes: 1, + }, + |size| *size, + ); + let parts: Vec<_> = parts + .into_iter() + .map(|part| match part { + Partition::Batch(items) => (false, items), + Partition::Oversized(item) => (true, vec![item]), + }) + .collect(); + assert_eq!( + parts, + vec![ + (false, vec![2, 3]), + (true, vec![50]), + (false, vec![1, 1]), + (false, vec![1]) + ] + ); + } + + #[test] + fn weights_cannot_wrap_and_empty_input_has_no_batches() { + let limits = BatchLimits { + max_items: 0, + max_bytes: usize::MAX, + container_bytes: 2, + separator_bytes: 1, + }; + assert!(Batcher::partition(Vec::::new(), limits, |size| *size).is_empty()); + assert!(matches!( + Batcher::partition(vec![usize::MAX], limits, |size| *size)[0], + Partition::Oversized(_) + )); + } +} diff --git a/src/internal/ingest.rs b/src/internal/ingest.rs new file mode 100644 index 0000000..27e2ec5 --- /dev/null +++ b/src/internal/ingest.rs @@ -0,0 +1,247 @@ +//! Pure logs3 encoding and upload decisions. No transport or retry policy. +use super::batching::{BatchLimits, Batcher, Partition}; +use crate::types::{Logs3OverflowInputRow, Logs3OverflowInputRowMeta, LOGS_API_VERSION}; +use bytes::Bytes; +use serde::Serialize; +use serde_json::{Map, Value}; + +#[derive(Debug, Clone)] +pub(crate) struct EncodedRow { + data: Bytes, + overflow: Logs3OverflowInputRow, +} + +impl EncodedRow { + pub fn encode( + row: &impl Serialize, + mut overflow: Logs3OverflowInputRow, + ) -> Result { + let data = serde_json::to_vec(row)?; + overflow.input_row.byte_size = data.len(); + Ok(Self { + data: data.into(), + overflow, + }) + } + + pub fn from_map(row: &Map) -> Result { + let object_ids = [ + "experiment_id", + "dataset_id", + "prompt_session_id", + "project_id", + "log_id", + "function_data", + ] + .into_iter() + .filter_map(|key| row.get(key).map(|value| (key.to_owned(), value.clone()))) + .collect(); + Self::encode( + row, + Logs3OverflowInputRow { + object_ids, + is_delete: row.get("_object_delete").and_then(Value::as_bool), + input_row: Logs3OverflowInputRowMeta { byte_size: 0 }, + }, + ) + } +} + +#[derive(Debug, Clone)] +pub(crate) struct EncodedBatch { + pub data: Bytes, + pub overflow_rows: Vec, + rows: Vec, +} + +impl EncodedBatch { + fn encode(rows: Vec) -> Self { + let mut data = Vec::new(); + data.extend_from_slice(b"{\"rows\":["); + for (index, row) in rows.iter().enumerate() { + if index > 0 { + data.push(b','); + } + data.extend_from_slice(&row.data); + } + data.extend_from_slice(format!("],\"api_version\":{LOGS_API_VERSION}}}").as_bytes()); + Self { + data: data.into(), + overflow_rows: rows.iter().map(|row| row.overflow.clone()).collect(), + rows, + } + } + + pub fn partition(rows: Vec, max_items: usize, max_bytes: usize) -> Vec { + let container_bytes = Self::encode(Vec::new()).data.len(); + Batcher::partition( + rows, + BatchLimits { + max_items, + max_bytes, + container_bytes, + separator_bytes: 1, + }, + |row| row.data.len(), + ) + .into_iter() + .map(|part| { + Self::encode(match part { + Partition::Batch(rows) => rows, + // Keep the oversized row as a singleton for the planner to handle. + Partition::Oversized(row) => vec![row], + }) + }) + .collect() + } + + pub fn len(&self) -> usize { + self.rows.len() + } + + fn split(mut self) -> (Self, Self) { + let right = self.rows.split_off(self.rows.len() / 2); + (Self::encode(self.rows), Self::encode(right)) + } +} + +pub(crate) enum UploadPlan { + Direct(EncodedBatch), + Overflow(EncodedBatch), + Split(EncodedBatch, EncodedBatch), +} + +#[derive(Clone, Copy)] +pub(crate) struct UploadPlanner { + pub max_request_bytes: usize, + pub overflow: bool, +} + +impl UploadPlanner { + pub fn plan(&self, batch: EncodedBatch) -> Result { + self.decide(batch, false) + } + + /// A rejected direct request must change shape, never retry identical bytes. + pub fn rejected( + &self, + batch: EncodedBatch, + ) -> Result { + self.decide(batch, true) + } + + fn decide( + &self, + batch: EncodedBatch, + rejected: bool, + ) -> Result { + if !rejected && batch.data.len() <= self.max_request_bytes { + Ok(UploadPlan::Direct(batch)) + } else if self.overflow { + Ok(UploadPlan::Overflow(batch)) + } else if batch.len() > 1 { + let (left, right) = batch.split(); + Ok(UploadPlan::Split(left, right)) + } else { + Err(crate::error::BraintrustError::Api { + status: 413, + message: "single row exceeds server payload limit and overflow is unavailable" + .into(), + }) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn rows() -> Vec { + ["é\\\"", "second", "third"] + .into_iter() + .map(|id| { + EncodedRow::from_map( + json!({"id":id,"dataset_id":"ds","_object_delete":true}) + .as_object() + .unwrap(), + ) + .unwrap() + }) + .collect() + } + + #[test] + fn encoded_limits_include_envelope_separators_and_utf8() { + let expected = EncodedBatch::encode(rows()); + let exact = EncodedBatch::partition(rows(), 3, expected.data.len()); + assert_eq!(exact.len(), 1); + assert_eq!(exact[0].data, expected.data); + let smaller = EncodedBatch::partition(rows(), 3, expected.data.len() - 1); + assert_eq!(smaller.len(), 2); + for batch in smaller { + assert!(batch.data.len() < expected.data.len()); + let value: Value = serde_json::from_slice(&batch.data).unwrap(); + for (row, meta) in value["rows"] + .as_array() + .unwrap() + .iter() + .zip(batch.overflow_rows) + { + assert_eq!( + meta.input_row.byte_size, + serde_json::to_vec(row).unwrap().len() + ); + assert_eq!(meta.object_ids["dataset_id"], "ds"); + assert_eq!(meta.is_delete, Some(true)); + } + } + } + + #[test] + fn planner_handles_oversized_and_rejected_batches_without_reordering() { + let batch = EncodedBatch::encode(rows()); + let planner = UploadPlanner { + max_request_bytes: batch.data.len(), + overflow: false, + }; + assert!(matches!( + planner.plan(batch.clone()).unwrap(), + UploadPlan::Direct(_) + )); + let UploadPlan::Split(left, right) = planner.rejected(batch).unwrap() else { + panic!("expected split") + }; + let ids: Vec<_> = [left.clone(), right] + .iter() + .flat_map(|batch| { + serde_json::from_slice::(&batch.data).unwrap()["rows"] + .as_array() + .unwrap() + .iter() + .map(|row| row["id"].clone()) + .collect::>() + }) + .collect(); + assert_eq!(ids, vec![json!("é\\\""), json!("second"), json!("third")]); + assert!(planner.rejected(left.clone()).is_err()); + assert!(matches!( + UploadPlanner { + overflow: true, + ..planner + } + .rejected(left) + .unwrap(), + UploadPlan::Overflow(_) + )); + assert!(matches!( + UploadPlanner { + max_request_bytes: 1, + overflow: true + } + .plan(EncodedBatch::encode(rows())) + .unwrap(), + UploadPlan::Overflow(_) + )); + } +} diff --git a/src/internal/mod.rs b/src/internal/mod.rs index bb52cea..f24492a 100644 --- a/src/internal/mod.rs +++ b/src/internal/mod.rs @@ -1 +1,3 @@ +pub(crate) mod batching; +pub(crate) mod ingest; pub(crate) mod retry; diff --git a/src/log_queue/batching.rs b/src/log_queue/batching.rs index 0aaae99..338e6f6 100644 --- a/src/log_queue/batching.rs +++ b/src/log_queue/batching.rs @@ -1,134 +1,30 @@ use super::config::LogQueueConfig; -use crate::types::{Logs3OverflowInputRow, Logs3OverflowInputRowMeta, Logs3Row, LOGS_API_VERSION}; +use crate::types::{Logs3OverflowInputRow, Logs3OverflowInputRowMeta, Logs3Row}; use tracing::warn; -/// Generic batching function — matches TypeScript SDK's `batchItems()`. -/// -/// Splits `items` into batches respecting `max_items` and `max_bytes` limits. -/// `byte_size` is called once per item to determine its size. -/// -/// Flush condition (matches TypeScript): a new batch is started when the current -/// batch is non-empty AND (`batch.len() >= max_items` OR `batch_bytes + item_size >= max_bytes`). -/// This means a single item that exceeds `max_bytes` on its own is always placed in its own batch. -/// -/// Note: the byte comparison uses `>=` (not `>`), so two items that sum to exactly `max_bytes` -/// are split into separate batches. This matches the TypeScript SDK's strict `<` check in the -/// negated condition (`!(size < max && len < max)`). The test -/// `test_byte_boundary_flushes_at_exact_limit` verifies this invariant explicitly. -pub(crate) fn batch_items( - items: Vec, - max_items: usize, - max_bytes: usize, - byte_size: impl Fn(&T) -> usize, -) -> Vec> { - let max_items = max_items.max(1); - let max_bytes = max_bytes.max(1); - let mut output = Vec::new(); - let mut batch: Vec = Vec::new(); - let mut batch_bytes = 0usize; - - for item in items { - let item_size = byte_size(&item); - if !batch.is_empty() && (batch.len() >= max_items || batch_bytes + item_size >= max_bytes) { - output.push(batch); - batch = Vec::new(); - batch_bytes = 0; - } - batch_bytes += item_size; - batch.push(item); - } - if !batch.is_empty() { - output.push(batch); - } - output -} - -/// A serialized batch with per-row overflow metadata. -pub(crate) struct SerializedBatch { - /// The full batch serialized as a Logs3Request JSON payload. - pub data: Vec, - /// Per-row overflow metadata (for `POST logs3/overflow` requests). - pub overflow_rows: Vec, -} +pub(crate) use crate::internal::ingest::EncodedBatch as SerializedBatch; +use crate::internal::ingest::EncodedRow; -/// Batch rows respecting size limits, then serialize each batch to JSON. -/// -/// Each row is serialized once directly to bytes, and overflow metadata is -/// extracted from the `Logs3Row` struct without intermediate `Value` representation. -/// -/// `batch_max_bytes` is the effective byte limit for this flush, which may be -/// lower than `config.batch_max_bytes()` when the server advertises a tighter -/// limit via `GET /version` (callers should use `min(config, server_limit / 2)`). +/// Encode rows once, then partition using the complete request size. pub(crate) fn batch_and_serialize_rows( rows: Vec, config: &LogQueueConfig, batch_max_bytes: usize, ) -> Result, anyhow::Error> { - // Serialize each row once directly to bytes, extract overflow metadata from struct - let prepared: Vec<(Vec, Logs3OverflowInputRow)> = rows + let rows = rows .into_iter() .map(|row| { - let overflow_row = build_overflow_row_from_logs3row(&row); - let row_bytes = match serde_json::to_vec(&row) { - Ok(bytes) => bytes, - Err(error) => { - warn!(error = %error, "failed to serialize row"); - return Err(error); - } - }; - let overflow_row = Logs3OverflowInputRow { - input_row: Logs3OverflowInputRowMeta { - byte_size: row_bytes.len(), - }, - ..overflow_row - }; - Ok((row_bytes, overflow_row)) - }) - .collect::>()?; - - // Split into batches using the shared batch_items logic. - let batches = batch_items( - prepared, - config.batch_max_items(), - batch_max_bytes, - |(row_bytes, _)| row_bytes.len(), - ); - - // Assemble each batch by concatenating the pre-serialized row bytes. - batches - .into_iter() - .map(|batch| { - let (row_bytes, overflow_rows): (Vec<_>, Vec<_>) = batch.into_iter().unzip(); - assemble_logs3_request(&row_bytes, overflow_rows).map_err(|error| { - warn!(error = %error, "failed to assemble batch"); + EncodedRow::encode(&row, build_overflow_row_from_logs3row(&row)).map_err(|error| { + warn!(%error, "failed to serialize row"); error }) }) - .collect() -} - -/// Assemble a `Logs3Request` payload by concatenating pre-serialized row bytes. -/// -/// Produces `{"rows":[,,...],"api_version":}` without re-serializing rows. -fn assemble_logs3_request( - row_bytes: &[Vec], - overflow_rows: Vec, -) -> Result { - let mut data = Vec::new(); - data.extend_from_slice(b"{\"rows\":["); - for (i, row) in row_bytes.iter().enumerate() { - if i > 0 { - data.push(b','); - } - data.extend_from_slice(row); - } - data.extend_from_slice(b"],\"api_version\":"); - data.extend_from_slice(LOGS_API_VERSION.to_string().as_bytes()); - data.push(b'}'); - Ok(SerializedBatch { - data, - overflow_rows, - }) + .collect::, _>>()?; + Ok(SerializedBatch::partition( + rows, + config.batch_max_items(), + batch_max_bytes, + )) } /// Build overflow metadata directly from a Logs3Row struct. @@ -182,6 +78,7 @@ mod tests { use super::*; use crate::log_queue::config::LogQueueConfig; use crate::types::LogDestination; + use crate::types::LOGS_API_VERSION; use chrono::Utc; use std::collections::HashMap; @@ -277,21 +174,6 @@ mod tests { assert!(overflow_row.input_row.byte_size > 0); } - #[test] - fn test_byte_boundary_flushes_at_exact_limit() { - // A batch should flush when adding the next item would reach the limit (>=), - // not just exceed it (>). With max_bytes=100, two items of 50 bytes each - // should NOT be in the same batch (50 + 50 >= 100). - let items: Vec> = vec![vec![0u8; 50], vec![0u8; 50]]; - let batches = batch_items(items, 1000, 100, |b| b.len()); - - assert_eq!( - batches.len(), - 2, - "Items summing to exactly max_bytes should be split" - ); - } - #[test] fn test_overflow_metadata_is_delete_set_from_object_delete_field() { let config = LogQueueConfig::default(); @@ -322,15 +204,4 @@ mod tests { "is_delete should be None when _object_delete is not set" ); } - - #[test] - fn test_single_oversized_item_gets_its_own_batch() { - // An item larger than max_bytes should still be placed in a batch alone. - let items: Vec> = vec![vec![0u8; 200], vec![0u8; 10]]; - let batches = batch_items(items, 1000, 100, |b| b.len()); - - assert_eq!(batches.len(), 2); - assert_eq!(batches[0].len(), 1); // oversized item alone - assert_eq!(batches[1].len(), 1); - } } diff --git a/src/log_queue/http.rs b/src/log_queue/http.rs index 49ffeea..782664a 100644 --- a/src/log_queue/http.rs +++ b/src/log_queue/http.rs @@ -1,5 +1,6 @@ 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 std::path::PathBuf; @@ -105,18 +106,61 @@ pub(crate) async fn send_batch_with_retry( max_request_size: usize, can_use_overflow: bool, ) -> Result<(), anyhow::Error> { - let json_bytes = batch.data; - let overflow_rows = batch.overflow_rows; + 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(); + 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}"))?; - let use_overflow = can_use_overflow && json_bytes.len() > max_request_size; - // 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; @@ -132,12 +176,12 @@ pub(crate) async fn send_batch_with_retry( // 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) + 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 { + match upload_overflow_payload(client, &upload, json_bytes).await { Ok(()) => overflow_upload = Some(upload), Err(e) => { if attempt < config.num_retries() { @@ -145,7 +189,7 @@ pub(crate) async fn send_batch_with_retry( continue; } save_payload_debug( - &json_bytes, + json_bytes, config.failed_publish_payloads_dir(), "failed", ) @@ -160,7 +204,7 @@ pub(crate) async fn send_batch_with_retry( continue; } save_payload_debug( - &json_bytes, + json_bytes, config.failed_publish_payloads_dir(), "failed", ) @@ -209,7 +253,7 @@ pub(crate) async fn send_batch_with_retry( warn!(status = %status, attempt, "overflow ref send failed, retrying"); continue; } - save_payload_debug(&json_bytes, config.failed_publish_payloads_dir(), "failed") + save_payload_debug(json_bytes, config.failed_publish_payloads_dir(), "failed") .ok(); return Err(crate::error::BraintrustError::Api { status: status.as_u16(), @@ -222,7 +266,7 @@ pub(crate) async fn send_batch_with_retry( warn!(error = %e, attempt, "overflow ref send network error, retrying"); continue; } - save_payload_debug(&json_bytes, config.failed_publish_payloads_dir(), "failed") + save_payload_debug(json_bytes, config.failed_publish_payloads_dir(), "failed") .ok(); return Err(e.into()); } @@ -256,7 +300,7 @@ pub(crate) async fn send_batch_with_retry( warn!(status = %status, attempt, "batch send failed with retryable error, retrying"); continue; } - save_payload_debug(&json_bytes, config.failed_publish_payloads_dir(), "failed") + save_payload_debug(json_bytes, config.failed_publish_payloads_dir(), "failed") .ok(); return Err(crate::error::BraintrustError::Api { status: status.as_u16(), @@ -269,7 +313,7 @@ pub(crate) async fn send_batch_with_retry( warn!(error = %e, attempt, "batch send network error, retrying"); continue; } - save_payload_debug(&json_bytes, config.failed_publish_payloads_dir(), "failed") + save_payload_debug(json_bytes, config.failed_publish_payloads_dir(), "failed") .ok(); return Err(e.into()); } @@ -395,11 +439,14 @@ mod tests { use super::*; use crate::log_queue::config::LogQueueConfig; - fn make_batch(data: &[u8]) -> SerializedBatch { - SerializedBatch { - data: data.to_vec(), - overflow_rows: vec![], - } + fn make_batch() -> SerializedBatch { + SerializedBatch::partition( + vec![crate::internal::ingest::EncodedRow::from_map(&serde_json::Map::new()).unwrap()], + 1, + usize::MAX, + ) + .pop() + .unwrap() } #[test] @@ -437,6 +484,114 @@ mod tests { assert!(result.is_ok()); } + #[tokio::test] + async fn rejected_direct_batch_splits_in_order_without_retrying_identical_bytes() { + use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, + }; + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(|request: &wiremock::Request| { + let body: serde_json::Value = serde_json::from_slice(&request.body).unwrap(); + ResponseTemplate::new(if body["rows"].as_array().unwrap().len() > 1 { + 413 + } else { + 200 + }) + }) + .expect(5) + .mount(&server) + .await; + let rows = ["first", "second", "third"] + .into_iter() + .map(|id| { + crate::internal::ingest::EncodedRow::from_map( + serde_json::json!({"id":id}).as_object().unwrap(), + ) + .unwrap() + }) + .collect(); + let batch = SerializedBatch::partition(rows, 100, usize::MAX) + .pop() + .unwrap(); + send_batch_with_retry( + &reqwest::Client::new(), + &Url::parse(&server.uri()).unwrap(), + "token", + None, + batch, + &LogQueueConfig::builder().num_retries(2).build(), + usize::MAX, + false, + ) + .await + .unwrap(); + let requests = server.received_requests().await.unwrap(); + let delivered: Vec<_> = requests + .iter() + .filter_map(|request| { + let body: serde_json::Value = serde_json::from_slice(&request.body).unwrap(); + let rows = body["rows"].as_array().unwrap(); + (rows.len() == 1).then(|| rows[0]["id"].clone()) + }) + .collect(); + assert_eq!( + delivered, + vec![ + serde_json::json!("first"), + serde_json::json!("second"), + serde_json::json!("third") + ] + ); + } + + #[tokio::test] + async fn rejected_singleton_uses_overflow_once_when_supported() { + use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, + }; + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(413)) + .up_to_n_times(1) + .with_priority(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; + Mock::given(method("POST")).and(path("/logs3/overflow")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_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(200)) + .expect(1) + .mount(&server) + .await; + send_batch_with_retry( + &reqwest::Client::new(), + &Url::parse(&server.uri()).unwrap(), + "token", + None, + make_batch(), + &LogQueueConfig::builder().num_retries(0).build(), + usize::MAX, + true, + ) + .await + .unwrap(); + } + #[tokio::test] async fn exhausted_queue_retries_send_exactly_the_configured_attempts() { use wiremock::{ @@ -455,7 +610,7 @@ mod tests { &Url::parse(&server.uri()).unwrap(), "token", None, - make_batch(b"{}"), + make_batch(), &LogQueueConfig::builder().num_retries(2).build(), usize::MAX, false, @@ -506,7 +661,7 @@ mod tests { &Url::parse(&server.uri()).unwrap(), "token", None, - make_batch(b"{}"), + make_batch(), &LogQueueConfig::builder().num_retries(1).build(), 1, true, @@ -545,7 +700,7 @@ mod tests { &api_url, "token", None, - make_batch(b"{}"), + make_batch(), &config, usize::MAX, false, @@ -581,7 +736,7 @@ mod tests { &api_url, "token", None, - make_batch(b"{}"), + make_batch(), &config, usize::MAX, false, @@ -616,7 +771,7 @@ mod tests { &api_url, "token", None, - make_batch(b"{}"), + make_batch(), &config, usize::MAX, false, @@ -645,7 +800,7 @@ mod tests { &api_url, "token", None, - make_batch(b"{}"), + make_batch(), &config, usize::MAX, false, @@ -675,7 +830,7 @@ mod tests { &api_url, "token", None, - make_batch(b"{}"), + make_batch(), &config, usize::MAX, false, @@ -710,7 +865,7 @@ mod tests { &api_url, "token", None, - make_batch(b"{}"), + make_batch(), &config, usize::MAX, false, diff --git a/src/log_queue/mod.rs b/src/log_queue/mod.rs index 32ffd50..a5dd237 100644 --- a/src/log_queue/mod.rs +++ b/src/log_queue/mod.rs @@ -7,6 +7,5 @@ mod row_key; pub(crate) mod worker; // Public API exports -pub(crate) use batching::batch_items; pub use config::LogQueueConfig; pub use queue::LogQueue; diff --git a/src/logs3.rs b/src/logs3.rs index fbf249d..f2c6669 100644 --- a/src/logs3.rs +++ b/src/logs3.rs @@ -3,13 +3,14 @@ use std::time::Duration; use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; use reqwest::{multipart, Client, Url}; -use serde::{Deserialize, Serialize}; +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::log_queue::batch_items; +use crate::types::Logs3OverflowInputRow; const DEFAULT_MAX_REQUEST_SIZE: usize = 6 * 1024 * 1024; const LOGS_API_VERSION: u8 = 2; @@ -17,15 +18,6 @@ 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; -const OBJECT_ID_KEYS: [&str; 6] = [ - "experiment_id", - "dataset_id", - "prompt_session_id", - "project_id", - "log_id", - "function_data", -]; - #[derive(Debug, Clone, Default)] pub struct Logs3UploadResult { pub rows_uploaded: usize, @@ -49,26 +41,6 @@ struct PayloadLimits { can_use_overflow: bool, } -#[derive(Debug, Clone)] -struct RowPayload { - row_json: String, - row_bytes: usize, - overflow_meta: Logs3OverflowInputRow, -} - -#[derive(Debug, Clone, Serialize)] -struct Logs3OverflowInputRow { - object_ids: Map, - #[serde(skip_serializing_if = "Option::is_none")] - is_delete: Option, - input_row: OverflowInputRowInfo, -} - -#[derive(Debug, Clone, Serialize)] -struct OverflowInputRowInfo { - byte_size: usize, -} - #[derive(Debug, Deserialize)] struct VersionInfo { #[serde(default)] @@ -157,20 +129,22 @@ impl Logs3BatchUploader { } let limits = self.get_payload_limits().await; - let mut row_payloads = Vec::with_capacity(rows.len()); - for row in rows { - row_payloads.push(RowPayload::from_row(row)?); - } - - let batch_limit_rows = batch_max_num_items.max(1); - let batch_limit_bytes = (limits.max_request_size / 2).max(1); - let batches = batch_items(row_payloads, batch_limit_rows, batch_limit_bytes, |item| { - item.row_bytes - }); + let row_payloads = rows + .iter() + .map(EncodedRow::from_map) + .collect::, _>>() + .map_err(|error| { + BraintrustError::InvalidConfig(format!("failed to serialize logs3 row: {error}")) + })?; + let batches = EncodedBatch::partition( + row_payloads, + batch_max_num_items.max(1), + (limits.max_request_size / 2).max(1), + ); let mut result = Logs3UploadResult::default(); for batch in batches { - let batch_result = self.submit_row_batch(&batch, &limits).await?; + let batch_result = self.submit_row_batch(batch, &limits).await?; result.bytes_processed += batch_result.bytes_processed; result.requests_sent += batch_result.requests_sent; result.rows_uploaded += batch_result.rows_uploaded; @@ -222,103 +196,70 @@ impl Logs3BatchUploader { async fn submit_row_batch( &self, - items: &[RowPayload], + batch: EncodedBatch, limits: &PayloadLimits, ) -> Result { - if items.is_empty() { - return Ok(Logs3UploadResult::default()); - } - let mut pending: Vec> = vec![items.to_vec()]; + 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(chunk) = pending.pop() { - if chunk.is_empty() { - continue; - } - let payload = construct_logs3_payload(&chunk); - let payload_bytes = payload.len(); - - if limits.can_use_overflow && payload_bytes > limits.max_request_size { - let overflow_result = self.submit_overflow(&chunk, payload, payload_bytes).await?; - result.rows_uploaded += overflow_result.rows_uploaded; - result.bytes_processed += overflow_result.bytes_processed; - result.requests_sent += overflow_result.requests_sent; - continue; - } - - let (post_result, post_attempts) = - run_with_backoff(|| self.post_logs3_raw(&payload)).await; - result.requests_sent += post_attempts; - match post_result { - Ok(()) => { - result.rows_uploaded += chunk.len(); - result.bytes_processed += payload_bytes; + 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; } - Err(BraintrustError::Api { status, .. }) if status == 413 => { - if limits.can_use_overflow { - let overflow_result = self - .submit_overflow(&chunk, payload.clone(), payload_bytes) - .await?; - result.rows_uploaded += overflow_result.rows_uploaded; - result.bytes_processed += overflow_result.bytes_processed; - result.requests_sent += overflow_result.requests_sent; - } else if chunk.len() > 1 { - let mid = chunk.len() / 2; - pending.push(chunk[mid..].to_vec()); - pending.push(chunk[..mid].to_vec()); - } else { - return Err(BraintrustError::Api { - status, - message: - "single row exceeds server payload limit and overflow is unavailable" - .to_string(), - }); + 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), } } - Err(err) => return Err(err), } } - Ok(result) } - async fn submit_overflow( - &self, - items: &[RowPayload], - payload: String, - payload_bytes: usize, - ) -> Result { - let overflow_rows: Vec = items - .iter() - .map(|item| item.overflow_meta.clone()) - .collect(); - let mut requests_sent = 0usize; - let (upload_result, upload_attempts) = - run_with_backoff(|| self.request_overflow_upload(&overflow_rows, payload_bytes)).await; - requests_sent += upload_attempts; + 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_payload_result, upload_payload_attempts) = - run_with_backoff(|| self.upload_overflow_payload(&upload, &payload)).await; - requests_sent += upload_payload_attempts; - upload_payload_result?; - - let overflow_reference = json!({ - "rows": { - "type": LOGS3_OVERFLOW_REFERENCE_TYPE, - "key": upload.key, - }, + 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, - }) - .to_string(); - let (post_result, post_attempts) = - run_with_backoff(|| self.post_logs3_raw(&overflow_reference)).await; - requests_sent += post_attempts; - post_result?; - + })) + .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: items.len(), - bytes_processed: payload_bytes, + rows_uploaded: batch.len(), + bytes_processed: batch.data.len(), requests_sent, }) } @@ -365,7 +306,7 @@ impl Logs3BatchUploader { async fn upload_overflow_payload( &self, upload: &Logs3OverflowUpload, - payload: &str, + payload: &[u8], ) -> Result<()> { let signed_url = Url::parse(&upload.signed_url).map_err(|e| { BraintrustError::InvalidConfig(format!("invalid logs3 overflow signed URL: {e}")) @@ -385,7 +326,7 @@ impl Logs3BatchUploader { for (key, value) in &fields { form = form.text(key.clone(), value.clone()); } - let file_part = multipart::Part::text(payload.to_string()) + let file_part = multipart::Part::bytes(payload.to_vec()) .mime_str(&content_type) .map_err(|e| { BraintrustError::InvalidConfig(format!( @@ -415,7 +356,7 @@ impl Logs3BatchUploader { headers.insert("x-ms-blob-type", HeaderValue::from_static("BlockBlob")); } - let mut request = self.client.put(signed_url).body(payload.to_string()); + let mut request = self.client.put(signed_url).body(payload.to_vec()); if !headers.is_empty() { request = request.headers(headers); } @@ -430,7 +371,7 @@ impl Logs3BatchUploader { Ok(()) } - async fn post_logs3_raw(&self, payload: &str) -> Result<()> { + async fn post_logs3_raw(&self, payload: &[u8]) -> Result<()> { let url = self .api_url .join("logs3") @@ -440,7 +381,7 @@ impl Logs3BatchUploader { .post(url) .bearer_auth(&self.api_key) .header("content-type", "application/json") - .body(payload.to_string()); + .body(payload.to_vec()); if let Some(org_name) = &self.org_name { request = request.header("x-bt-org-name", org_name); } @@ -454,42 +395,6 @@ impl Logs3BatchUploader { } } -impl RowPayload { - fn from_row(row: &Map) -> Result { - let row_json = serde_json::to_string(row).map_err(|e| { - BraintrustError::InvalidConfig(format!("failed to serialize logs3 row: {e}")) - })?; - let row_bytes = row_json.len(); - let mut object_ids = Map::new(); - for key in OBJECT_ID_KEYS { - if let Some(value) = row.get(key) { - object_ids.insert(key.to_string(), value.clone()); - } - } - let is_delete = row.get("_object_delete").and_then(Value::as_bool); - Ok(Self { - row_json, - row_bytes, - overflow_meta: Logs3OverflowInputRow { - object_ids, - is_delete, - input_row: OverflowInputRowInfo { - byte_size: row_bytes, - }, - }, - }) - } -} - -fn construct_logs3_payload(items: &[RowPayload]) -> String { - let rows = items - .iter() - .map(|item| item.row_json.as_str()) - .collect::>() - .join(","); - format!("{{\"rows\":[{rows}],\"api_version\":{LOGS_API_VERSION}}}") -} - fn header_map_from_pairs(headers: Option<&HashMap>) -> Result { let mut out = HeaderMap::new(); if let Some(headers) = headers {