feat: Add helper functions to ApiKeyFactory struct by lym953 · Pull Request #24 · DataDog/serverless-components

lym953 · 2025-06-25T15:31:31Z

What does this PR do?

Add helper functions to ApiKeyFactory struct, including:
1. new_from_resolver(), which takes in a resolver function and enables lazy resolution
2. new_from_static_key(), which takes in a static api key string slice
3. get_api_key()
Move ApiKeyFactory related code to a separate module api_key

Motivation

As a continuation of #21.

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

Additional Notes

Describe how to test/QA your changes

Passed the added automated tests
Used it in Lambda extension, which did improve start time. See feat: Lazily resolve api key datadog-lambda-extension#717

DarcyRaynerDD · 2025-06-27T17:35:30Z

+
+#[derive(Clone)]
+pub struct ApiKeyFactory {
+    inner: Arc<ApiKeyFactoryInner>,


Do you need this Arc? the ApiKeyResolverFn is already wrapped in an Arc, and the OnceCell should already be Send + Sync. Can you just derive Clone on ApiKeyFactoryInner? And then if you can do that, you no longer need the wrapper struct. You might even be able to get away with not having the Arc around the ApiKeyResolverFn, since you are already required it to be Send+Sync

I might be wrong on OnceCell, it looks like it's memory is copied. You might be able to get away with just wrapping the OnceCell in an Arc, and that way at least you don't need the wrapper struct, and in the default case there is no overhead for using the static string.

There is also LazyLock which might give the correct behaviour with less overhead https://doc.rust-lang.org/std/sync/struct.LazyLock.html

1.

Do you need this Arc? the ApiKeyResolverFn is already wrapped in an Arc, and the OnceCell should already be Send + Sync. Can you just derive Clone on ApiKeyFactoryInner? And then if you can do that, you no longer need the wrapper struct.

Great! This works.

2.

You might even be able to get away with not having the Arc around the ApiKeyResolverFn, since you are already required it to be Send+Sync

Do you mean this?

pub type ApiKeyResolverFn = dyn Fn() -> Pin<Box<dyn Future<Output = String> + Send>> + Send + Sync; #[derive(Clone)] pub enum ApiKeyFactory { Static(String), Dynamic { resolver_fn: ApiKeyResolverFn, api_key: OnceCell<String>, }, }

This doesn't work because resolver_fn has unknown size.

3.

I might be wrong on OnceCell, it looks like it's memory is copied. You might be able to get away with just wrapping the OnceCell in an Arc

You are right. Will add Arc.

4.

There is also LazyLock which might give the correct behaviour with less overhead

This doesn't work because std::sync::LazyLock doesn't allow async functions. OnceCell is the best choice I found that works inside tokio.

Thanks for the comments! I also just learned that enum can have impl.

Copilot

Pull Request Overview

This PR introduces a dedicated api_key module to encapsulate API key resolution logic and updates existing components to use the new ApiKeyFactory helpers.

Moved ApiKeyFactory logic into its own module with static and dynamic constructors.
Refactored Flusher to depend on ApiKeyFactory instead of raw closures.
Updated integration tests and serverless-compat usage to call the new helpers.

Reviewed Changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
crates/dogstatsd/tests/integration_test.rs	Swapped manual API key closure for `ApiKeyFactory::new_from_static_key`
crates/dogstatsd/src/lib.rs	Added `pub mod api_key`
crates/dogstatsd/src/flusher.rs	Updated `Flusher` to accept `Arc<ApiKeyFactory>`
crates/dogstatsd/src/api_key.rs	Added `ApiKeyFactory` enum, helper constructors, and unit tests
crates/datadog-serverless-compat/src/main.rs	Replaced inline API key closures with `ApiKeyFactory::new_from_static_key`

Comments suppressed due to low confidence (2)

crates/dogstatsd/src/api_key.rs:37

Add a unit test to verify that the dynamic resolver is only invoked once (i.e., that the OnceCell caching works as intended).

                    .get_or_init(|| async { (resolver_fn)().await })

crates/dogstatsd/src/api_key.rs:17

[nitpick] Add module- or item-level doc comments for ApiKeyFactory, new_from_resolver, new_from_static_key, and get_api_key to clarify intended usage.

impl ApiKeyFactory {

Copilot · 2025-06-27T18:32:42Z

 pub struct Flusher {
-    // Accept a future so the API key resolution is deferred until the flush happens
-    api_key_factory: ApiKeyFactory,
+    // Allow accepting a future so the API key resolution is deferred until the flush happens


[nitpick] Update this comment to reflect that api_key_factory now uses ApiKeyFactory rather than accepting a raw future-producing closure.

Suggested change

// Allow accepting a future so the API key resolution is deferred until the flush happens

// Uses ApiKeyFactory to defer API key resolution until the flush happens

duncanista · 2025-06-27T18:43:42Z

+}
+
+impl ApiKeyFactory {
+    pub fn new_from_resolver(resolver_fn: ApiKeyResolverFn) -> Self {


Maybe some docs on what is considered a resolver

duncanista · 2025-06-27T18:46:18Z

+
+#[cfg(test)]
+pub mod tests {
+    use crate::api_key::ApiKeyFactory;


in tests. it's common to just use use super::*, which would give you all the imports from the file you're in

duncanista · 2025-06-27T18:47:46Z

+    use std::sync::Arc;
+
+    #[tokio::test]
+    async fn new_from_resolver() {


It's a good practice to name tests in the form of test_new_from_resolver to not confuse user to the usage above, sometimes, if you wanna be more descriptive and add more context, you can also do test_new_from_resolver_<EXPECTED-ACTION>

This reverts commit b29bba8.

@astuyve

# 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%) <img width="461" alt="image" src="https://github.com/user-attachments/assets/b2973aae-d8f2-4003-a37f-6af05a42e059" /> Both use 5 samples. # Notes https://datadoghq.atlassian.net/issues/SVLS-6996 https://datadoghq.atlassian.net/issues/SVLS-6998

lym953 requested a review from Copilot June 25, 2025 17:50

This comment was marked as outdated.

Sign in to view

lym953 added 4 commits June 27, 2025 10:44

feat: Move key resolution logic to ApiKeyFactory struct

5d2dcdb

Avoid async closure

d3f95e3

Use Arc<ApiKeyFactory>

2252249

impl Debug for ApiKeyFactory

f6c7a65

lym953 force-pushed the yiming.luo/api-key-future-cursor-2 branch from f606cce to f6c7a65 Compare June 27, 2025 14:44

lym953 added 4 commits June 27, 2025 10:46

Change unwrap to expect with msg

689429e

Add an enum ApiKeyFactoryInner

0e06eb2

Rename functions

e94c813

Format

6bdd328

lym953 requested a review from Copilot June 27, 2025 15:52

lym953 changed the title ~~feat: Move key resolution logic to ApiKeyFactory struct~~ feat: Add helper functions to ApiKeyFactory struct Jun 27, 2025

This comment was marked as outdated.

Sign in to view

Add #[allow(clippy::expect_used)]

8987283

lym953 commented Jun 27, 2025

View reviewed changes

Comment thread crates/dogstatsd/src/api_key.rs

lym953 marked this pull request as ready for review June 27, 2025 16:08

DarcyRaynerDD reviewed Jun 27, 2025

View reviewed changes

Address comments

8d99785

lym953 requested a review from Copilot June 27, 2025 18:31

Format

aa17d9a

Copilot AI reviewed Jun 27, 2025

View reviewed changes

duncanista reviewed Jun 27, 2025

View reviewed changes

Comment thread crates/dogstatsd/src/api_key.rs Outdated

duncanista reviewed Jun 27, 2025

View reviewed changes

lym953 added 4 commits June 27, 2025 14:53

Rename: new_from_static_key() -> new()

74f8ab1

Add doc for new() functions

98eac71

Use super::* in tests

07c343d

Rename tests

1822e69

duncanista approved these changes Jun 27, 2025

View reviewed changes

lym953 merged commit b29bba8 into main Jun 27, 2025
26 checks passed

lym953 deleted the yiming.luo/api-key-future-cursor-2 branch June 27, 2025 20:25

lym953 mentioned this pull request Jul 3, 2025

feat: Lazily resolve api key DataDog/datadog-lambda-extension#717

Merged

lym953 added a commit that referenced this pull request Jul 10, 2025

Revert "feat: Add helper functions to ApiKeyFactory struct (#24)"

e783068

This reverts commit b29bba8.

lym953 mentioned this pull request Jul 10, 2025

Revert ApiKeyFactory commits #28

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Add helper functions to ApiKeyFactory struct#24

feat: Add helper functions to ApiKeyFactory struct#24
lym953 merged 15 commits intomainfrom
yiming.luo/api-key-future-cursor-2

lym953 commented Jun 25, 2025 •

edited

Loading

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

Uh oh!

DarcyRaynerDD Jun 27, 2025

Uh oh!

DarcyRaynerDD Jun 27, 2025

Uh oh!

DarcyRaynerDD Jun 27, 2025

Uh oh!

lym953 Jun 27, 2025

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Jun 27, 2025

Uh oh!

Uh oh!

duncanista Jun 27, 2025

Uh oh!

duncanista Jun 27, 2025 •

edited

Loading

Uh oh!

duncanista Jun 27, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

	// Allow accepting a future so the API key resolution is deferred until the flush happens
	// Uses ApiKeyFactory to defer API key resolution until the flush happens

Conversation

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

Uh oh!

What does this PR do?

Motivation

Additional Notes

Describe how to test/QA your changes

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

Uh oh!

DarcyRaynerDD Jun 27, 2025

Choose a reason for hiding this comment

Uh oh!

DarcyRaynerDD Jun 27, 2025

Choose a reason for hiding this comment

Uh oh!

DarcyRaynerDD Jun 27, 2025

Choose a reason for hiding this comment

Uh oh!

lym953 Jun 27, 2025

Choose a reason for hiding this comment

1.

2.

3.

4.

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Copilot AI Jun 27, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

duncanista Jun 27, 2025

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

duncanista Jun 27, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

lym953 commented Jun 25, 2025 •

edited

Loading

duncanista Jun 27, 2025 •

edited

Loading