feat(data-pipeline): expose API to send spans without going through deserialization (#992)

# What does this PR do?

This PR adds the following features:
* A public API to send trace chunks through the trace exporter. This function takes a `Vec<Vec<Span<T>>` so we can rapidly iterate in ddtrace-rs on which datastructure we want to use to pass spans
* A public API to wait for the /info endpoint to be ready. This is needed for deterministic tests with the test agent, because otherwise it's possible to not have stats during the first submission
* A way to specify env vars when instantiating an apm test agent. We need to skip some meta keys in dd-trace-rs using the `SNAPSHOT_IGNORED_ATTRS` env, and to recreate the snapshots when running outside of the CI with the `SNAPSHOT_CI` env parameters

# Motivation

Working on integrating the trace-exporter in dd-trace-rs

Author: paullegranddc

data-pipeline/Cargo.toml

@@ -61,3 +61,6 @@ tokio = { version = "1.23", features = [
     "time",
     "test-util",
 ], default-features = false }
+
+[features]
+test-utils = []

data-pipeline/src/trace_exporter/mod.rs

@@ -13,7 +13,7 @@ use arc_swap::{ArcSwap, ArcSwapOption};
 use bytes::Bytes;
 use datadog_trace_utils::msgpack_decoder::{self, decode::error::DecodeError};
 use datadog_trace_utils::send_with_retry::{send_with_retry, RetryStrategy, SendWithRetryError};
-use datadog_trace_utils::span::SpanSlice;
+use datadog_trace_utils::span::{Span, SpanText};
 use datadog_trace_utils::trace_utils::{self, TracerHeaderTags};
 use datadog_trace_utils::tracer_payload;
 use ddcommon::header::{
@@ -199,51 +199,47 @@ pub struct TraceExporter {
     telemetry: Option<TelemetryClient>,
 }
 
+enum DeserInputFormat {
+    V04,
+    V05,
+}
+
 impl TraceExporter {
     #[allow(missing_docs)]
     pub fn builder() -> TraceExporterBuilder {
         TraceExporterBuilder::default()
     }
 
     /// Send msgpack serialized traces to the agent
-    #[allow(missing_docs)]
+    ///
+    /// # Arguments
+    ///
+    /// * data: A slice containing the serialized traces. This slice should be encoded following the
+    ///   input_format passed to the TraceExporter on creating.
+    /// * trace_count: The number of traces in the data
+    ///
+    /// # Returns
+    /// * Ok(AgentResponse): The response from the agent
+    /// * Err(TraceExporterError): An error detailling what went wrong in the process
     pub fn send(
         &self,
         data: &[u8],
         trace_count: usize,
     ) -> Result<AgentResponse, TraceExporterError> {
         self.check_agent_info();
 
-        match self.input_format {
+        let res = match self.input_format {
             TraceExporterInputFormat::Proxy => self.send_proxy(data.as_ref(), trace_count),
-            TraceExporterInputFormat::V04 => match msgpack_decoder::v04::from_slice(data) {
-                Ok((traces, _)) => self.send_deser_ser(traces),
-                Err(e) => Err(TraceExporterError::Deserialization(e)),
-            },
-            TraceExporterInputFormat::V05 => match msgpack_decoder::v05::from_slice(data) {
-                Ok((traces, _)) => self.send_deser_ser(traces),
-                Err(e) => Err(TraceExporterError::Deserialization(e)),
-            },
+            TraceExporterInputFormat::V04 => self.send_deser(data, DeserInputFormat::V04),
+            TraceExporterInputFormat::V05 => self.send_deser(data, DeserInputFormat::V05),
+        }?;
+        if res.is_empty() {
+            return Err(TraceExporterError::Agent(
+                error::AgentErrorKind::EmptyResponse,
+            ));
         }
-        .and_then(|res| {
-            if res.is_empty() {
-                return Err(TraceExporterError::Agent(
-                    error::AgentErrorKind::EmptyResponse,
-                ));
-            }
 
-            Ok(AgentResponse::from(res))
-        })
-        .map_err(|err| {
-            if let TraceExporterError::Deserialization(ref e) = err {
-                error!("Error deserializing trace from request body: {e}");
-                self.emit_metric(
-                    HealthMetric::Count(health_metrics::STAT_DESER_TRACES_ERRORS, 1),
-                    None,
-                );
-            }
-            err
-        })
+        Ok(AgentResponse::from(res))
     }
 
     /// Safely shutdown the TraceExporter and all related tasks
@@ -425,6 +421,35 @@ impl TraceExporter {
         }
     }
 
+    /// !!! This function is only for testing purposes !!!
+    ///
+    /// Waits the agent info to be ready by checking the agent_info state.
+    /// It will only return Ok after the agent info has been fetched at least once or Err if timeout
+    /// has been reached
+    ///
+    /// In production:
+    /// 1) We should not synchronously wait for this to be ready before sending traces
+    /// 2) It's not guaranteed to not block forever, since the /info endpoint might not be
+    ///    available.
+    ///
+    /// The `send`` function will check agent_info when running, which will only be available if the
+    /// fetcher had time to reach to the agent.
+    /// Since agent_info can enable CSS computation, waiting for this during testing can make
+    /// snapshots non-determinitic.
+    #[cfg(feature = "test-utils")]
+    pub fn wait_agent_info_ready(&self, timeout: Duration) -> anyhow::Result<()> {
+        let start = std::time::Instant::now();
+        loop {
+            if std::time::Instant::now().duration_since(start) > timeout {
+                anyhow::bail!("Timeout waiting for agent info to be ready",);
+            }
+            if self.agent_info.load().is_some() {
+                return Ok(());
+            }
+            std::thread::sleep(Duration::from_millis(10));
+        }
+    }
+
     fn send_proxy(&self, data: &[u8], trace_count: usize) -> Result<String, TraceExporterError> {
         self.send_data_to_url(
             data,
@@ -546,42 +571,79 @@ impl TraceExporter {
     /// Add all spans from the given iterator into the stats concentrator
     /// # Panic
     /// Will panic if another thread panicked will holding the lock on `stats_concentrator`
-    fn add_spans_to_stats(&self, traces: &[Vec<SpanSlice>]) {
-        if let StatsComputationStatus::Enabled {
-            stats_concentrator,
-            cancellation_token: _,
-            exporter_handle: _,
-        } = &**self.client_side_stats.load()
-        {
-            #[allow(clippy::unwrap_used)]
-            let mut stats_concentrator = stats_concentrator.lock().unwrap();
-
-            let spans = traces.iter().flat_map(|trace| trace.iter());
-            for span in spans {
-                stats_concentrator.add_span(span);
-            }
+    fn add_spans_to_stats<T: SpanText>(
+        &self,
+        stats_concentrator: &Mutex<SpanConcentrator>,
+        traces: &[Vec<Span<T>>],
+    ) {
+        #[allow(clippy::unwrap_used)]
+        let mut stats_concentrator = stats_concentrator.lock().unwrap();
+
+        let spans = traces.iter().flat_map(|trace| trace.iter());
+        for span in spans {
+            stats_concentrator.add_span(span);
         }
     }
 
-    fn send_deser_ser(
+    /// Send a list of trace chunks to the agent
+    ///
+    /// # Arguments
+    /// * trace_chunks: A list of trace chunks. Each trace chunk is a list of spans.
+    ///
+    /// # Returns
+    /// * Ok(String): The response from the agent
+    /// * Err(TraceExporterError): An error detailing what went wrong in the process
+    pub fn send_trace_chunks<T: SpanText>(
+        &self,
+        trace_chunks: Vec<Vec<Span<T>>>,
+    ) -> Result<String, TraceExporterError> {
+        self.check_agent_info();
+        self.send_trace_chunks_inner(trace_chunks)
+    }
+
+    /// Deserializes, processes and sends trace chunks to the agent
+    fn send_deser(
         &self,
-        mut traces: Vec<Vec<SpanSlice>>,
+        data: &[u8],
+        format: DeserInputFormat,
     ) -> Result<String, TraceExporterError> {
+        let (traces, _) = match format {
+            DeserInputFormat::V04 => msgpack_decoder::v04::from_slice(data),
+            DeserInputFormat::V05 => msgpack_decoder::v05::from_slice(data),
+        }
+        .map_err(|e| {
+            error!("Error deserializing trace from request body: {e}");
+            self.emit_metric(
+                HealthMetric::Count(health_metrics::STAT_DESER_TRACES_ERRORS, 1),
+                None,
+            );
+            TraceExporterError::Deserialization(e)
+        })?;
         self.emit_metric(
             HealthMetric::Count(health_metrics::STAT_DESER_TRACES, traces.len() as i64),
             None,
         );
 
+        self.send_trace_chunks_inner(traces)
+    }
+
+    fn send_trace_chunks_inner<T: SpanText>(
+        &self,
+        mut traces: Vec<Vec<Span<T>>>,
+    ) -> Result<String, TraceExporterError> {
         let mut header_tags: TracerHeaderTags = self.metadata.borrow().into();
 
         // Stats computation
-        if let StatsComputationStatus::Enabled { .. } = &**self.client_side_stats.load() {
+        if let StatsComputationStatus::Enabled {
+            stats_concentrator, ..
+        } = &**self.client_side_stats.load()
+        {
             if !self.client_computed_top_level {
                 for chunk in traces.iter_mut() {
                     datadog_trace_utils::span::trace_utils::compute_top_level_span(chunk);
                 }
             }
-            self.add_spans_to_stats(&traces);
+            self.add_spans_to_stats(stats_concentrator, &traces);
             // Once stats have been computed we can drop all chunks that are not going to be
             // sampled by the agent
             let datadog_trace_utils::span::trace_utils::DroppedP0Stats {
@@ -597,18 +659,9 @@ impl TraceExporter {
             header_tags.dropped_p0_spans = dropped_p0_spans;
         }
 
-        let use_v05_format = match (self.input_format, self.output_format) {
-            (TraceExporterInputFormat::V04, TraceExporterOutputFormat::V04) => false,
-            (TraceExporterInputFormat::V04, TraceExporterOutputFormat::V05)
-            | (TraceExporterInputFormat::V05, TraceExporterOutputFormat::V05) => true,
-            (TraceExporterInputFormat::V05, TraceExporterOutputFormat::V04) => {
-                // TODO: Properly handle non-OK states to prevent possible panics (APMSP-18190).
-                unreachable!("Conversion from v05 to v04 not implemented")
-            }
-            (TraceExporterInputFormat::Proxy, _) => {
-                // TODO: Properly handle non-OK states to prevent possible panics (APMSP-18190).
-                unreachable!("Codepath invalid for proxy mode",)
-            }
+        let use_v05_format = match self.output_format {
+            TraceExporterOutputFormat::V05 => true,
+            TraceExporterOutputFormat::V04 => false,
         };
         let payload = trace_utils::collect_trace_chunks(traces, use_v05_format).map_err(|e| {
             TraceExporterError::Deserialization(DecodeError::InvalidFormat(e.to_string()))

data-pipeline/tests/test_fetch_info.rs

@@ -12,7 +12,7 @@ mod tracing_integration_tests {
     #[cfg_attr(miri, ignore)]
     #[tokio::test]
     async fn test_fetch_info_from_test_agent() {
-        let test_agent = DatadogTestAgent::new(None, None).await;
+        let test_agent = DatadogTestAgent::new(None, None, &[]).await;
         let endpoint = Endpoint::from_url(test_agent.get_uri_for_endpoint("info", None).await);
         let info = fetch_info(&endpoint)
             .await
@@ -28,7 +28,7 @@ mod tracing_integration_tests {
     #[cfg_attr(miri, ignore)]
     #[tokio::test]
     async fn test_agent_info_fetcher_with_test_agent() {
-        let test_agent = DatadogTestAgent::new(None, None).await;
+        let test_agent = DatadogTestAgent::new(None, None, &[]).await;
         let endpoint = Endpoint::from_url(test_agent.get_uri_for_endpoint("info", None).await);
         let fetcher = AgentInfoFetcher::new(endpoint, Duration::from_secs(1));
         let info_arc = fetcher.get_info();

data-pipeline/tests/test_trace_exporter.rs

@@ -9,10 +9,6 @@ mod tracing_integration_tests {
     use datadog_trace_utils::test_utils::datadog_test_agent::DatadogTestAgent;
     use datadog_trace_utils::test_utils::{create_test_json_span, create_test_v05_span};
     use serde_json::json;
-    #[cfg(target_os = "linux")]
-    use std::fs::Permissions;
-    #[cfg(target_os = "linux")]
-    use std::os::unix::fs::PermissionsExt;
     use tokio::task;
 
     fn get_v04_trace_snapshot_test_payload(name_prefix: &str) -> Vec<u8> {
@@ -105,7 +101,7 @@ mod tracing_integration_tests {
     async fn compare_v04_trace_snapshot_test() {
         let relative_snapshot_path = "data-pipeline/tests/snapshots/";
         let snapshot_name = "compare_exporter_v04_trace_snapshot_test";
-        let test_agent = DatadogTestAgent::new(Some(relative_snapshot_path), None).await;
+        let test_agent = DatadogTestAgent::new(Some(relative_snapshot_path), None, &[]).await;
         let url = test_agent.get_base_uri().await;
         let rate_param = "{\"service:test,env:test_env\": 0.5, \"service:test2,env:prod\": 0.2}";
         test_agent
@@ -154,7 +150,7 @@ mod tracing_integration_tests {
     async fn compare_v04_to_v05_trace_snapshot_test() {
         let relative_snapshot_path = "data-pipeline/tests/snapshots/";
         let snapshot_name = "compare_exporter_v04_to_v05_trace_snapshot_test";
-        let test_agent = DatadogTestAgent::new(Some(relative_snapshot_path), None).await;
+        let test_agent = DatadogTestAgent::new(Some(relative_snapshot_path), None, &[]).await;
         let url = test_agent.get_base_uri().await;
         let rate_param = "{\"service:test,env:test_env\": 0.5, \"service:test2,env:prod\": 0.2}";
         test_agent
@@ -197,7 +193,7 @@ mod tracing_integration_tests {
     async fn compare_v05_trace_snapshot_test() {
         let relative_snapshot_path = "data-pipeline/tests/snapshots/";
         let snapshot_name = "compare_exporter_v05_trace_snapshot_test";
-        let test_agent = DatadogTestAgent::new(Some(relative_snapshot_path), None).await;
+        let test_agent = DatadogTestAgent::new(Some(relative_snapshot_path), None, &[]).await;
         let url = test_agent.get_base_uri().await;
         let rate_param = "{\"service:test,env:test_env\": 0.5, \"service:test2,env:prod\": 0.2}";
         test_agent
@@ -243,6 +239,9 @@ mod tracing_integration_tests {
     // assign unique names to the spans and instantiate a unique session for the test to avoid flaky
     // behavior when running on CI
     async fn uds_snapshot_test() {
+        use std::fs::Permissions;
+        use std::os::unix::fs::PermissionsExt;
+
         let relative_snapshot_path = "data-pipeline/tests/snapshots/";
         let snapshot_name = "compare_exporter_v04_trace_snapshot_uds_test";
         // Create a temporary directory for the socket to be mounted in the test agent container
@@ -266,6 +265,7 @@ mod tracing_integration_tests {
         let test_agent = DatadogTestAgent::new(
             Some(relative_snapshot_path),
             Some(&absolute_socket_dir_path),
+            &[],
         )
         .await;
 

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

@@ -63,9 +63,21 @@ impl FromStr for SpanKey {
 /// Trait representing the requirements for a type to be used as a Span "string" type.
 /// Note: Borrow<str> is not required by the derived traits, but allows to access HashMap elements
 /// from a static str and check if the string is empty.
-pub trait SpanText: Eq + Hash + Borrow<str> + Serialize + Default + Clone {}
-/// Implement the SpanText trait for any type which satisfies the sub traits.
-impl<T: Eq + Hash + Borrow<str> + Serialize + Default + Clone> SpanText for T {}
+pub trait SpanText: Eq + Hash + Borrow<str> + Serialize + Default + Clone {
+    fn from_static_str(value: &'static str) -> Self;
+}
+
+impl SpanText for &str {
+    fn from_static_str(value: &'static str) -> Self {
+        value
+    }
+}
+
+impl SpanText for BytesString {
+    fn from_static_str(value: &'static str) -> Self {
+        BytesString::from_static(value)
+    }
+}
 
 /// Checks if the `value` represents an empty string. Used to skip serializing empty strings
 /// with serde.

datadog-trace-utils/src/span/trace_utils.rs

@@ -13,12 +13,12 @@ const TRACER_TOP_LEVEL_KEY: &str = "_dd.top_level";
 const MEASURED_KEY: &str = "_dd.measured";
 const PARTIAL_VERSION_KEY: &str = "_dd.partial_version";
 
-fn set_top_level_span<'a, T>(span: &mut Span<T>, is_top_level: bool)
+fn set_top_level_span<T>(span: &mut Span<T>, is_top_level: bool)
 where
-    T: SpanText + From<&'a str>,
+    T: SpanText,
 {
     if is_top_level {
-        span.metrics.insert(TOP_LEVEL_KEY.into(), 1.0);
+        span.metrics.insert(T::from_static_str(TOP_LEVEL_KEY), 1.0);
     } else {
         span.metrics.remove(TOP_LEVEL_KEY);
     }
@@ -30,9 +30,9 @@ where
 ///   - OR its parent is unknown (other part of the code, distributed trace)
 ///   - OR its parent belongs to another service (in that case it's a "local root" being the highest
 ///     ancestor of other spans belonging to this service and attached to it).
-pub fn compute_top_level_span<'a, T>(trace: &mut [Span<T>])
+pub fn compute_top_level_span<T>(trace: &mut [Span<T>])
 where
-    T: SpanText + Clone + From<&'a str>,
+    T: SpanText,
 {
     let mut span_id_to_service: HashMap<u64, T> = HashMap::new();
     for span in trace.iter() {
@@ -102,7 +102,7 @@ const SAMPLING_ANALYTICS_RATE_KEY: &str = "_dd1.sr.eausr";
 /// dropped and the latter to the spans dropped.
 pub fn drop_chunks<T>(traces: &mut Vec<Vec<Span<T>>>) -> DroppedP0Stats
 where
-    T: SpanText + Default,
+    T: SpanText,
 {
     let mut dropped_p0_traces = 0;
     let mut dropped_p0_spans = 0;