feat(trace_utils): support generic connector in send_data (#1340)

feat(trace_utils): support generic connector in send_data

define ddcommon connect trait

fix import

Co-authored-by: vianney.ruhlmann <vianney.ruhlmann@datadoghq.com>

Author: VianneyRuhlmann

libdd-common/src/connector/mod.rs

@@ -211,14 +211,14 @@ mod tests {
     /// Verify that the Connector type implements the correct bound Connect + Clone
     /// to be able to use the hyper::Client
     fn test_hyper_client_from_connector() {
-        let _: hyper_migration::HttpClient = hyper_migration::new_default_client();
+        let _ = hyper_migration::new_default_client();
     }
 
     #[test]
     #[cfg_attr(miri, ignore)]
     #[cfg(feature = "use_webpki_roots")]
     fn test_hyper_client_from_connector_with_webpki_roots() {
-        let _: hyper_migration::HttpClient = hyper_migration::new_default_client();
+        let _ = hyper_migration::new_default_client();
     }
 
     #[tokio::test]

libdd-common/src/hyper_migration.rs

@@ -4,7 +4,7 @@
 use core::fmt;
 use std::{convert::Infallible, task::Poll};
 
-use crate::connector::{self, Connector};
+use crate::connector::Connector;
 use http_body_util::BodyExt;
 use hyper::body::Incoming;
 use pin_project::pin_project;
@@ -17,7 +17,7 @@ use hyper::Request as HyperRequest;
 /// every second connection because of low keep alive in the agent.
 ///
 /// This is on general not a problem if we use the client once every tens of seconds.
-pub fn new_client_periodic() -> HttpClient {
+pub fn new_client_periodic() -> GenericHttpClient<Connector> {
     hyper_util::client::legacy::Client::builder(hyper_util::rt::TokioExecutor::default())
         .pool_max_idle_per_host(0)
         .build(Connector::default())
@@ -26,7 +26,7 @@ pub fn new_client_periodic() -> HttpClient {
 /// Create a new default configuration hyper client.
 ///
 /// It will keep connections open for a longer time and reuse them.
-pub fn new_default_client() -> HttpClient {
+pub fn new_default_client() -> GenericHttpClient<Connector> {
     hyper_util::client::legacy::Client::builder(hyper_util::rt::TokioExecutor::default())
         .build(Connector::default())
 }
@@ -209,7 +209,7 @@ impl hyper::body::Body for Body {
     }
 }
 
-pub type HttpClient = hyper_util::client::legacy::Client<connector::Connector, Body>;
+pub type GenericHttpClient<C> = hyper_util::client::legacy::Client<C, Body>;
 
 pub fn client_builder() -> hyper_util::client::legacy::Builder {
     hyper_util::client::legacy::Client::builder(hyper_util::rt::TokioExecutor::default())

libdd-common/src/lib.rs

@@ -108,9 +108,18 @@ pub mod header {
         HeaderName::from_static("x-datadog-test-session-token");
 }
 
-pub type HttpClient = hyper_migration::HttpClient;
+pub type HttpClient = hyper_migration::GenericHttpClient<connector::Connector>;
+pub type GenericHttpClient<C> = hyper_migration::GenericHttpClient<C>;
 pub type HttpResponse = hyper_migration::HttpResponse;
 pub type HttpRequestBuilder = hyper::http::request::Builder;
+pub trait Connect:
+    hyper_util::client::legacy::connect::Connect + Clone + Send + Sync + 'static
+{
+}
+impl<C: hyper_util::client::legacy::connect::Connect + Clone + Send + Sync + 'static> Connect
+    for C
+{
+}
 
 // Used by tag! macro
 use crate::entity_id::DD_EXTERNAL_ENV;

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

@@ -11,13 +11,12 @@ use anyhow::{anyhow, Context};
 use futures::stream::FuturesUnordered;
 use futures::StreamExt;
 use hyper::header::CONTENT_TYPE;
-use libdd_common::HttpClient;
 use libdd_common::{
     header::{
         APPLICATION_MSGPACK_STR, APPLICATION_PROTOBUF_STR, DATADOG_SEND_REAL_HTTP_STATUS_STR,
         DATADOG_TRACE_COUNT_STR,
     },
-    Endpoint,
+    Connect, Endpoint, GenericHttpClient,
 };
 use libdd_trace_protobuf::pb::{AgentPayload, TracerPayload};
 use send_data_result::SendDataResult;
@@ -237,24 +236,27 @@ impl SendData {
     /// # Returns
     ///
     /// A `SendDataResult` instance containing the result of the operation.
-    pub async fn send(&self, http_client: &HttpClient) -> SendDataResult {
+    pub async fn send<C: Connect>(&self, http_client: &GenericHttpClient<C>) -> SendDataResult {
         self.send_internal(http_client).await
     }
 
-    async fn send_internal(&self, http_client: &HttpClient) -> SendDataResult {
+    async fn send_internal<C: Connect>(
+        &self,
+        http_client: &GenericHttpClient<C>,
+    ) -> SendDataResult {
         if self.use_protobuf() {
             self.send_with_protobuf(http_client).await
         } else {
             self.send_with_msgpack(http_client).await
         }
     }
 
-    async fn send_payload(
+    async fn send_payload<C: Connect>(
         &self,
         chunks: u64,
         payload: Vec<u8>,
         headers: HashMap<&'static str, String>,
-        http_client: &HttpClient,
+        http_client: &GenericHttpClient<C>,
     ) -> (SendWithRetryResult, u64, u64) {
         #[allow(clippy::unwrap_used)]
         let payload_len = u64::try_from(payload.len()).unwrap();
@@ -298,7 +300,10 @@ impl SendData {
         }
     }
 
-    async fn send_with_protobuf(&self, http_client: &HttpClient) -> SendDataResult {
+    async fn send_with_protobuf<C: Connect>(
+        &self,
+        http_client: &GenericHttpClient<C>,
+    ) -> SendDataResult {
         let mut result = SendDataResult::default();
 
         #[allow(clippy::unwrap_used)]
@@ -336,7 +341,10 @@ impl SendData {
         }
     }
 
-    async fn send_with_msgpack(&self, http_client: &HttpClient) -> SendDataResult {
+    async fn send_with_msgpack<C: Connect>(
+        &self,
+        http_client: &GenericHttpClient<C>,
+    ) -> SendDataResult {
         let mut result = SendDataResult::default();
         let mut futures = FuturesUnordered::new();
 

libdd-trace-utils/src/send_with_retry/mod.rs

@@ -9,7 +9,7 @@ pub use retry_strategy::{RetryBackoffType, RetryStrategy};
 
 use bytes::Bytes;
 use hyper::Method;
-use libdd_common::{hyper_migration, Endpoint, HttpClient, HttpRequestBuilder};
+use libdd_common::{hyper_migration, Connect, Endpoint, GenericHttpClient, HttpRequestBuilder};
 use std::{collections::HashMap, time::Duration};
 use tracing::{debug, error};
 
@@ -111,8 +111,8 @@ impl std::error::Error for RequestError {}
 /// send_with_retry(&client, &target, payload, &headers, &retry_strategy).await
 /// # }
 /// ```
-pub async fn send_with_retry(
-    client: &HttpClient,
+pub async fn send_with_retry<C: Connect>(
+    client: &GenericHttpClient<C>,
     target: &Endpoint,
     payload: Vec<u8>,
     headers: &HashMap<&'static str, String>,
@@ -222,8 +222,8 @@ pub async fn send_with_retry(
     }
 }
 
-async fn send_request(
-    client: &HttpClient,
+async fn send_request<C: Connect>(
+    client: &GenericHttpClient<C>,
     timeout: Duration,
     req: HttpRequestBuilder,
     payload: Bytes,