feat: Lazily resolve api key

[lym953]

Motivation

From @astuyve:

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

[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]

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]

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]

[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]

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

[lym953]

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

[lym953]

Lazily initialize headers, which includes api key.

[lym953]

A core change: added RwLock

[astuyve]

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

[lym953]

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

[lym953]

Lazily resolve endpoint, which contains api key.

[lym953]

Any better way to handle this?

[duncanista]

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

[lym953]

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]

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

[lym953]

Done! Extracted a function create_api_key_factory()

[duncanista]

Great PR @lym953 !