websocket otel trace handling correction

common/bin-common/src/opentelemetry/context.rs

@@ -3,6 +3,7 @@ use opentelemetry::propagation::{Injector, Extractor, TextMapPropagator};
 use opentelemetry::trace::{SpanContext, TraceContextExt, TraceId};
 use opentelemetry_sdk::propagation::TraceContextPropagator;
 use opentelemetry_sdk::trace::IdGenerator;
+use tracing_opentelemetry::OpenTelemetrySpanExt;
 use std::collections::HashMap;
 
 /// Make a Carrier for context propagation
@@ -63,44 +64,31 @@ impl Extractor for ContextCarrier {
     }
 }
 
-pub struct ManualSpanContextExt {
-    pub context_carrier: ContextCarrier,
+pub struct ManualContextPropagator {
     pub root_span: tracing::Span,
-    pub trace_id: Option<TraceId>,
-    pub is_valid: bool,
+    pub trace_id: TraceId,
 }
 
-impl ManualSpanContextExt {
-    pub fn new() -> Self {
-        ManualSpanContextExt {
-            context_carrier: ContextCarrier::new_empty(),
-            root_span: tracing::Span::current(),
-            trace_id: None,
-            is_valid: false,
+impl ManualContextPropagator {
+    pub fn new(name: &str, context: HashMap<String, String>) -> Self {
+        let carrier = ContextCarrier::new_with_data(context);
+        let trace_id = match carrier.extract_trace_id() {
+            Some(id) => id,
+            None => Context::current().span().span_context().trace_id(),
+        };
+
+        let root_span_builder = new_span_context_with_id(trace_id.clone());
+        let _guard = root_span_builder.clone().attach();
+
+        let root_span = tracing::info_span!("trace_root", name = %name, trace_id = %trace_id);
+        root_span.set_parent(root_span_builder);
+
+        ManualContextPropagator {
+            root_span,
+            trace_id,
         }
     }
-
-    pub fn with_context_carrier(mut self, carrier: &ContextCarrier) -> Self {
-        self.context_carrier = ContextCarrier::new_with_data(carrier.data.clone());
-        self.trace_id = self.context_carrier.extract_trace_id();
-        self
-    }
-
-    pub fn with_extracted_context(mut self, external_context: HashMap<String, String> ) -> Self {
-        self.context_carrier = ContextCarrier::new_with_data(external_context);
-        self.trace_id = self.context_carrier.extract_trace_id();
-        self
-    }
-
-    pub fn set_root_span(mut self, span: tracing::Span) -> Self {
-        self.root_span = span;
-        self
-    }
-
-    pub fn is_valid(&self) -> bool {
-        self.is_valid
-    }
- }
+}
 
 pub fn new_span_context_with_id(trace_id: TraceId) -> Context {
     let id_gen = opentelemetry_sdk::trace::RandomIdGenerator::default();

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

@@ -14,12 +14,11 @@ use futures::{
     future::{FusedFuture, OptionFuture},
     FutureExt, StreamExt,
 };
-use nym_bin_common::opentelemetry::context::{new_span_context_with_id, ManualSpanContextExt};
+use nym_bin_common::opentelemetry::context::ManualContextPropagator;
 use nym_credential_verification::CredentialVerifier;
 use nym_credential_verification::{
     bandwidth_storage_manager::BandwidthStorageManager, ClientBandwidth,
 };
-use nym_crypto::shared_key::new_ephemeral_shared_key;
 use nym_gateway_requests::{
     types::{BinaryRequest, ServerResponse},
     ClientControlRequest, ClientRequest, GatewayRequestsError, SensitiveServerResponse,
@@ -33,10 +32,9 @@ use nym_sphinx::forwarding::packet::MixPacket;
 use nym_statistics_common::{gateways::GatewaySessionEvent, types::SessionType};
 use nym_validator_client::coconut::EcashApiError;
 use rand::{random, CryptoRng, Rng};
-use tracing_opentelemetry::OpenTelemetrySpanExt;
-use std::{process, time::Duration};
+use std::{collections::HashMap, process, time::Duration};
 use thiserror::Error;
-use tokio::{io::{AsyncRead, AsyncWrite}, time::{error::Elapsed, interval}};
+use tokio::io::{AsyncRead, AsyncWrite};
 use tokio_tungstenite::tungstenite::{protocol::Message, Error as WsError};
 use tracing::*;
 
@@ -150,7 +148,7 @@ pub(crate) struct AuthenticatedHandler<R, S> {
     // senders that are used to return the result of the ping to the handler requesting the ping.
     is_active_request_receiver: IsActiveRequestReceiver,
     is_active_ping_pending_reply: Option<(u64, IsActiveResultSender)>,
-    root_span: Option<tracing::Span>,
+    otel_propagator: Option<ManualContextPropagator>,
 }
 
 // explicitly remove handle from the global store upon being dropped
@@ -193,6 +191,17 @@ impl<R, S> AuthenticatedHandler<R, S> {
                 client_address: client.address.as_base58_string(),
             })?;
 
+        let context = match client.otel_context {
+            Some(ref ctx) => ctx.clone(),
+            None => HashMap::new(),
+        };
+
+        let manual_ctx_propagator = if !context.is_empty() {
+            Some(ManualContextPropagator::new("upgrading_fresh_to_authenticated", context))
+        } else {
+            None
+        };
+
         let handler = AuthenticatedHandler {
             bandwidth_storage_manager: BandwidthStorageManager::new(
                 Box::new(fresh.shared_state.storage.clone()),
@@ -206,7 +215,7 @@ impl<R, S> AuthenticatedHandler<R, S> {
             mix_receiver,
             is_active_request_receiver,
             is_active_ping_pending_reply: None,
-            root_span: None,
+            otel_propagator: manual_ctx_propagator,
         };
         handler.send_metrics(GatewaySessionEvent::new_session_start(
             handler.client.address,
@@ -227,10 +236,6 @@ impl<R, S> AuthenticatedHandler<R, S> {
         self.inner.send_metrics(event)
     }
 
-    pub fn get_trace_id(&self) -> Option<opentelemetry::trace::TraceId> {
-        self.client.get_extracted_trace_id()
-    }
-
     /// Forwards the received mix packet from the client into the mix network.
     ///
     /// # Arguments
@@ -238,8 +243,10 @@ impl<R, S> AuthenticatedHandler<R, S> {
     /// * `mix_packet`: packet received from the client that should get forwarded into the network.
     #[instrument(skip_all)]
     fn forward_packet(&self, mix_packet: MixPacket) {
-        let herited_span = self.root_span.as_ref().cloned().unwrap_or_else(|| tracing::Span::none());
-        let span = info_span!(parent: &herited_span, "forwarding_packet");
+        let span = match &self.otel_propagator {
+            Some(propagator) => info_span!(parent: &propagator.root_span, "forwarding_mix_packet"),
+            None => info_span!("forwarding_mix_packet_no_otel"),
+        };
         let _enter = span.enter();
 
         if let Err(err) = self
@@ -301,8 +308,10 @@ impl<R, S> AuthenticatedHandler<R, S> {
         &mut self,
         mix_packet: MixPacket,
     ) -> Result<ServerResponse, RequestHandlingError> {
-        let herited_span = self.root_span.as_ref().cloned().unwrap_or_else(|| tracing::Span::none());
-        let span = info_span!(parent: &herited_span, "forwarding_sphinx_packet");
+        let span = match &self.otel_propagator {
+            Some(propagator) => info_span!(parent: &propagator.root_span, "handling_forward_sphinx"),
+            None => info_span!("handling_forward_sphinx_no_otel"),
+        };
         let _enter = span.enter();
 
         let required_bandwidth = mix_packet.packet().len() as i64;
@@ -326,8 +335,10 @@ impl<R, S> AuthenticatedHandler<R, S> {
     #[instrument(skip_all)]
     async fn handle_binary(&mut self, bin_msg: Vec<u8>) -> Message {
         trace!("binary request");
-        let herited_span = self.root_span.as_ref().cloned().unwrap_or_else(|| tracing::Span::none());
-        let span = info_span!(parent: &herited_span, "handling_binary");
+        let span = match &self.otel_propagator {
+            Some(propagator) => info_span!(parent: &propagator.root_span, "handling_binary_request"),
+            None => info_span!("handling_binary_request_no_otel"),
+        };
         let _enter = span.enter();
 
         // this function decrypts the request and checks the MAC
@@ -618,32 +629,6 @@ impl<R, S> AuthenticatedHandler<R, S> {
         S: AsyncRead + AsyncWrite + Unpin,
     {
         trace!("Started listening for ALL incoming requests...");
-        let context_ext: ManualSpanContextExt = match &self.client.otel_context {
-            Some(otel_context) => {
-                let context_ext = ManualSpanContextExt::new()
-                    .with_context_carrier(&otel_context.context_carrier);
-                let span = if let Some(trace_id) = self.client.get_extracted_trace_id() {
-                    let cx = new_span_context_with_id(trace_id);
-                    let _context_guard = cx.clone().attach();
-
-                    let span = info_span!("client handling with otel", %trace_id);
-                    span.set_parent(cx.clone());
-                    span
-                } else {
-                    info_span!("client handling without otel context")
-                };
-                let context_ext = context_ext.set_root_span(span);
-                context_ext
-            }
-            None => {
-                warn!("No otel context associated with the client - this should never happen if otel has been set up!");
-                ManualSpanContextExt::new()
-            }
-        };
-
-        let _guard = context_ext.root_span.enter();
-        let handling_span = info_span!(parent: &context_ext.root_span, "client_handling_loop");
-
         // Ping timeout future used to check if the client responded to our ping request
         let mut ping_timeout: OptionFuture<_> = None.into();
 
@@ -668,7 +653,10 @@ impl<R, S> AuthenticatedHandler<R, S> {
                     self.handle_ping_timeout().await;
                 },
                 socket_msg = self.inner.read_websocket_message() => {
-                    let span = info_span!(parent: &handling_span, "client_message_received");
+                    let span = match &self.otel_propagator {
+                        Some(propagator) => info_span!(parent: &propagator.root_span, "client_message_received"),
+                        None => info_span!("client_message_received_no_otel"),
+                    };
                     let _enter = span.enter();
                     let socket_msg = match socket_msg {
                         None => break,
@@ -693,7 +681,10 @@ impl<R, S> AuthenticatedHandler<R, S> {
                     }
                 },
                 mix_messages = self.mix_receiver.next() => {
-                    let span = info_span!(parent: &handling_span, "mix_message_received");
+                    let span = match &self.otel_propagator {
+                        Some(propagator) => info_span!(parent: &propagator.root_span, "mix_message_received"),
+                        None => info_span!("mix_message_received_no_otel"),
+                    };
                     let _enter = span.enter();
                     let mix_messages = match mix_messages {
                         None => {

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

@@ -13,7 +13,7 @@ use futures::{
     channel::{mpsc, oneshot},
     SinkExt, StreamExt,
 };
-use nym_bin_common::opentelemetry::context::{new_span_context_with_id, ManualSpanContextExt};
+use nym_bin_common::opentelemetry::context::ManualContextPropagator;
 use nym_credentials_interface::AvailableBandwidth;
 use nym_crypto::aes::cipher::crypto_common::rand_core::RngCore;
 use nym_crypto::asymmetric::ed25519;
@@ -28,21 +28,14 @@ use nym_gateway_requests::{
     INITIAL_PROTOCOL_VERSION,
 };
 use nym_gateway_storage::error::GatewayStorageError;
-use nym_gateway_storage::models::Client;
 use nym_gateway_storage::traits::BandwidthGatewayStorage;
 use nym_gateway_storage::traits::InboxGatewayStorage;
 use nym_gateway_storage::traits::SharedKeyGatewayStorage;
 use nym_node_metrics::events::MetricsEvent;
 use nym_sphinx::DestinationAddressBytes;
 use nym_task::ShutdownToken;
-use opentelemetry::trace::{SpanContext, TraceContextExt, TraceFlags};
-use opentelemetry::TraceId;
-use opentelemetry::Context;
-use opentelemetry_sdk::logs::TraceContext;
-use opentelemetry_sdk::trace::{IdGenerator, RandomIdGenerator};
-use opentelemetry::propagation::TextMapPropagator;
-use opentelemetry_sdk::propagation::TraceContextPropagator;
 use rand::CryptoRng;
+use std::collections::HashMap;
 use std::net::SocketAddr;
 use std::time::Duration;
 use thiserror::Error;
@@ -50,8 +43,7 @@ use time::OffsetDateTime;
 use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::time::timeout;
 use tokio_tungstenite::tungstenite::{protocol::Message, Error as WsError};
-use tracing::{debug, error, info, info_span, instrument, span, warn};
-use tracing_opentelemetry::OpenTelemetrySpanExt;
+use tracing::{debug, error, info, info_span, instrument, warn};
 
 
 
@@ -663,13 +655,12 @@ impl<R, S> FreshHandler<R, S> {
     async fn handle_authenticate_v2(
         &mut self,
         request: Box<AuthenticateRequest>,
-        otel_context: Option<ManualSpanContextExt>,
+        otel_context: Option<HashMap<String, String>>,
     ) -> Result<InitialAuthResult, InitialAuthenticationError>
     where
         S: AsyncRead + AsyncWrite + Unpin,
     {
         debug!("handling client authentication (v2)");
-        tracing::info_span!("Authenticate v2");
 
         let negotiated_protocol =
             self.negotiate_client_protocol(Some(request.content.protocol_version))?;
@@ -879,45 +870,26 @@ impl<R, S> FreshHandler<R, S> {
         R: CryptoRng + RngCore + Send,
     {
         // extract and set up opentelemetry context if provided
-        let context_ext = if let ClientControlRequest::AuthenticateV2(ref auth_req) = request {
+        let (context_propagator, otel_ctx) = if let ClientControlRequest::AuthenticateV2(ref auth_req) = request {
             if let Some(otel_context) = &auth_req.otel_context {
-                // Extract OpenTelemetry context
-                let context_ext = ManualSpanContextExt::new()
-                    .with_extracted_context(otel_context.clone());
-                info!("Extracted trace id: {:?}", context_ext.trace_id);
-
-                // Build imported context and set it as parent
-                let extractor = TraceContextPropagator::new();
-                let extracted_context = extractor.extract(&context_ext.context_carrier);
-                let trace_id = if let Some(trace_id) = &context_ext.trace_id {
-                    *trace_id
-                } else {
-                    warn!("No trace id provided in the request, falling back to extracted context");
-                    extracted_context.span().span_context().trace_id()
-                };
-                let span_cx = new_span_context_with_id(trace_id);
-                let _context_guard = span_cx.clone().attach();
-
-                // Build root_span with extracted context as parent
-                let span = info_span!("=== Manual context propagation starting point ===", %trace_id);
-                span.set_parent(span_cx.clone());
-                let context_ext = context_ext.set_root_span(span);
-                warn!("==== Context propagation successful ====");
-                context_ext
+                (Some(ManualContextPropagator::new("handling_initial_client_request_with_otel", otel_context.clone())), Some(otel_context.clone()))
             } else {
                 warn!("No OpenTelemetry context provided in the request");
-                ManualSpanContextExt::new()
+                (None, None)
             }
         } else {
             warn!("No OpenTelemetry context provided in the request");
-            ManualSpanContextExt::new()
+            (None, None)
         };
 
-        let child_span = if context_ext.is_valid() {
-            info_span!(parent: &context_ext.root_span, "handling initial client request with otel context")
-        } else {
-            info_span!("handling_initial_client_request")
+        let child_span = match context_propagator {
+            Some(ref propagator) => {
+                let span = info_span!(parent: &propagator.root_span, "=== Handling initial client request with otel context ===");
+                span
+            }
+            None => info_span!("=== Handling initial client request without otel context ==="),
         };
+        
         let _enter = child_span.enter();
 
         let auth_result = match request {
@@ -931,7 +903,7 @@ impl<R, S> FreshHandler<R, S> {
                 self.handle_legacy_authenticate(protocol_version, address, enc_address, iv)
                     .await
             }
-            ClientControlRequest::AuthenticateV2(req) => self.handle_authenticate_v2(req, Some(context_ext)).await,
+            ClientControlRequest::AuthenticateV2(req) => self.handle_authenticate_v2(req, otel_ctx).await,
             ClientControlRequest::RegisterHandshakeInitRequest {
                 protocol_version,
                 data,

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

@@ -2,19 +2,17 @@
 // SPDX-License-Identifier: GPL-3.0-only
 
 use crate::config::Config;
-use nym_bin_common::opentelemetry::context::{new_span_context_with_id, ManualSpanContextExt};
 use nym_credential_verification::BandwidthFlushingBehaviourConfig;
 use nym_gateway_requests::shared_key::SharedGatewayKey;
 use nym_gateway_requests::ServerResponse;
 use nym_sphinx::DestinationAddressBytes;
-use opentelemetry::TraceId;
 use rand::{CryptoRng, Rng};
+use std::collections::HashMap;
 use std::time::Duration;
 use time::OffsetDateTime;
 use tokio::io::{AsyncRead, AsyncWrite};
 use tokio_tungstenite::WebSocketStream;
-use tracing::{debug, instrument, trace, warn, info_span};
-use tracing_opentelemetry::OpenTelemetrySpanExt;
+use tracing::{debug, instrument, trace, warn};
 
 pub(crate) use self::authenticated::AuthenticatedHandler;
 pub(crate) use self::fresh::FreshHandler;
@@ -51,7 +49,7 @@ pub(crate) struct ClientDetails {
     // note, this does **NOT ALWAYS** indicate timestamp of when client connected
     // it is (for v2 auth) timestamp the client **signed** when it created the request
     pub(crate) session_request_timestamp: OffsetDateTime,
-    pub(crate) otel_context: Option<ManualSpanContextExt>,
+    pub(crate) otel_context: Option<HashMap<String, String>>,
 }
 
 impl ClientDetails {
@@ -60,7 +58,7 @@ impl ClientDetails {
         address: DestinationAddressBytes,
         shared_keys: SharedGatewayKey,
         session_request_timestamp: OffsetDateTime,
-        otel_context: Option<ManualSpanContextExt>,
+        otel_context: Option<HashMap<String, String>>,
     ) -> Self {
         ClientDetails {
             address,
@@ -70,10 +68,6 @@ impl ClientDetails {
             otel_context,
         }
     }
-
-    pub(crate) fn get_extracted_trace_id(&self) -> Option<TraceId> {
-        self.otel_context.as_ref().and_then(|ctx| ctx.trace_id)
-    }
 }
 
 pub(crate) struct InitialAuthResult {
@@ -128,18 +122,6 @@ where
     trace!("managed to perform websocket handshake!");
 
     if let Some(auth_handle) = handle.handle_until_authenticated_or_failure().await {
-        let span = if let Some(trace_id) = auth_handle.get_trace_id() {
-            let cx = new_span_context_with_id(trace_id);
-            let _context_guard = cx.clone().attach();
-
-            let span = info_span!("authentication handler with otel", %trace_id);
-            span.set_parent(cx.clone());
-            span
-        } else {
-            info_span!("authentication handler without otel context")
-        };
-        let _guard = span.enter();
-
         auth_handle.listen_for_requests().await
     }