feat: Add helper functions to ApiKeyFactory struct

[lym953]

What does this PR do?

1. 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()
2. Move ApiKeyFactory related code to a separate module api_key

Motivation

As a continuation of #21.

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

Additional Notes

Describe how to test/QA your changes

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

[DarcyRaynerDD]

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

[DarcyRaynerDD]

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.

[DarcyRaynerDD]

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

[lym953]

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]

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

Maybe some docs on what is considered a resolver

[duncanista]

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

[duncanista]

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>