Refactor data-pipeline to hold a client pool so the connection with the agent is reused.

Author: hoolioh

data-pipeline/src/stats_exporter.rs

@@ -14,7 +14,8 @@ use std::{
 use crate::{span_concentrator::SpanConcentrator, trace_exporter::TracerMetadata};
 use datadog_trace_protobuf::pb;
 use datadog_trace_utils::send_with_retry::{send_with_retry, RetryStrategy};
-use ddcommon::{worker::Worker, Endpoint};
+use ddcommon::hyper_migration::new_default_client;
+use ddcommon::{worker::Worker, Endpoint, HttpClient};
 use hyper;
 use tokio::select;
 use tokio_util::sync::CancellationToken;
@@ -31,6 +32,7 @@ pub struct StatsExporter {
     meta: TracerMetadata,
     sequence_id: AtomicU64,
     cancellation_token: CancellationToken,
+    client: HttpClient,
 }
 
 impl StatsExporter {
@@ -48,6 +50,7 @@ impl StatsExporter {
         meta: TracerMetadata,
         endpoint: Endpoint,
         cancellation_token: CancellationToken,
+        client: Option<HttpClient>,
     ) -> Self {
         Self {
             flush_interval,
@@ -56,6 +59,7 @@ impl StatsExporter {
             meta,
             sequence_id: AtomicU64::new(0),
             cancellation_token,
+            client: client.unwrap_or(new_default_client()),
         }
     }
 
@@ -89,11 +93,11 @@ impl StatsExporter {
         );
 
         let result = send_with_retry(
+            &self.client,
             &self.endpoint,
             body,
             &headers,
             &RetryStrategy::default(),
-            None,
         )
         .await;
 
@@ -266,6 +270,7 @@ mod tests {
             get_test_metadata(),
             Endpoint::from_url(stats_url_from_agent_url(&server.url("/")).unwrap()),
             CancellationToken::new(),
+            None,
         );
 
         let send_status = stats_exporter.send(true).await;
@@ -293,6 +298,7 @@ mod tests {
             get_test_metadata(),
             Endpoint::from_url(stats_url_from_agent_url(&server.url("/")).unwrap()),
             CancellationToken::new(),
+            None,
         );
 
         let send_status = stats_exporter.send(true).await;
@@ -325,6 +331,7 @@ mod tests {
             get_test_metadata(),
             Endpoint::from_url(stats_url_from_agent_url(&server.url("/")).unwrap()),
             CancellationToken::new(),
+            None,
         );
 
         tokio::time::pause();
@@ -365,6 +372,7 @@ mod tests {
             get_test_metadata(),
             Endpoint::from_url(stats_url_from_agent_url(&server.url("/")).unwrap()),
             cancellation_token.clone(),
+            None,
         );
 
         tokio::spawn(async move {

data-pipeline/src/trace_exporter/builder.rs

@@ -12,6 +12,7 @@ use crate::trace_exporter::{
     INFO_ENDPOINT,
 };
 use arc_swap::ArcSwap;
+use ddcommon::hyper_migration::new_default_client;
 use ddcommon::{parse_uri, tag, Endpoint};
 use dogstatsd_client::new;
 use std::sync::{Arc, Mutex};
@@ -343,6 +344,7 @@ impl TraceExporterBuilder {
             agent_payload_response_version: self
                 .agent_rates_payload_version_enabled
                 .then(AgentResponsePayloadVersion::new),
+            http_client: new_default_client(),
         })
     }
 

data-pipeline/src/trace_exporter/mod.rs

@@ -36,9 +36,9 @@ use datadog_trace_utils::send_with_retry::{
 };
 use datadog_trace_utils::span::{Span, SpanText};
 use datadog_trace_utils::trace_utils::TracerHeaderTags;
-use ddcommon::MutexExt;
 use ddcommon::{hyper_migration, Endpoint};
 use ddcommon::{tag, tag::Tag};
+use ddcommon::{HttpClient, MutexExt};
 use ddtelemetry::worker::TelemetryWorker;
 use dogstatsd_client::Client;
 use http_body_util::BodyExt;
@@ -208,6 +208,7 @@ pub struct TraceExporter {
     health_metrics_enabled: bool,
     workers: Arc<Mutex<TraceExporterWorkers>>,
     agent_payload_response_version: Option<AgentResponsePayloadVersion>,
+    http_client: HttpClient,
 }
 
 impl TraceExporter {
@@ -424,6 +425,7 @@ impl TraceExporter {
                             &agent_info,
                             &self.client_side_stats,
                             &self.workers,
+                            Some(self.http_client.clone()),
                         );
                     }
                     StatsComputationStatus::Enabled {
@@ -627,7 +629,8 @@ impl TraceExporter {
         let payload_len = mp_payload.len();
 
         // Send traces to the agent
-        let result = send_with_retry(endpoint, mp_payload, &headers, &strategy, None).await;
+        let result =
+            send_with_retry(&self.http_client, endpoint, mp_payload, &headers, &strategy).await;
 
         // Emit http.requests health metric based on number of attempts
         let requests_count = match &result {

data-pipeline/src/trace_exporter/stats.rs

@@ -11,7 +11,7 @@ use crate::agent_info::schema::AgentInfo;
 use crate::span_concentrator::SpanConcentrator;
 use crate::stats_exporter;
 use arc_swap::ArcSwap;
-use ddcommon::{Endpoint, MutexExt};
+use ddcommon::{Endpoint, HttpClient, MutexExt};
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
 use tokio::runtime::Runtime;
@@ -64,6 +64,7 @@ pub(crate) fn start_stats_computation(
     workers: &Arc<Mutex<super::TraceExporterWorkers>>,
     span_kinds: Vec<String>,
     peer_tags: Vec<String>,
+    client: Option<HttpClient>,
 ) -> anyhow::Result<()> {
     if let StatsComputationStatus::DisabledByAgent { bucket_size } = **client_side_stats.load() {
         let stats_concentrator = Arc::new(Mutex::new(SpanConcentrator::new(
@@ -80,6 +81,7 @@ pub(crate) fn start_stats_computation(
             &cancellation_token,
             workers,
             client_side_stats,
+            client,
         )?;
     }
     Ok(())
@@ -93,13 +95,15 @@ fn create_and_start_stats_worker(
     cancellation_token: &CancellationToken,
     workers: &Arc<Mutex<super::TraceExporterWorkers>>,
     client_side_stats: &ArcSwap<StatsComputationStatus>,
+    client: Option<HttpClient>,
 ) -> anyhow::Result<()> {
     let stats_exporter = stats_exporter::StatsExporter::new(
         bucket_size,
         stats_concentrator.clone(),
         ctx.metadata.clone(),
         Endpoint::from_url(add_path(ctx.endpoint_url, STATS_ENDPOINT)),
         cancellation_token.clone(),
+        client,
     );
     let mut stats_worker = crate::pausable_worker::PausableWorker::new(stats_exporter);
 
@@ -160,6 +164,7 @@ pub(crate) fn handle_stats_disabled_by_agent(
     agent_info: &Arc<AgentInfo>,
     client_side_stats: &ArcSwap<StatsComputationStatus>,
     workers: &Arc<Mutex<super::TraceExporterWorkers>>,
+    client: Option<HttpClient>,
 ) {
     if agent_info.info.client_drop_p0s.is_some_and(|v| v) {
         // Client-side stats is supported by the agent
@@ -169,6 +174,7 @@ pub(crate) fn handle_stats_disabled_by_agent(
             workers,
             get_span_kinds_for_stats(agent_info),
             agent_info.info.peer_tags.clone().unwrap_or_default(),
+            client,
         );
         match status {
             Ok(()) => info!("Client-side stats enabled"),