chore: refactor logs flusher (#562)

* fix clippy holding lock warning

* refactor logs flusher

* use futures instead of task joinSet

Author: duncanista

bottlecap/Cargo.lock

@@ -41,6 +41,7 @@ rustls = { version = "0.23.18", default-features = false, features = ["aws-lc-rs
 rand = { version = "0.8", default-features = false }
 prost = { version = "0.11.6", default-features = false }
 zstd = { version = "0.13.2", default-features = false }
+futures = { version = "0.3.31", default-features = false}
 
 [dev-dependencies]
 figment = { version = "0.10", default-features = false, features = ["yaml", "env", "test"] }

bottlecap/Cargo.toml

@@ -3,6 +3,7 @@ use tracing::warn;
 
 use crate::logs::constants;
 
+#[derive(Debug, Clone)]
 pub struct Aggregator {
     messages: VecDeque<String>,
     max_batch_entries_size: usize,

bottlecap/src/logs/aggregator.rs

@@ -1,22 +1,24 @@
 use crate::config;
 use crate::http_client;
 use crate::logs::aggregator::Aggregator;
+use futures::future::join_all;
+use reqwest::header::HeaderMap;
 use std::time::Instant;
 use std::{
     error::Error,
     io::Write,
     sync::{Arc, Mutex},
 };
-use tokio::task::JoinSet;
 use tracing::{debug, error};
 use zstd::stream::write::Encoder;
 
+#[derive(Debug, Clone)]
 pub struct Flusher {
-    api_key: String,
     fqdn_site: String,
     client: reqwest::Client,
     aggregator: Arc<Mutex<Aggregator>>,
     config: Arc<config::Config>,
+    headers: HeaderMap,
 }
 
 #[inline]
@@ -25,7 +27,6 @@ pub fn build_fqdn_logs(site: String) -> String {
     format!("https://http-intake.logs.{site}")
 }
 
-#[allow(clippy::await_holding_lock)]
 impl Flusher {
     pub fn new(
         api_key: String,
@@ -34,113 +35,118 @@ impl Flusher {
         config: Arc<config::Config>,
     ) -> Self {
         let client = http_client::get_client(config.clone());
+        let mut headers = HeaderMap::new();
+        headers.insert(
+            "DD-API-KEY",
+            api_key.clone().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 config.logs_config_use_compression {
+            headers.insert(
+                "Content-Encoding",
+                "zstd".parse().expect("failed to parse header"),
+            );
+        }
+
         Flusher {
-            api_key,
             fqdn_site: site,
             client,
             aggregator,
             config,
+            headers,
         }
     }
     pub async fn flush(&self) {
-        let mut guard = self.aggregator.lock().expect("lock poisoned");
-        let mut set = JoinSet::new();
-
-        let mut logs = guard.get_batch();
-        while !logs.is_empty() {
-            let api_key = self.api_key.clone();
-            let site = self.fqdn_site.clone();
-            let cloned_client = self.client.clone();
-            let cloned_use_compression = self.config.logs_config_use_compression;
-            let cloned_compression_level = self.config.logs_config_compression_level;
-            set.spawn(async move {
-                Self::send(
-                    cloned_client,
-                    api_key,
-                    site,
-                    logs,
-                    cloned_use_compression,
-                    cloned_compression_level,
-                )
-                .await;
-            });
-            logs = guard.get_batch();
-        }
-        drop(guard);
-        while let Some(res) = set.join_next().await {
-            match res {
-                Ok(()) => (),
-                Err(e) => {
-                    error!("Failed to wait for request sending {}", e);
-                }
+        let logs_batches = {
+            let mut guard = self.aggregator.lock().expect("lock poisoned");
+            let mut batches = Vec::new();
+            let mut current_batch = guard.get_batch();
+
+            while !current_batch.is_empty() {
+                batches.push(current_batch);
+                current_batch = guard.get_batch();
+            }
+
+            batches
+        };
+
+        let futures = logs_batches.into_iter().filter(|b| !b.is_empty()).map(|b| {
+            let req = self.create_request(b);
+            Self::send(req)
+        });
+
+        let results = join_all(futures).await;
+
+        for result in results {
+            if let Err(e) = result {
+                debug!("Failed to send logs: {}", e);
             }
         }
     }
 
-    #[allow(clippy::unwrap_used)]
-    async fn send(
-        client: reqwest::Client,
-        api_key: String,
-        fqdn: String,
-        data: Vec<u8>,
-        compression_enabled: bool,
-        compression_level: i32,
-    ) {
-        if !data.is_empty() {
-            let url = format!("{fqdn}/api/v2/logs");
-            let start = Instant::now();
-            let body = if compression_enabled {
-                let result = (|| -> Result<Vec<u8>, Box<dyn Error>> {
-                    let mut encoder = Encoder::new(Vec::new(), compression_level)
-                        .map_err(|e| Box::new(e) as Box<dyn Error>)?;
-                    encoder
-                        .write_all(&data)
-                        .map_err(|e| Box::new(e) as Box<dyn Error>)?;
-
-                    encoder.finish().map_err(|e| Box::new(e) as Box<dyn Error>)
-                })();
-
-                if let Ok(compressed_data) = result {
-                    compressed_data
-                } else {
-                    debug!("Failed to compress data, sending uncompressed data");
-                    data
-                }
-            } else {
-                data
-            };
-            let req = client
-                .post(&url)
-                .header("DD-API-KEY", api_key)
-                .header("DD-PROTOCOL", "agent-json")
-                .header("Content-Type", "application/json");
-            let req = if compression_enabled {
-                req.header("Content-Encoding", "zstd")
-            } else {
-                req
-            };
-            let resp: Result<reqwest::Response, reqwest::Error> = req.body(body).send().await;
-
-            let elapsed = start.elapsed();
-
-            match resp {
-                Ok(resp) => {
-                    if resp.status() != 202 {
-                        debug!(
-                            "Failed to send logs to datadog after {}ms: {}",
-                            elapsed.as_millis(),
-                            resp.status()
-                        );
-                    }
-                }
-                Err(e) => {
-                    error!(
+    fn create_request(&self, data: Vec<u8>) -> reqwest::RequestBuilder {
+        let url = format!("{}/api/v2/logs", self.fqdn_site);
+        let body = self.compress(data);
+        self.client
+            .post(&url)
+            .headers(self.headers.clone())
+            .body(body)
+    }
+
+    async fn send(req: reqwest::RequestBuilder) -> Result<(), Box<dyn Error>> {
+        let time = Instant::now();
+        let resp = req.send().await;
+        let elapsed = time.elapsed();
+
+        match resp {
+            Ok(resp) => {
+                if resp.status() != 202 {
+                    debug!(
                         "Failed to send logs to datadog after {}ms: {}",
                         elapsed.as_millis(),
-                        e
+                        resp.status()
                     );
                 }
             }
+            Err(e) => {
+                error!(
+                    "Failed to send logs to datadog after {}ms: {}",
+                    elapsed.as_millis(),
+                    e
+                );
+
+                return Err(Box::new(e));
+            }
         }
+
+        Ok(())
+    }
+
+    fn compress(&self, data: Vec<u8>) -> Vec<u8> {
+        if !self.config.logs_config_use_compression {
+            return data;
+        }
+
+        match self.encode(&data) {
+            Ok(compressed_data) => compressed_data,
+            Err(e) => {
+                debug!("Failed to compress data: {}", e);
+                data
+            }
+        }
+    }
+
+    fn encode(&self, data: &[u8]) -> Result<Vec<u8>, Box<dyn Error>> {
+        let mut encoder = Encoder::new(Vec::new(), self.config.logs_config_compression_level)?;
+        encoder.write_all(data)?;
+        encoder.finish().map_err(|e| Box::new(e) as Box<dyn Error>)
     }
 }