add snapstart span

Author: jchrostek-dd

bottlecap/src/bin/bottlecap/main.rs

@@ -803,6 +803,7 @@ async fn blocking_flush_all(
 }
 
 #[allow(clippy::too_many_arguments)]
+#[allow(clippy::too_many_lines)]
 async fn handle_event_bus_event(
     event: Event,
     invocation_processor_handle: InvocationProcessorHandle,
@@ -848,6 +849,24 @@ async fn handle_event_bus_event(
                         error!("Failed to send platform init report to processor: {}", e);
                     }
                 }
+                TelemetryRecord::PlatformRestoreStart { .. } => {
+                    if let Err(e) = invocation_processor_handle
+                        .on_platform_restore_start(event.time)
+                        .await
+                    {
+                        error!("Failed to send platform restore start to processor: {}", e);
+                    }
+                }
+                TelemetryRecord::PlatformRestoreReport { metrics, .. } => {
+                    if let Some(m) = metrics {
+                        if let Err(e) = invocation_processor_handle
+                            .on_platform_restore_report(m.duration_ms, event.time.timestamp())
+                            .await
+                        {
+                            error!("Failed to send platform restore report to processor: {}", e);
+                        }
+                    }
+                }
                 TelemetryRecord::PlatformStart { request_id, .. } => {
                     if let Err(e) = invocation_processor_handle
                         .on_platform_start(request_id, event.time)

bottlecap/src/extension/telemetry/events.rs

@@ -142,6 +142,8 @@ pub enum TelemetryRecord {
         status: Status,
         /// When unsuccessful, the `error_type` describes what kind of error occurred
         error_type: Option<String>,
+        /// Metrics about the restore phase
+        metrics: Option<RestoreReportMetrics>,
     },
     #[serde(rename = "platform.restoreRuntimeDone", rename_all = "camelCase")]
     PlatformRestoreRuntimeDone {
@@ -203,6 +205,14 @@ pub struct InitReportMetrics {
     pub duration_ms: f64,
 }
 
+/// Restore report metrics
+#[derive(Clone, Copy, Debug, Deserialize, PartialEq)]
+#[serde(rename_all = "camelCase")]
+pub struct RestoreReportMetrics {
+    /// Duration of restore phase in milliseconds
+    pub duration_ms: f64,
+}
+
 /// Runtime done metrics
 #[derive(Clone, Copy, Debug, Deserialize, PartialEq)]
 #[serde(rename_all = "camelCase")]

bottlecap/src/lifecycle/invocation/context.rs

@@ -35,6 +35,10 @@ pub struct Context {
     ///
     /// This span is only present if the function is being invoked for the first time.
     pub cold_start_span: Option<Span>,
+    /// The span representing the `SnapStart` restore phase.
+    ///
+    /// This span is only present if the function is using `SnapStart` and being invoked for the first time.
+    pub snapstart_restore_span: Option<Span>,
     /// The extracted span context from the incoming request, used for distributed
     /// tracing.
     ///
@@ -87,6 +91,7 @@ impl Default for Context {
             invocation_span: Span::default(),
             runtime_done_received: false,
             cold_start_span: None,
+            snapstart_restore_span: None,
             tracer_span: None,
             extracted_span_context: None,
         }

bottlecap/src/lifecycle/invocation/processor.rs

@@ -262,6 +262,53 @@ impl Processor {
         }
     }
 
+    /// Called when the `SnapStart` restore phase starts.
+    ///
+    /// This is used to create a `snapstart_restore` span, since this telemetry event does not
+    /// provide a `request_id`, we try to guess which invocation is the restore similar to init.
+    pub fn on_platform_restore_start(&mut self, time: DateTime<Utc>) {
+        let start_time: i64 = SystemTime::from(time)
+            .duration_since(UNIX_EPOCH)
+            .expect("time went backwards")
+            .as_nanos()
+            .try_into()
+            .unwrap_or_default();
+
+        // Get the closest context
+        let Some(context) = self.context_buffer.get_closest_mut(start_time) else {
+            debug!("Cannot process on platform restore start, no invocation context found");
+            return;
+        };
+
+        // Create a SnapStart restore span
+        let mut snapstart_restore_span = create_empty_span(
+            String::from("aws.lambda.snapstart_restore"),
+            &self.resource,
+            &self.service,
+        );
+        snapstart_restore_span.span_id = generate_span_id();
+        snapstart_restore_span.start = start_time;
+        context.snapstart_restore_span = Some(snapstart_restore_span);
+    }
+
+    /// Given the duration of the platform restore report, set the snapstart restore duration.
+    ///
+    #[allow(clippy::cast_possible_truncation)]
+    pub fn on_platform_restore_report(&mut self, duration_ms: f64, timestamp: i64) {
+        self.enhanced_metrics
+            .set_snapstart_restore_duration_metric(duration_ms, timestamp);
+
+        let Some(context) = self.context_buffer.get_closest_mut(timestamp) else {
+            debug!("Cannot process on platform restore report, no invocation context found");
+            return;
+        };
+
+        if let Some(snapstart_restore_span) = &mut context.snapstart_restore_span {
+            // `round` is intentionally meant to be a whole integer
+            snapstart_restore_span.duration = (duration_ms * MS_TO_NS) as i64;
+        }
+    }
+
     /// Given a `request_id` and the time of the platform start, add the start time to the context buffer.
     ///
     pub fn on_platform_start(&mut self, request_id: String, time: DateTime<Utc>) {
@@ -409,6 +456,14 @@ impl Processor {
                 cold_start_span.parent_id = context.invocation_span.parent_id;
             }
         }
+
+        // Handle snapstart restore span if present
+        if let Some(snapstart_restore_span) = &mut context.snapstart_restore_span {
+            if context.invocation_span.trace_id != 0 {
+                snapstart_restore_span.trace_id = context.invocation_span.trace_id;
+                snapstart_restore_span.parent_id = context.invocation_span.parent_id;
+            }
+        }
         Some(context.clone())
     }
 
@@ -431,7 +486,14 @@ impl Processor {
             traces.push(ws.clone());
         }
 
-        if let Some(cold_start_span) = &context.cold_start_span {
+        // SnapStart includes telemetry events from Init (Cold Start).
+        // However, these Init events are from when the snapshot was created and
+        // not when the lambda sandbox is actually created.
+        // So, if we have a snapstart restore span, use it instead of cold start span.
+        if let Some(snapstart_restore_span) = &context.snapstart_restore_span {
+            body_size += std::mem::size_of_val(snapstart_restore_span);
+            traces.push(snapstart_restore_span.clone());
+        } else if let Some(cold_start_span) = &context.cold_start_span {
             body_size += std::mem::size_of_val(cold_start_span);
             traces.push(cold_start_span.clone());
         }

bottlecap/src/lifecycle/invocation/processor_service.rs

@@ -44,6 +44,13 @@ pub enum ProcessorCommand {
         duration_ms: f64,
         timestamp: i64,
     },
+    PlatformRestoreStart {
+        time: DateTime<Utc>,
+    },
+    PlatformRestoreReport {
+        duration_ms: f64,
+        timestamp: i64,
+    },
     PlatformStart {
         request_id: String,
         time: DateTime<Utc>,
@@ -144,6 +151,28 @@ impl InvocationProcessorHandle {
             .await
     }
 
+    pub async fn on_platform_restore_start(
+        &self,
+        time: DateTime<Utc>,
+    ) -> Result<(), mpsc::error::SendError<ProcessorCommand>> {
+        self.sender
+            .send(ProcessorCommand::PlatformRestoreStart { time })
+            .await
+    }
+
+    pub async fn on_platform_restore_report(
+        &self,
+        duration_ms: f64,
+        timestamp: i64,
+    ) -> Result<(), mpsc::error::SendError<ProcessorCommand>> {
+        self.sender
+            .send(ProcessorCommand::PlatformRestoreReport {
+                duration_ms,
+                timestamp,
+            })
+            .await
+    }
+
     pub async fn on_platform_start(
         &self,
         request_id: String,
@@ -405,6 +434,16 @@ impl InvocationProcessorService {
                     self.processor
                         .on_platform_init_report(init_type, duration_ms, timestamp);
                 }
+                ProcessorCommand::PlatformRestoreStart { time } => {
+                    self.processor.on_platform_restore_start(time);
+                }
+                ProcessorCommand::PlatformRestoreReport {
+                    duration_ms,
+                    timestamp,
+                } => {
+                    self.processor
+                        .on_platform_restore_report(duration_ms, timestamp);
+                }
                 ProcessorCommand::PlatformStart { request_id, time } => {
                     self.processor.on_platform_start(request_id, time);
                 }

bottlecap/src/logs/lambda/processor.rs

@@ -222,6 +222,18 @@ impl LambdaProcessor {
                     None,
                 ))
             },
+            TelemetryRecord::PlatformRestoreStart { .. } => {
+                if let Err(e) = self.event_bus.send(Event::Telemetry(event)).await {
+                    error!("Failed to send PlatformRestoreStart to the main event bus: {}", e);
+                }
+                Err("Unsupported event type".into())
+            }
+            TelemetryRecord::PlatformRestoreReport { .. } => {
+                if let Err(e) = self.event_bus.send(Event::Telemetry(event)).await {
+                    error!("Failed to send PlatformRestoreReport to the main event bus: {}", e);
+                }
+                Err("Unsupported event type".into())
+            }
             // TODO: PlatformInitRuntimeDone
             // TODO: PlatformExtension
             // TODO: PlatformTelemetrySubscription

bottlecap/src/metrics/enhanced/constants.rs

@@ -17,6 +17,8 @@ pub const DURATION_METRIC: &str = "aws.lambda.enhanced.duration";
 pub const POST_RUNTIME_DURATION_METRIC: &str = "aws.lambda.enhanced.post_runtime_duration";
 pub const ESTIMATED_COST_METRIC: &str = "aws.lambda.enhanced.estimated_cost";
 pub const INIT_DURATION_METRIC: &str = "aws.lambda.enhanced.init_duration";
+pub const SNAPSTART_RESTORE_DURATION_METRIC: &str =
+    "aws.lambda.enhanced.snapstart_restore_duration";
 pub const RESPONSE_LATENCY_METRIC: &str = "aws.lambda.enhanced.response_latency";
 pub const RESPONSE_DURATION_METRIC: &str = "aws.lambda.enhanced.response_duration";
 pub const PRODUCED_BYTES_METRIC: &str = "aws.lambda.enhanced.produced_bytes";

bottlecap/src/metrics/enhanced/lambda.rs

@@ -118,6 +118,26 @@ impl Lambda {
         }
     }
 
+    pub fn set_snapstart_restore_duration_metric(
+        &mut self,
+        restore_duration_ms: f64,
+        timestamp: i64,
+    ) {
+        if !self.config.enhanced_metrics {
+            return;
+        }
+        let metric = Metric::new(
+            constants::SNAPSTART_RESTORE_DURATION_METRIC.into(),
+            MetricValue::distribution(restore_duration_ms * constants::MS_TO_SEC),
+            self.get_dynamic_value_tags(),
+            Some(timestamp),
+        );
+
+        if let Err(e) = self.aggr_handle.insert_batch(vec![metric]) {
+            error!("failed to insert metric: {}", e);
+        }
+    }
+
     pub fn set_invoked_received(&mut self) {
         self.invoked_received = true;
     }