Refactor trace-utils to pass an http client so the connections are reused.

Author: hoolioh

datadog-trace-utils/Cargo.toml

@@ -38,11 +38,6 @@ tinybytes = { path = "../tinybytes", features = [
     "serialization",
 ] }
 
-# Proxy feature
-hyper-http-proxy = { version = "1.1.0", default-features = false, features = [
-    "rustls-tls-webpki-roots",
-], optional = true }
-
 # Compression feature
 flate2 = { version = "1.0", optional = true }
 zstd = { version = "0.13.3", default-features = false, optional = true }
@@ -66,15 +61,14 @@ tempfile = "3.3.0"
 [features]
 default = ["https"]
 https = ["ddcommon/https"]
-mini_agent = ["proxy", "compression", "ddcommon/use_webpki_roots"]
+mini_agent = ["compression", "ddcommon/use_webpki_roots"]
 test-utils = [
     "hyper/server",
     "httpmock",
     "cargo_metadata",
     "cargo-platform",
     "urlencoding",
 ]
-proxy = ["hyper-http-proxy"]
 compression = ["zstd", "flate2"]
 # FIPS mode uses the FIPS-compliant cryptographic provider (Unix only)
 fips = ["ddcommon/fips"]

datadog-trace-utils/src/send_data/mod.rs

@@ -9,6 +9,7 @@ use crate::trace_utils::TracerHeaderTags;
 use crate::tracer_payload::TracerPayloadCollection;
 use anyhow::{anyhow, Context};
 use datadog_trace_protobuf::pb::{AgentPayload, TracerPayload};
+use ddcommon::HttpClient;
 use ddcommon::{
     header::{
         APPLICATION_MSGPACK_STR, APPLICATION_PROTOBUF_STR, DATADOG_SEND_REAL_HTTP_STATUS_STR,
@@ -42,6 +43,7 @@ use zstd::stream::write::Encoder;
 /// use datadog_trace_utils::trace_utils::TracerHeaderTags;
 /// use datadog_trace_utils::tracer_payload::TracerPayloadCollection;
 /// use ddcommon::Endpoint;
+/// use ddcommon::hyper_migration::new_default_client;
 ///
 /// #[cfg_attr(miri, ignore)]
 /// async fn update_send_results_example() {
@@ -58,8 +60,9 @@ use zstd::stream::write::Encoder;
 ///
 ///     send_data.set_retry_strategy(retry_strategy);
 ///
+///     let client = new_default_client();
 ///     // Send the data
-///     let result = send_data.send().await;
+///     let result = send_data.send(&client).await;
 /// }
 /// ```
 pub struct SendData {
@@ -234,24 +237,15 @@ impl SendData {
     /// # Returns
     ///
     /// A `SendDataResult` instance containing the result of the operation.
-    pub async fn send(&self) -> SendDataResult {
-        self.send_internal(None).await
+    pub async fn send(&self, http_client: &HttpClient) -> SendDataResult {
+        self.send_internal(http_client).await
     }
 
-    /// Sends the data to the target endpoint.
-    ///
-    /// # Returns
-    ///
-    /// A `SendDataResult` instance containing the result of the operation.
-    pub async fn send_proxy(&self, http_proxy: Option<&str>) -> SendDataResult {
-        self.send_internal(http_proxy).await
-    }
-
-    async fn send_internal(&self, http_proxy: Option<&str>) -> SendDataResult {
+    async fn send_internal(&self, http_client: &HttpClient) -> SendDataResult {
         if self.use_protobuf() {
-            self.send_with_protobuf(http_proxy).await
+            self.send_with_protobuf(http_client).await
         } else {
-            self.send_with_msgpack(http_proxy).await
+            self.send_with_msgpack(http_client).await
         }
     }
 
@@ -260,17 +254,17 @@ impl SendData {
         chunks: u64,
         payload: Vec<u8>,
         headers: HashMap<&'static str, String>,
-        http_proxy: Option<&str>,
+        http_client: &HttpClient,
     ) -> (SendWithRetryResult, u64, u64) {
         #[allow(clippy::unwrap_used)]
         let payload_len = u64::try_from(payload.len()).unwrap();
         (
             send_with_retry(
+                http_client,
                 &self.target,
                 payload,
                 &headers,
                 &self.retry_strategy,
-                http_proxy,
             )
             .await,
             payload_len,
@@ -304,7 +298,7 @@ impl SendData {
         }
     }
 
-    async fn send_with_protobuf(&self, http_proxy: Option<&str>) -> SendDataResult {
+    async fn send_with_protobuf(&self, http_client: &HttpClient) -> SendDataResult {
         let mut result = SendDataResult::default();
 
         #[allow(clippy::unwrap_used)]
@@ -331,7 +325,7 @@ impl SendData {
                 request_headers.insert(CONTENT_TYPE.as_str(), APPLICATION_PROTOBUF_STR.to_string());
 
                 let (response, bytes_sent, chunks) = self
-                    .send_payload(chunks, final_payload, request_headers, http_proxy)
+                    .send_payload(chunks, final_payload, request_headers, http_client)
                     .await;
 
                 result.update(response, bytes_sent, chunks);
@@ -342,7 +336,7 @@ impl SendData {
         }
     }
 
-    async fn send_with_msgpack(&self, http_proxy: Option<&str>) -> SendDataResult {
+    async fn send_with_msgpack(&self, http_client: &HttpClient) -> SendDataResult {
         let mut result = SendDataResult::default();
         let mut futures = FuturesUnordered::new();
 
@@ -360,7 +354,7 @@ impl SendData {
                         Err(e) => return result.error(anyhow!(e)),
                     };
 
-                    futures.push(self.send_payload(chunks, payload, headers, http_proxy));
+                    futures.push(self.send_payload(chunks, payload, headers, http_client));
                 }
             }
             TracerPayloadCollection::V04(payload) => {
@@ -372,7 +366,7 @@ impl SendData {
 
                 let payload = msgpack_encoder::v04::to_vec(payload);
 
-                futures.push(self.send_payload(chunks, payload, headers, http_proxy));
+                futures.push(self.send_payload(chunks, payload, headers, http_client));
             }
             TracerPayloadCollection::V05(payload) => {
                 #[allow(clippy::unwrap_used)]
@@ -386,7 +380,7 @@ impl SendData {
                     Err(e) => return result.error(anyhow!(e)),
                 };
 
-                futures.push(self.send_payload(chunks, payload, headers, http_proxy));
+                futures.push(self.send_payload(chunks, payload, headers, http_client));
             }
         }
 
@@ -592,7 +586,8 @@ mod tests {
         );
 
         let data_payload_len = compute_payload_len(&data.tracer_payloads);
-        let res = data.send().await;
+        let client = ddcommon::hyper_migration::new_default_client();
+        let res = data.send(&client).await;
 
         mock.assert_async().await;
 
@@ -637,7 +632,8 @@ mod tests {
         );
 
         let data_payload_len = compute_payload_len(&data.tracer_payloads);
-        let res = data.send().await;
+        let client = ddcommon::hyper_migration::new_default_client();
+        let res = data.send(&client).await;
 
         mock.assert_async().await;
 
@@ -696,7 +692,8 @@ mod tests {
         );
 
         let data_payload_len = rmp_compute_payload_len(&data.tracer_payloads);
-        let res = data.send().await;
+        let client = ddcommon::hyper_migration::new_default_client();
+        let res = data.send(&client).await;
 
         mock.assert_async().await;
 
@@ -754,7 +751,8 @@ mod tests {
         );
 
         let data_payload_len = rmp_compute_payload_len(&data.tracer_payloads);
-        let res = data.send().await;
+        let client = ddcommon::hyper_migration::new_default_client();
+        let res = data.send(&client).await;
 
         mock.assert_async().await;
 
@@ -798,7 +796,8 @@ mod tests {
         );
 
         let data_payload_len = rmp_compute_payload_len(&data.tracer_payloads);
-        let res = data.send().await;
+        let client = ddcommon::hyper_migration::new_default_client();
+        let res = data.send(&client).await;
 
         mock.assert_calls_async(2).await;
 
@@ -839,7 +838,8 @@ mod tests {
             },
         );
 
-        let res = data.send().await;
+        let client = ddcommon::hyper_migration::new_default_client();
+        let res = data.send(&client).await;
 
         mock.assert_calls_async(5).await;
 
@@ -871,7 +871,8 @@ mod tests {
             },
         );
 
-        let res = data.send().await;
+        let client = ddcommon::hyper_migration::new_default_client();
+        let res = data.send(&client).await;
 
         assert!(res.last_result.is_err());
         match std::env::consts::OS {
@@ -938,7 +939,8 @@ mod tests {
             },
         );
 
-        let res = data.send().await;
+        let client = ddcommon::hyper_migration::new_default_client();
+        let res = data.send(&client).await;
 
         mock.assert_calls_async(5).await;
 
@@ -980,7 +982,8 @@ mod tests {
             },
         );
 
-        let res = data.send().await;
+        let client = ddcommon::hyper_migration::new_default_client();
+        let res = data.send(&client).await;
 
         mock.assert_calls_async(10).await;