fix(shutdown): send Tombstone after telemetry is drained

.gitlab/datasources/test-suites.yaml

@@ -10,3 +10,4 @@ test_suites:
   - name: payload-size-e2e
   - name: dsm
   - name: durable-cold-start
+  - name: failure-accounting

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,20 +806,17 @@ 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
             drop(event_bus_tx);
 
             // 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<TelemetryEvent>,
-    event_bus_tx: Sender<Event>,
     logs_enabled: bool,
     managed_instance_mode: bool,
 ) -> anyhow::Result<CancellationToken> {
-    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) => {

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<TelemetryEvent>,
-    event_bus_tx: Sender<event_bus::Event>,
 }
 
 #[derive(Clone)]
@@ -41,19 +39,13 @@ struct HandlerState {
 
 impl TelemetryListener {
     #[must_use]
-    pub fn new(
-        host: [u8; 4],
-        port: u16,
-        logs_tx: Sender<TelemetryEvent>,
-        event_bus_tx: Sender<event_bus::Event>,
-    ) -> Self {
+    pub fn new(host: [u8; 4], port: u16, logs_tx: Sender<TelemetryEvent>) -> 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<event_bus::Event>,
-    ) {
+    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<HandlerState>, 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() {

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<DurableContextUpdate>,
     processor: LogsProcessor,
     aggregator_handle: AggregatorHandle,
+    event_bus: Sender<Event>,
     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<TelemetryEvent>, 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");
+    }
+}