diff --git a/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs b/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs index 22e9d340e3..49d2644417 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs +++ b/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 AuthenticatedHandler { 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 AuthenticatedHandler { 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 AuthenticatedHandler { 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 AuthenticatedHandler { } }, 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.");