use bounded channel for stats concentrator

Author: duncanpharvey

crates/datadog-trace-agent/src/stats_concentrator_service.rs

@@ -12,6 +12,10 @@ use tracing::error;
 const S_TO_NS: u64 = 1_000_000_000;
 const BUCKET_DURATION_NS: u64 = 10 * S_TO_NS; // 10 seconds
 
+/// A bounded channel applies backpressure on the trace request path when the concentrator
+/// cannot keep up, instead of growing without limit as an unbounded channel would.
+const CONCENTRATOR_COMMAND_CHANNEL_CAPACITY: usize = 8192;
+
 #[derive(Debug, thiserror::Error)]
 pub enum StatsError {
     #[error("Failed to send command to concentrator: {0}")]
@@ -39,7 +43,7 @@ pub enum ConcentratorCommand {
 }
 
 pub struct StatsConcentratorHandle {
-    tx: mpsc::UnboundedSender<ConcentratorCommand>,
+    tx: mpsc::Sender<ConcentratorCommand>,
     is_tracer_metadata_set: AtomicBool,
 }
 
@@ -58,14 +62,14 @@ impl Clone for StatsConcentratorHandle {
 
 impl StatsConcentratorHandle {
     #[must_use]
-    pub fn new(tx: mpsc::UnboundedSender<ConcentratorCommand>) -> Self {
+    pub fn new(tx: mpsc::Sender<ConcentratorCommand>) -> Self {
         Self {
             tx,
             is_tracer_metadata_set: AtomicBool::new(false),
         }
     }
 
-    pub fn set_tracer_metadata(&self, trace: &TracerPayload) -> Result<(), StatsError> {
+    pub async fn set_tracer_metadata(&self, trace: &TracerPayload) -> Result<(), StatsError> {
         // Set tracer metadata only once for the first trace because
         // it is the same for all traces.
         if !self.is_tracer_metadata_set.load(Ordering::Acquire) {
@@ -78,14 +82,16 @@ impl StatsConcentratorHandle {
             };
             self.tx
                 .send(ConcentratorCommand::SetTracerMetadata(tracer_metadata))
+                .await
                 .map_err(StatsError::SendError)?;
         }
         Ok(())
     }
 
-    pub fn add(&self, span: &pb::Span) -> Result<(), StatsError> {
+    pub async fn add(&self, span: &pb::Span) -> Result<(), StatsError> {
         self.tx
             .send(ConcentratorCommand::Add(Box::new(span.clone())))
+            .await
             .map_err(StatsError::SendError)?;
         Ok(())
     }
@@ -94,14 +100,15 @@ impl StatsConcentratorHandle {
         let (response_tx, response_rx) = oneshot::channel();
         self.tx
             .send(ConcentratorCommand::Flush(force_flush, response_tx))
+            .await
             .map_err(StatsError::SendError)?;
         response_rx.await.map_err(StatsError::RecvError)
     }
 }
 
 pub struct StatsConcentratorService {
     concentrator: SpanConcentrator,
-    rx: mpsc::UnboundedReceiver<ConcentratorCommand>,
+    rx: mpsc::Receiver<ConcentratorCommand>,
     tracer_metadata: TracerMetadata,
     config: Arc<Config>,
 }
@@ -111,7 +118,7 @@ pub struct StatsConcentratorService {
 impl StatsConcentratorService {
     #[must_use]
     pub fn new(config: Arc<Config>) -> (Self, StatsConcentratorHandle) {
-        let (tx, rx) = mpsc::unbounded_channel();
+        let (tx, rx) = mpsc::channel(CONCENTRATOR_COMMAND_CHANNEL_CAPACITY);
         let handle = StatsConcentratorHandle::new(tx);
         // TODO: set span_kinds_stats_computed and peer_tag_keys
         let concentrator = SpanConcentrator::new(

crates/datadog-trace-agent/src/stats_generator.rs

@@ -21,19 +21,19 @@ impl StatsGenerator {
         Self { stats_concentrator }
     }
 
-    pub fn send(&self, traces: &TracerPayloadCollection) -> Result<(), StatsGeneratorError> {
+    pub async fn send(&self, traces: &TracerPayloadCollection) -> Result<(), StatsGeneratorError> {
         if let TracerPayloadCollection::V07(traces) = traces {
             for trace in traces {
                 // Set tracer metadata
-                if let Err(err) = self.stats_concentrator.set_tracer_metadata(trace) {
+                if let Err(err) = self.stats_concentrator.set_tracer_metadata(trace).await {
                     error!("Failed to set tracer metadata: {err}");
                     return Err(StatsGeneratorError::ConcentratorCommandError(err));
                 }
 
                 // Generate stats for each span in the trace
                 for chunk in &trace.chunks {
                     for span in &chunk.spans {
-                        if let Err(err) = self.stats_concentrator.add(span) {
+                        if let Err(err) = self.stats_concentrator.add(span).await {
                             error!("Failed to send trace stats: {err}");
                             return Err(StatsGeneratorError::ConcentratorCommandError(err));
                         }

crates/datadog-trace-agent/src/trace_processor.rs

@@ -145,7 +145,7 @@ impl TraceProcessor for ServerlessTraceProcessor {
 
         if let Some(stats_generator) = self.stats_generator.as_ref()
             && !tracer_header_tags.client_computed_stats
-            && let Err(e) = stats_generator.send(&payload)
+            && let Err(e) = stats_generator.send(&payload).await
         {
             error!("Stats generator error: {e}");
         }