From 1ab7ff657eee2f02e9b920a2bdda78f4fd31c152 Mon Sep 17 00:00:00 2001 From: "shreya.malpani" Date: Fri, 28 Aug 2026 14:57:20 -0400 Subject: [PATCH 1/8] bump logs maxBytes limit to 1MiB --- bottlecap/src/extension/telemetry/mod.rs | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/bottlecap/src/extension/telemetry/mod.rs b/bottlecap/src/extension/telemetry/mod.rs index 2e4913e18..ce703653b 100644 --- a/bottlecap/src/extension/telemetry/mod.rs +++ b/bottlecap/src/extension/telemetry/mod.rs @@ -1,5 +1,5 @@ use reqwest::{Client, Response}; -use tracing::debug; +use tracing::{debug, error}; use crate::extension::{EXTENSION_ID_HEADER, base_url}; @@ -62,15 +62,25 @@ pub async fn subscribe( "URI": format!("http://sandbox:{}/", destination_port), }, "types": get_subscription_event_types(logs_enabled), - "buffering": { // TODO: re evaluate using default values + "buffering": { "maxItems": 1000, - "maxBytes": 256 * 1024, + // 1 MiB is the AWS maximum. A log line larger than this limit is split across POSTs. + "maxBytes": 1024 * 1024, "timeoutMs": 25 } })) .send() .await?; - debug!("EXTENSION | Subscribed to Telemetry API: {:?}", response); + if response.status().is_success() { + debug!("EXTENSION | Subscribed to Telemetry API: {:?}", response); + } else { + // A rejected subscription means no telemetry at all, including platform.runtimeDone, so + // surface it rather than leaving it at debug level. + error!( + "EXTENSION | Telemetry API rejected subscription with status {}", + response.status() + ); + } Ok(response) } From 564e63afde7ce5f59bf6592253cc0d4528ee5b46 Mon Sep 17 00:00:00 2001 From: "shreya.malpani" Date: Thu, 3 Sep 2026 10:59:45 -0700 Subject: [PATCH 2/8] stitch together large payload that was split to recover telemetry events --- bottlecap/src/extension/telemetry/listener.rs | 346 +++++++++++++++++- 1 file changed, 326 insertions(+), 20 deletions(-) diff --git a/bottlecap/src/extension/telemetry/listener.rs b/bottlecap/src/extension/telemetry/listener.rs index 68349374f..f6ba0546d 100644 --- a/bottlecap/src/extension/telemetry/listener.rs +++ b/bottlecap/src/extension/telemetry/listener.rs @@ -6,16 +6,41 @@ use crate::{ use axum::{ Router, - extract::{Request, State}, + extract::{DefaultBodyLimit, Request, State}, http::StatusCode, response::{IntoResponse, Response}, routing::post, }; -use std::net::SocketAddr; +use serde_json::error::Category; +use std::{ + net::SocketAddr, + sync::{Arc, Mutex, PoisonError}, + time::{Duration, Instant}, +}; use tokio::{net::TcpListener, sync::mpsc::Sender}; use tokio_util::sync::CancellationToken; use tracing::debug; +/// Ceiling on a held fragment. The Telemetry API can POST up to +/// `2 * maxBytes + metadataBytes`, so this fits a full-size head at the 1 MiB `maxBytes` we +/// subscribe with. +const MAX_FRAGMENT_BYTES: usize = 2 * 1024 * 1024; + +/// Body ceiling, replacing axum's 2 MiB default so a full-size POST isn't rejected before +/// the handler sees it. +const MAX_BODY_BYTES: usize = 2 * MAX_FRAGMENT_BYTES; + +/// A continuation POST follows its head immediately, so anything older belongs to a payload +/// whose continuation never arrived. +const FRAGMENT_TTL: Duration = Duration::from_secs(1); + +const RECORD_KEY: &[u8] = b"\"record\":"; +const RECORD_START: &[u8] = b"{\"time\":"; + +/// The Telemetry API writes a record's framing even after cutting its value short, so a head +/// fragment ends with the envelope's `}` and the array's `]`. +const APPENDED_CLOSERS: &[u8] = b"}]"; + #[allow(clippy::module_name_repetitions)] #[derive(Debug, Clone)] pub struct TelemetryListener { @@ -26,6 +51,12 @@ pub struct TelemetryListener { event_bus_tx: Sender, } +#[derive(Clone)] +struct HandlerState { + logs_tx: Sender, + fragments: FragmentBuffer, +} + impl TelemetryListener { #[must_use] pub fn new( @@ -70,12 +101,16 @@ impl TelemetryListener { } fn make_router(&self) -> Router { - let logs_tx: Sender = self.logs_tx.clone(); + let state = HandlerState { + logs_tx: self.logs_tx.clone(), + fragments: FragmentBuffer::default(), + }; Router::new() .route("/", post(Self::handle)) .fallback(handler_not_found) - .with_state(logs_tx) + .with_state(state) + .layer(DefaultBodyLimit::max(MAX_BODY_BYTES)) } async fn graceful_shutdown( @@ -93,7 +128,7 @@ impl TelemetryListener { debug!("TELEMETRY API | Shutting down"); } - async fn handle(State(logs_tx): State>, request: Request) -> Response { + async fn handle(State(state): State, request: Request) -> Response { let (_, body) = match extract_request_body(request).await { Ok(r) => r, Err(e) => { @@ -105,29 +140,159 @@ impl TelemetryListener { } }; - let body = std::str::from_utf8(&body).expect("infallible"); - - let mut telemetry_events: Vec = match serde_json::from_str(body) { + let mut telemetry_events: Vec = match serde_json::from_slice(&body) { Ok(events) => events, - Err(e) => { - // If we can't parse the event, we will receive it again in a new batch - // causing an infinite loop and resource contention. - // Instead, log it and move on. - // This will result in a dropped payload, but may be from - // events we haven't added support for yet - debug!("Failed to parse telemetry events `{body}`, failed with: {e}"); - return (StatusCode::OK, "Failed to parse telemetry events").into_response(); - } + Err(e) => match state.fragments.stitch(&body, &e) { + Stitch::Complete(events) => { + debug!( + "TELEMETRY API | Reassembled a split payload, recovered {} events", + events.len() + ); + events + } + Stitch::Pending => { + return (StatusCode::OK, "Holding split telemetry payload").into_response(); + } + Stitch::Discarded => { + // If we can't parse the event, we will receive it again in a new batch + // causing an infinite loop and resource contention. + // Instead, log it and move on. + // This will result in a dropped payload, but may be from + // events we haven't added support for yet + let body = String::from_utf8_lossy(&body); + debug!("Failed to parse telemetry events `{body}`, failed with: {e}"); + return (StatusCode::OK, "Failed to parse telemetry events").into_response(); + } + }, }; for event in telemetry_events.drain(..) { - logs_tx.send(event).await.expect("infallible"); + state.logs_tx.send(event).await.expect("infallible"); } (StatusCode::OK, "OK").into_response() } } +/// Outcome of trying to pair an unparseable body with a held fragment. +#[derive(Debug)] +enum Stitch { + /// The body completed a split payload. + Complete(Vec), + /// The body opens a split payload; it is held for its continuation. + Pending, + /// The body is not part of a split payload, or its continuation never came. + Discarded, +} + +/// Reassembles telemetry payloads that the Telemetry API split across consecutive POSTs. +/// +/// A record larger than the subscription's `maxBytes` is cut mid-value and resumed in the +/// next POST, which repeats the cut record's envelope (`[{"time":..,"type":..,"record":`) +/// ahead of the resumed bytes. Neither half parses on its own, and the second half is where +/// the rest of the batch lands — including the `platform.runtimeDone` that the on-demand +/// loop waits on before calling `/next`. +#[derive(Clone, Default)] +struct FragmentBuffer { + held: Arc>>, +} + +impl FragmentBuffer { + /// Joins `body` onto a held fragment, or holds it if it opens a split payload. + /// + /// Assumes the continuation is the next POST: the Telemetry API sends fragments back to + /// back and gives no sequence number to correlate on. The repeated envelope, which + /// carries the cut record's timestamp, is what pairs the two halves. + fn stitch(&self, body: &[u8], error: &serde_json::Error) -> Stitch { + let mut held = self.held.lock().unwrap_or_else(PoisonError::into_inner); + + if let Some(Fragment { + body: mut joined, + continuation_prefix, + received, + }) = held.take() + { + if received.elapsed() > FRAGMENT_TTL { + debug!( + "TELEMETRY API | Dropping {} held bytes, no continuation arrived", + joined.len() + ); + } else if let Some(tail) = body.strip_prefix(continuation_prefix.as_slice()) { + // Left in place, the head's closers land inside the resumed value, where + // they parse but corrupt the record. + if joined.ends_with(APPENDED_CLOSERS) { + joined.truncate(joined.len() - APPENDED_CLOSERS.len()); + } + joined.extend_from_slice(tail); + + return match serde_json::from_slice(&joined) { + Ok(events) => Stitch::Complete(events), + // A record can be cut more than once, so keep going. + Err(e) => Self::hold(&mut held, joined, &e), + }; + } else { + debug!( + "TELEMETRY API | Dropping {} held bytes, next payload does not continue it", + joined.len() + ); + } + } + + Self::hold(&mut held, body.to_vec(), error) + } + + fn hold(held: &mut Option, body: Vec, error: &serde_json::Error) -> Stitch { + *held = Fragment::opening(body, error); + if held.is_some() { + Stitch::Pending + } else { + Stitch::Discarded + } + } +} + +struct Fragment { + body: Vec, + /// What the continuation POST repeats before the resumed bytes: `[` followed by the + /// keys of the cut record up to and including `"record":`. + continuation_prefix: Vec, + received: Instant, +} + +impl Fragment { + /// Holds `body` only if it opens a split payload: an array that ran out of input inside + /// the value of its last record. Any other parse failure won't be fixed by joining, and + /// holding it would poison the next stitch. + fn opening(body: Vec, error: &serde_json::Error) -> Option { + if error.classify() != Category::Eof + || body.len() > MAX_FRAGMENT_BYTES + || body.first() != Some(&b'[') + { + return None; + } + + // The cut record is the last one in the body. + let key_end = rfind(&body, RECORD_KEY)? + RECORD_KEY.len(); + let record_start = rfind(body.get(..key_end)?, RECORD_START)?; + + let mut continuation_prefix = vec![b'[']; + continuation_prefix.extend_from_slice(body.get(record_start..key_end)?); + + Some(Self { + body, + continuation_prefix, + received: Instant::now(), + }) + } +} + +/// Byte offset of the last occurrence of `needle` in `haystack`. +fn rfind(haystack: &[u8], needle: &[u8]) -> Option { + haystack + .windows(needle.len()) + .rposition(|window| window == needle) +} + #[cfg(test)] mod tests { use super::*; @@ -135,7 +300,148 @@ mod tests { use axum::http::Request; use chrono::DateTime; - use crate::extension::telemetry::events::{InitPhase, InitType, TelemetryRecord}; + use crate::extension::telemetry::events::{ + InitPhase, InitType, RuntimeDoneMetrics, Status, TelemetryRecord, + }; + + /// Leading half of a split payload: a `function` record cut inside its `message` value, + /// with the framing the Telemetry API writes anyway. + const HEAD: &str = + r#"[{"time":"2026-09-03T14:29:52.929Z","type":"function","record":{"message":"AAAA}]"#; + + /// Its continuation, which repeats the cut record's envelope before the resumed bytes + /// and carries the rest of the batch. + const TAIL: &str = r#"[{"time":"2026-09-03T14:29:52.929Z","type":"function","record":BBBB"}},{"time":"2026-09-03T14:29:52.930Z","type":"platform.runtimeDone","record":{"requestId":"abc123","status":"success","metrics":{"durationMs":18.074,"producedBytes":329814}}}]"#; + + /// Mirrors the handler: a body only reaches the buffer once it has failed to parse. + fn stitch(fragments: &FragmentBuffer, body: &str) -> Stitch { + let error = serde_json::from_slice::>(body.as_bytes()) + .expect_err("fixture must not parse on its own"); + fragments.stitch(body.as_bytes(), &error) + } + + fn state(logs_tx: Sender) -> HandlerState { + HandlerState { + logs_tx, + fragments: FragmentBuffer::default(), + } + } + + fn post(body: &str) -> Request { + Request::builder() + .method("POST") + .uri("http://localhost:8080") + .body(Body::from(body.to_string())) + .expect("failed to build request") + } + + #[test] + fn test_stitch_split_payload() { + let fragments = FragmentBuffer::default(); + + assert!(matches!(stitch(&fragments, HEAD), Stitch::Pending)); + + let Stitch::Complete(events) = stitch(&fragments, TAIL) else { + panic!("expected the continuation to complete the payload"); + }; + + assert_eq!(events.len(), 2); + // `AAAA}]BBBB` would mean the head's framing was left in the resumed value. + assert_eq!( + events[0].record, + TelemetryRecord::Function(serde_json::json!({"message": "AAAABBBB"})) + ); + assert_eq!( + events[1].record, + TelemetryRecord::PlatformRuntimeDone { + request_id: "abc123".to_string(), + status: Status::Success, + error_type: None, + metrics: Some(RuntimeDoneMetrics { + duration_ms: 18.074, + produced_bytes: Some(329_814), + }), + } + ); + + // The pair is consumed, so a following payload starts from nothing. + assert!(matches!(stitch(&fragments, TAIL), Stitch::Discarded)); + } + + #[test] + fn test_complete_payload_is_not_held() { + let fragments = FragmentBuffer::default(); + + // Parses as JSON, so it failed for a reason joining won't fix. Holding it would + // poison the next stitch. + let unsupported = + r#"[{"time":"2026-09-03T14:29:52.929Z","type":"platform.brandNew","record":{}}]"#; + assert!(matches!(stitch(&fragments, unsupported), Stitch::Discarded)); + + assert!(matches!(stitch(&fragments, HEAD), Stitch::Pending)); + assert!(matches!(stitch(&fragments, TAIL), Stitch::Complete(_))); + } + + #[test] + fn test_head_dropped_when_next_payload_is_unrelated() { + let fragments = FragmentBuffer::default(); + + assert!(matches!(stitch(&fragments, HEAD), Stitch::Pending)); + + // A different record's envelope: not this head's continuation, so the head goes and + // this one is held in its place. + let other = + r#"[{"time":"2026-09-03T14:30:11.001Z","type":"function","record":{"message":"CCCC}]"#; + assert!(matches!(stitch(&fragments, other), Stitch::Pending)); + + // Proof the original head is gone: its continuation no longer stitches. + assert!(matches!(stitch(&fragments, TAIL), Stitch::Discarded)); + } + + #[test] + fn test_stitch_head_holding_complete_records() { + let fragments = FragmentBuffer::default(); + + // The cut record is the last of several, so the envelope the continuation repeats is + // in the middle of the head. + let head = format!( + r#"[{{"time":"2026-09-03T14:29:52.900Z","type":"extension","record":"ready"}},{}"#, + HEAD.trim_start_matches('[') + ); + assert!(matches!(stitch(&fragments, &head), Stitch::Pending)); + + let Stitch::Complete(events) = stitch(&fragments, TAIL) else { + panic!("expected the continuation to complete the payload"); + }; + assert_eq!(events.len(), 3); + assert_eq!( + events[1].record, + TelemetryRecord::Function(serde_json::json!({"message": "AAAABBBB"})) + ); + } + + #[tokio::test] + async fn test_handle_split_payload() { + let (tx, mut rx) = tokio::sync::mpsc::channel(2); + let state = state(tx); + + let response = + TelemetryListener::handle(axum::extract::State(state.clone()), post(HEAD)).await; + assert_eq!(response.status(), StatusCode::OK); + assert!(rx.try_recv().is_err(), "held fragment must not be emitted"); + + let response = TelemetryListener::handle(axum::extract::State(state), post(TAIL)).await; + assert_eq!(response.status(), StatusCode::OK); + + assert!(matches!( + rx.try_recv().expect("function record").record, + TelemetryRecord::Function(_) + )); + assert!(matches!( + rx.try_recv().expect("runtimeDone record").record, + TelemetryRecord::PlatformRuntimeDone { .. } + )); + } #[tokio::test] #[allow(clippy::unwrap_used)] @@ -155,7 +461,7 @@ mod tests { let (parts, body) = req.into_parts(); let req = Request::from_parts(parts, body); - let response = TelemetryListener::handle(axum::extract::State(tx), req).await; + let response = TelemetryListener::handle(axum::extract::State(state(tx)), req).await; // Check that the response is OK assert_eq!(response.status(), axum::http::StatusCode::OK); From 7e615a5af8c52e4c07bd5644ff47784c6a3e301d Mon Sep 17 00:00:00 2001 From: "shreya.malpani" Date: Thu, 3 Sep 2026 14:33:20 -0700 Subject: [PATCH 3/8] separate payload stitching into separate module --- bottlecap/src/extension/telemetry/listener.rs | 308 ++---------------- bottlecap/src/extension/telemetry/mod.rs | 1 + bottlecap/src/extension/telemetry/stitch.rs | 286 ++++++++++++++++ 3 files changed, 323 insertions(+), 272 deletions(-) create mode 100644 bottlecap/src/extension/telemetry/stitch.rs diff --git a/bottlecap/src/extension/telemetry/listener.rs b/bottlecap/src/extension/telemetry/listener.rs index f6ba0546d..693767df5 100644 --- a/bottlecap/src/extension/telemetry/listener.rs +++ b/bottlecap/src/extension/telemetry/listener.rs @@ -1,6 +1,9 @@ use crate::{ event_bus, - extension::telemetry::events::TelemetryEvent, + extension::telemetry::{ + events::TelemetryEvent, + stitch::{FragmentBuffer, Stitch}, + }, http::{extract_request_body, handler_not_found}, }; @@ -11,35 +14,14 @@ use axum::{ response::{IntoResponse, Response}, routing::post, }; -use serde_json::error::Category; -use std::{ - net::SocketAddr, - sync::{Arc, Mutex, PoisonError}, - time::{Duration, Instant}, -}; +use std::net::SocketAddr; use tokio::{net::TcpListener, sync::mpsc::Sender}; use tokio_util::sync::CancellationToken; use tracing::debug; -/// Ceiling on a held fragment. The Telemetry API can POST up to -/// `2 * maxBytes + metadataBytes`, so this fits a full-size head at the 1 MiB `maxBytes` we -/// subscribe with. -const MAX_FRAGMENT_BYTES: usize = 2 * 1024 * 1024; - -/// Body ceiling, replacing axum's 2 MiB default so a full-size POST isn't rejected before -/// the handler sees it. -const MAX_BODY_BYTES: usize = 2 * MAX_FRAGMENT_BYTES; - -/// A continuation POST follows its head immediately, so anything older belongs to a payload -/// whose continuation never arrived. -const FRAGMENT_TTL: Duration = Duration::from_secs(1); - -const RECORD_KEY: &[u8] = b"\"record\":"; -const RECORD_START: &[u8] = b"{\"time\":"; - -/// The Telemetry API writes a record's framing even after cutting its value short, so a head -/// fragment ends with the envelope's `}` and the array's `]`. -const APPENDED_CLOSERS: &[u8] = b"}]"; +/// Body ceiling, replacing axum's 2 MiB default so a full-size POST — up to +/// `2 * maxBytes + metadataBytes` — isn't rejected before the handler sees it. +const MAX_BODY_BYTES: usize = 4 * 1024 * 1024; #[allow(clippy::module_name_repetitions)] #[derive(Debug, Clone)] @@ -142,6 +124,8 @@ impl TelemetryListener { let mut telemetry_events: Vec = match serde_json::from_slice(&body) { Ok(events) => events, + // The Telemetry API splits an oversized record across two POSTs, and neither half + // parses alone. See `stitch`. Err(e) => match state.fragments.stitch(&body, &e) { Stitch::Complete(events) => { debug!( @@ -174,125 +158,6 @@ impl TelemetryListener { } } -/// Outcome of trying to pair an unparseable body with a held fragment. -#[derive(Debug)] -enum Stitch { - /// The body completed a split payload. - Complete(Vec), - /// The body opens a split payload; it is held for its continuation. - Pending, - /// The body is not part of a split payload, or its continuation never came. - Discarded, -} - -/// Reassembles telemetry payloads that the Telemetry API split across consecutive POSTs. -/// -/// A record larger than the subscription's `maxBytes` is cut mid-value and resumed in the -/// next POST, which repeats the cut record's envelope (`[{"time":..,"type":..,"record":`) -/// ahead of the resumed bytes. Neither half parses on its own, and the second half is where -/// the rest of the batch lands — including the `platform.runtimeDone` that the on-demand -/// loop waits on before calling `/next`. -#[derive(Clone, Default)] -struct FragmentBuffer { - held: Arc>>, -} - -impl FragmentBuffer { - /// Joins `body` onto a held fragment, or holds it if it opens a split payload. - /// - /// Assumes the continuation is the next POST: the Telemetry API sends fragments back to - /// back and gives no sequence number to correlate on. The repeated envelope, which - /// carries the cut record's timestamp, is what pairs the two halves. - fn stitch(&self, body: &[u8], error: &serde_json::Error) -> Stitch { - let mut held = self.held.lock().unwrap_or_else(PoisonError::into_inner); - - if let Some(Fragment { - body: mut joined, - continuation_prefix, - received, - }) = held.take() - { - if received.elapsed() > FRAGMENT_TTL { - debug!( - "TELEMETRY API | Dropping {} held bytes, no continuation arrived", - joined.len() - ); - } else if let Some(tail) = body.strip_prefix(continuation_prefix.as_slice()) { - // Left in place, the head's closers land inside the resumed value, where - // they parse but corrupt the record. - if joined.ends_with(APPENDED_CLOSERS) { - joined.truncate(joined.len() - APPENDED_CLOSERS.len()); - } - joined.extend_from_slice(tail); - - return match serde_json::from_slice(&joined) { - Ok(events) => Stitch::Complete(events), - // A record can be cut more than once, so keep going. - Err(e) => Self::hold(&mut held, joined, &e), - }; - } else { - debug!( - "TELEMETRY API | Dropping {} held bytes, next payload does not continue it", - joined.len() - ); - } - } - - Self::hold(&mut held, body.to_vec(), error) - } - - fn hold(held: &mut Option, body: Vec, error: &serde_json::Error) -> Stitch { - *held = Fragment::opening(body, error); - if held.is_some() { - Stitch::Pending - } else { - Stitch::Discarded - } - } -} - -struct Fragment { - body: Vec, - /// What the continuation POST repeats before the resumed bytes: `[` followed by the - /// keys of the cut record up to and including `"record":`. - continuation_prefix: Vec, - received: Instant, -} - -impl Fragment { - /// Holds `body` only if it opens a split payload: an array that ran out of input inside - /// the value of its last record. Any other parse failure won't be fixed by joining, and - /// holding it would poison the next stitch. - fn opening(body: Vec, error: &serde_json::Error) -> Option { - if error.classify() != Category::Eof - || body.len() > MAX_FRAGMENT_BYTES - || body.first() != Some(&b'[') - { - return None; - } - - // The cut record is the last one in the body. - let key_end = rfind(&body, RECORD_KEY)? + RECORD_KEY.len(); - let record_start = rfind(body.get(..key_end)?, RECORD_START)?; - - let mut continuation_prefix = vec![b'[']; - continuation_prefix.extend_from_slice(body.get(record_start..key_end)?); - - Some(Self { - body, - continuation_prefix, - received: Instant::now(), - }) - } -} - -/// Byte offset of the last occurrence of `needle` in `haystack`. -fn rfind(haystack: &[u8], needle: &[u8]) -> Option { - haystack - .windows(needle.len()) - .rposition(|window| window == needle) -} - #[cfg(test)] mod tests { use super::*; @@ -300,25 +165,8 @@ mod tests { use axum::http::Request; use chrono::DateTime; - use crate::extension::telemetry::events::{ - InitPhase, InitType, RuntimeDoneMetrics, Status, TelemetryRecord, - }; - - /// Leading half of a split payload: a `function` record cut inside its `message` value, - /// with the framing the Telemetry API writes anyway. - const HEAD: &str = - r#"[{"time":"2026-09-03T14:29:52.929Z","type":"function","record":{"message":"AAAA}]"#; - - /// Its continuation, which repeats the cut record's envelope before the resumed bytes - /// and carries the rest of the batch. - const TAIL: &str = r#"[{"time":"2026-09-03T14:29:52.929Z","type":"function","record":BBBB"}},{"time":"2026-09-03T14:29:52.930Z","type":"platform.runtimeDone","record":{"requestId":"abc123","status":"success","metrics":{"durationMs":18.074,"producedBytes":329814}}}]"#; - - /// Mirrors the handler: a body only reaches the buffer once it has failed to parse. - fn stitch(fragments: &FragmentBuffer, body: &str) -> Stitch { - let error = serde_json::from_slice::>(body.as_bytes()) - .expect_err("fixture must not parse on its own"); - fragments.stitch(body.as_bytes(), &error) - } + use crate::extension::telemetry::events::{InitPhase, InitType, TelemetryRecord}; + use crate::extension::telemetry::stitch::fixtures::{HEAD, TAIL}; fn state(logs_tx: Sender) -> HandlerState { HandlerState { @@ -335,114 +183,6 @@ mod tests { .expect("failed to build request") } - #[test] - fn test_stitch_split_payload() { - let fragments = FragmentBuffer::default(); - - assert!(matches!(stitch(&fragments, HEAD), Stitch::Pending)); - - let Stitch::Complete(events) = stitch(&fragments, TAIL) else { - panic!("expected the continuation to complete the payload"); - }; - - assert_eq!(events.len(), 2); - // `AAAA}]BBBB` would mean the head's framing was left in the resumed value. - assert_eq!( - events[0].record, - TelemetryRecord::Function(serde_json::json!({"message": "AAAABBBB"})) - ); - assert_eq!( - events[1].record, - TelemetryRecord::PlatformRuntimeDone { - request_id: "abc123".to_string(), - status: Status::Success, - error_type: None, - metrics: Some(RuntimeDoneMetrics { - duration_ms: 18.074, - produced_bytes: Some(329_814), - }), - } - ); - - // The pair is consumed, so a following payload starts from nothing. - assert!(matches!(stitch(&fragments, TAIL), Stitch::Discarded)); - } - - #[test] - fn test_complete_payload_is_not_held() { - let fragments = FragmentBuffer::default(); - - // Parses as JSON, so it failed for a reason joining won't fix. Holding it would - // poison the next stitch. - let unsupported = - r#"[{"time":"2026-09-03T14:29:52.929Z","type":"platform.brandNew","record":{}}]"#; - assert!(matches!(stitch(&fragments, unsupported), Stitch::Discarded)); - - assert!(matches!(stitch(&fragments, HEAD), Stitch::Pending)); - assert!(matches!(stitch(&fragments, TAIL), Stitch::Complete(_))); - } - - #[test] - fn test_head_dropped_when_next_payload_is_unrelated() { - let fragments = FragmentBuffer::default(); - - assert!(matches!(stitch(&fragments, HEAD), Stitch::Pending)); - - // A different record's envelope: not this head's continuation, so the head goes and - // this one is held in its place. - let other = - r#"[{"time":"2026-09-03T14:30:11.001Z","type":"function","record":{"message":"CCCC}]"#; - assert!(matches!(stitch(&fragments, other), Stitch::Pending)); - - // Proof the original head is gone: its continuation no longer stitches. - assert!(matches!(stitch(&fragments, TAIL), Stitch::Discarded)); - } - - #[test] - fn test_stitch_head_holding_complete_records() { - let fragments = FragmentBuffer::default(); - - // The cut record is the last of several, so the envelope the continuation repeats is - // in the middle of the head. - let head = format!( - r#"[{{"time":"2026-09-03T14:29:52.900Z","type":"extension","record":"ready"}},{}"#, - HEAD.trim_start_matches('[') - ); - assert!(matches!(stitch(&fragments, &head), Stitch::Pending)); - - let Stitch::Complete(events) = stitch(&fragments, TAIL) else { - panic!("expected the continuation to complete the payload"); - }; - assert_eq!(events.len(), 3); - assert_eq!( - events[1].record, - TelemetryRecord::Function(serde_json::json!({"message": "AAAABBBB"})) - ); - } - - #[tokio::test] - async fn test_handle_split_payload() { - let (tx, mut rx) = tokio::sync::mpsc::channel(2); - let state = state(tx); - - let response = - TelemetryListener::handle(axum::extract::State(state.clone()), post(HEAD)).await; - assert_eq!(response.status(), StatusCode::OK); - assert!(rx.try_recv().is_err(), "held fragment must not be emitted"); - - let response = TelemetryListener::handle(axum::extract::State(state), post(TAIL)).await; - assert_eq!(response.status(), StatusCode::OK); - - assert!(matches!( - rx.try_recv().expect("function record").record, - TelemetryRecord::Function(_) - )); - assert!(matches!( - rx.try_recv().expect("runtimeDone record").record, - TelemetryRecord::PlatformRuntimeDone { .. } - )); - } - #[tokio::test] #[allow(clippy::unwrap_used)] async fn test_handle() { @@ -477,4 +217,28 @@ mod tests { runtime_version_arn: Some("arn:aws:lambda:us-east-1::runtime:da57c20c4b965d5b75540f6865a35fc8030358e33ec44ecfed33e90901a27a72".to_string()), }); } + + /// A split payload is held on the first POST and forwarded whole on the second. + #[tokio::test] + async fn test_handle_split_payload() { + let (tx, mut rx) = tokio::sync::mpsc::channel(2); + let state = state(tx); + + let response = + TelemetryListener::handle(axum::extract::State(state.clone()), post(HEAD)).await; + assert_eq!(response.status(), StatusCode::OK); + assert!(rx.try_recv().is_err(), "held fragment must not be emitted"); + + let response = TelemetryListener::handle(axum::extract::State(state), post(TAIL)).await; + assert_eq!(response.status(), StatusCode::OK); + + assert!(matches!( + rx.try_recv().expect("function record").record, + TelemetryRecord::Function(_) + )); + assert!(matches!( + rx.try_recv().expect("runtimeDone record").record, + TelemetryRecord::PlatformRuntimeDone { .. } + )); + } } diff --git a/bottlecap/src/extension/telemetry/mod.rs b/bottlecap/src/extension/telemetry/mod.rs index ce703653b..8d504e7eb 100644 --- a/bottlecap/src/extension/telemetry/mod.rs +++ b/bottlecap/src/extension/telemetry/mod.rs @@ -5,6 +5,7 @@ use crate::extension::{EXTENSION_ID_HEADER, base_url}; pub mod events; pub mod listener; +mod stitch; pub const TELEMETRY_SUBSCRIPTION_ROUTE: &str = "2022-07-01/telemetry"; diff --git a/bottlecap/src/extension/telemetry/stitch.rs b/bottlecap/src/extension/telemetry/stitch.rs new file mode 100644 index 000000000..89c1f9c96 --- /dev/null +++ b/bottlecap/src/extension/telemetry/stitch.rs @@ -0,0 +1,286 @@ +//! Reassembles telemetry payloads that the Telemetry API split across two POSTs. +//! +//! A record larger than the subscription's `maxBytes` is cut mid-value. The rest of it — plus +//! the rest of the batch — arrives in the next POST, which repeats the cut record's envelope +//! ahead of the resumed bytes. Neither half parses on its own: +//! +//! ```text +//! POST 1 [{"time":"..929Z","type":"function","record":{"message":"iVBORw0KGgoAAA}] +//! `------------- envelope --------------------'`--- cut here ---'`framing' +//! +//! POST 2 [{"time":"..929Z","type":"function","record":QICAgIfAhkiAAA"}},{..},{..runtimeDone..}] +//! `------- the same envelope, repeated -------'`-- resumed --'`- rest of the batch -' +//! ``` +//! +//! So the two are joined by dropping POST 1's framing and POST 2's repeated envelope. That +//! byte-identical envelope, which carries the cut record's timestamp, is the only thing +//! pairing them: the API sends no sequence number. +//! +//! Recovering the batch matters beyond the log line itself. `platform.runtimeDone` lands in +//! the second half, and the on-demand loop waits for it before calling `/next` — so dropping +//! the batch holds the invocation open until Lambda times out the sandbox. + +use serde_json::error::Category; +use std::{ + sync::{Arc, Mutex, PoisonError}, + time::{Duration, Instant}, +}; +use tracing::debug; + +use crate::extension::telemetry::events::TelemetryEvent; + +/// Ceiling on a held fragment. The API can POST up to `2 * maxBytes + metadataBytes`, which +/// fits here at the 1 MiB `maxBytes` we subscribe with. +const MAX_FRAGMENT_BYTES: usize = 2 * 1024 * 1024; + +/// Fragments arrive back to back, so one held this long is waiting on a continuation that +/// never came. +const FRAGMENT_TTL: Duration = Duration::from_secs(1); + +/// Precedes a record's value, so everything up to and including it is the envelope. +const RECORD_KEY: &[u8] = b"\"record\":"; + +/// Opens a record. +const RECORD_START: &[u8] = b"{\"time\":"; + +/// The API writes the record's closing `}` and the array's `]` even after cutting the +/// record's value short, so a fragment ends with framing that belongs to neither half. +const FRAMING: &[u8] = b"}]"; + +/// What came of pairing an unparseable body with a held fragment. +#[derive(Debug)] +pub(crate) enum Stitch { + /// The body completed a split payload. + Complete(Vec), + /// The body opens a split payload, and is held for its continuation. + Pending, + /// The body is not part of a split payload. + Discarded, +} + +/// Holds the leading half of a split payload until its continuation arrives. +#[derive(Clone, Default)] +pub(crate) struct FragmentBuffer { + held: Arc>>, +} + +impl FragmentBuffer { + /// Joins `body` onto the held fragment, or holds `body` if it opens a split payload. + /// + /// `error` is the failure `body` produced on its own; it is how a payload cut mid-record + /// is told apart from one we simply can't interpret. + pub(crate) fn stitch(&self, body: &[u8], error: &serde_json::Error) -> Stitch { + let mut slot = self.held.lock().unwrap_or_else(PoisonError::into_inner); + + if let Some(stale) = slot.take_if(|held| held.received.elapsed() > FRAGMENT_TTL) { + debug!( + "TELEMETRY API | Dropping {} held bytes, no continuation arrived", + stale.body.len() + ); + } + + if let Some(head) = slot.take() { + if let Some(resumed) = head.resumed_bytes(body) { + let joined = head.join(resumed); + return match serde_json::from_slice(&joined) { + Ok(events) => Stitch::Complete(events), + // A record can be cut more than once, so keep accumulating. + Err(e) => hold_or_discard(&mut slot, joined, &e), + }; + } + + debug!( + "TELEMETRY API | Dropping {} held bytes, the next payload does not continue it", + head.body.len() + ); + } + + hold_or_discard(&mut slot, body.to_vec(), error) + } +} + +/// Holds `body` for its continuation, or reports that nothing can be recovered from it. +fn hold_or_discard( + slot: &mut Option, + body: Vec, + error: &serde_json::Error, +) -> Stitch { + *slot = Fragment::from_cut_payload(body, error); + if slot.is_some() { + Stitch::Pending + } else { + Stitch::Discarded + } +} + +/// The leading half of a split payload. +struct Fragment { + body: Vec, + /// The envelope the continuation repeats before the resumed bytes. + repeated_envelope: Vec, + received: Instant, +} + +impl Fragment { + /// A fragment, if `body` is the leading half of a split payload: an array that ran out of + /// input inside its last record's value. Any other parse failure won't be fixed by + /// joining, and holding such a body would poison the next stitch. + fn from_cut_payload(body: Vec, error: &serde_json::Error) -> Option { + if error.classify() != Category::Eof + || body.len() > MAX_FRAGMENT_BYTES + || body.first() != Some(&b'[') + { + return None; + } + + // Rebuild the cut record's envelope as the continuation will send it: `[` then the + // record's keys, up to the value that got cut. The cut record is the last one here. + let value_start = rfind(&body, RECORD_KEY)? + RECORD_KEY.len(); + let record_start = rfind(body.get(..value_start)?, RECORD_START)?; + + let mut repeated_envelope = vec![b'[']; + repeated_envelope.extend_from_slice(body.get(record_start..value_start)?); + + Some(Self { + body, + repeated_envelope, + received: Instant::now(), + }) + } + + /// The bytes that resume this fragment, if `body` is its continuation. + fn resumed_bytes<'a>(&self, body: &'a [u8]) -> Option<&'a [u8]> { + body.strip_prefix(self.repeated_envelope.as_slice()) + } + + /// Joins the resumed bytes on, dropping the framing: left in place it would land inside + /// the resumed value, where it parses but corrupts the record. + fn join(mut self, resumed: &[u8]) -> Vec { + if self.body.ends_with(FRAMING) { + self.body.truncate(self.body.len() - FRAMING.len()); + } + self.body.extend_from_slice(resumed); + self.body + } +} + +/// Offset of the last occurrence of `needle` in `haystack`. +fn rfind(haystack: &[u8], needle: &[u8]) -> Option { + haystack + .windows(needle.len()) + .rposition(|window| window == needle) +} + +/// The two halves of a real split payload, trimmed to the bytes that matter. +#[cfg(test)] +pub(crate) mod fixtures { + /// A `function` record cut inside its `message` value, plus the framing. + pub(crate) const HEAD: &str = + r#"[{"time":"2026-09-03T14:29:52.929Z","type":"function","record":{"message":"AAAA}]"#; + + /// The continuation: the same envelope, the resumed bytes, then the rest of the batch. + pub(crate) const TAIL: &str = r#"[{"time":"2026-09-03T14:29:52.929Z","type":"function","record":BBBB"}},{"time":"2026-09-03T14:29:52.930Z","type":"platform.runtimeDone","record":{"requestId":"abc123","status":"success","metrics":{"durationMs":18.074,"producedBytes":329814}}}]"#; +} + +#[cfg(test)] +mod tests { + use super::fixtures::{HEAD, TAIL}; + use super::*; + use crate::extension::telemetry::events::{RuntimeDoneMetrics, Status, TelemetryRecord}; + + /// Mirrors the handler: a body only reaches the buffer once it has failed to parse. + fn stitch(fragments: &FragmentBuffer, body: &str) -> Stitch { + let error = serde_json::from_slice::>(body.as_bytes()) + .expect_err("fixture must not parse on its own"); + fragments.stitch(body.as_bytes(), &error) + } + + /// The events of a stitch that should have completed, reporting what came back if it did not. + fn completed(stitch: Stitch) -> Vec { + match stitch { + Stitch::Complete(events) => events, + other => panic!("expected the continuation to complete the payload, got {other:?}"), + } + } + + #[test] + fn joins_a_split_payload() { + let fragments = FragmentBuffer::default(); + + assert!(matches!(stitch(&fragments, HEAD), Stitch::Pending)); + + let events = completed(stitch(&fragments, TAIL)); + + assert_eq!(events.len(), 2); + // `AAAA}]BBBB` would mean the head's framing was left in the resumed value. + assert_eq!( + events[0].record, + TelemetryRecord::Function(serde_json::json!({"message": "AAAABBBB"})) + ); + assert_eq!( + events[1].record, + TelemetryRecord::PlatformRuntimeDone { + request_id: "abc123".to_string(), + status: Status::Success, + error_type: None, + metrics: Some(RuntimeDoneMetrics { + duration_ms: 18.074, + produced_bytes: Some(329_814), + }), + } + ); + + // The pair is consumed, so a following payload starts from nothing. + assert!(matches!(stitch(&fragments, TAIL), Stitch::Discarded)); + } + + #[test] + fn joins_a_payload_cut_after_whole_records() { + let fragments = FragmentBuffer::default(); + + // The cut record is the last of several, so the envelope the continuation repeats is + // in the middle of the fragment. + let head = format!( + r#"[{{"time":"2026-09-03T14:29:52.900Z","type":"extension","record":"ready"}},{}"#, + HEAD.trim_start_matches('[') + ); + assert!(matches!(stitch(&fragments, &head), Stitch::Pending)); + + let events = completed(stitch(&fragments, TAIL)); + assert_eq!(events.len(), 3); + assert_eq!( + events[1].record, + TelemetryRecord::Function(serde_json::json!({"message": "AAAABBBB"})) + ); + } + + #[test] + fn does_not_hold_a_payload_that_arrived_whole() { + let fragments = FragmentBuffer::default(); + + // Parses as JSON, so it failed for a reason joining won't fix. Holding it would + // poison the next stitch. + let unsupported = + r#"[{"time":"2026-09-03T14:29:52.929Z","type":"platform.brandNew","record":{}}]"#; + assert!(matches!(stitch(&fragments, unsupported), Stitch::Discarded)); + + assert!(matches!(stitch(&fragments, HEAD), Stitch::Pending)); + assert!(matches!(stitch(&fragments, TAIL), Stitch::Complete(_))); + } + + #[test] + fn drops_a_fragment_the_next_payload_does_not_continue() { + let fragments = FragmentBuffer::default(); + + assert!(matches!(stitch(&fragments, HEAD), Stitch::Pending)); + + // A different record's envelope, so the held fragment goes and this one takes its + // place. + let other = + r#"[{"time":"2026-09-03T14:30:11.001Z","type":"function","record":{"message":"CCCC}]"#; + assert!(matches!(stitch(&fragments, other), Stitch::Pending)); + + // Proof the first fragment is gone: its own continuation no longer joins. + assert!(matches!(stitch(&fragments, TAIL), Stitch::Discarded)); + } +} From 7e3ac18784ab0215897b9c0136bce6ce72e51b7d Mon Sep 17 00:00:00 2001 From: "shreya.malpani" Date: Fri, 4 Sep 2026 07:41:22 -0700 Subject: [PATCH 4/8] fixes --- bottlecap/src/extension/telemetry/listener.rs | 6 ++++-- bottlecap/src/extension/telemetry/mod.rs | 12 +++++++----- bottlecap/src/extension/telemetry/stitch.rs | 7 ++++--- 3 files changed, 15 insertions(+), 10 deletions(-) diff --git a/bottlecap/src/extension/telemetry/listener.rs b/bottlecap/src/extension/telemetry/listener.rs index 693767df5..56e96647d 100644 --- a/bottlecap/src/extension/telemetry/listener.rs +++ b/bottlecap/src/extension/telemetry/listener.rs @@ -143,8 +143,10 @@ impl TelemetryListener { // Instead, log it and move on. // This will result in a dropped payload, but may be from // events we haven't added support for yet - let body = String::from_utf8_lossy(&body); - debug!("Failed to parse telemetry events `{body}`, failed with: {e}"); + debug!( + "TELEMETRY API | Failed to parse telemetry events ({} bytes), failed with: {e}", + body.len() + ); return (StatusCode::OK, "Failed to parse telemetry events").into_response(); } }, diff --git a/bottlecap/src/extension/telemetry/mod.rs b/bottlecap/src/extension/telemetry/mod.rs index 8d504e7eb..93b362917 100644 --- a/bottlecap/src/extension/telemetry/mod.rs +++ b/bottlecap/src/extension/telemetry/mod.rs @@ -73,15 +73,17 @@ pub async fn subscribe( .send() .await?; - if response.status().is_success() { - debug!("EXTENSION | Subscribed to Telemetry API: {:?}", response); - } else { - // A rejected subscription means no telemetry at all, including platform.runtimeDone, so - // surface it rather than leaving it at debug level. + if let Err(e) = response.error_for_status_ref() { + // A rejected subscription means no telemetry at all, including platform.runtimeDone. + // Fail rather than run blind: the caller falls back to the idle loop, which calls + // /next straight through instead of waiting for a runtimeDone that never arrives. error!( "EXTENSION | Telemetry API rejected subscription with status {}", response.status() ); + return Err(e.into()); } + + debug!("EXTENSION | Subscribed to Telemetry API: {:?}", response); Ok(response) } diff --git a/bottlecap/src/extension/telemetry/stitch.rs b/bottlecap/src/extension/telemetry/stitch.rs index 89c1f9c96..d6c40d349 100644 --- a/bottlecap/src/extension/telemetry/stitch.rs +++ b/bottlecap/src/extension/telemetry/stitch.rs @@ -29,9 +29,10 @@ use tracing::debug; use crate::extension::telemetry::events::TelemetryEvent; -/// Ceiling on a held fragment. The API can POST up to `2 * maxBytes + metadataBytes`, which -/// fits here at the 1 MiB `maxBytes` we subscribe with. -const MAX_FRAGMENT_BYTES: usize = 2 * 1024 * 1024; +/// Ceiling on a held fragment. Sized above the largest single POST — the API can send up to +/// `2 * maxBytes + metadataBytes`, so over 2 MiB at the 1 MiB `maxBytes` we subscribe with — +/// while still bounding the accumulation when one record is cut repeatedly. +const MAX_FRAGMENT_BYTES: usize = 4 * 1024 * 1024; /// Fragments arrive back to back, so one held this long is waiting on a continuation that /// never came. From b7e216004ff47f223a95e155e20889001be104c6 Mon Sep 17 00:00:00 2001 From: "shreya.malpani" Date: Fri, 4 Sep 2026 10:41:53 -0700 Subject: [PATCH 5/8] locate the record envelope with a depth-aware scan --- bottlecap/src/extension/telemetry/stitch.rs | 99 ++++++++++++++++++--- 1 file changed, 89 insertions(+), 10 deletions(-) diff --git a/bottlecap/src/extension/telemetry/stitch.rs b/bottlecap/src/extension/telemetry/stitch.rs index d6c40d349..5e37a4ff9 100644 --- a/bottlecap/src/extension/telemetry/stitch.rs +++ b/bottlecap/src/extension/telemetry/stitch.rs @@ -41,9 +41,6 @@ const FRAGMENT_TTL: Duration = Duration::from_secs(1); /// Precedes a record's value, so everything up to and including it is the envelope. const RECORD_KEY: &[u8] = b"\"record\":"; -/// Opens a record. -const RECORD_START: &[u8] = b"{\"time\":"; - /// The API writes the record's closing `}` and the array's `]` even after cutting the /// record's value short, so a fragment ends with framing that belongs to neither half. const FRAMING: &[u8] = b"}]"; @@ -136,8 +133,7 @@ impl Fragment { // Rebuild the cut record's envelope as the continuation will send it: `[` then the // record's keys, up to the value that got cut. The cut record is the last one here. - let value_start = rfind(&body, RECORD_KEY)? + RECORD_KEY.len(); - let record_start = rfind(body.get(..value_start)?, RECORD_START)?; + let (record_start, value_start) = envelope_bounds(&body)?; let mut repeated_envelope = vec![b'[']; repeated_envelope.extend_from_slice(body.get(record_start..value_start)?); @@ -165,11 +161,53 @@ impl Fragment { } } -/// Offset of the last occurrence of `needle` in `haystack`. -fn rfind(haystack: &[u8], needle: &[u8]) -> Option { - haystack - .windows(needle.len()) - .rposition(|window| window == needle) +/// Where the last record opens, and where its `record` value begins. +/// +/// Depth- and string-aware, because a structured record can carry a `record` key of its own. +/// A plain search would find that one instead — always at a later offset, so the envelope +/// would come out too long, fail to match the continuation, and drop the batch. +fn envelope_bounds(body: &[u8]) -> Option<(usize, usize)> { + let mut depth = 0usize; + let mut in_string = false; + let mut escaped = false; + let mut record_start = None; + let mut value_start = None; + + for (i, &byte) in body.iter().enumerate() { + if in_string { + match byte { + _ if escaped => escaped = false, + b'\\' => escaped = true, + b'"' => in_string = false, + _ => {} + } + continue; + } + + match byte { + b'"' => { + // A key at depth 2 sits directly on a record, so this is the envelope's. + if depth == 2 && body[i..].starts_with(RECORD_KEY) { + value_start = Some(i + RECORD_KEY.len()); + } + in_string = true; + } + b'{' => { + if depth == 1 { + record_start = Some(i); + } + depth += 1; + } + b'[' => depth += 1, + b'}' | b']' => depth = depth.saturating_sub(1), + _ => {} + } + } + + // Out of order means the payload was cut before the last record reached its `record` key, + // leaving nothing to pair the continuation against. + let (record_start, value_start) = (record_start?, value_start?); + (record_start < value_start).then_some((record_start, value_start)) } /// The two halves of a real split payload, trimmed to the bytes that matter. @@ -255,6 +293,47 @@ mod tests { ); } + /// A structured record can nest a `record` key of its own, which must not be mistaken for + /// the envelope's. + #[test] + fn joins_a_structured_record_that_nests_a_record_key() { + let fragments = FragmentBuffer::default(); + + let head = r#"[{"time":"2026-09-03T14:29:52.929Z","type":"function","record":{"level":"INFO","message":{"record":"AAAA}]"#; + let tail = r#"[{"time":"2026-09-03T14:29:52.929Z","type":"function","record":BBBB"}}},{"time":"2026-09-03T14:29:52.930Z","type":"platform.runtimeDone","record":{"requestId":"abc123","status":"success","metrics":{"durationMs":18.074,"producedBytes":329814}}}]"#; + + assert!(matches!(stitch(&fragments, head), Stitch::Pending)); + + let events = completed(stitch(&fragments, tail)); + assert_eq!(events.len(), 2); + assert_eq!( + events[0].record, + TelemetryRecord::Function( + serde_json::json!({"level": "INFO", "message": {"record": "AAAABBBB"}}) + ) + ); + } + + /// The same key inside a string is escaped, so it never looked like the envelope's. + #[test] + fn joins_a_message_whose_text_looks_like_a_record_key() { + let fragments = FragmentBuffer::default(); + + let head = r#"[{"time":"2026-09-03T14:29:52.929Z","type":"function","record":{"message":"{\"record\":\"AAAA}]"#; + let tail = r#"[{"time":"2026-09-03T14:29:52.929Z","type":"function","record":BBBB\"}"}},{"time":"2026-09-03T14:29:52.930Z","type":"platform.runtimeDone","record":{"requestId":"abc123","status":"success","metrics":{"durationMs":18.074,"producedBytes":329814}}}]"#; + + assert!(matches!(stitch(&fragments, head), Stitch::Pending)); + + let events = completed(stitch(&fragments, tail)); + assert_eq!(events.len(), 2); + assert_eq!( + events[0].record, + TelemetryRecord::Function( + serde_json::json!({"message": r#"{"record":"AAAABBBB"}"#}) + ) + ); + } + #[test] fn does_not_hold_a_payload_that_arrived_whole() { let fragments = FragmentBuffer::default(); From bb0251d1a410607fbfe34de4e748e76dd153aa77 Mon Sep 17 00:00:00 2001 From: "shreya.malpani" Date: Fri, 4 Sep 2026 11:35:02 -0700 Subject: [PATCH 6/8] replace rfind search --- bottlecap/src/extension/telemetry/stitch.rs | 73 +++++---------------- 1 file changed, 15 insertions(+), 58 deletions(-) diff --git a/bottlecap/src/extension/telemetry/stitch.rs b/bottlecap/src/extension/telemetry/stitch.rs index 5e37a4ff9..ae9d07f7d 100644 --- a/bottlecap/src/extension/telemetry/stitch.rs +++ b/bottlecap/src/extension/telemetry/stitch.rs @@ -114,8 +114,6 @@ fn hold_or_discard( /// The leading half of a split payload. struct Fragment { body: Vec, - /// The envelope the continuation repeats before the resumed bytes. - repeated_envelope: Vec, received: Instant, } @@ -131,23 +129,24 @@ impl Fragment { return None; } - // Rebuild the cut record's envelope as the continuation will send it: `[` then the - // record's keys, up to the value that got cut. The cut record is the last one here. - let (record_start, value_start) = envelope_bounds(&body)?; - - let mut repeated_envelope = vec![b'[']; - repeated_envelope.extend_from_slice(body.get(record_start..value_start)?); - Some(Self { body, - repeated_envelope, received: Instant::now(), }) } /// The bytes that resume this fragment, if `body` is its continuation. + /// + /// A continuation opens with the cut record's envelope, so its `record` key is the first + /// one in the payload — anything the customer nested sits inside the value that follows. + /// Finding the same envelope in the fragment is what pairs the two. fn resumed_bytes<'a>(&self, body: &'a [u8]) -> Option<&'a [u8]> { - body.strip_prefix(self.repeated_envelope.as_slice()) + let payload = body.strip_prefix(b"[")?; + let value_start = find(payload, RECORD_KEY)? + RECORD_KEY.len(); + let envelope = payload.get(..value_start)?; + + find(&self.body, envelope)?; + payload.get(value_start..) } /// Joins the resumed bytes on, dropping the framing: left in place it would land inside @@ -161,53 +160,11 @@ impl Fragment { } } -/// Where the last record opens, and where its `record` value begins. -/// -/// Depth- and string-aware, because a structured record can carry a `record` key of its own. -/// A plain search would find that one instead — always at a later offset, so the envelope -/// would come out too long, fail to match the continuation, and drop the batch. -fn envelope_bounds(body: &[u8]) -> Option<(usize, usize)> { - let mut depth = 0usize; - let mut in_string = false; - let mut escaped = false; - let mut record_start = None; - let mut value_start = None; - - for (i, &byte) in body.iter().enumerate() { - if in_string { - match byte { - _ if escaped => escaped = false, - b'\\' => escaped = true, - b'"' => in_string = false, - _ => {} - } - continue; - } - - match byte { - b'"' => { - // A key at depth 2 sits directly on a record, so this is the envelope's. - if depth == 2 && body[i..].starts_with(RECORD_KEY) { - value_start = Some(i + RECORD_KEY.len()); - } - in_string = true; - } - b'{' => { - if depth == 1 { - record_start = Some(i); - } - depth += 1; - } - b'[' => depth += 1, - b'}' | b']' => depth = depth.saturating_sub(1), - _ => {} - } - } - - // Out of order means the payload was cut before the last record reached its `record` key, - // leaving nothing to pair the continuation against. - let (record_start, value_start) = (record_start?, value_start?); - (record_start < value_start).then_some((record_start, value_start)) +/// Offset of the first occurrence of `needle` in `haystack`. +fn find(haystack: &[u8], needle: &[u8]) -> Option { + haystack + .windows(needle.len()) + .position(|window| window == needle) } /// The two halves of a real split payload, trimmed to the bytes that matter. From 7055b50d096f867c5a09cc949d5d88a34e91431b Mon Sep 17 00:00:00 2001 From: "shreya.malpani" Date: Tue, 8 Sep 2026 07:43:51 -0700 Subject: [PATCH 7/8] cargo fmt fix --- bottlecap/src/extension/telemetry/stitch.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/bottlecap/src/extension/telemetry/stitch.rs b/bottlecap/src/extension/telemetry/stitch.rs index ae9d07f7d..a375584c2 100644 --- a/bottlecap/src/extension/telemetry/stitch.rs +++ b/bottlecap/src/extension/telemetry/stitch.rs @@ -285,9 +285,7 @@ mod tests { assert_eq!(events.len(), 2); assert_eq!( events[0].record, - TelemetryRecord::Function( - serde_json::json!({"message": r#"{"record":"AAAABBBB"}"#}) - ) + TelemetryRecord::Function(serde_json::json!({"message": r#"{"record":"AAAABBBB"}"#})) ); } From d884824d36e1d9e536936a0a1053bbd99cc68826 Mon Sep 17 00:00:00 2001 From: "shreya.malpani" Date: Tue, 8 Sep 2026 09:20:23 -0700 Subject: [PATCH 8/8] fix position of runtime_id and extension_id --- bottlecap/src/bin/bottlecap/main.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bottlecap/src/bin/bottlecap/main.rs b/bottlecap/src/bin/bottlecap/main.rs index 70739f716..d293cbc5b 100644 --- a/bottlecap/src/bin/bottlecap/main.rs +++ b/bottlecap/src/bin/bottlecap/main.rs @@ -289,7 +289,7 @@ async fn extension_loop_idle( aws_config: &AwsConfig, ) -> anyhow::Result<()> { loop { - match extension::next_event(client, &r.extension_id, &aws_config.runtime_api).await { + match extension::next_event(client, &aws_config.runtime_api, &r.extension_id).await { Ok(_) => { debug!("Extension is idle, skipping next event"); }