Improve trace_id test

apollographql · Dec 15, 2024 · 0632ee6 · 0632ee6
1 parent 3b5d256
commit 0632ee6
Show file tree

Hide file tree

Showing 8 changed files with 112 additions and 36 deletions.
diff --git a/apollo-router/tests/integration/telemetry/datadog.rs b/apollo-router/tests/integration/telemetry/datadog.rs
@@ -12,6 +12,7 @@ use crate::integration::common::graph_os_enabled;
 use crate::integration::common::Query;
 use crate::integration::common::Telemetry;
 use crate::integration::telemetry::verifier::Verifier;
+use crate::integration::telemetry::DatadogId;
 use crate::integration::telemetry::TraceSpec;
 use crate::integration::IntegrationTest;
 use crate::integration::ValueExt;
@@ -357,6 +358,69 @@ async fn test_priority_sampling_parent_sampler_very_small() -> Result<(), BoxErr
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn test_priority_sampling_parent_sampler_very_small_no_parent() -> Result<(), BoxError> {
+    // Note that there is a very small chance this test will fail. We are trying to test a non-zero sampler.
+
+    if !graph_os_enabled() {
+        return Ok(());
+    }
+    let mut router = IntegrationTest::builder()
+        .telemetry(Telemetry::Datadog)
+        .config(include_str!(
+            "fixtures/datadog_parent_sampler_very_small_no_parent.router.yaml"
+        ))
+        .build()
+        .await;
+
+    router.start().await;
+    router.assert_started().await;
+
+    // // The router should respect upstream but also almost never sample if left to its own devices.
+    TraceSpec::builder()
+        .services(["client", "router"].into())
+        .priority_sampled("0")
+        .subgraph_sampled(false)
+        .build()
+        .validate_datadog_trace(&mut router, Query::builder().psr("-1").traced(true).build())
+        .await?;
+    TraceSpec::builder()
+        .services(["client", "router"].into())
+        .priority_sampled("0")
+        .subgraph_sampled(false)
+        .build()
+        .validate_datadog_trace(&mut router, Query::builder().psr("0").traced(true).build())
+        .await?;
+
+    TraceSpec::builder()
+        .services(["client", "router"].into())
+        .priority_sampled("0")
+        .subgraph_sampled(false)
+        .build()
+        .validate_datadog_trace(&mut router, Query::builder().psr("1").traced(true).build())
+        .await?;
+
+    TraceSpec::builder()
+        .services(["client", "router"].into())
+        .priority_sampled("0")
+        .subgraph_sampled(false)
+        .build()
+        .validate_datadog_trace(&mut router, Query::builder().psr("2").traced(true).build())
+        .await?;
+
+    TraceSpec::builder()
+        .services(["router"].into())
+        .priority_sampled("0")
+        .subgraph_sampled(false)
+        .build()
+        .validate_datadog_trace(&mut router, Query::builder().psr("2").traced(false).build())
+        .await?;
+
+    router.graceful_shutdown().await;
+
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn test_untraced_request() -> Result<(), BoxError> {
     if !graph_os_enabled() {
@@ -524,17 +588,19 @@ async fn test_header_propagator_override() -> Result<(), BoxError> {
         .build()
         .await;
 
+    let trace_id = opentelemetry::trace::TraceId::from_u128(uuid::Uuid::new_v4().as_u128());
+
     router.start().await;
     router.assert_started().await;
     TraceSpec::builder()
         .services(["router", "subgraph"].into())
         .subgraph_sampled(true)
-        .trace_id("00000000000000000000000000000001")
+        .trace_id(format!("{:032x}", trace_id.to_datadog()))
         .build()
         .validate_datadog_trace(
             &mut router,
             Query::builder()
-                .header("trace-id", "00000000000000000000000000000001")
+                .header("trace-id", trace_id.to_string())
                 .header("x-datadog-trace-id", "2")
                 .traced(false)
                 .build(),
@@ -779,16 +845,6 @@ async fn test_span_metrics() -> Result<(), BoxError> {
     Ok(())
 }
 
-pub(crate) trait DatadogId {
-    fn to_datadog(&self) -> String;
-}
-impl DatadogId for TraceId {
-    fn to_datadog(&self) -> String {
-        let bytes = &self.to_bytes()[std::mem::size_of::<u64>()..std::mem::size_of::<u128>()];
-        u64::from_be_bytes(bytes.try_into().unwrap()).to_string()
-    }
-}
-
 struct DatadogTraceSpec {
     trace_spec: TraceSpec,
 }

diff --git a/apollo-router/tests/integration/telemetry/fixtures/datadog.router.yaml b/apollo-router/tests/integration/telemetry/fixtures/datadog.router.yaml
@@ -5,9 +5,6 @@ telemetry:
         enabled: true
         header_name: apollo-custom-trace-id
         format: datadog
-      propagation:
-        trace_context: true
-        jaeger: true
       common:
         service_name: router
         resource:

diff --git a/apollo-router/tests/integration/telemetry/fixtures/datadog_no_parent_sampler.router.yaml b/apollo-router/tests/integration/telemetry/fixtures/datadog_no_parent_sampler.router.yaml
@@ -5,9 +5,6 @@ telemetry:
         enabled: true
         header_name: apollo-custom-trace-id
         format: datadog
-      propagation:
-        trace_context: true
-        jaeger: true
       common:
         service_name: router
         parent_based_sampler: false

diff --git a/...router/tests/integration/telemetry/fixtures/datadog_parent_sampler_very_small.router.yaml b/...router/tests/integration/telemetry/fixtures/datadog_parent_sampler_very_small.router.yaml
@@ -5,12 +5,9 @@ telemetry:
         enabled: true
         header_name: apollo-custom-trace-id
         format: datadog
-      propagation:
-        trace_context: true
-        jaeger: true
       common:
         service_name: router
-        sampler: 0.000000001
+        sampler: 0.00001
         parent_based_sampler: true
         resource:
           env: local1

diff --git a/...ts/integration/telemetry/fixtures/datadog_parent_sampler_very_small_no_parent.router.yaml b/...ts/integration/telemetry/fixtures/datadog_parent_sampler_very_small_no_parent.router.yaml
@@ -0,0 +1,25 @@
+telemetry:
+  exporters:
+    tracing:
+      experimental_response_trace_id:
+        enabled: true
+        header_name: apollo-custom-trace-id
+        format: datadog
+      common:
+        service_name: router
+        sampler: 0.00001
+        parent_based_sampler: false
+        resource:
+          env: local1
+          service.version: router_version_override
+        preview_datadog_agent_sampling: true
+      datadog:
+        enabled: true
+        batch_processor:
+          scheduled_delay: 100ms
+  instrumentation:
+    spans:
+      mode: spec_compliant
+      supergraph:
+        attributes:
+          graphql.operation.name: true
diff --git a/apollo-router/tests/integration/telemetry/mod.rs b/apollo-router/tests/integration/telemetry/mod.rs
@@ -1,6 +1,8 @@
 use std::collections::HashMap;
 use std::collections::HashSet;
 
+use opentelemetry_api::trace::TraceId;
+
 #[cfg(any(not(feature = "ci"), all(target_arch = "x86_64", target_os = "linux")))]
 mod datadog;
 #[cfg(any(not(feature = "ci"), all(target_arch = "x86_64", target_os = "linux")))]
@@ -22,7 +24,7 @@ struct TraceSpec {
     unmeasured_spans: HashSet<&'static str>,
     priority_sampled: Option<&'static str>,
     subgraph_sampled: Option<bool>,
-    trace_id: Option<&'static str>,
+    trace_id: Option<String>,
     span_attributes: HashMap<&'static str, Vec<(&'static str, &'static str)>>,
 }
 
@@ -39,7 +41,7 @@ impl TraceSpec {
         unmeasured_spans: HashSet<&'static str>,
         priority_sampled: Option<&'static str>,
         subgraph_sampled: Option<bool>,
-        trace_id: Option<&'static str>,
+        trace_id: Option<String>,
         span_attributes: HashMap<&'static str, Vec<(&'static str, &'static str)>>,
     ) -> Self {
         Self {
@@ -56,3 +58,14 @@ impl TraceSpec {
         }
     }
 }
+
+#[allow(dead_code)]
+pub trait DatadogId {
+    fn to_datadog(&self) -> u64;
+}
+impl DatadogId for TraceId {
+    fn to_datadog(&self) -> u64 {
+        let bytes = &self.to_bytes()[std::mem::size_of::<u64>()..std::mem::size_of::<u128>()];
+        u64::from_be_bytes(bytes.try_into().unwrap())
+    }
+}
diff --git a/apollo-router/tests/integration/telemetry/otlp.rs b/apollo-router/tests/integration/telemetry/otlp.rs
@@ -20,6 +20,7 @@ use crate::integration::common::graph_os_enabled;
 use crate::integration::common::Query;
 use crate::integration::common::Telemetry;
 use crate::integration::telemetry::verifier::Verifier;
+use crate::integration::telemetry::DatadogId;
 use crate::integration::telemetry::TraceSpec;
 use crate::integration::IntegrationTest;
 use crate::integration::ValueExt;
@@ -769,16 +770,6 @@ async fn mock_otlp_server() -> MockServer {
     mock_server
 }
 
-pub(crate) trait DatadogId {
-    fn to_datadog(&self) -> u64;
-}
-impl DatadogId for TraceId {
-    fn to_datadog(&self) -> u64 {
-        let bytes = &self.to_bytes()[std::mem::size_of::<u64>()..std::mem::size_of::<u128>()];
-        u64::from_be_bytes(bytes.try_into().unwrap())
-    }
-}
-
 impl TraceSpec {
     async fn validate_otlp_trace(
         self,

diff --git a/apollo-router/tests/integration/telemetry/verifier.rs b/apollo-router/tests/integration/telemetry/verifier.rs
@@ -18,8 +18,8 @@ pub trait Verifier {
         query: Query,
     ) -> Result<(), BoxError> {
         let (id, response) = router.execute_query(query).await;
-        if let Some(spec_id) = self.spec().trace_id {
-            assert_eq!(id.to_string(), spec_id, "trace id");
+        if let Some(spec_id) = &self.spec().trace_id {
+            assert_eq!(id.to_string(), *spec_id, "trace id");
         }
         for _ in 0..20 {
             if self.find_valid_trace(id).await.is_ok() {