Fix rebase conflicts

Author: lym953

bottlecap/src/logs/flusher.rs

@@ -2,8 +2,8 @@ use crate::config;
 use crate::http::get_client;
 use crate::logs::aggregator::Aggregator;
 use crate::FLUSH_RETRY_COUNT;
-use futures::future::join_all;
 use dogstatsd::api_key::ApiKeyFactory;
+use futures::future::join_all;
 use reqwest::header::HeaderMap;
 use std::error::Error;
 use std::time::Instant;
@@ -95,7 +95,6 @@ impl Flusher {
 
     async fn create_request(&self, data: Vec<u8>) -> reqwest::RequestBuilder {
         let url = format!("{}/api/v2/logs", self.endpoint);
-        let body = self.compress(data);
         let headers = self.get_headers().await;
         self.client
             .post(&url)
@@ -146,6 +145,35 @@ impl Flusher {
             }
         }
     }
+
+    async fn get_headers(&self) -> &HeaderMap {
+        self.headers
+            .get_or_init(move || async move {
+                let api_key = self.api_key_factory.get_api_key().await;
+                let mut headers = HeaderMap::new();
+                headers.insert(
+                    "DD-API-KEY",
+                    api_key.parse().expect("failed to parse header"),
+                );
+                headers.insert(
+                    "DD-PROTOCOL",
+                    "agent-json".parse().expect("failed to parse header"),
+                );
+                headers.insert(
+                    "Content-Type",
+                    "application/json".parse().expect("failed to parse header"),
+                );
+
+                if self.config.logs_config_use_compression {
+                    headers.insert(
+                        "Content-Encoding",
+                        "zstd".parse().expect("failed to parse header"),
+                    );
+                }
+                headers
+            })
+            .await
+    }
 }
 
 #[allow(clippy::module_name_repetitions)]
@@ -157,15 +185,15 @@ pub struct LogsFlusher {
 
 impl LogsFlusher {
     pub fn new(
-        api_key: String,
+        api_key_factory: Arc<ApiKeyFactory>,
         aggregator: Arc<Mutex<Aggregator>>,
         config: Arc<config::Config>,
     ) -> Self {
         let mut flushers = Vec::new();
 
         // Create primary flusher
         flushers.push(Flusher::new(
-            api_key.clone(),
+            Arc::clone(&api_key_factory),
             config.logs_config_logs_dd_url.clone(),
             aggregator.clone(),
             config.clone(),
@@ -175,7 +203,7 @@ impl LogsFlusher {
         for endpoint in &config.logs_config_additional_endpoints {
             let endpoint_url = format!("https://{}:{}", endpoint.host, endpoint.port);
             flushers.push(Flusher::new(
-                endpoint.api_key.clone(),
+                Arc::clone(&api_key_factory),
                 endpoint_url,
                 aggregator.clone(),
                 config.clone(),
@@ -252,33 +280,4 @@ impl LogsFlusher {
         encoder.write_all(data)?;
         encoder.finish().map_err(|e| Box::new(e) as Box<dyn Error>)
     }
-
-    async fn get_headers(&self) -> &HeaderMap {
-        self.headers
-            .get_or_init(move || async move {
-                let api_key = self.api_key_factory.get_api_key().await;
-                let mut headers = HeaderMap::new();
-                headers.insert(
-                    "DD-API-KEY",
-                    api_key.parse().expect("failed to parse header"),
-                );
-                headers.insert(
-                    "DD-PROTOCOL",
-                    "agent-json".parse().expect("failed to parse header"),
-                );
-                headers.insert(
-                    "Content-Type",
-                    "application/json".parse().expect("failed to parse header"),
-                );
-
-                if self.config.logs_config_use_compression {
-                    headers.insert(
-                        "Content-Encoding",
-                        "zstd".parse().expect("failed to parse header"),
-                    );
-                }
-                headers
-            })
-            .await
-    }
 }

bottlecap/src/traces/trace_agent.rs

@@ -76,7 +76,7 @@ pub struct ProxyState {
     pub config: Arc<config::Config>,
     pub tags_provider: Arc<provider::Provider>,
     pub http_client: reqwest::Client,
-    pub api_key: String,
+    pub api_key_factory: Arc<ApiKeyFactory>,
 }
 
 pub struct TraceAgent {
@@ -197,7 +197,7 @@ impl TraceAgent {
             config: Arc::clone(&self.config),
             tags_provider: Arc::clone(&self.tags_provider),
             http_client: self.http_client.clone(),
-            api_key: self.api_key.clone(),
+            api_key_factory: Arc::clone(&self.api_key_factory),
         };
 
         let trace_router = Router::new()
@@ -289,7 +289,7 @@ impl TraceAgent {
         Self::handle_proxy(
             state.config,
             state.http_client,
-            state.api_key,
+            state.api_key_factory,
             state.tags_provider,
             request,
             "trace.agent",
@@ -303,7 +303,7 @@ impl TraceAgent {
         Self::handle_proxy(
             state.config,
             state.http_client,
-            state.api_key,
+            state.api_key_factory,
             state.tags_provider,
             request,
             "intake.profile",
@@ -320,7 +320,7 @@ impl TraceAgent {
         Self::handle_proxy(
             state.config,
             state.http_client,
-            state.api_key,
+            state.api_key_factory,
             state.tags_provider,
             request,
             "api",
@@ -337,7 +337,7 @@ impl TraceAgent {
         Self::handle_proxy(
             state.config,
             state.http_client,
-            state.api_key,
+            state.api_key_factory,
             state.tags_provider,
             request,
             "api",
@@ -351,7 +351,7 @@ impl TraceAgent {
         Self::handle_proxy(
             state.config,
             state.http_client,
-            state.api_key,
+            state.api_key_factory,
             state.tags_provider,
             request,
             "llmobs-intake",
@@ -365,7 +365,7 @@ impl TraceAgent {
         Self::handle_proxy(
             state.config,
             state.http_client,
-            state.api_key,
+            state.api_key_factory,
             state.tags_provider,
             request,
             "http-intake.logs",