use otel span to parent tracing span

gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs

@@ -31,8 +31,8 @@ use nym_node_metrics::events::MetricsEvent;
 use nym_sphinx::forwarding::packet::MixPacket;
 use nym_statistics_common::{gateways::GatewaySessionEvent, types::SessionType};
 use nym_validator_client::coconut::EcashApiError;
-use opentelemetry_sdk::propagation::TraceContextPropagator;
-use opentelemetry::propagation::TextMapPropagator;
+use opentelemetry_sdk::{propagation::TraceContextPropagator, trace::IdGenerator};
+use opentelemetry::{propagation::TextMapPropagator, trace::{SpanContext, TraceContextExt, Tracer}, TraceFlags};
 use rand::{random, CryptoRng, Rng};
 use tracing_opentelemetry::OpenTelemetrySpanExt;
 use std::{process, time::Duration};
@@ -603,12 +603,26 @@ impl<R, S> AuthenticatedHandler<R, S> {
             let extracted_trace_id = carrier.extract_trace_id();
             warn!("=== Listen for requests: Extracted trace id from client: {:?} ===", extracted_trace_id);
 
-            let propagator = TraceContextPropagator::new();
-            let extracted_cx = propagator.extract(&carrier);
+            let id_generator = opentelemetry_sdk::trace::RandomIdGenerator::default();
+            let span_id = id_generator.new_span_id();
+            let trace_id = extracted_trace_id.unwrap_or_else(|| {
+                warn!("Failed to extract trace id from the provided otel context carrier");
+                let context = opentelemetry::Context::current();
+                context.span().span_context().trace_id()
+            });
+
+            let span_context = SpanContext::new(
+                trace_id,
+                span_id,
+                TraceFlags::SAMPLED,
+                true,
+                Default::default(),
+            );
+            let cx = opentelemetry::Context::current().with_remote_span_context(span_context);
+            let _context_guard = cx.clone().attach();
+            let span = info_span!("starting point with otel", trace_id = %trace_id);
+            span.set_parent(cx);
 
-            tracing::Span::current().set_parent(extracted_cx);
-            let span = info_span!("starting point with otel");
-            warn!("=== Context propagated to authenticated client handler listen_for_requests ===");
             Some(span)
         } else {
             None
@@ -620,6 +634,37 @@ impl<R, S> AuthenticatedHandler<R, S> {
             info_span!("listen_for_requests without otel")
         };
         let _child_span_guard = child_span.enter();
+
+        // Note: the commented out code below is an alternative way of propagating the context
+        // it does not work because of the way we spawn the task 
+        // i.e. `tokio::spawn(async move { ... })` prevents the context from being propagated
+        // since the `async move` creates a new future that does not capture the current context
+        // so for now we just manually extract the trace id from the client and create a new span with it
+        
+        // if we ever change the way we spawn the task, we can try using the code below instead
+        // let span = if let Some(remote_ctx) = &self.client.otel_context {
+        //     let carrier = ContextCarrier::from_map(remote_ctx.clone());
+        //     let extracted_trace_id = carrier.extract_trace_id().unwrap_or_else(|| {
+        //         warn!("Failed to extract trace id from the provided otel context carrier");
+        //         let context = opentelemetry::Context::current();
+        //         context.span().span_context().trace_id()
+        //     });
+        //     warn!("=== Listen for requests: Extracted trace id from client: {:?} ===", extracted_trace_id);
+
+        //     let tracer = opentelemetry::global::tracer("nym-ws-listener");
+        //     tracer
+        //         .span_builder("websocket listen_for_requests")
+        //         .with_trace_id(extracted_trace_id)
+        //         .start(&tracer)
+        // } else {
+        //     let tracer = opentelemetry::global::tracer("nym-ws-listener");
+        //     tracer
+        //         .span_builder("websocket listen_for_requests")
+        //         .start(&tracer)
+        // };
+        // let context = opentelemetry::Context::current_with_span(span);
+        // let _guard = context.clone().attach();
+        // warn!("=== Context propagated to authenticated client handler listen_for_requests ===");
         
         // Ping timeout future used to check if the client responded to our ping request
         let mut ping_timeout: OptionFuture<_> = None.into();
@@ -645,6 +690,8 @@ impl<R, S> AuthenticatedHandler<R, S> {
                     self.handle_ping_timeout().await;
                 },
                 socket_msg = self.inner.read_websocket_message() => {
+                    let span = info_span!(parent: &child_span, "client_message_received");
+                    let _enter = span.enter();
                     let socket_msg = match socket_msg {
                         None => break,
                         Some(Ok(socket_msg)) => socket_msg,
@@ -668,6 +715,8 @@ impl<R, S> AuthenticatedHandler<R, S> {
                     }
                 },
                 mix_messages = self.mix_receiver.next() => {
+                    let span = info_span!(parent: &child_span, "mix_message_received");
+                    let _enter = span.enter();
                     let mix_messages = match mix_messages {
                         None => {
                             debug!("mix receiver was closed! Assuming the connection is dead.");