Aj/fix shutdown deadlock (#746)

We added this check for the platform.Report item because we'd
occasionally race past it when shutting down, leading to missing metrics
(namely post-runtime duration for the last request). This predominantly
affected hourly/infrequent functions.

But that change introduced a race condition where sometimes Lambda would
respond to /next with the shutdown event at the same time we'd receive
the last invocation's report event. In some cases, Tokio would run the
report event first followed by the response to /next, where then we'd
hit the shutdown event bus loop and wait for a report event that was
never coming.

This change also adds a 300ms timeout in case of a deadlock.

Jira: https://datadoghq.atlassian.net/browse/SLES-2151

Author: astuyve

bottlecap/src/bin/bottlecap/main.rs

@@ -669,28 +669,26 @@ async fn extension_loop_active(
             }
         }
 
-        if let NextEventResponse::Shutdown {
-            shutdown_reason, ..
-        } = maybe_shutdown_event
-        {
+        if let NextEventResponse::Shutdown { .. } = maybe_shutdown_event {
             // Redrive/block on any failed payloads
             let tf = trace_flusher.clone();
             pending_flush_handles
                 .await_flush_handles(&logs_flusher.clone(), &tf, &metrics_flushers)
                 .await;
-            // The Shutdown event we get during a timeout will
-            // never include a report log
-            if shutdown_reason != "timeout" {
-                'shutdown: loop {
-                    tokio::select! {
-                        Some(event) = event_bus.rx.recv() => {
-                            if let Some(telemetry_event) = handle_event_bus_event(event, invocation_processor.clone(), tags_provider.clone(), trace_processor.clone(), trace_agent_channel.clone()).await {
-                                if let TelemetryRecord::PlatformReport{ .. } = telemetry_event.record {
-                                    // Wait for the report event before shutting down
-                                    break 'shutdown;
-                                }
-                            }
+            // Wait for tombstone event from telemetry listener to ensure all events are processed
+            'shutdown: loop {
+                tokio::select! {
+                    Some(event) = event_bus.rx.recv() => {
+                    if let Event::Telemetry(TelemetryEvent { record: TelemetryRecord::PlatformTombstone, .. }) = event {
+                            debug!("Received tombstone event, proceeding with shutdown");
+                            break 'shutdown;
                         }
+                    handle_event_bus_event(event, invocation_processor.clone(), tags_provider.clone(), trace_processor.clone(), trace_agent_channel.clone()).await;
+                    }
+                    // Add timeout to prevent hanging indefinitely
+                    () = tokio::time::sleep(tokio::time::Duration::from_millis(300)) => {
+                        debug!("Timeout waiting for tombstone event, proceeding with shutdown");
+                        break 'shutdown;
                     }
                 }
             }

bottlecap/src/telemetry/events.rs

@@ -150,6 +150,10 @@ pub enum TelemetryRecord {
         /// When unsuccessful, the `error_type` describes what kind of error occurred
         error_type: Option<String>,
     },
+
+    /// Tombstone event to signal shutdown
+    #[serde(rename = "platform.tombstone")]
+    PlatformTombstone,
 }
 
 /// Type of Initialization

bottlecap/src/telemetry/listener.rs

@@ -1,6 +1,6 @@
 use crate::{
     http::{extract_request_body, handler_not_found},
-    telemetry::events::TelemetryEvent,
+    telemetry::events::{TelemetryEvent, TelemetryRecord},
 };
 
 use axum::{
@@ -10,6 +10,7 @@ use axum::{
     routing::post,
     Router,
 };
+use chrono::Utc;
 use std::net::SocketAddr;
 use tokio::{net::TcpListener, sync::mpsc::Sender};
 use tokio_util::sync::CancellationToken;
@@ -46,13 +47,17 @@ impl TelemetryListener {
         let router = self.make_router();
 
         let cancel_token_clone = self.cancel_token();
+        let event_bus_clone = self.event_bus.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))
+                .with_graceful_shutdown(Self::graceful_shutdown(
+                    cancel_token_clone,
+                    event_bus_clone,
+                ))
                 .await
                 .expect("Failed to start telemetry listener");
         });
@@ -69,9 +74,21 @@ impl TelemetryListener {
             .with_state(event_bus)
     }
 
-    async fn graceful_shutdown(cancel_token: CancellationToken) {
+    async fn graceful_shutdown(cancel_token: CancellationToken, event_bus: Sender<TelemetryEvent>) {
         cancel_token.cancelled().await;
-        debug!("Telemetry API | Shutdown signal received, shutting down");
+        debug!("Telemetry API | Shutdown signal received, sending tombstone event");
+
+        // Send tombstone event to signal shutdown
+        let tombstone_event = TelemetryEvent {
+            time: Utc::now(),
+            record: TelemetryRecord::PlatformTombstone,
+        };
+
+        if let Err(e) = event_bus.send(tombstone_event).await {
+            debug!("Failed to send tombstone event: {:?}", e);
+        }
+
+        debug!("Telemetry API | Shutting down");
     }
 
     async fn handle(State(event_bus): State<Sender<TelemetryEvent>>, request: Request) -> Response {