[SVLS-9883] feat: log the duration of each cold start step

bottlecap/src/bin/bottlecap/main.rs

@@ -44,7 +44,10 @@ use bottlecap::{
         invocation::processor_service::{InvocationProcessorHandle, InvocationProcessorService},
         listener::Listener as LifecycleListener,
     },
-    logger,
+    logger::{
+        self,
+        cold_start::{ColdStartDurations, elapsed_ms},
+    },
     logs::{
         agent::LogsAgent,
         aggregator_service::{
@@ -80,7 +83,9 @@ use bottlecap::{
     },
 };
 use datadog_fips::reqwest_adapter::create_reqwest_client_builder;
-use decrypt::{resolve_additional_endpoints_secrets, resolve_secrets};
+use decrypt::{
+    has_additional_endpoints_secrets, resolve_additional_endpoints_secrets, resolve_secrets,
+};
 use dogstatsd::{
     aggregator::{
         AggregatorHandle as MetricsAggregatorHandle, AggregatorService as MetricsAggregatorService,
@@ -117,11 +122,14 @@ fn resolve_dsm_service(
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
-    let start_time = Instant::now();
+    let mut cold_start_durations = ColdStartDurations::new(Instant::now());
     init_ustr();
     enable_logging_subsystem();
-    let aws_config = AwsConfig::from_env(start_time);
+    cold_start_durations.enable_logging_ms = elapsed_ms(cold_start_durations.start_time);
+    let step_start_time = Instant::now();
+    let aws_config = AwsConfig::from_env(cold_start_durations.start_time);
     log_fips_status(&aws_config.region);
+    cold_start_durations.load_aws_config_ms = elapsed_ms(step_start_time);
     let version_without_next = EXTENSION_VERSION.split('-').next().unwrap_or("NA");
     debug!("Starting Datadog Extension v{version_without_next}");
 
@@ -135,32 +143,43 @@ async fn main() -> anyhow::Result<()> {
         debug!("DD_DEBUG_WAIT_FOR_ATTACH: Continuing execution...");
     }
 
+    let step_start_time = Instant::now();
     prepare_client_provider()?;
     let client = create_reqwest_client_builder()
         .map_err(|e| anyhow::anyhow!("Failed to create client builder: {e:?}"))?
         .no_proxy()
         .build()
         .map_err(|e| anyhow::anyhow!("Failed to create client: {e:?}"))?;
+    cold_start_durations.create_extensions_api_client_ms = elapsed_ms(step_start_time);
 
     let cloned_client = client.clone();
     let runtime_api = aws_config.runtime_api.clone();
     let managed_instance_mode = aws_config.is_managed_instance_mode();
     let response = tokio::task::spawn(async move {
-        extension::register(
+        let step_start_time = Instant::now();
+        let response = extension::register(
             &cloned_client,
             &runtime_api,
             extension::EXTENSION_NAME,
             managed_instance_mode,
         )
-        .await
+        .await;
+        (response, elapsed_ms(step_start_time))
     });
     // First load the AWS configuration
     let lambda_directory: String =
         env::var("LAMBDA_TASK_ROOT").unwrap_or_else(|_| "/var/task".to_string());
     let aws_config = Arc::new(aws_config);
 
+    let step_start_time = Instant::now();
     let mut lambda_config = config::get_config(Path::new(&lambda_directory));
-    resolve_additional_endpoints_secrets(&mut lambda_config, &aws_config).await;
+    cold_start_durations.load_config_ms = elapsed_ms(step_start_time);
+    if has_additional_endpoints_secrets(&lambda_config) {
+        let step_start_time = Instant::now();
+        resolve_additional_endpoints_secrets(&mut lambda_config, &aws_config).await;
+        cold_start_durations.resolve_additional_endpoint_secrets_ms =
+            Some(elapsed_ms(step_start_time));
+    }
     let config = Arc::new(lambda_config);
 
     // Build one shared reqwest::Client for metrics, logs, trace proxy flushing, and calls to
@@ -169,10 +188,13 @@ async fn main() -> anyhow::Result<()> {
     let shared_client = bottlecap::http::get_client(&config);
     let api_key_factory = create_api_key_factory(&config, &aws_config, &shared_client);
 
-    let r = response
+    let step_start_time = Instant::now();
+    let (response, register_extension_ms) = response
         .await
-        .map_err(|e| anyhow::anyhow!("Failed to join task: {e:?}"))?
-        .map_err(|e| anyhow::anyhow!("Failed to register extension: {e:?}"))?;
+        .map_err(|e| anyhow::anyhow!("Failed to join task: {e:?}"))?;
+    let r = response.map_err(|e| anyhow::anyhow!("Failed to register extension: {e:?}"))?;
+    cold_start_durations.register_extension_ms = register_extension_ms;
+    cold_start_durations.wait_for_register_ms = elapsed_ms(step_start_time);
 
     match extension_loop_active(
         Arc::clone(&aws_config),
@@ -181,7 +203,7 @@ async fn main() -> anyhow::Result<()> {
         shared_client,
         &r,
         Arc::clone(&api_key_factory),
-        start_time,
+        cold_start_durations,
     )
     .await
     {
@@ -312,7 +334,7 @@ async fn extension_loop_active(
     shared_client: reqwest::Client,
     r: &RegisterResponse,
     api_key_factory: Arc<ApiKeyFactory>,
-    start_time: Instant,
+    mut cold_start_durations: ColdStartDurations,
 ) -> anyhow::Result<()> {
     let (mut event_bus, event_bus_tx) = EventBus::run();
 
@@ -334,13 +356,15 @@ async fn extension_loop_active(
         &shared_client,
     );
 
+    let step_start_time = Instant::now();
     let (metrics_flushers, metrics_aggregator_handle, dogstatsd_cancel_token) = start_dogstatsd(
         tags_provider.clone(),
         Arc::clone(&api_key_factory),
         config,
         &shared_client,
     )
     .await;
+    cold_start_durations.start_dogstatsd_ms = elapsed_ms(step_start_time);
 
     let propagator = Arc::new(DatadogCompositePropagator::new(Arc::clone(config)));
 
@@ -406,10 +430,15 @@ async fn extension_loop_active(
     });
 
     // AppSec processor (if enabled)
+    let step_start_time = Instant::now();
     let appsec_processor = match AppSecProcessor::new(config) {
-        Ok(p) => Some(Arc::new(TokioMutex::new(p))),
+        Ok(p) => {
+            cold_start_durations.create_appsec_processor_ms = Some(elapsed_ms(step_start_time));
+            Some(Arc::new(TokioMutex::new(p)))
+        }
         Err(AppSecFeatureDisabled) => None,
         Err(e) => {
+            cold_start_durations.create_appsec_processor_ms = Some(elapsed_ms(step_start_time));
             error!(
                 "AAP | error creating App & API Protection processor, the feature will be disabled: {e}"
             );
@@ -454,6 +483,7 @@ async fn extension_loop_active(
         }
     });
 
+    let step_start_time = Instant::now();
     let telemetry_listener_cancel_token = setup_telemetry_client(
         client,
         &r.extension_id,
@@ -464,6 +494,7 @@ async fn extension_loop_active(
         aws_config.is_managed_instance_mode(),
     )
     .await?;
+    cold_start_durations.start_telemetry_listener_ms = elapsed_ms(step_start_time);
 
     let otlp_cancel_token = start_otlp_agent(
         config,
@@ -481,8 +512,13 @@ async fn extension_loop_active(
 
     debug!(
         "Datadog Next-Gen Extension ready in {:}ms",
-        start_time.elapsed().as_millis().to_string()
+        cold_start_durations
+            .start_time
+            .elapsed()
+            .as_millis()
+            .to_string()
     );
+    cold_start_durations.log();
 
     if aws_config.is_managed_instance_mode() {
         // Clone Arc references for the background flusher task

bottlecap/src/logger/cold_start.rs

@@ -0,0 +1,70 @@
+use tokio::time::Instant;
+
+use super::COLD_START_DURATION;
+
+/// Duration of each cold start step in milliseconds. The steps go in one log line as numeric
+/// fields, so that Datadog can graph each step.
+#[derive(Debug, Clone, Copy)]
+pub struct ColdStartDurations {
+    /// When the extension process started.
+    pub start_time: Instant,
+    pub enable_logging_ms: f64,
+    pub load_aws_config_ms: f64,
+    pub create_extensions_api_client_ms: f64,
+    pub register_extension_ms: f64,
+    pub load_config_ms: f64,
+    /// Set only when the config has an additional endpoints secret ARN.
+    pub resolve_additional_endpoint_secrets_ms: Option<f64>,
+    pub wait_for_register_ms: f64,
+    pub start_dogstatsd_ms: f64,
+    /// Set only when App and API Protection is enabled.
+    pub create_appsec_processor_ms: Option<f64>,
+    pub start_telemetry_listener_ms: f64,
+}
+
+impl ColdStartDurations {
+    #[must_use]
+    pub fn new(start_time: Instant) -> Self {
+        Self {
+            start_time,
+            enable_logging_ms: 0.0,
+            load_aws_config_ms: 0.0,
+            create_extensions_api_client_ms: 0.0,
+            register_extension_ms: 0.0,
+            load_config_ms: 0.0,
+            resolve_additional_endpoint_secrets_ms: None,
+            wait_for_register_ms: 0.0,
+            start_dogstatsd_ms: 0.0,
+            create_appsec_processor_ms: None,
+            start_telemetry_listener_ms: 0.0,
+        }
+    }
+
+    /// Logs all durations in one line, with the time since `start_time` as the total.
+    pub fn log(&self) {
+        // `debug!` cannot parse dotted field names after `target:`, so this uses `event!`.
+        tracing::event!(
+            target: COLD_START_DURATION,
+            tracing::Level::DEBUG,
+            cold_start.total_ms = elapsed_ms(self.start_time),
+            cold_start.enable_logging_ms = self.enable_logging_ms,
+            cold_start.load_aws_config_ms = self.load_aws_config_ms,
+            cold_start.create_extensions_api_client_ms = self.create_extensions_api_client_ms,
+            cold_start.register_extension_ms = self.register_extension_ms,
+            cold_start.load_config_ms = self.load_config_ms,
+            cold_start.resolve_additional_endpoint_secrets_ms =
+                self.resolve_additional_endpoint_secrets_ms,
+            cold_start.wait_for_register_ms = self.wait_for_register_ms,
+            cold_start.start_dogstatsd_ms = self.start_dogstatsd_ms,
+            cold_start.create_appsec_processor_ms = self.create_appsec_processor_ms,
+            cold_start.start_telemetry_listener_ms = self.start_telemetry_listener_ms,
+            "Cold start step durations"
+        );
+    }
+}
+
+/// Milliseconds since `start_time`, rounded to 2 decimals.
+#[must_use]
+pub fn elapsed_ms(start_time: Instant) -> f64 {
+    (start_time.elapsed().as_secs_f64() * 100_000.0).round() / 100.0
+}