diff --git a/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs b/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs index a5bdb0f479..9e2c842095 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs @@ -34,13 +34,12 @@ use nym_gateway_storage::traits::SharedKeyGatewayStorage; use nym_node_metrics::events::MetricsEvent; use nym_sphinx::DestinationAddressBytes; use nym_task::ShutdownToken; -use opentelemetry::trace::TraceContextExt; +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::trace::TraceContextExt; use opentelemetry_sdk::propagation::TraceContextPropagator; use rand::CryptoRng; use std::collections::HashMap; @@ -867,12 +866,12 @@ impl FreshHandler { pub(crate) async fn handle_initial_client_request( &mut self, request: ClientControlRequest, - ) -> Result<(Option, Option>), InitialAuthenticationError> + ) -> Result<(Option, Option), InitialAuthenticationError> where S: AsyncRead + AsyncWrite + Unpin + Send, R: CryptoRng + RngCore + Send, { - let remote_cx_span = if let ClientControlRequest::AuthenticateV2(ref auth_req) = request { + let (remote_cx_span, trace_id) = if let ClientControlRequest::AuthenticateV2(ref auth_req) = request { if let Some(otel_context) = &auth_req.otel_context { let carrier = ContextCarrier::from_map(otel_context.clone()); @@ -886,13 +885,13 @@ impl FreshHandler { let span = info_span!("extracted_otel_context"); warn!("==== Context propagation successful ===="); - Some(span) + (Some(span), extracted_trace_id) } else { warn!("No OpenTelemetry context provided in the request"); - None + (None, None) } } else { - None + (None, None) }; let child_span = if let Some(ref parent_span) = remote_cx_span { @@ -902,18 +901,6 @@ impl FreshHandler { }; let _enter = child_span.enter(); - let gateway_context = opentelemetry::Context::current(); - let mut carrier = ContextCarrier::new(); - let propagator = TraceContextPropagator::new(); - propagator.inject_context(&gateway_context, &mut carrier); - let extracted_trace_id = carrier.extract_trace_id(); - info!("handle_initial_client_request: Extracted trace id after child_span: {:?}", extracted_trace_id); - let gw_carrier = carrier.into_map(); - - let context = opentelemetry::Context::current(); - let trace_id = context.span().span_context().trace_id(); - error!("=== [DEBUG SPAN] Current trace id in handle_initial_client_request: {:?}", trace_id); - let auth_result = match request { ClientControlRequest::Authenticate { protocol_version, @@ -932,7 +919,7 @@ impl FreshHandler { } => self.handle_register(protocol_version, data).await, ClientControlRequest::SupportedProtocol { .. } => { self.handle_reply_supported_protocol_request().await; - return Ok((None, Some(gw_carrier))); + return Ok((None, trace_id)); } _ => { debug!("received an invalid client request"); @@ -971,8 +958,7 @@ impl FreshHandler { warn!("could not establish client details"); return Err(InitialAuthenticationError::EmptyClientDetails); }; - - Ok((Some(client_details), Some(gw_carrier))) + Ok((Some(client_details), trace_id)) } #[instrument(skip_all)] @@ -994,8 +980,8 @@ impl FreshHandler { }; // see if we managed to register the client through this request - let (maybe_auth_res, context_carrier) = match self.handle_initial_client_request(initial_request).await { - Ok((maybe_auth_res, context_carrier)) => (maybe_auth_res, context_carrier), + let (maybe_auth_res, maybe_trace_id) = match self.handle_initial_client_request(initial_request).await { + Ok((maybe_auth_res, trace_id)) => (maybe_auth_res, trace_id), Err(err) => { debug!("initial client request handling error: {err}"); self.send_and_forget_error_response(err).await; @@ -1003,14 +989,26 @@ impl FreshHandler { } }; - if let (Some(registration_details), Some(context_carrier)) = (maybe_auth_res, context_carrier) { - let new_carrier = ContextCarrier::from_map(context_carrier); - let propagator = TraceContextPropagator::new(); - let extracted_context = propagator.extract(&new_carrier); - tracing::Span::current().set_parent(extracted_context); - let span = tracing::info_span!("handle_authenticated_client"); - let _entered_span = span.enter(); - warn!("==== Context propagation successful to handle until authenticated client handler ===="); + if let (Some(registration_details), Some(trace_id)) = (maybe_auth_res, maybe_trace_id) { + let span = { + let id_gen = opentelemetry_sdk::trace::RandomIdGenerator::default(); + let span_id = id_gen.new_span_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!("authentication handler with otel", %trace_id); + span.set_parent(cx.clone()); + span + }; + + let _guard = span.enter(); let (mix_sender, mix_receiver) = mpsc::unbounded(); // Channel for handlers to ask other handlers if they are still active. diff --git a/gateway/src/node/client_handling/websocket/connection_handler/mod.rs b/gateway/src/node/client_handling/websocket/connection_handler/mod.rs index 1784af9e92..661003d49b 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/mod.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/mod.rs @@ -9,13 +9,15 @@ use nym_gateway_requests::ServerResponse; use nym_sphinx::DestinationAddressBytes; use opentelemetry::propagation::TextMapPropagator; use opentelemetry_sdk::propagation::TraceContextPropagator; +use opentelemetry_sdk::trace::IdGenerator; use rand::{CryptoRng, Rng}; use std::time::Duration; use time::OffsetDateTime; use tokio::io::{AsyncRead, AsyncWrite}; use tokio_tungstenite::WebSocketStream; -use tracing::{debug, instrument, trace, warn}; +use tracing::{debug, instrument, trace, warn, info_span}; use tracing_opentelemetry::OpenTelemetrySpanExt; +use opentelemetry::trace::{SpanContext, TraceContextExt, TraceFlags}; pub(crate) use self::authenticated::AuthenticatedHandler; pub(crate) use self::fresh::FreshHandler;