From 3ea012074b9becd6644f3d4adec1dce9c17f5da4 Mon Sep 17 00:00:00 2001 From: Darcy Rayner Date: Thu, 22 Oct 2026 13:52:36 -0400 Subject: [PATCH 1/3] fix(shutdown): send Tombstone after telemetry is drained At shutdown the Telemetry API listener pushed the Tombstone straight onto the event bus, while telemetry events still had to pass through the logs agent. The main loop stops reading at the first Tombstone, so a platform.runtimeDone, platform.report or OOM log line still in the logs agent was dropped. For an invocation that timed out or crashed, that lost its errors/timeouts/OOM enhanced metrics. The logs agent now sends the Tombstone after its channel disconnects (the listener has shut down and dropped its sender) or a 250ms drain timeout elapses. This also replaces the busy-spin drain with an awaited recv. Co-Authored-By: Claude Opus 5.5 --- bottlecap/src/bin/bottlecap/main.rs | 17 +-- bottlecap/src/extension/telemetry/listener.rs | 30 ++--- bottlecap/src/logs/agent.rs | 110 ++++++++++++++---- 3 files changed, 103 insertions(+), 54 deletions(-) diff --git a/bottlecap/src/bin/bottlecap/main.rs b/bottlecap/src/bin/bottlecap/main.rs index 8124e4a87..c197efcff 100644 --- a/bottlecap/src/bin/bottlecap/main.rs +++ b/bottlecap/src/bin/bottlecap/main.rs @@ -478,7 +478,6 @@ async fn extension_loop_active( &r.extension_id, &aws_config.runtime_api, logs_agent_channel, - event_bus_tx.clone(), config.ext.serverless_logs_enabled, aws_config.is_managed_instance_mode(), ) @@ -807,12 +806,9 @@ async fn extension_loop_active( } if let NextEventResponse::Shutdown { .. } = maybe_shutdown_event { - // Cancel Telemetry API listener - // Important to do this first, so we can receive the Tombstone event which signals - // that there are no more Telemetry events to process + // Stop the Telemetry API listener, then let the logs agent drain. The agent sends the + // Tombstone only after it has forwarded every telemetry event to the event bus. telemetry_listener_cancel_token.cancel(); - - // Cancel Logs Agent which might have Telemetry API events to process logs_agent_cancel_token.cancel(); // Drop the event bus sender to allow the channel to close properly @@ -820,7 +816,7 @@ async fn extension_loop_active( // Redrive/block on any failed payloads flushing_service.await_handles().await; - // Wait for tombstone event from telemetry listener to ensure all events are processed + // Wait for the Tombstone from the logs agent to ensure all events are processed wait_for_tombstone_event( &mut event_bus, &invocation_processor_handle, @@ -862,7 +858,7 @@ async fn extension_loop_active( } } -/// Wait for the `Tombstone` event from telemetry listener to ensure all events are processed. +/// Wait for the `Tombstone` event from the logs agent to ensure all events are processed. /// This function will timeout after the specified duration to prevent hanging indefinitely. #[allow(clippy::too_many_arguments)] async fn wait_for_tombstone_event( @@ -1338,16 +1334,15 @@ async fn setup_telemetry_client( extension_id: &str, runtime_api: &str, logs_tx: Sender, - event_bus_tx: Sender, logs_enabled: bool, managed_instance_mode: bool, ) -> anyhow::Result { - let listener = TelemetryListener::new(EXTENSION_HOST_IP, TELEMETRY_PORT, logs_tx, event_bus_tx); + let listener = TelemetryListener::new(EXTENSION_HOST_IP, TELEMETRY_PORT, logs_tx); let cancel_token = listener.cancel_token(); match listener.start() { Ok(()) => { - // Drop the listener, so event_bus_tx is closed + // Drop the listener, so its `logs_tx` doesn't keep the logs agent's channel open drop(listener); } Err(e) => { diff --git a/bottlecap/src/extension/telemetry/listener.rs b/bottlecap/src/extension/telemetry/listener.rs index 7e8884e05..75ba987f2 100644 --- a/bottlecap/src/extension/telemetry/listener.rs +++ b/bottlecap/src/extension/telemetry/listener.rs @@ -1,5 +1,4 @@ use crate::{ - event_bus, extension::telemetry::{ events::TelemetryEvent, stitch::{FragmentBuffer, Stitch}, @@ -30,7 +29,6 @@ pub struct TelemetryListener { port: u16, cancel_token: CancellationToken, logs_tx: Sender, - event_bus_tx: Sender, } #[derive(Clone)] @@ -41,19 +39,13 @@ struct HandlerState { impl TelemetryListener { #[must_use] - pub fn new( - host: [u8; 4], - port: u16, - logs_tx: Sender, - event_bus_tx: Sender, - ) -> Self { + pub fn new(host: [u8; 4], port: u16, logs_tx: Sender) -> Self { let cancel_token = CancellationToken::new(); Self { host, port, cancel_token, logs_tx, - event_bus_tx, } } @@ -67,16 +59,18 @@ impl TelemetryListener { let router = self.make_router(); let cancel_token_clone = self.cancel_token(); - let event_bus_tx = self.event_bus_tx.clone(); tokio::spawn(async move { let listener = TcpListener::bind(&socket) .await .expect("Failed to bind socket"); debug!("TELEMETRY API | Starting listener on {}", socket); axum::serve(listener, router) - .with_graceful_shutdown(Self::graceful_shutdown(cancel_token_clone, event_bus_tx)) + .with_graceful_shutdown(Self::graceful_shutdown(cancel_token_clone)) .await .expect("Failed to start telemetry listener"); + // The router, and with it the last `logs_tx`, drops here. That disconnects the + // logs agent's channel, which then drains and sends the Tombstone. + debug!("TELEMETRY API | Shut down"); }); Ok(()) @@ -95,19 +89,9 @@ impl TelemetryListener { .layer(DefaultBodyLimit::max(MAX_BODY_BYTES)) } - async fn graceful_shutdown( - cancel_token: CancellationToken, - event_bus_tx: Sender, - ) { + async fn graceful_shutdown(cancel_token: CancellationToken) { cancel_token.cancelled().await; - debug!("TELEMETRY API | Shutdown signal received, sending tombstone event"); - - // Send tombstone event to signal shutdown - if let Err(e) = event_bus_tx.send(event_bus::Event::Tombstone).await { - debug!("TELEMETRY API |Failed to send tombstone event: {:?}", e); - } - - debug!("TELEMETRY API | Shutting down"); + debug!("TELEMETRY API | Shutdown signal received, shutting down"); } async fn handle(State(state): State, request: Request) -> Response { diff --git a/bottlecap/src/logs/agent.rs b/bottlecap/src/logs/agent.rs index 31eb78930..7bfb9b67e 100644 --- a/bottlecap/src/logs/agent.rs +++ b/bottlecap/src/logs/agent.rs @@ -1,5 +1,5 @@ use std::sync::Arc; -use std::time::{Duration, Instant}; +use std::time::Duration; use tokio::sync::mpsc::{self, Sender}; use tokio_util::sync::CancellationToken; use tracing::debug; @@ -10,7 +10,8 @@ use crate::logs::{aggregator_service::AggregatorHandle, processor::LogsProcessor use crate::tags; use crate::{LAMBDA_RUNTIME_SLUG, config}; -const DRAIN_LOG_INTERVAL: Duration = Duration::from_millis(100); +/// Upper bound on waiting for the Telemetry API listener to close the channel at shutdown. +const DRAIN_TIMEOUT: Duration = Duration::from_millis(250); use crate::logs::lambda::DurableContextUpdate; @@ -20,6 +21,7 @@ pub struct LogsAgent { durable_context_rx: mpsc::Receiver, processor: LogsProcessor, aggregator_handle: AggregatorHandle, + event_bus: Sender, cancel_token: CancellationToken, } @@ -35,7 +37,7 @@ impl LogsAgent { let processor = LogsProcessor::new( Arc::clone(&datadog_config), tags_provider, - event_bus, + event_bus.clone(), LAMBDA_RUNTIME_SLUG.to_string(), is_managed_instance_mode, ); @@ -49,6 +51,7 @@ impl LogsAgent { durable_context_rx, processor, aggregator_handle, + event_bus, cancel_token, }; @@ -67,26 +70,22 @@ impl LogsAgent { () = self.cancel_token.cancelled() => { debug!("LOGS_AGENT | Received shutdown signal, draining remaining events"); - // Drain remaining telemetry events - let mut last_drain_log_time = Instant::now().checked_sub(DRAIN_LOG_INTERVAL).expect("Failed to subtract interval from now"); - 'drain_logs_loop: loop { - match self.rx.try_recv() { - Ok(event) => { + // Drain until the listener has shut down and dropped its sender, so every + // telemetry event reaches the event bus ahead of the Tombstone. + let deadline = tokio::time::Instant::now() + DRAIN_TIMEOUT; + loop { + match tokio::time::timeout_at(deadline, self.rx.recv()).await { + Ok(Some(event)) => { self.processor.process(event, &self.aggregator_handle).await; } - Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => { + Ok(None) => { debug!("LOGS_AGENT | Channel disconnected, finished draining"); - break 'drain_logs_loop; - }, - // Empty signals there are still outstanding senders - Err(tokio::sync::mpsc::error::TryRecvError::Empty) => { - // Log at most once every 100ms to avoid spamming the logs - let now = Instant::now(); - if now.duration_since(last_drain_log_time) >= DRAIN_LOG_INTERVAL { - debug!("LOGS_AGENT | No more events to process but still have senders, continuing to drain..."); - last_drain_log_time = now; - } - }, + break; + } + Err(_) => { + debug!("LOGS_AGENT | Timed out waiting for channel to disconnect"); + break; + } } } @@ -99,6 +98,10 @@ impl LogsAgent { // Drain remaining held logs without durable context tags so no logs are lost. self.processor.drain_held_logs(&self.aggregator_handle); + if let Err(e) = self.event_bus.send(Event::Tombstone).await { + debug!("LOGS_AGENT | Failed to send tombstone event: {e}"); + } + break; } } @@ -123,3 +126,70 @@ impl LogsAgent { self.cancel_token.clone() } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::event_bus::EventBus; + use crate::extension::telemetry::events::TelemetryRecord; + use crate::logs::aggregator_service::AggregatorService; + use std::collections::HashMap; + + const RUNTIME_DONE: &str = r#"{"time":"2022-10-21T14:05:05.764Z","type":"platform.runtimeDone","record":{"requestId":"459921b5-681c-4a96-beb0-81e0aa586026","status":"timeout","metrics":{"durationMs":3000.0,"producedBytes":0}}}"#; + + fn agent() -> (LogsAgent, Sender, EventBus) { + let config = Arc::new(config::Config::default()); + let tags_provider = Arc::new(tags::provider::Provider::new( + Arc::clone(&config), + LAMBDA_RUNTIME_SLUG.to_string(), + &HashMap::from([("function_arn".to_string(), "test-arn".to_string())]), + )); + let (event_bus, event_bus_tx) = EventBus::run(); + let (aggregator_service, aggregator_handle) = AggregatorService::default(); + tokio::spawn(aggregator_service.run()); + let (agent, logs_tx, _) = LogsAgent::new( + tags_provider, + config, + event_bus_tx, + aggregator_handle, + false, + ); + (agent, logs_tx, event_bus) + } + + #[tokio::test] + async fn test_tombstone_follows_events_sent_after_cancel() { + let (mut agent, logs_tx, mut event_bus) = agent(); + let cancel_token = agent.cancel_token(); + let handle = tokio::spawn(async move { agent.spin().await }); + + // A telemetry POST that is still in flight when shutdown starts. + cancel_token.cancel(); + tokio::task::yield_now().await; + logs_tx + .send(serde_json::from_str(RUNTIME_DONE).expect("valid event")) + .await + .expect("agent is draining"); + drop(logs_tx); + + let first = event_bus.rx.recv().await.expect("event"); + assert!(matches!( + first, + Event::Telemetry(TelemetryEvent { + record: TelemetryRecord::PlatformRuntimeDone { .. }, + .. + }) + )); + assert!(matches!(event_bus.rx.recv().await, Some(Event::Tombstone))); + handle.await.expect("agent task"); + } + + #[tokio::test] + async fn test_tombstone_sent_when_sender_never_closes() { + let (mut agent, _logs_tx, mut event_bus) = agent(); + agent.cancel_token().cancel(); + agent.spin().await; + + assert!(matches!(event_bus.rx.try_recv(), Ok(Event::Tombstone))); + } +} From 3f7585cf72d6fd7995be53e347fe800c2dcd62fa Mon Sep 17 00:00:00 2001 From: Darcy Rayner Date: Thu, 22 Oct 2026 13:58:04 -0400 Subject: [PATCH 2/3] fix(telemetry): don't panic when the logs agent stops receiving at shutdown The logs agent's shutdown drain can now time out and drop its receiver while a Telemetry API request is still in flight. Return 503 instead of panicking, which aborts the process in release builds. Also correct the Managed Instance shutdown comment. Co-Authored-By: Claude Opus 5.5 --- bottlecap/src/bin/bottlecap/main.rs | 4 ++-- bottlecap/src/extension/telemetry/listener.rs | 22 ++++++++++++++++++- 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/bottlecap/src/bin/bottlecap/main.rs b/bottlecap/src/bin/bottlecap/main.rs index c197efcff..f47bcd26f 100644 --- a/bottlecap/src/bin/bottlecap/main.rs +++ b/bottlecap/src/bin/bottlecap/main.rs @@ -675,8 +675,8 @@ async fn extension_loop_active( // Shutdown sequence debug!("Initiating shutdown sequence"); - // Wait for tombstone event from telemetry listener to ensure all events are processed - // This is the result of code refactoring which is shared by OnDemand mode as well. + // Process remaining events. The logs agent isn't cancelled in this mode, so no Tombstone + // arrives and this ends on the idle timeout. wait_for_tombstone_event( &mut event_bus, &invocation_processor_handle, diff --git a/bottlecap/src/extension/telemetry/listener.rs b/bottlecap/src/extension/telemetry/listener.rs index 75ba987f2..e9b6a6e71 100644 --- a/bottlecap/src/extension/telemetry/listener.rs +++ b/bottlecap/src/extension/telemetry/listener.rs @@ -137,7 +137,10 @@ impl TelemetryListener { }; for event in telemetry_events.drain(..) { - state.logs_tx.send(event).await.expect("infallible"); + // The logs agent stops receiving once its shutdown drain times out. + if state.logs_tx.send(event).await.is_err() { + return (StatusCode::SERVICE_UNAVAILABLE, "Shutting down").into_response(); + } } (StatusCode::OK, "OK").into_response() @@ -204,6 +207,23 @@ mod tests { }); } + #[tokio::test] + async fn test_handle_after_receiver_dropped() { + let (tx, rx) = tokio::sync::mpsc::channel(1); + drop(rx); + + let response = TelemetryListener::handle( + axum::extract::State(state(tx)), + post(r#"[{"time":"2024-04-25T17:35:59.944Z","type":"platform.initStart","record":{"initializationType":"on-demand","phase":"init"}}]"#), + ) + .await; + + assert_eq!( + response.status(), + axum::http::StatusCode::SERVICE_UNAVAILABLE + ); + } + /// A split payload is held on the first POST and forwarded whole on the second. #[tokio::test] async fn test_handle_split_payload() { From beb8006b6ab85208045024abddee8d3265c1fae1 Mon Sep 17 00:00:00 2001 From: Darcy Rayner Date: Fri, 23 Oct 2026 12:50:09 -0400 Subject: [PATCH 3/3] test(logs): cover queued events at the drain deadline; register failure-accounting suite Events still queued when the logs agent's shutdown drain deadline passes are forwarded before the Tombstone, because timeout_at polls recv first. Add a test that pins this down. Register the failure-accounting integration suite now that the fix is in place. Co-Authored-By: Claude Opus 5.5 --- .gitlab/datasources/test-suites.yaml | 1 + bottlecap/src/logs/agent.rs | 43 ++++++++++++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/.gitlab/datasources/test-suites.yaml b/.gitlab/datasources/test-suites.yaml index 54a176c41..7da1e67fe 100644 --- a/.gitlab/datasources/test-suites.yaml +++ b/.gitlab/datasources/test-suites.yaml @@ -10,3 +10,4 @@ test_suites: - name: payload-size-e2e - name: dsm - name: durable-cold-start + - name: failure-accounting diff --git a/bottlecap/src/logs/agent.rs b/bottlecap/src/logs/agent.rs index 7bfb9b67e..0cf4bc60f 100644 --- a/bottlecap/src/logs/agent.rs +++ b/bottlecap/src/logs/agent.rs @@ -136,6 +136,8 @@ mod tests { use std::collections::HashMap; const RUNTIME_DONE: &str = r#"{"time":"2022-10-21T14:05:05.764Z","type":"platform.runtimeDone","record":{"requestId":"459921b5-681c-4a96-beb0-81e0aa586026","status":"timeout","metrics":{"durationMs":3000.0,"producedBytes":0}}}"#; + const FUNCTION_LOG: &str = + r#"{"time":"2022-10-21T14:05:04.000Z","type":"function","record":"hello"}"#; fn agent() -> (LogsAgent, Sender, EventBus) { let config = Arc::new(config::Config::default()); @@ -192,4 +194,45 @@ mod tests { assert!(matches!(event_bus.rx.try_recv(), Ok(Event::Tombstone))); } + + #[tokio::test] + async fn test_queued_events_forwarded_after_drain_timeout() { + const FILL: usize = 150; + let (mut agent, logs_tx, mut event_bus) = agent(); + let runtime_done: TelemetryEvent = serde_json::from_str(RUNTIME_DONE).expect("valid event"); + let function_log: TelemetryEvent = serde_json::from_str(FUNCTION_LOG).expect("valid event"); + // Events still queued when the drain deadline passes must be forwarded before the + // Tombstone: `timeout_at` polls `recv` first, and `recv` is ready while events are + // queued. The run of function logs, which aren't forwarded, makes the agent drain past + // the deadline without yielding. + for _ in 0..FILL { + logs_tx + .send(runtime_done.clone()) + .await + .expect("channel open"); + } + for _ in 0..500 { + logs_tx + .send(function_log.clone()) + .await + .expect("channel open"); + } + logs_tx.send(runtime_done).await.expect("channel open"); + agent.cancel_token().cancel(); + let handle = tokio::spawn(async move { agent.spin().await }); + + // Leave the event bus full until the drain deadline has passed. + tokio::time::sleep(DRAIN_TIMEOUT + Duration::from_millis(50)).await; + + let mut forwarded = 0; + loop { + match event_bus.rx.recv().await.expect("event") { + Event::Telemetry(_) => forwarded += 1, + Event::Tombstone => break, + Event::OutOfMemory { .. } => {} + } + } + assert_eq!(forwarded, FILL + 1); + handle.await.expect("agent task"); + } }