Parent based sampling tests (#3136)

This extracts the tests for opentelemetry's parent based sampling that were introduced in #3109 --------- Co-authored-by: bryn <[email protected]> Co-authored-by: Bryn Cooke <[email protected]> Co-authored-by: Jeremy Lempereur <[email protected]>
apollographql · Aug 9, 2023 · 2627143 · 2627143
1 parent dd7bac0
commit 2627143
Show file tree

Hide file tree

Showing 7 changed files with 271 additions and 35 deletions.
diff --git a/.changesets/maint_geal_parent_based_sampling_test.md b/.changesets/maint_geal_parent_based_sampling_test.md
@@ -0,0 +1,5 @@
+### Parent based sampling tests ([PR #3136](https://github.com/apollographql/router/pull/3136))
+
+This adds test for OpenTelemetry sampling defined either in the configuration or in headers carried by the request
+
+By [@Geal](https://github.com/Geal) in https://github.com/apollographql/router/pull/3136
diff --git a/apollo-router/src/tracer.rs b/apollo-router/src/tracer.rs
@@ -4,7 +4,6 @@
 use std::fmt;
 
 use opentelemetry::trace::TraceContextExt;
-use opentelemetry::trace::TraceId as OtelTraceId;
 use serde::Deserialize;
 use serde::Serialize;
 use tracing::Span;
@@ -15,15 +14,16 @@ use tracing_opentelemetry::OpenTelemetrySpanExt;
 pub struct TraceId([u8; 16]);
 
 impl TraceId {
-    /// Create a TraceId. If called from an invalid context
-    /// (e.g.: not in a span, or in a disabled span), then
-    /// None is returned.
+    /// Create a TraceId. If the span is not sampled then return None.
     pub fn maybe_new() -> Option<Self> {
-        let trace_id = Span::current().context().span().span_context().trace_id();
-        if trace_id == OtelTraceId::INVALID {
-            None
+        let span = Span::current();
+        let context = span.context();
+        let span_ref = context.span();
+        let span_context = span_ref.span_context();
+        if span_context.is_sampled() {
+            Some(Self(span_context.trace_id().to_bytes()))
         } else {
-            Some(Self(trace_id.to_bytes()))
+            None
         }
     }
 

diff --git a/apollo-router/tests/common.rs b/apollo-router/tests/common.rs
@@ -11,12 +11,14 @@ use std::time::SystemTime;
 use buildstructor::buildstructor;
 use http::header::ACCEPT;
 use http::header::CONTENT_TYPE;
+use http::HeaderValue;
 use jsonpath_lib::Selector;
 use mime::APPLICATION_JSON;
 use once_cell::sync::OnceCell;
 use opentelemetry::global;
 use opentelemetry::propagation::TextMapPropagator;
 use opentelemetry::trace::Span;
+use opentelemetry::trace::TraceContextExt;
 use opentelemetry::trace::Tracer;
 use opentelemetry::trace::TracerProvider;
 use serde_json::json;
@@ -343,11 +345,11 @@ impl IntegrationTest {
         );
         let default_query = &json!({"query":"query {topProducts{name}}","variables":{}});
         let query = query.unwrap_or(default_query).clone();
-        let id = Uuid::new_v4().to_string();
         let dispatch = self.subscriber.clone();
 
         async move {
-            let span = info_span!("client_request", unit_test = id.as_str());
+            let span = info_span!("client_request");
+            let span_id = span.context().span().span_context().trace_id().to_string();
 
             async move {
                 let client = reqwest::Client::new();
@@ -368,7 +370,7 @@ impl IntegrationTest {
                 });
                 request.headers_mut().remove(ACCEPT);
                 match client.execute(request).await {
-                    Ok(response) => (id, response),
+                    Ok(response) => (span_id, response),
                     Err(err) => {
                         panic!("unable to send successful request to router, {err}")
                     }
@@ -380,6 +382,51 @@ impl IntegrationTest {
         .with_subscriber(dispatch.unwrap_or_default())
     }
 
+    #[allow(dead_code)]
+    pub fn execute_untraced_query(
+        &self,
+        query: &Value,
+    ) -> impl std::future::Future<Output = (String, reqwest::Response)> {
+        assert!(
+            self.router.is_some(),
+            "router was not started, call `router.start().await; router.assert_started().await`"
+        );
+        let query = query.clone();
+        let dispatch = self.subscriber.clone();
+
+        async move {
+            let client = reqwest::Client::new();
+
+            let mut request = client
+                .post("http://localhost:4000")
+                .header(CONTENT_TYPE, APPLICATION_JSON.essence_str())
+                .header("apollographql-client-name", "custom_name")
+                .header("apollographql-client-version", "1.0")
+                .json(&query)
+                .build()
+                .unwrap();
+
+            request.headers_mut().remove(ACCEPT);
+            match client.execute(request).await {
+                Ok(response) => (
+                    response
+                        .headers()
+                        .get("apollo-custom-trace-id")
+                        .cloned()
+                        .unwrap_or(HeaderValue::from_static("no-trace-id"))
+                        .to_str()
+                        .unwrap_or_default()
+                        .to_string(),
+                    response,
+                ),
+                Err(err) => {
+                    panic!("unable to send successful request to router, {err}")
+                }
+            }
+        }
+        .with_subscriber(dispatch.unwrap_or_default())
+    }
+
     #[allow(dead_code)]
     pub async fn run_subscription(&self, subscription: &str) -> (String, reqwest::Response) {
         assert!(

diff --git a/apollo-router/tests/fixtures/jaeger-no-sample.router.yaml b/apollo-router/tests/fixtures/jaeger-no-sample.router.yaml
@@ -0,0 +1,32 @@
+telemetry:
+  apollo:
+    field_level_instrumentation_sampler: always_off
+  tracing:
+
+    experimental_response_trace_id:
+      enabled: true
+      header_name: apollo-custom-trace-id
+    propagation:
+      jaeger: true
+    trace_config:
+      service_name: router
+      sampler: always_off
+    jaeger:
+      batch_processor:
+        scheduled_delay: 100ms
+      agent:
+        endpoint: default
+  experimental_logging:
+    when_header:
+      - name: apollo-router-log-request
+        value: test
+        headers: true # default: false
+        body: true # default: false
+      # log request for all requests coming from Iphones
+      - name: custom-header
+        match: ^foo.*
+        headers: true
+override_subgraph_url:
+  products: http://localhost:4005
+include_subgraph_errors:
+  all: true
diff --git a/apollo-router/tests/fixtures/jaeger.router.yaml b/apollo-router/tests/fixtures/jaeger.router.yaml
@@ -7,6 +7,7 @@ telemetry:
       jaeger: true
     trace_config:
       service_name: router
+      sampler: always_on
     jaeger:
       batch_processor:
         scheduled_delay: 100ms

diff --git a/apollo-router/tests/fixtures/no-telemetry.router.yaml b/apollo-router/tests/fixtures/no-telemetry.router.yaml
@@ -0,0 +1,22 @@
+telemetry:
+  tracing:
+    experimental_response_trace_id:
+      enabled: true
+      header_name: apollo-custom-trace-id
+    trace_config:
+      service_name: router
+      sampler: always_on
+  experimental_logging:
+    when_header:
+      - name: apollo-router-log-request
+        value: test
+        headers: true # default: false
+        body: true # default: false
+      # log request for all requests coming from Iphones
+      - name: custom-header
+        match: ^foo.*
+        headers: true
+override_subgraph_url:
+  products: http://localhost:4005
+include_subgraph_errors:
+  all: true