fix(shutdown): handle runtimeDone without metrics; wait for the processor before the final flush

bottlecap/src/bin/bottlecap/main.rs

@@ -829,6 +829,12 @@ async fn extension_loop_active(
             )
             .await;
 
+            // Make sure the processor has handled every earlier command, including the
+            // fire-and-forget shutdown and OOM events, before the final flush.
+            if let Err(e) = invocation_processor_handle.barrier().await {
+                error!("Failed to wait for the invocation processor: {e}");
+            }
+
             // Cancel background services
             cancel_background_services(
                 api_runtime_proxy_shutdown_signal.as_ref(),
@@ -1000,7 +1006,7 @@ async fn handle_event_bus_event(
                 }
                 TelemetryRecord::PlatformRuntimeDone {
                     ref request_id,
-                    metrics: Some(metrics),
+                    metrics,
                     status,
                     ref error_type,
                     ..

bottlecap/src/lifecycle/invocation/processor.rs

@@ -580,16 +580,18 @@ impl Processor {
     pub async fn on_platform_runtime_done(
         &mut self,
         request_id: &String,
-        metrics: RuntimeDoneMetrics,
+        metrics: Option<RuntimeDoneMetrics>,
         status: Status,
         error_type: Option<String>,
         tags_provider: Arc<provider::Provider>,
         trace_sender: Arc<SendingTraceProcessor>,
         timestamp: i64,
     ) {
         // Set the runtime duration metric
-        self.enhanced_metrics
-            .set_runtime_done_metrics(&metrics, timestamp);
+        if let Some(metrics) = &metrics {
+            self.enhanced_metrics
+                .set_runtime_done_metrics(metrics, timestamp);
+        }
 
         if status != Status::Success {
             // Increment the error metric
@@ -614,6 +616,11 @@ impl Processor {
             self.enhanced_metrics.set_usage_enhanced_metrics();
         }
 
+        // Without a duration the invocation span would be sent with zero length, so skip span
+        // processing. `runtime_done_received` stays false so reparenting doesn't send it either.
+        let Some(metrics) = metrics else {
+            return;
+        };
         self.context_buffer
             .add_runtime_duration(request_id, metrics.duration_ms);
 
@@ -1014,7 +1021,7 @@ impl Processor {
         if self.context_buffer.get(request_id).is_some() {
             self.on_platform_runtime_done(
                 request_id,
-                runtime_done_metrics,
+                Some(runtime_done_metrics),
                 status,
                 error_type,
                 tags_provider,
@@ -1735,6 +1742,7 @@ mod tests {
     use super::*;
     use crate::LAMBDA_RUNTIME_SLUG;
     use crate::extension::telemetry::events::ManagedInstanceReportMetrics;
+    use crate::metrics::enhanced::constants::{ERRORS_METRIC, TIMEOUTS_METRIC};
     use crate::traces::stats_concentrator_service::StatsConcentratorService;
     use crate::traces::stats_generator::StatsGenerator;
     use crate::traces::trace_processor;
@@ -3195,10 +3203,10 @@ mod tests {
 
         p.on_platform_runtime_done(
             &request_id,
-            RuntimeDoneMetrics {
+            Some(RuntimeDoneMetrics {
                 duration_ms: 100.0,
                 produced_bytes: None,
-            },
+            }),
             Status::Timeout,
             None,
             tags_provider,
@@ -3714,4 +3722,59 @@ mod tests {
             );
         }
     }
+
+    async fn metric_sum(p: &Processor, metric: &str, timestamp: i64) -> Option<f64> {
+        p.enhanced_metrics
+            .aggr_handle
+            .get_entry_by_id(metric.into(), None, (timestamp / 10) * 10)
+            .await
+            .unwrap()
+            .map(|entry| {
+                entry
+                    .value
+                    .get_sketch()
+                    .expect("distribution sketch")
+                    .sum()
+                    .expect("sketch sum")
+            })
+    }
+
+    #[tokio::test]
+    async fn test_runtime_done_without_metrics_counts_timeout_without_span() {
+        let request_id = String::from("req-no-metrics");
+        let mut p = setup();
+        p.context_buffer.start_context(&request_id, Span::default());
+        p.tracer_detected = true;
+        let config = Arc::clone(&p.config);
+        let tags_provider = Arc::new(provider::Provider::new(
+            Arc::clone(&config),
+            LAMBDA_RUNTIME_SLUG.to_string(),
+            &HashMap::from([("function_arn".to_string(), "test-arn".to_string())]),
+        ));
+        let (trace_sender, mut trace_rx) = make_trace_sender_with_rx(config);
+        let now = chrono::Utc::now().timestamp();
+
+        p.on_platform_runtime_done(
+            &request_id,
+            None,
+            Status::Timeout,
+            None,
+            tags_provider,
+            Arc::clone(&trace_sender),
+            now,
+        )
+        .await;
+
+        assert_eq!(metric_sum(&p, ERRORS_METRIC, now).await, Some(1.0));
+        assert_eq!(metric_sum(&p, TIMEOUTS_METRIC, now).await, Some(1.0));
+        // Without a duration no zero-length invocation span is sent, and the span isn't marked
+        // complete, so reparenting won't send it either.
+        assert!(trace_rx.try_recv().is_err());
+        assert!(
+            !p.context_buffer
+                .get(&request_id)
+                .expect("context")
+                .runtime_done_received
+        );
+    }
 }

bottlecap/src/lifecycle/invocation/processor_service.rs

@@ -59,7 +59,7 @@ pub enum ProcessorCommand {
     },
     PlatformRuntimeDone {
         request_id: String,
-        metrics: RuntimeDoneMetrics,
+        metrics: Option<RuntimeDoneMetrics>,
         status: Status,
         error_type: Option<String>,
         tags_provider: Arc<provider::Provider>,
@@ -127,6 +127,9 @@ pub enum ProcessorCommand {
         timestamp: i64,
     },
     OnShutdownEvent,
+    Barrier {
+        response: oneshot::Sender<()>,
+    },
     SendCtxSpans {
         tags_provider: Arc<provider::Provider>,
         trace_sender: Arc<SendingTraceProcessor>,
@@ -180,7 +183,8 @@ impl InvocationProcessorHandle {
                         let _ = response.send(Ok(None));
                     }
                     ProcessorCommand::PlatformRuntimeDone { response, .. }
-                    | ProcessorCommand::PlatformReport { response, .. } => {
+                    | ProcessorCommand::PlatformReport { response, .. }
+                    | ProcessorCommand::Barrier { response } => {
                         let _ = response.send(());
                     }
                     // Fire-and-forget commands: drop silently.
@@ -279,7 +283,7 @@ impl InvocationProcessorHandle {
     pub async fn on_platform_runtime_done(
         &self,
         request_id: String,
-        metrics: RuntimeDoneMetrics,
+        metrics: Option<RuntimeDoneMetrics>,
         status: Status,
         error_type: Option<String>,
         tags_provider: Arc<provider::Provider>,
@@ -508,6 +512,18 @@ impl InvocationProcessorHandle {
         self.sender.send(ProcessorCommand::OnShutdownEvent).await
     }
 
+    /// Resolves once the processor has handled this and every earlier command.
+    pub async fn barrier(&self) -> Result<(), ProcessorError> {
+        let (tx, rx) = oneshot::channel();
+        self.sender
+            .send(ProcessorCommand::Barrier { response: tx })
+            .await
+            .map_err(|e| ProcessorError::ChannelSend(e.to_string()))?;
+        rx.await
+            .map_err(|e| ProcessorError::ChannelReceive(e.to_string()))?;
+        Ok(())
+    }
+
     pub async fn send_ctx_spans(
         &self,
         tags_provider: &Arc<provider::Provider>,
@@ -746,6 +762,9 @@ impl InvocationProcessorService {
                 ProcessorCommand::OnShutdownEvent => {
                     self.processor.on_shutdown_event();
                 }
+                ProcessorCommand::Barrier { response } => {
+                    let _ = response.send(());
+                }
                 ProcessorCommand::SendCtxSpans {
                     tags_provider,
                     trace_sender,
@@ -878,10 +897,10 @@ mod tests {
             timeout_dur,
             handle.on_platform_runtime_done(
                 "rid".to_string(),
-                RuntimeDoneMetrics {
+                Some(RuntimeDoneMetrics {
                     duration_ms: 0.0,
                     produced_bytes: None,
-                },
+                }),
                 Status::Success,
                 None,
                 Arc::clone(&tags_provider),
@@ -893,6 +912,11 @@ mod tests {
         .expect("on_platform_runtime_done timed out")
         .expect("on_platform_runtime_done");
 
+        timeout(timeout_dur, handle.barrier())
+            .await
+            .expect("barrier timed out")
+            .expect("barrier");
+
         timeout(
             timeout_dur,
             handle.on_platform_report(