feat: Lazily resolve api key #717

lym953 · 2025-06-24T19:55:20Z

Motivation

today we basically block/await on that decrypt call before we can call /next
so if we can instead make that async and then resolve the future only when we need to flush data, that can be a big win for many customers.

https://datadoghq.atlassian.net/browse/SVLS-6995

Previous work

DataDog/serverless-components#21, DataDog/serverless-components#24 created ApiKeyFactory, which is a util to enable lazy API key resolution.

This PR

Updates Bottlecap code to use ApiKeyFactory to lazily resolve API key, i.e. instead of resolving it by querying Secret Manager or KMS during init phase, do it at flushing time when api key is actually needed.

Note

This PR changes the behavior when key resolution fails, i.e. when resolve_secrets() returns None.

Before: run extension_loop_idle(), which does not stop the runtime
After: panic, which will stop the runtime (if I understand correctly). Of course it's not ideal. Any better idea?
- It's harder now to run extension_loop_idle() because api key resolution code is not in the main loop anymore, but in various consumer code of api key
- Is there a way to gracefully shut down the extension without affecting the runtime?

Update: Added a PR to address resolution failure: #732
These two PRs should be merged together. Keeping them separate PRs just to make review easier.

Testing

Setup

Runtime: Go1 on Amazon Linux 2
Architecture: arm64
An app with empty implementation code

Result

Below is the Datadog Next-Gen Extension ready in: time logged.

Before: (prod extension arn:aws:lambda:us-east-1:464622532012:layer:Datadog-Extension-ARM:82)
- 88.6 ± 1.8 (ms)
After: (test extension arn:aws:lambda:us-east-1:425362996713:layer:Datadog-Bottlecap-Beta-ARM-yiming:2)
- 35.4 ± 5.1 (ms)
- (-60.0%)

Both use 5 samples.

Notes

https://datadoghq.atlassian.net/issues/SVLS-6996
https://datadoghq.atlassian.net/issues/SVLS-6998

# Problem Right now `AwsConfig` has a lot of fields, including the ones related to credential: ``` pub aws_access_key_id: String, pub aws_secret_access_key: String, pub aws_session_token: String, pub aws_container_credentials_full_uri: String, pub aws_container_authorization_token: String, ``` The next PR #717 wants to lazily load API key and the credentials. To do that, for the resolver function `resolve_secrets()`, I need to change the param `aws_config` from `&AwsConfig` to `Arc<RwLock<AwsConfig>>`. Because `aws_config` is passed to many places, this change involves updating lots of functions, which is formidable. # This PR Separates these credential-related fields out from `AwsConfig` and creates a new struct `AwsCredentials` Thus, the next PR will only need to change the param `aws_credentials` from `&AwsCredentials` to `Arc<RwLock<AwsCredentials>>`. Because `aws_credentials` is not used in lots of places, the next PR becomes easier. https://datadoghq.atlassian.net/issues/SVLS-6996 https://datadoghq.atlassian.net/issues/SVLS-6998

Copilot

Pull Request Overview

This PR integrates ApiKeyFactory across Bottlecap to defer DD-API-KEY resolution until flush/send time, reducing init latency. Key changes include:

Replace direct API key strings with Arc<ApiKeyFactory> in all flushers, agents, and tests
Refactor trace/stat flusher and log flusher to initialize endpoints and headers lazily via OnceCell
Update resolve_secrets to use an async RwLock for AWS credentials and adjust related helper signatures

Reviewed Changes

Copilot reviewed 13 out of 15 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
tests/metrics_integration_test.rs	Switched `FlusherConfig` to use `ApiKeyFactory`
tests/logs_integration_test.rs	Switched `LogsFlusher` instantiation to use factory
src/traces/trace_processor.rs	Made `process_traces` async and use `ApiKeyFactory`
src/traces/trace_agent.rs	Replaced stored key with factory; await per request
src/traces/stats_flusher.rs	Swapped in factory and lazily build `Endpoint`
src/secrets/decrypt.rs	Converted credentials to `Arc<RwLock<_>>` and updated calls
src/proxy/mod.rs	Changed `should_start_proxy` to take `Arc<AwsConfig>`
src/proxy/interceptor.rs	Updated interceptor to use `Arc<AwsConfig>`
src/otlp/agent.rs	Updated OTLP agent to await `process_traces`
src/logs/flusher.rs	Introduced factory plus lazy `HeaderMap` caching
src/lifecycle/invocation/span_inferrer.rs	Updated tests to pass `Arc<AwsConfig>`
src/lifecycle/invocation/processor.rs	Refactored processor to use `Arc<AwsConfig>`
Cargo.toml	Bumped `dogstatsd` revision

Copilot · 2025-07-03T15:55:20Z

bottlecap/src/traces/stats_flusher.rs

+    endpoint: OnceCell<Endpoint>,
 }



Caching the Endpoint with an initial API key may lead to stale keys if ApiKeyFactory rotates or refreshes credentials. Consider regenerating or invalidating the cell when the underlying key changes.

Suggested change

endpoint: OnceCell<Endpoint>,

}

}

impl ServerlessStatsFlusher {

async fn construct_endpoint(&self) -> Endpoint {

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()),

timeout_ms: self.config.flush_timeout * 1_000,

test_token: None,

}

}

}

Copilot · 2025-07-03T15:55:20Z

bottlecap/src/logs/flusher.rs

@@ -26,42 +28,24 @@ pub struct Flusher {
    client: reqwest::Client,
    aggregator: Arc<Mutex<Aggregator>>,
    config: Arc<config::Config>,
-    headers: HeaderMap,
+    api_key_factory: Arc<ApiKeyFactory>,
+    headers: OnceCell<HeaderMap>,


Storing headers including DD-API-KEY in a OnceCell will cache the first resolved key indefinitely. If the key can change at runtime, you might emit outdated headers; consider refreshing per flush or using a time-to-live.

Suggested change

headers: OnceCell<HeaderMap>,

headers: Arc<Mutex<HeaderMap>>,

Copilot · 2025-07-03T15:55:21Z

bottlecap/src/traces/trace_processor.rs

@@ -323,7 +329,7 @@ mod tests {
        };

        let received_payload =
-            if let TracerPayloadCollection::V07(payload) = tracer_payload.get_payloads() {
+            if let TracerPayloadCollection::V07(payload) = tracer_payload.await.get_payloads() {


[nitpick] Awaiting the future inline in the match expression can reduce readability. It may help to .await process_traces into a local send_data variable first, then call get_payloads() on it.

lym953 · 2025-07-03T15:48:08Z

bottlecap/src/bin/bottlecap/main.rs

+    let aws_config = Arc::new(aws_config);
+    let aws_credentials = Arc::new(RwLock::new(aws_credentials));
+    let api_key_factory = {
+        let config = Arc::clone(&config);
+        let aws_config = Arc::clone(&aws_config);
+        let aws_credentials = Arc::clone(&aws_credentials);
+
+        Arc::new(ApiKeyFactory::new_from_resolver(Arc::new(move || {
+            let config = Arc::clone(&config);
+            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
+                    .unwrap_or_else(|| {
+                        error!("Failed to resolve API key");
+                        String::new()
+                    })
+            })
+        })))
+    };


During INIT phase, instead of resolving API key, just initialize an API key factory.

lym953 · 2025-07-03T15:50:19Z

bottlecap/src/bin/bottlecap/main.rs

@@ -415,11 +430,11 @@ async fn extension_loop_idle(client: &Client, r: &RegisterResponse) -> Result<()

 #[allow(clippy::too_many_lines)]
 async fn extension_loop_active(
-    aws_config: &AwsConfig,
+    aws_config: Arc<AwsConfig>,


Changing it to Arc so it can be passed to ApiKeyFactory and shared across threads.

lym953 · 2025-07-03T15:51:12Z

bottlecap/src/logs/flusher.rs

@@ -26,42 +28,24 @@ pub struct Flusher {
    client: reqwest::Client,
    aggregator: Arc<Mutex<Aggregator>>,
    config: Arc<config::Config>,
-    headers: HeaderMap,
+    api_key_factory: Arc<ApiKeyFactory>,
+    headers: OnceCell<HeaderMap>,


Lazily initialize headers, which includes api key.

lym953 · 2025-07-03T15:52:04Z

bottlecap/src/secrets/decrypt.rs

-    aws_config: &AwsConfig,
-    aws_credentials: &mut AwsCredentials,
+    aws_config: Arc<AwsConfig>,
+    aws_credentials: Arc<RwLock<AwsCredentials>>,


A core change: added RwLock

ah and we need an RwLock here because the factory will lazily write/update this struct member?

Yes. The factory fills in aws_access_key_id, aws_secret_access_key and aws_session_token if we are in snap start.

bottlecap/src/secrets/decrypt.rs

lym953 · 2025-07-03T15:53:40Z

bottlecap/src/traces/stats_flusher.rs

@@ -34,36 +36,47 @@ pub struct ServerlessStatsFlusher {
    // pub buffer: Arc<Mutex<Vec<pb::ClientStatsPayload>>>,
    aggregator: Arc<Mutex<StatsAggregator>>,
    config: Arc<config::Config>,
-    endpoint: Endpoint,
+    api_key_factory: Arc<ApiKeyFactory>,
+    endpoint: OnceCell<Endpoint>,


Lazily resolve endpoint, which contains api key.

lym953 · 2025-07-03T17:07:16Z

bottlecap/src/bin/bottlecap/main.rs

+            Box::pin(async move {
+                resolve_secrets(config, aws_config, aws_credentials)
+                    .await
+                    .expect("Failed to resolve API key")


Any better way to handle this?

Mmm ideally, if failing on resolve you'd enter the noop loop, what would happen if you added that here?

It will panic at flush time and stop the runtime, so I wonder if there's a way for the extension to stop gracefully at that time without stopping the runtime.

astuyve

Nice work Yiming! Let's get this onto self monitoring to test for a few days

duncanista · 2025-07-09T16:43:09Z

bottlecap/src/bin/bottlecap/main.rs

+    let api_key_factory = {
+        let config = Arc::clone(&config);
+        let aws_config = Arc::clone(&aws_config);
+        let aws_credentials = Arc::clone(&aws_credentials);
+
+        Arc::new(ApiKeyFactory::new_from_resolver(Arc::new(move || {
+            let config = Arc::clone(&config);
+            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")
+            })
+        })))
+    };


Wondering if we could make this a method here on in the api key resolver, really don't like the pattern of having the declaration of the same variable 3 times just because we're nesting it, given how main code has increasingly grown, it would be good to hide this in some way and document it

Done! Extracted a function create_api_key_factory()

duncanista

Great PR @lym953 !

…745) # Background Right now `SendData` is passed around across channels. # This PR Instead of passing `SendData`, pass `SendDataBuilderInfo`, which bundles `SendDataBuilder` and payload size. Just before flush, call `SendDataBuilder.build()` to build `SendData`. # Motivation DataDog/libdatadog#1140 (comment) It is suggested that the function `set_api_key()` shouldn't be added on `SendData`, but should be added on `SendDataBuilder`. Because need to call `set_api_key()` just before flush, we need to make sure the object is `SendDataBuilder` instead of `SendData` until flush time. And because we need payload size in Trace Aggregator, and `SendDataBuilder` doesn't expose this field, we need to pass it explicitly along with `SendDataBuilder`. # Next steps Update #717 #732 so that `get_api_key()` is called just before flush. # Dependency DataDog/libdatadog#1140

Fix clippy

lym953 mentioned this pull request Jun 27, 2025

feat: Add helper functions to ApiKeyFactory struct DataDog/serverless-components#24

Merged

lym953 force-pushed the yiming.luo/separate-aws-creds branch from f749307 to 11c1c77 Compare June 27, 2025 17:35

lym953 mentioned this pull request Jun 30, 2025

chore: Separate AwsCredentials from AwsConfig #716

Merged

Base automatically changed from yiming.luo/separate-aws-creds to main July 2, 2025 21:11

lym953 force-pushed the yiming.luo/lazy-api-key-3 branch from fdd93a4 to 4ac16ad Compare July 2, 2025 21:31

lym953 requested a review from Copilot July 3, 2025 15:52

Copilot AI reviewed Jul 3, 2025

View reviewed changes

lym953 commented Jul 3, 2025

View reviewed changes

lym953 marked this pull request as ready for review July 3, 2025 17:09

lym953 requested a review from a team as a code owner July 3, 2025 17:09

lym953 force-pushed the yiming.luo/lazy-api-key-3 branch from d623d67 to 3724235 Compare July 7, 2025 17:02

lym953 mentioned this pull request Jul 8, 2025

feat: Handle API key resolution failure #732

Open

astuyve approved these changes Jul 9, 2025

View reviewed changes

lym953 force-pushed the yiming.luo/lazy-api-key-3 branch from cec3247 to 727d04f Compare July 9, 2025 15:47

duncanista reviewed Jul 9, 2025

View reviewed changes

duncanista approved these changes Jul 9, 2025

View reviewed changes

lym953 force-pushed the yiming.luo/lazy-api-key-3 branch from 3be2928 to e820be0 Compare July 10, 2025 20:18

lym953 changed the base branch from main to jordan.gonzalez/trace-agent/aggregation-for-proxy July 10, 2025 20:19

duncanista approved these changes Jul 11, 2025

View reviewed changes

duncanista force-pushed the jordan.gonzalez/trace-agent/aggregation-for-proxy branch from 4c1162d to 07aa341 Compare July 14, 2025 19:11

lym953 mentioned this pull request Jul 14, 2025

chore: pass SendDataBuilderInfo instead of SendData until flush time #745

Merged

Base automatically changed from jordan.gonzalez/trace-agent/aggregation-for-proxy to main July 16, 2025 19:47

lym953 force-pushed the yiming.luo/lazy-api-key-3 branch from 2ef2a9a to 37caca4 Compare July 16, 2025 20:51

feat: Lazily resolve api key

ef63759

Fix clippy

lym953 force-pushed the yiming.luo/lazy-api-key-3 branch from 37caca4 to ef63759 Compare July 16, 2025 21:46

-    endpoint: OnceCell<Endpoint>,
-}
+}
+impl ServerlessStatsFlusher {
+    async fn construct_endpoint(&self) -> Endpoint {
+        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()),
+            timeout_ms: self.config.flush_timeout * 1_000,
+            test_token: None,
+        }
+    }
+}

	headers: OnceCell<HeaderMap>,
	headers: Arc<Mutex<HeaderMap>>,

feat: Lazily resolve api key #717

Are you sure you want to change the base?

feat: Lazily resolve api key #717

Conversation

lym953 commented Jun 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Previous work

This PR

Note

Testing

Setup

Result

Notes

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Copilot AI Jul 3, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 3, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 3, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

astuyve left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

duncanista left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

lym953 commented Jun 24, 2025 •

edited

Loading