diff --git a/crates/datadog-serverless-compat/src/main.rs b/crates/datadog-serverless-compat/src/main.rs index 24f7f5fb..270d03a5 100644 --- a/crates/datadog-serverless-compat/src/main.rs +++ b/crates/datadog-serverless-compat/src/main.rs @@ -27,7 +27,6 @@ use datadog_trace_utils::{config_utils::read_cloud_env, trace_utils::Environment use dogstatsd::{ aggregator::Aggregator as MetricsAggregator, - api_key::ApiKeyFactory, constants::CONTEXTS, datadog::{MetricsIntakeUrlPrefix, RetryStrategy, Site}, dogstatsd::{DogStatsD, DogStatsDConfig}, @@ -204,7 +203,7 @@ async fn start_dogstatsd( Some(dd_api_key) => { #[allow(clippy::expect_used)] let metrics_flusher = Flusher::new(FlusherConfig { - api_key_factory: Arc::new(ApiKeyFactory::new(&dd_api_key)), + api_key: dd_api_key, aggregator: Arc::clone(&metrics_aggr), metrics_intake_url_prefix: MetricsIntakeUrlPrefix::new( Some(Site::new(dd_site).expect("Failed to parse site")), diff --git a/crates/dogstatsd/src/api_key.rs b/crates/dogstatsd/src/api_key.rs deleted file mode 100644 index 6a32c806..00000000 --- a/crates/dogstatsd/src/api_key.rs +++ /dev/null @@ -1,71 +0,0 @@ -use std::fmt::Debug; -use std::sync::Arc; -use std::{future::Future, pin::Pin}; -use tokio::sync::OnceCell; - -pub type ApiKeyResolverFn = - Arc Pin + Send>> + Send + Sync>; - -#[derive(Clone)] -pub enum ApiKeyFactory { - Static(String), - Dynamic { - resolver_fn: ApiKeyResolverFn, - api_key: Arc>, - }, -} - -impl ApiKeyFactory { - /// Create a new `ApiKeyFactory` with a static API key. - pub fn new(api_key: &str) -> Self { - Self::Static(api_key.to_string()) - } - - /// Create a new `ApiKeyFactory` with a dynamic API key resolver function. - pub fn new_from_resolver(resolver_fn: ApiKeyResolverFn) -> Self { - Self::Dynamic { - resolver_fn, - api_key: Arc::new(OnceCell::new()), - } - } - - pub async fn get_api_key(&self) -> &str { - match self { - Self::Static(api_key) => api_key, - Self::Dynamic { - resolver_fn, - api_key, - } => { - api_key - .get_or_init(|| async { (resolver_fn)().await }) - .await - } - } - } -} - -impl Debug for ApiKeyFactory { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "ApiKeyFactory") - } -} - -#[cfg(test)] -pub mod tests { - use super::*; - - #[tokio::test] - async fn test_new() { - let api_key_factory = ApiKeyFactory::new("mock-api-key"); - assert_eq!(api_key_factory.get_api_key().await, "mock-api-key"); - } - - #[tokio::test] - async fn test_new_from_resolver() { - let api_key_factory = Arc::new(ApiKeyFactory::new_from_resolver(Arc::new(move || { - let api_key = "mock-api-key".to_string(); - Box::pin(async move { api_key }) - }))); - assert_eq!(api_key_factory.get_api_key().await, "mock-api-key"); - } -} diff --git a/crates/dogstatsd/src/flusher.rs b/crates/dogstatsd/src/flusher.rs index 2947ad7b..b4bc771c 100644 --- a/crates/dogstatsd/src/flusher.rs +++ b/crates/dogstatsd/src/flusher.rs @@ -2,7 +2,6 @@ // SPDX-License-Identifier: Apache-2.0 use crate::aggregator::Aggregator; -use crate::api_key::ApiKeyFactory; use crate::datadog::{DdApi, MetricsIntakeUrlPrefix, RetryStrategy}; use reqwest::{Response, StatusCode}; use std::sync::{Arc, Mutex}; @@ -11,18 +10,12 @@ use tracing::{debug, error}; #[derive(Clone)] pub struct Flusher { - // Allow accepting a future so the API key resolution is deferred until the flush happens - api_key_factory: Arc, - metrics_intake_url_prefix: MetricsIntakeUrlPrefix, - https_proxy: Option, - timeout: Duration, - retry_strategy: RetryStrategy, + dd_api: DdApi, aggregator: Arc>, - dd_api: Option, } pub struct FlusherConfig { - pub api_key_factory: Arc, + pub api_key: String, pub aggregator: Arc>, pub metrics_intake_url_prefix: MetricsIntakeUrlPrefix, pub https_proxy: Option, @@ -33,35 +26,19 @@ pub struct FlusherConfig { #[allow(clippy::await_holding_lock)] impl Flusher { pub fn new(config: FlusherConfig) -> Self { + let dd_api = DdApi::new( + config.api_key, + config.metrics_intake_url_prefix, + config.https_proxy, + config.timeout, + config.retry_strategy, + ); Flusher { - api_key_factory: Arc::clone(&config.api_key_factory), - metrics_intake_url_prefix: config.metrics_intake_url_prefix, - https_proxy: config.https_proxy, - timeout: config.timeout, - retry_strategy: config.retry_strategy, + dd_api, aggregator: config.aggregator, - dd_api: None, } } - async fn get_dd_api(&mut self) -> &DdApi { - if self.dd_api.is_none() { - let api_key = self.api_key_factory.get_api_key().await; - self.dd_api = Some(DdApi::new( - api_key.to_string(), - self.metrics_intake_url_prefix.clone(), - self.https_proxy.clone(), - self.timeout, - self.retry_strategy.clone(), - )); - } - - #[allow(clippy::expect_used)] - self.dd_api - .as_ref() - .expect("dd_api should have been initialized") - } - /// Flush metrics from the aggregator pub async fn flush( &mut self, @@ -98,9 +75,7 @@ impl Flusher { let series_copy = series.clone(); let distributions_copy = distributions.clone(); - let dd_api = self.get_dd_api().await; - - let dd_api_clone = dd_api.clone(); + let dd_api_clone = self.dd_api.clone(); let series_handle = tokio::spawn(async move { let mut failed = Vec::new(); let mut had_shipping_error = false; @@ -118,7 +93,7 @@ impl Flusher { (failed, had_shipping_error) }); - let dd_api_clone = dd_api.clone(); + let dd_api_clone = self.dd_api.clone(); let distributions_handle = tokio::spawn(async move { let mut failed = Vec::new(); let mut had_shipping_error = false; diff --git a/crates/dogstatsd/src/lib.rs b/crates/dogstatsd/src/lib.rs index 59cbd5b9..968ee549 100644 --- a/crates/dogstatsd/src/lib.rs +++ b/crates/dogstatsd/src/lib.rs @@ -8,7 +8,6 @@ #![cfg_attr(not(test), deny(clippy::unimplemented))] pub mod aggregator; -pub mod api_key; pub mod constants; pub mod datadog; pub mod dogstatsd; diff --git a/crates/dogstatsd/tests/integration_test.rs b/crates/dogstatsd/tests/integration_test.rs index 10fbc783..2ddbed1a 100644 --- a/crates/dogstatsd/tests/integration_test.rs +++ b/crates/dogstatsd/tests/integration_test.rs @@ -4,7 +4,6 @@ use dogstatsd::metric::SortedTags; use dogstatsd::{ aggregator::Aggregator as MetricsAggregator, - api_key::ApiKeyFactory, constants::CONTEXTS, datadog::{DdDdUrl, MetricsIntakeUrlPrefix, MetricsIntakeUrlPrefixOverride}, dogstatsd::{DogStatsD, DogStatsDConfig}, @@ -40,10 +39,8 @@ async fn dogstatsd_server_ships_series() { let _ = start_dogstatsd(&metrics_aggr).await; - let api_key_factory = ApiKeyFactory::new("mock-api-key"); - let mut metrics_flusher = Flusher::new(FlusherConfig { - api_key_factory: Arc::new(api_key_factory), + api_key: "mock-api-key".to_string(), aggregator: Arc::clone(&metrics_aggr), metrics_intake_url_prefix: MetricsIntakeUrlPrefix::new( None,