From daa6ca6609ee5c2aff05d5d1837a39ab1fa56b33 Mon Sep 17 00:00:00 2001
From: Paul Scoropan <1paulscoropan@gmail.com>
Date: Fri, 15 Nov 2024 12:35:34 -0500
Subject: [PATCH] more tracing logic fixes - traces are being split due to
 async instrumentation logic

Signed-off-by: Paul Scoropan <1paulscoropan@gmail.com>
---
 src/clients.rs      | 6 ++++--
 src/clients/http.rs | 7 +++++--
 src/utils/trace.rs  | 9 ++++-----
 3 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/src/clients.rs b/src/clients.rs
index 3cdd5888..a2c14080 100644
--- a/src/clients.rs
+++ b/src/clients.rs
@@ -29,7 +29,8 @@ use futures::Stream;
 use ginepro::LoadBalancedChannel;
 use hyper_util::rt::TokioExecutor;
 use tonic::{metadata::MetadataMap, Request};
-use tracing::{debug, instrument};
+use tracing::{debug, instrument, Span};
+use tracing_opentelemetry::OpenTelemetrySpanExt;
 use url::Url;
 
 use crate::{
@@ -335,7 +336,8 @@ pub fn is_valid_hostname(hostname: &str) -> bool {
 /// Turns a gRPC client request body of type `T` and header map into a `tonic::Request<T>`.
 /// Will also inject the current `traceparent` header into the request based on the current span.
 fn grpc_request_with_headers<T>(request: T, headers: HeaderMap) -> Request<T> {
-    let headers = with_traceparent_header(headers);
+    let ctx = Span::current().context();
+    let headers = with_traceparent_header(&ctx, headers);
     let metadata = MetadataMap::from_headers(headers);
     Request::from_parts(metadata, Extensions::new(), request)
 }
diff --git a/src/clients/http.rs b/src/clients/http.rs
index d18f4fd1..70a25592 100644
--- a/src/clients/http.rs
+++ b/src/clients/http.rs
@@ -26,6 +26,7 @@ use hyper_rustls::HttpsConnector;
 use hyper_util::client::legacy::connect::HttpConnector;
 use serde::{de::DeserializeOwned, Serialize};
 use tracing::{debug, error, instrument, Instrument, Span};
+use tracing_opentelemetry::OpenTelemetrySpanExt;
 use url::Url;
 
 use super::{Client, Error};
@@ -130,7 +131,8 @@ impl HttpClient {
         headers: HeaderMap,
         body: impl Serialize,
     ) -> Result<Response, Error> {
-        let headers = trace::with_traceparent_header(headers.to_owned());
+        let ctx = Span::current().context();
+        let headers = trace::with_traceparent_header(&ctx, headers.to_owned());
         let mut builder = hyper::http::request::Builder::new()
             .method(method)
             .uri(url.as_uri());
@@ -160,7 +162,8 @@ impl HttpClient {
                     .await
                     .map_err(|e| Error::internal("sending client request failed", e))?
                     .into();
-                trace_context_from_http_response(&response);
+                let span = Span::current();
+                trace_context_from_http_response(&span, &response);
                 Ok(response)
             }
             None => Err(builder.body(body).err().map_or_else(
diff --git a/src/utils/trace.rs b/src/utils/trace.rs
index 8629cb35..415283dc 100644
--- a/src/utils/trace.rs
+++ b/src/utils/trace.rs
@@ -340,12 +340,11 @@ pub fn on_outgoing_eos(trailers: Option<&HeaderMap>, stream_duration: Duration,
 /// vendor-specific trace context.
 /// Used by both gRPC and HTTP requests since `tonic::Metadata` uses `http::HeaderMap`.
 /// See https://www.w3.org/TR/trace-context/#trace-context-http-headers-format.
-pub fn with_traceparent_header(headers: HeaderMap) -> HeaderMap {
+pub fn with_traceparent_header(ctx: &opentelemetry::Context, headers: HeaderMap) -> HeaderMap {
     let mut headers = headers.clone();
-    let ctx = Span::current().context();
     global::get_text_map_propagator(|propagator| {
         // Injects current `traceparent` (and by default empty `tracestate`)
-        propagator.inject_context(&ctx, &mut HeaderInjector(&mut headers))
+        propagator.inject_context(ctx, &mut HeaderInjector(&mut headers))
     });
     headers
 }
@@ -354,12 +353,12 @@ pub fn with_traceparent_header(headers: HeaderMap) -> HeaderMap {
 /// tracing span context (i.e. use `traceparent` as parent to the current span).
 /// Defaults to using the current context when no `traceparent` is found.
 /// See https://www.w3.org/TR/trace-context/#trace-context-http-headers-format.
-pub fn trace_context_from_http_response(response: &http::Response) {
+pub fn trace_context_from_http_response(span: &Span, response: &http::Response) {
     let ctx = global::get_text_map_propagator(|propagator| {
         // Returns the current context if no `traceparent` is found
         propagator.extract(&HeaderExtractor(response.headers()))
     });
-    Span::current().set_parent(ctx);
+    span.set_parent(ctx);
 }
 
 /// Extracts the `traceparent` header from a gRPC response's metadata and uses it to set the current