Use TraceParent::to_remote_context

prisma · Oct 7, 2024 · f90f448 · f90f448
1 parent dffe32b
commit f90f448
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 10 deletions.
diff --git a/libs/telemetry/src/helpers.rs b/libs/telemetry/src/helpers.rs
@@ -68,12 +68,12 @@ impl TraceParent {
 
     /// Returns a remote `opentelemetry::Context`. By "remote" we mean that it wasn't emitted in the
     /// current process.
-    pub fn as_remote_context(&self) -> opentelemetry::Context {
+    pub fn to_remote_context(&self) -> opentelemetry::Context {
         // This relies on the fact that global text map propagator was installed that
         // can handle `traceparent` field (for example, `TraceContextPropagator`).
-        let mut extractor = HashMap::new();
-        extractor.insert("traceparent".to_string(), self.to_string());
-        opentelemetry::global::get_text_map_propagator(|propagator| propagator.extract(&extractor))
+        opentelemetry::global::get_text_map_propagator(|propagator| {
+            propagator.extract(&TraceParentExtractor::new(self))
+        })
     }
 }
 
@@ -82,7 +82,7 @@ impl TraceParent {
 pub struct TraceParentExtractor(String);
 
 impl TraceParentExtractor {
-    pub fn new(traceparent: TraceParent) -> Self {
+    pub fn new(traceparent: &TraceParent) -> Self {
         Self(traceparent.to_string())
     }
 }

diff --git a/query-engine/query-engine/src/server/mod.rs b/query-engine/query-engine/src/server/mod.rs
@@ -14,7 +14,7 @@ use std::net::SocketAddr;
 use std::sync::Arc;
 use std::time::Instant;
 use telemetry::capturing::Capturer;
-use telemetry::helpers::{TraceParent, TraceParentExtractor};
+use telemetry::helpers::TraceParent;
 use tracing::{Instrument, Span};
 use tracing_opentelemetry::OpenTelemetrySpanExt;
 
@@ -426,10 +426,7 @@ async fn setup_telemetry(span: Span, headers: &HeaderMap) -> (Span, Option<Trace
             #[allow(deprecated)]
             TraceParent::new_random()
         };
-        let context = opentelemetry::global::get_text_map_propagator(|propagator| {
-            propagator.extract(&TraceParentExtractor::new(traceparent))
-        });
-        span.set_parent(context);
+        span.set_parent(traceparent.to_remote_context());
         Some(traceparent)
     } else {
         None