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/bin/bottlecap/main.rs b/bottlecap/src/bin/bottlecap/main.rs index 8124e4a87..f47bcd26f 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(), ) @@ -676,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, @@ -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..e9b6a6e71 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 { @@ -153,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() @@ -220,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() { diff --git a/bottlecap/src/logs/agent.rs b/bottlecap/src/logs/agent.rs index 31eb78930..0cf4bc60f 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,113 @@ 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}}}"#; + 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()); + 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))); + } + + #[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"); + } +}