Merge branch 'main' into zarir/OTLP-operation-name-change

Author: zarirhamza

bottlecap/Cargo.lock

@@ -28,6 +28,8 @@ reqwest = { version = "0.12.11", features = ["json", "http2"], default-features
 serde = { version = "1.0", default-features = false, features = ["derive"] }
 serde_json = { version = "1.0", default-features = false, features = ["alloc"] }
 thiserror = { version = "1.0", default-features = false }
+# Transitive dependency (pulled in via cookie). Pinned to >=0.3.47 so cargo audit / CI passes (RUSTSEC-2026-0009).
+time = { version = "0.3.47", default-features = false }
 tokio = { version = "1.47", default-features = false, features = ["macros", "rt-multi-thread", "time"] }
 tokio-util = { version = "0.7", default-features = false }
 tracing = { version = "0.1", default-features = false }
@@ -71,8 +73,8 @@ libdd-trace-utils = { git = "https://github.com/DataDog/libdatadog", rev = "158b
 libdd-trace-normalization = { git = "https://github.com/DataDog/libdatadog",  rev = "158b59471f1132e3cb36023fa3c46ccb2dd0eda1" }
 libdd-trace-obfuscation = { git = "https://github.com/DataDog/libdatadog", rev = "158b59471f1132e3cb36023fa3c46ccb2dd0eda1"  }
 libdd-trace-stats = { git = "https://github.com/DataDog/libdatadog", rev = "158b59471f1132e3cb36023fa3c46ccb2dd0eda1"  }
-dogstatsd = { git = "https://github.com/DataDog/serverless-components", rev = "18b49baba8bfef97060d7edd8b830584d0da3373", default-features = false }
-datadog-fips = { git = "https://github.com/DataDog/serverless-components", rev = "18b49baba8bfef97060d7edd8b830584d0da3373", default-features = false }
+dogstatsd = { git = "https://github.com/DataDog/serverless-components", rev = "2eb009a59ed07ffcaf174db1c31af413365e9bc6", default-features = false }
+datadog-fips = { git = "https://github.com/DataDog/serverless-components", rev = "2eb009a59ed07ffcaf174db1c31af413365e9bc6", default-features = false }
 libddwaf = { version = "1.28.1", git = "https://github.com/DataDog/libddwaf-rust", rev = "d1534a158d976bd4f747bf9fcc58e0712d2d17fc", default-features = false, features = ["serde"] }
 
 [dev-dependencies]

bottlecap/Cargo.toml

@@ -202,7 +202,7 @@ impl Context {
                 self.process_result(result);
             }
             Err(e) => log_waf_run_error(e),
-        };
+        }
     }
 
     /// Obtain the information about the endpoint that was invoked, if available.

bottlecap/src/appsec/processor/context.rs

@@ -34,7 +34,7 @@ pub struct Processor {
     context_buffer: VecDeque<Context>,
 }
 impl Processor {
-    const CONTEXT_BUFFER_DEFAULT_CAPACITY: NonZero<usize> = unsafe { NonZero::new_unchecked(5) };
+    const CONTEXT_BUFFER_DEFAULT_CAPACITY: NonZero<usize> = NonZero::new(5).unwrap();
 
     /// Creates a new [`Processor`] instance using the provided [`Config`].
     ///
@@ -578,9 +578,9 @@ impl InvocationPayload for IdentifiedTrigger {
                 .unwrap_or_default()
                 .iter()
                 .flat_map(parse_cookie)
-                .chunk_by(|(k, _)| (*k).to_string())
+                .chunk_by(|(k, _)| (*k).clone())
                 .into_iter()
-                .map(|(k, v)| (k, v.into_iter().map(|(_, v)| v.to_string()).collect()))
+                .map(|(k, v)| (k, v.into_iter().map(|(_, v)| v.clone()).collect()))
                 .collect(),
             Self::APIGatewayWebSocketEvent(t) => t
                 .multi_value_headers
@@ -589,9 +589,9 @@ impl InvocationPayload for IdentifiedTrigger {
                 .unwrap_or_default()
                 .iter()
                 .flat_map(parse_cookie)
-                .chunk_by(|(k, _)| k.to_string())
+                .chunk_by(|(k, _)| k.clone())
                 .into_iter()
-                .map(|(k, v)| (k, v.into_iter().map(|(_, v)| v.to_string()).collect()))
+                .map(|(k, v)| (k, v.into_iter().map(|(_, v)| v.clone()).collect()))
                 .collect(),
             Self::ALBEvent(t) => t
                 .multi_value_headers
@@ -601,9 +601,9 @@ impl InvocationPayload for IdentifiedTrigger {
                 .unwrap_or_default()
                 .iter()
                 .flat_map(parse_cookie)
-                .chunk_by(|(k, _)| (*k).to_string())
+                .chunk_by(|(k, _)| (*k).clone())
                 .into_iter()
-                .map(|(k, v)| (k, v.into_iter().map(|(_, v)| v.to_string()).collect()))
+                .map(|(k, v)| (k, v.into_iter().map(|(_, v)| v.clone()).collect()))
                 .collect(),
             Self::LambdaFunctionUrlEvent(t) => {
                 t.cookies.as_ref().map(list_to_map).unwrap_or_default()

bottlecap/src/appsec/processor/mod.rs

@@ -7,7 +7,7 @@ use crate::appsec::processor::InvocationPayload;
 use crate::lifecycle::invocation::triggers::body::Body;
 
 /// The expected payload of a response. This is different from trigger to trigger.
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
 pub enum ExpectedResponseFormat {
     /// API Gateway style integration responses (REST and HTTP)
     ApiGatewayResponse,
@@ -16,6 +16,7 @@ pub enum ExpectedResponseFormat {
     Raw,
 
     /// Unknown or unsupported response format
+    #[default]
     Unknown,
 }
 impl ExpectedResponseFormat {
@@ -39,11 +40,6 @@ impl ExpectedResponseFormat {
         }
     }
 }
-impl Default for ExpectedResponseFormat {
-    fn default() -> Self {
-        Self::Unknown
-    }
-}
 
 #[derive(Debug, Default, Deserialize)]
 #[serde(default)]

bottlecap/src/appsec/processor/response.rs

@@ -113,13 +113,13 @@ async fn main() -> anyhow::Result<()> {
     debug!("Starting Datadog Extension v{version_without_next}");
 
     // Debug: Wait for debugger to attach if DD_DEBUG_WAIT_FOR_ATTACH is set
-    if let Ok(wait_secs) = env::var("DD_DEBUG_WAIT_FOR_ATTACH") {
-        if let Ok(secs) = wait_secs.parse::<u64>() {
-            debug!("DD_DEBUG_WAIT_FOR_ATTACH: Waiting {secs} seconds for debugger to attach...");
-            debug!("Connect your debugger to port 2345 now!");
-            tokio::time::sleep(tokio::time::Duration::from_secs(secs)).await;
-            debug!("DD_DEBUG_WAIT_FOR_ATTACH: Continuing execution...");
-        }
+    if let Ok(wait_secs) = env::var("DD_DEBUG_WAIT_FOR_ATTACH")
+        && let Ok(secs) = wait_secs.parse::<u64>()
+    {
+        debug!("DD_DEBUG_WAIT_FOR_ATTACH: Waiting {secs} seconds for debugger to attach...");
+        debug!("Connect your debugger to port 2345 now!");
+        tokio::time::sleep(tokio::time::Duration::from_secs(secs)).await;
+        debug!("DD_DEBUG_WAIT_FOR_ATTACH: Continuing execution...");
     }
 
     prepare_client_provider()?;
@@ -273,7 +273,7 @@ async fn extension_loop_idle(
                 error!("Error getting next event: {e:?}");
                 return Err(e.into());
             }
-        };
+        }
     }
 }
 
@@ -288,11 +288,7 @@ async fn extension_loop_active(
 ) -> anyhow::Result<()> {
     let (mut event_bus, event_bus_tx) = EventBus::run();
 
-    let account_id = r
-        .account_id
-        .as_ref()
-        .unwrap_or(&"none".to_string())
-        .to_string();
+    let account_id = r.account_id.as_ref().unwrap_or(&"none".to_string()).clone();
     let tags_provider = setup_tag_provider(&Arc::clone(&aws_config), config, &account_id);
 
     let (logs_agent_channel, logs_flusher, logs_agent_cancel_token, logs_aggregator_handle) =
@@ -450,10 +446,8 @@ async fn extension_loop_active(
                         // Wait for any pending flushes
                         flushing_service.await_handles().await;
                         // Final flush to capture any data that accumulated since the last
-                        // spawn_non_blocking(). We pass force_stats=true since this is our
-                        // last opportunity to send data before shutdown.
-                        let mut locked_metrics = flushing_service.metrics_flushers().lock().await;
-                        flushing_service.flush_blocking(true, &mut locked_metrics).await;
+                        // spawn_non_blocking(). This is our last opportunity to send data.
+                        flushing_service.flush_blocking_final().await;
                         break;
                     }
                 }
@@ -519,7 +513,6 @@ async fn extension_loop_active(
                             "Transient network error waiting for shutdown event: {}. Retrying...",
                             e
                         );
-                        continue;
                     }
                     Err(e) => {
                         error!(
@@ -628,26 +621,19 @@ async fn extension_loop_active(
                     tokio::select! {
                     biased;
                         Some(event) = event_bus.rx.recv() => {
-                            if let Some(telemetry_event) = handle_event_bus_event(event, invocation_processor_handle.clone(), appsec_processor.clone(), tags_provider.clone(), trace_processor.clone(), trace_agent_channel.clone(), stats_concentrator.clone()).await {
-                                if let TelemetryRecord::PlatformRuntimeDone{ .. } = telemetry_event.record {
+                            if let Some(telemetry_event) = handle_event_bus_event(event, invocation_processor_handle.clone(), appsec_processor.clone(), tags_provider.clone(), trace_processor.clone(), trace_agent_channel.clone(), stats_concentrator.clone()).await
+                                && let TelemetryRecord::PlatformRuntimeDone{ .. } = telemetry_event.record {
                                     break 'flush_end;
                                 }
-                            }
                         }
                         _ = race_flush_interval.tick() => {
-                            let mut locked_metrics = metrics_flushers.lock().await;
-                            flushing_service
-                                .flush_blocking(false, &mut locked_metrics)
-                                .await;
+                            flushing_service.flush_blocking().await;
                             race_flush_interval.reset();
                         }
                     }
                 }
                 // flush
-                let mut locked_metrics = metrics_flushers.lock().await;
-                flushing_service
-                    .flush_blocking(false, &mut locked_metrics)
-                    .await;
+                flushing_service.flush_blocking().await;
                 race_flush_interval.reset();
                 let next_response =
                     extension::next_event(client, &aws_config.runtime_api, &r.extension_id).await;
@@ -664,10 +650,7 @@ async fn extension_loop_active(
                         }
                     }
                     FlushDecision::Periodic => {
-                        let mut locked_metrics = metrics_flushers.lock().await;
-                        flushing_service
-                            .flush_blocking(false, &mut locked_metrics)
-                            .await;
+                        flushing_service.flush_blocking().await;
                         race_flush_interval.reset();
                     }
                     _ => {
@@ -695,10 +678,7 @@ async fn extension_loop_active(
                         }
                         _ = race_flush_interval.tick() => {
                             if flush_control.flush_strategy == FlushStrategy::Default {
-                                let mut locked_metrics = metrics_flushers.lock().await;
-                                flushing_service
-                                    .flush_blocking(false, &mut locked_metrics)
-                                    .await;
+                                flushing_service.flush_blocking().await;
                                 race_flush_interval.reset();
                             }
                         }
@@ -744,11 +724,8 @@ async fn extension_loop_active(
                 &lifecycle_listener_shutdown_token,
             );
 
-            // Final flush with force_stats=true since this is our last opportunity
-            let mut locked_metrics = metrics_flushers.lock().await;
-            flushing_service
-                .flush_blocking(true, &mut locked_metrics)
-                .await;
+            // Final flush - this is our last opportunity to send data before shutdown
+            flushing_service.flush_blocking_final().await;
 
             // Even though we're shutting down, we need to reset the flush interval to prevent any future flushes
             race_flush_interval.reset();
@@ -1178,7 +1155,7 @@ async fn start_dogstatsd(
     api_key_factory: Arc<ApiKeyFactory>,
     config: &Arc<Config>,
 ) -> (
-    Arc<TokioMutex<Vec<MetricsFlusher>>>,
+    Arc<Vec<MetricsFlusher>>,
     MetricsAggregatorHandle,
     CancellationToken,
 ) {
@@ -1200,17 +1177,20 @@ async fn start_dogstatsd(
     });
 
     // Get flushers with aggregator handle
-    let flushers = Arc::new(TokioMutex::new(start_metrics_flushers(
+    let flushers = Arc::new(start_metrics_flushers(
         Arc::clone(&api_key_factory),
         &aggregator_handle,
         config,
-    )));
+    ));
 
     // Create Dogstatsd server
     let dogstatsd_config = DogStatsDConfig {
         host: EXTENSION_HOST.to_string(),
         port: DOGSTATSD_PORT,
         metric_namespace: config.statsd_metric_namespace.clone(),
+        so_rcvbuf: config.dogstatsd_so_rcvbuf,
+        buffer_size: config.dogstatsd_buffer_size,
+        queue_size: config.dogstatsd_queue_size,
     };
     let cancel_token = tokio_util::sync::CancellationToken::new();
     let dogstatsd_agent = DogStatsD::new(

bottlecap/src/bin/bottlecap/main.rs

@@ -277,6 +277,24 @@ pub struct EnvConfig {
     #[serde(deserialize_with = "deserialize_optional_string")]
     pub statsd_metric_namespace: Option<String>,
 
+    /// @env `DD_DOGSTATSD_SO_RCVBUF`
+    /// Size of the receive buffer for `DogStatsD` UDP packets, in bytes (`SO_RCVBUF`).
+    /// Increase to reduce packet loss under high-throughput metric bursts.
+    #[serde(deserialize_with = "deserialize_option_lossless")]
+    pub dogstatsd_so_rcvbuf: Option<usize>,
+
+    /// @env `DD_DOGSTATSD_BUFFER_SIZE`
+    /// Maximum size of a single read from any transport (UDP or named pipe), in bytes.
+    /// Defaults to 8192.
+    #[serde(deserialize_with = "deserialize_option_lossless")]
+    pub dogstatsd_buffer_size: Option<usize>,
+
+    /// @env `DD_DOGSTATSD_QUEUE_SIZE`
+    /// Internal queue capacity between the socket reader and metric processor.
+    /// Defaults to 1024. Increase if the processor can't keep up with burst traffic.
+    #[serde(deserialize_with = "deserialize_option_lossless")]
+    pub dogstatsd_queue_size: Option<usize>,
+
     // OTLP
     //
     // - APM / Traces
@@ -554,6 +572,11 @@ fn merge_config(config: &mut Config, env_config: &EnvConfig) {
         config.statsd_metric_namespace = parse_metric_namespace(namespace);
     }
 
+    // DogStatsD
+    merge_option!(config, env_config, dogstatsd_so_rcvbuf);
+    merge_option!(config, env_config, dogstatsd_buffer_size);
+    merge_option!(config, env_config, dogstatsd_queue_size);
+
     // OTLP
     merge_option_to_value!(config, env_config, otlp_config_traces_enabled);
     merge_option_to_value!(
@@ -830,6 +853,11 @@ mod tests {
             );
             jail.set_env("DD_OTLP_CONFIG_LOGS_ENABLED", "true");
 
+            // DogStatsD
+            jail.set_env("DD_DOGSTATSD_SO_RCVBUF", "1048576");
+            jail.set_env("DD_DOGSTATSD_BUFFER_SIZE", "65507");
+            jail.set_env("DD_DOGSTATSD_QUEUE_SIZE", "2048");
+
             // AWS Lambda
             jail.set_env(
                 "DD_API_KEY_SECRET_ARN",
@@ -984,6 +1012,9 @@ mod tests {
                 otlp_config_traces_probabilistic_sampler_sampling_percentage: Some(50),
                 otlp_config_logs_enabled: true,
                 statsd_metric_namespace: None,
+                dogstatsd_so_rcvbuf: Some(1_048_576),
+                dogstatsd_buffer_size: Some(65507),
+                dogstatsd_queue_size: Some(2048),
                 api_key_secret_arn: "arn:aws:secretsmanager:region:account:secret:datadog-api-key"
                     .to_string(),
                 kms_api_key: "test-kms-key".to_string(),
@@ -1170,4 +1201,43 @@ mod tests {
             Ok(())
         });
     }
+
+    #[test]
+    fn test_dogstatsd_config_from_env() {
+        figment::Jail::expect_with(|jail| {
+            jail.clear_env();
+            jail.set_env("DD_DOGSTATSD_SO_RCVBUF", "1048576");
+            jail.set_env("DD_DOGSTATSD_BUFFER_SIZE", "65507");
+            jail.set_env("DD_DOGSTATSD_QUEUE_SIZE", "2048");
+
+            let mut config = Config::default();
+            let env_config_source = EnvConfigSource;
+            env_config_source
+                .load(&mut config)
+                .expect("Failed to load config");
+
+            assert_eq!(config.dogstatsd_so_rcvbuf, Some(1_048_576));
+            assert_eq!(config.dogstatsd_buffer_size, Some(65507));
+            assert_eq!(config.dogstatsd_queue_size, Some(2048));
+            Ok(())
+        });
+    }
+
+    #[test]
+    fn test_dogstatsd_config_defaults_to_none() {
+        figment::Jail::expect_with(|jail| {
+            jail.clear_env();
+
+            let mut config = Config::default();
+            let env_config_source = EnvConfigSource;
+            env_config_source
+                .load(&mut config)
+                .expect("Failed to load config");
+
+            assert_eq!(config.dogstatsd_so_rcvbuf, None);
+            assert_eq!(config.dogstatsd_buffer_size, None);
+            assert_eq!(config.dogstatsd_queue_size, None);
+            Ok(())
+        });
+    }
 }