diff --git a/common/bin-common/src/opentelemetry/context.rs b/common/bin-common/src/opentelemetry/context.rs index a016e2a39c..f6f59a85b2 100644 --- a/common/bin-common/src/opentelemetry/context.rs +++ b/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, - 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) -> 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 ) -> 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(); 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 ebb7fe6594..27654fcd80 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs +++ b/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 { // 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, + otel_propagator: Option, } // explicitly remove handle from the global store upon being dropped @@ -193,6 +191,17 @@ impl AuthenticatedHandler { 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 AuthenticatedHandler { 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 AuthenticatedHandler { self.inner.send_metrics(event) } - pub fn get_trace_id(&self) -> Option { - self.client.get_extracted_trace_id() - } - /// Forwards the received mix packet from the client into the mix network. /// /// # Arguments @@ -238,8 +243,10 @@ impl AuthenticatedHandler { /// * `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 AuthenticatedHandler { &mut self, mix_packet: MixPacket, ) -> Result { - 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 AuthenticatedHandler { #[instrument(skip_all)] async fn handle_binary(&mut self, bin_msg: Vec) -> 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 AuthenticatedHandler { 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 AuthenticatedHandler { 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 AuthenticatedHandler { } }, 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 => { 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 fc74621a5f..df10d2dcdc 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs +++ b/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 FreshHandler { async fn handle_authenticate_v2( &mut self, request: Box, - otel_context: Option, + otel_context: Option>, ) -> Result 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 FreshHandler { 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 FreshHandler { 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, 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 bf65d6908d..232b2f4c74 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/mod.rs +++ b/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, + pub(crate) otel_context: Option>, } impl ClientDetails { @@ -60,7 +58,7 @@ impl ClientDetails { address: DestinationAddressBytes, shared_keys: SharedGatewayKey, session_request_timestamp: OffsetDateTime, - otel_context: Option, + otel_context: Option>, ) -> Self { ClientDetails { address, @@ -70,10 +68,6 @@ impl ClientDetails { otel_context, } } - - pub(crate) fn get_extracted_trace_id(&self) -> Option { - 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 }