feat: Handle API key resolution failure (#732)

# Context
The previous PR
#717 defers API
key resolution from extension init stage to flush time. However, that PR
doesn't well handle the failure case.
- Before that PR, if resolution fails in init stage, the extension will
run an idle loop.
- After that PR, the extension will crash at flush time, which will kill
the runtime as well, which is not desired.

# What does this PR do?
1. For traces, defer key resolution from
`TraceProcessor.process_traces()` to `TraceFlusher.flush()`.
- (This should ideally be in the previous PR, but since that is already
approved, let me add this change in this new PR.)
2. If resolution fails at flush time, then make flush a no-op, so the
extension can keep running and consume events without crashing.

# Dependencies
1. DataDog/serverless-components#25
2. DataDog/libdatadog#1140

# Manual Test

## Steps
1. Create a layer in sandbox
2. Apply the layer to a Lambda function
3. Set the env var `DD_API_KEY_SECRET_ARN` to an invalid value
5. Run the Lambda
6. Then set `DD_API_KEY_SECRET_ARN` to a valid value
7. Run the Lambda

## Result
1. The function was successful
<img width="319" alt="image"
src="https://github.com/user-attachments/assets/f8a5cb36-f678-4643-ba1c-85f41256ffa1"
/>

2. The extension printed some error logs
<img width="737" height="33" alt="image"
src="https://github.com/user-attachments/assets/22553d24-e1f5-4ee5-9a91-0d18e3e2f297"
/>

<img width="603" height="186" alt="image"
src="https://github.com/user-attachments/assets/e797f991-ecba-45f0-8f49-7b7b59dd9e7b"
/>

3. With valid secret ARN, the Lambda runs successfully and reports to
Datadog
<img width="678" height="150" alt="image"
src="https://github.com/user-attachments/assets/073089f8-1e9a-4728-b8d1-1db7aa85d031"
/>

<img width="533" height="96" alt="image"
src="https://github.com/user-attachments/assets/d5f2b81c-5e02-42bc-b3ef-85e611228fc6"
/>


# Automated Test

I didn't add any automated test because from what I see in the codebase,
existing tests are usually unit tests for short functions and not for
long functions that this PR touches. Please let me know if you think I
should add automated tests.

Author: lym953

bottlecap/Cargo.lock

@@ -57,7 +57,7 @@ datadog-trace-protobuf = { git = "https://github.com/DataDog/libdatadog", rev =
 datadog-trace-utils = { git = "https://github.com/DataDog/libdatadog", rev = "8a49c7df2d9cbf05118bfd5b85772676f71b34f2" , features = ["mini_agent"] }
 datadog-trace-normalization = { git = "https://github.com/DataDog/libdatadog",  rev = "8a49c7df2d9cbf05118bfd5b85772676f71b34f2" }
 datadog-trace-obfuscation = { git = "https://github.com/DataDog/libdatadog", rev = "8a49c7df2d9cbf05118bfd5b85772676f71b34f2"  }
-dogstatsd = { git = "https://github.com/DataDog/serverless-components", rev = "b29bba8b4178fc2089943fe28e853d529826888b", default-features = false }
+dogstatsd = { git = "https://github.com/DataDog/serverless-components", rev = "985120329d0ba96c1ec8d719cc38e1f7ce11a092", default-features = false }
 datadog-trace-agent = { git = "https://github.com/DataDog/serverless-components", rev = "c3d8ed4f90591c6958921145d485463860307f78" }
 datadog-fips = { git = "https://github.com/DataDog/serverless-components", rev = "c3d8ed4f90591c6958921145d485463860307f78", default-features = false }
 axum = { version = "0.8.4", default-features = false, features = ["default"] }

bottlecap/Cargo.toml

@@ -439,11 +439,7 @@ fn create_api_key_factory(
         let aws_config = Arc::clone(&aws_config);
         let aws_credentials = Arc::clone(&aws_credentials);
 
-        Box::pin(async move {
-            resolve_secrets(config, aws_config, aws_credentials)
-                .await
-                .expect("Failed to resolve API key")
-        })
+        Box::pin(async move { resolve_secrets(config, aws_config, aws_credentials).await })
     })))
 }
 
@@ -517,7 +513,7 @@ async fn extension_loop_active(
         trace_agent_shutdown_token,
     ) = start_trace_agent(
         config,
-        Arc::clone(&api_key_factory),
+        &api_key_factory,
         &tags_provider,
         Arc::clone(&invocation_processor),
         Arc::clone(&trace_aggregator),
@@ -1038,7 +1034,7 @@ fn start_metrics_flushers(
 #[allow(clippy::type_complexity)]
 fn start_trace_agent(
     config: &Arc<Config>,
-    api_key_factory: Arc<ApiKeyFactory>,
+    api_key_factory: &Arc<ApiKeyFactory>,
     tags_provider: &Arc<TagProvider>,
     invocation_processor: Arc<TokioMutex<InvocationProcessor>>,
     trace_aggregator: Arc<TokioMutex<trace_aggregator::TraceAggregator>>,
@@ -1064,6 +1060,7 @@ fn start_trace_agent(
     let trace_flusher = Arc::new(trace_flusher::ServerlessTraceFlusher {
         aggregator: trace_aggregator.clone(),
         config: Arc::clone(config),
+        api_key_factory: Arc::clone(api_key_factory),
     });
 
     let obfuscation_config = obfuscation_config::ObfuscationConfig {
@@ -1077,7 +1074,6 @@ fn start_trace_agent(
 
     let trace_processor = Arc::new(trace_processor::ServerlessTraceProcessor {
         obfuscation_config: Arc::new(obfuscation_config),
-        api_key_factory: api_key_factory.clone(),
     });
 
     // Proxy
@@ -1098,7 +1094,6 @@ fn start_trace_agent(
         proxy_aggregator,
         invocation_processor,
         Arc::clone(tags_provider),
-        api_key_factory,
     );
     let trace_agent_channel = trace_agent.get_sender_copy();
     let shutdown_token = trace_agent.shutdown_token();

bottlecap/src/bin/bottlecap/main.rs

@@ -42,7 +42,6 @@ impl Flusher {
         config: Arc<config::Config>,
     ) -> Self {
         let client = get_client(&config);
-
         Flusher {
             client,
             endpoint,
@@ -54,7 +53,10 @@ impl Flusher {
     }
 
     pub async fn flush(&self, batches: Option<Arc<Vec<Vec<u8>>>>) -> Vec<reqwest::RequestBuilder> {
-        let api_key = self.api_key_factory.get_api_key().await;
+        let Some(api_key) = self.api_key_factory.get_api_key().await else {
+            error!("Skipping flushing logs: Failed to resolve API key");
+            return vec![];
+        };
 
         let mut set = JoinSet::new();
 

bottlecap/src/logs/flusher.rs

@@ -399,7 +399,6 @@ mod tests {
         let metrics_aggregator = Arc::new(Mutex::new(
             MetricsAggregator::new(EMPTY_TAGS, 1024).unwrap(),
         ));
-
         let aws_config = Arc::new(AwsConfig {
             region: "us-east-1".to_string(),
             function_name: "arn:some-function".to_string(),

bottlecap/src/proxy/interceptor.rs

@@ -2,10 +2,9 @@ use dogstatsd::api_key::ApiKeyFactory;
 use reqwest::header::HeaderMap;
 use std::{error::Error, sync::Arc};
 use thiserror::Error as ThisError;
-use tokio::{
-    sync::{Mutex, OnceCell},
-    task::JoinSet,
-};
+use tokio::sync::OnceCell;
+use tokio::{sync::Mutex, task::JoinSet};
+
 use tracing::{debug, error};
 
 use crate::{
@@ -83,7 +82,10 @@ impl Flusher {
         &self,
         retry_requests: Option<Vec<reqwest::RequestBuilder>>,
     ) -> Option<Vec<reqwest::RequestBuilder>> {
-        let api_key = self.api_key_factory.get_api_key().await;
+        let Some(api_key) = self.api_key_factory.get_api_key().await else {
+            error!("Skipping flush in proxy flusher: Failed to resolve API key");
+            return None;
+        };
 
         let mut join_set = JoinSet::new();
         let mut requests = Vec::<reqwest::RequestBuilder>::new();

bottlecap/src/traces/proxy_flusher.rs

@@ -60,16 +60,21 @@ impl StatsFlusher for ServerlessStatsFlusher {
             return;
         }
 
+        let Some(api_key) = self.api_key_factory.get_api_key().await else {
+            error!("Skipping flushing stats: Failed to resolve API key");
+            return;
+        };
+
+        let api_key_clone = api_key.to_string();
         let endpoint = self
             .endpoint
             .get_or_init({
                 move || async move {
-                    let api_key = self.api_key_factory.get_api_key().await.to_string();
                     let stats_url = trace_stats_url(&self.config.site);
                     Endpoint {
                         url: hyper::Uri::from_str(&stats_url)
                             .expect("can't make URI from stats url, exiting"),
-                        api_key: Some(api_key.clone().into()),
+                        api_key: Some(api_key_clone.into()),
                         timeout_ms: self.config.flush_timeout * 1_000,
                         test_token: None,
                     }
@@ -95,12 +100,8 @@ impl StatsFlusher for ServerlessStatsFlusher {
 
         let start = std::time::Instant::now();
 
-        let resp = stats_utils::send_stats_payload(
-            serialized_stats_payload,
-            endpoint,
-            self.api_key_factory.get_api_key().await,
-        )
-        .await;
+        let resp =
+            stats_utils::send_stats_payload(serialized_stats_payload, endpoint, api_key).await;
         let elapsed = start.elapsed();
         debug!(
             "Stats request to {} took {}ms",

bottlecap/src/traces/stats_flusher.rs

@@ -37,7 +37,6 @@ use crate::{
 use datadog_trace_protobuf::pb;
 use datadog_trace_utils::trace_utils::{self};
 use ddcommon::hyper_migration;
-use dogstatsd::api_key::ApiKeyFactory;
 
 const TRACE_AGENT_PORT: usize = 8126;
 
@@ -86,7 +85,6 @@ pub struct StatsState {
 pub struct ProxyState {
     pub config: Arc<config::Config>,
     pub proxy_aggregator: Arc<Mutex<proxy_aggregator::Aggregator>>,
-    pub api_key_factory: Arc<ApiKeyFactory>,
 }
 
 pub struct TraceAgent {
@@ -99,7 +97,6 @@ pub struct TraceAgent {
     invocation_processor: Arc<Mutex<InvocationProcessor>>,
     shutdown_token: CancellationToken,
     tx: Sender<SendDataBuilderInfo>,
-    api_key_factory: Arc<ApiKeyFactory>,
 }
 
 #[derive(Clone, Copy)]
@@ -120,7 +117,6 @@ impl TraceAgent {
         proxy_aggregator: Arc<Mutex<proxy_aggregator::Aggregator>>,
         invocation_processor: Arc<Mutex<InvocationProcessor>>,
         tags_provider: Arc<provider::Provider>,
-        api_key_factory: Arc<ApiKeyFactory>,
     ) -> TraceAgent {
         // setup a channel to send processed traces to our flusher. tx is passed through each
         // endpoint_handler to the trace processor, which uses it to send de-serialized
@@ -147,7 +143,6 @@ impl TraceAgent {
             invocation_processor,
             tags_provider,
             tx: trace_tx,
-            api_key_factory,
             shutdown_token: CancellationToken::new(),
         }
     }
@@ -207,7 +202,6 @@ impl TraceAgent {
         let proxy_state = ProxyState {
             config: Arc::clone(&self.config),
             proxy_aggregator: Arc::clone(&self.proxy_aggregator),
-            api_key_factory: Arc::clone(&self.api_key_factory),
         };
 
         let trace_router = Router::new()
@@ -395,6 +389,8 @@ impl TraceAgent {
         (StatusCode::OK, response_json.to_string()).into_response()
     }
 
+    #[allow(clippy::too_many_arguments)]
+    #[allow(clippy::too_many_lines)]
     async fn handle_traces(
         config: Arc<config::Config>,
         request: Request,

bottlecap/src/traces/trace_agent.rs

@@ -10,13 +10,18 @@ use datadog_trace_utils::{
     send_data::SendDataBuilder,
     trace_utils::{self, SendData},
 };
+use dogstatsd::api_key::ApiKeyFactory;
 
 use crate::config::Config;
 use crate::traces::trace_aggregator::TraceAggregator;
 
 #[async_trait]
 pub trait TraceFlusher {
-    fn new(aggregator: Arc<Mutex<TraceAggregator>>, config: Arc<Config>) -> Self
+    fn new(
+        aggregator: Arc<Mutex<TraceAggregator>>,
+        config: Arc<Config>,
+        api_key_factory: Arc<ApiKeyFactory>,
+    ) -> Self
     where
         Self: Sized;
     /// Given a `Vec<SendData>`, a tracer payload, send it to the Datadog intake endpoint.
@@ -34,15 +39,29 @@ pub trait TraceFlusher {
 pub struct ServerlessTraceFlusher {
     pub aggregator: Arc<Mutex<TraceAggregator>>,
     pub config: Arc<Config>,
+    pub api_key_factory: Arc<ApiKeyFactory>,
 }
 
 #[async_trait]
 impl TraceFlusher for ServerlessTraceFlusher {
-    fn new(aggregator: Arc<Mutex<TraceAggregator>>, config: Arc<Config>) -> Self {
-        ServerlessTraceFlusher { aggregator, config }
+    fn new(
+        aggregator: Arc<Mutex<TraceAggregator>>,
+        config: Arc<Config>,
+        api_key_factory: Arc<ApiKeyFactory>,
+    ) -> Self {
+        ServerlessTraceFlusher {
+            aggregator,
+            config,
+            api_key_factory,
+        }
     }
 
     async fn flush(&self, failed_traces: Option<Vec<SendData>>) -> Option<Vec<SendData>> {
+        let Some(api_key) = self.api_key_factory.get_api_key().await else {
+            error!("Skipping flushing traces: Failed to resolve API key");
+            return None;
+        };
+
         let mut failed_batch: Option<Vec<SendData>> = None;
 
         if let Some(traces) = failed_traces {
@@ -64,6 +83,8 @@ impl TraceFlusher for ServerlessTraceFlusher {
         while !trace_builders.is_empty() {
             let traces: Vec<_> = trace_builders
                 .into_iter()
+                // Lazily set the API key
+                .map(|builder| builder.with_api_key(api_key))
                 .map(SendDataBuilder::build)
                 .collect();
             if let Some(failed) = self.send(traces).await {

bottlecap/src/traces/trace_flusher.rs

@@ -20,7 +20,6 @@ use datadog_trace_utils::trace_utils::{self};
 use datadog_trace_utils::tracer_header_tags;
 use datadog_trace_utils::tracer_payload::{TraceChunkProcessor, TracerPayloadCollection};
 use ddcommon::Endpoint;
-use dogstatsd::api_key::ApiKeyFactory;
 use std::str::FromStr;
 use std::sync::Arc;
 use tracing::error;
@@ -31,7 +30,6 @@ use super::trace_aggregator::SendDataBuilderInfo;
 #[allow(clippy::module_name_repetitions)]
 pub struct ServerlessTraceProcessor {
     pub obfuscation_config: Arc<obfuscation_config::ObfuscationConfig>,
-    pub api_key_factory: Arc<ApiKeyFactory>,
 }
 
 struct ChunkProcessor {
@@ -164,11 +162,11 @@ impl TraceProcessor for ServerlessTraceProcessor {
                 tracer_payload.tags.extend(tags.clone());
             }
         }
-        let api_key = self.api_key_factory.get_api_key().await.to_string();
         let endpoint = Endpoint {
             url: hyper::Uri::from_str(&config.apm_dd_url)
                 .expect("can't parse trace intake URL, exiting"),
-            api_key: Some(api_key.into()),
+            // Will be set at flush time
+            api_key: None,
             timeout_ms: config.flush_timeout * 1_000,
             test_token: None,
         };
@@ -194,7 +192,6 @@ mod tests {
     };
 
     use datadog_trace_obfuscation::obfuscation_config::ObfuscationConfig;
-    use dogstatsd::api_key::ApiKeyFactory;
 
     use crate::{config::Config, tags::provider::Provider, LAMBDA_RUNTIME_SLUG};
 
@@ -299,7 +296,6 @@ mod tests {
         };
 
         let trace_processor = ServerlessTraceProcessor {
-            api_key_factory: Arc::new(ApiKeyFactory::new("test-api-key")),
             obfuscation_config: Arc::new(ObfuscationConfig::new().unwrap()),
         };
         let config = create_test_config();