diff --git a/src/config/src/ider.rs b/src/config/src/ider.rs index 53d1da09a6..e2150b12f7 100644 --- a/src/config/src/ider.rs +++ b/src/config/src/ider.rs @@ -20,6 +20,7 @@ use std::{ use once_cell::sync::Lazy; use parking_lot::Mutex; +use rand::Rng; use svix_ksuid::{Ksuid, KsuidLike}; static IDER: Lazy> = Lazy::new(|| { @@ -49,6 +50,21 @@ pub fn generate_trace_id() -> String { opentelemetry::trace::TraceId::from(trace_id).to_string() } +/// Generate a new span_id. +pub fn generate_span_id() -> String { + let mut rng = rand::thread_rng(); + let mut bytes = [0u8; 8]; + rng.fill(&mut bytes); + + let hex = hex::encode(bytes); + + if hex == "0000000000000000" { + return generate_span_id(); + } + + hex +} + /// Convert a snowflake id to a timestamp in milliseconds. pub fn to_timestamp_millis(id: i64) -> i64 { id >> 22 diff --git a/src/handler/http/request/search/mod.rs b/src/handler/http/request/search/mod.rs index 48651b337e..1d6bd542ed 100644 --- a/src/handler/http/request/search/mod.rs +++ b/src/handler/http/request/search/mod.rs @@ -681,6 +681,7 @@ pub type FieldName = String; /// The `Request` contains the SQL query, from, size, start_time, end_time, etc. /// The `StreamType` is the type of the stream to search on. /// The `FieldName` is the name of the field to search on. +#[tracing::instrument(name = "handler:search:build_search_request_per_field", skip_all)] pub async fn build_search_request_per_field( req: &config::meta::search::ValuesRequest, org_id: &str, diff --git a/src/handler/http/request/ws_v2/session.rs b/src/handler/http/request/ws_v2/session.rs index aaed4284d4..9b58651319 100644 --- a/src/handler/http/request/ws_v2/session.rs +++ b/src/handler/http/request/ws_v2/session.rs @@ -31,6 +31,7 @@ use o2_enterprise::enterprise::common::{ }; use rand::prelude::SliceRandom; use tokio::sync::mpsc; +use tracing::Instrument; #[cfg(feature = "enterprise")] use crate::common::infra::cluster::get_cached_online_router_nodes; @@ -47,7 +48,7 @@ use crate::{ }, service::websocket_events::{ WsClientEvents, WsServerEvents, handle_search_request, handle_values_request, - search_registry_utils::SearchState, sessions_cache_utils, + search_registry_utils::SearchState, sessions_cache_utils, setup_tracing_with_trace_id, }, }; @@ -208,7 +209,8 @@ pub async fn run( cfg.common.instance_name, msg ); - handle_text_message(&user_id, &req_id, msg.to_string(), path.clone()).await; + handle_text_message(&user_id, &req_id, msg.to_string(), path.clone()) + .await; } Ok(actix_ws::AggregatedMessage::Close(reason)) => { if let Some(reason) = reason.as_ref() { @@ -311,6 +313,7 @@ async fn resolve_enterprise_user_id( /// Text message is parsed into `WsClientEvents` and processed accordingly /// Depending on each event type, audit must be done /// Currently audit is done only for the search event +#[tracing::instrument(name = "service:search:websocket::handle_text_message", skip_all)] pub async fn handle_text_message(user_id: &str, req_id: &str, msg: String, path: String) { match serde_json::from_str::(&msg) { Ok(client_msg) => { @@ -327,6 +330,15 @@ pub async fn handle_text_message(user_id: &str, req_id: &str, msg: String, path: return; } + // Setup tracing + let ws_span = setup_tracing_with_trace_id( + &client_msg.get_trace_id(), + tracing::info_span!( + "src::handler::http::request::websocket::ws_v2::session::handle_text_message" + ), + ) + .await; + match client_msg { WsClientEvents::Search(ref search_req) => { #[allow(unused_mut)] @@ -363,6 +375,7 @@ pub async fn handle_text_message(user_id: &str, req_id: &str, msg: String, path: req_id, path.clone(), ) + .instrument(ws_span) .await; } WsClientEvents::Values(ref values_req) => { @@ -400,6 +413,7 @@ pub async fn handle_text_message(user_id: &str, req_id: &str, msg: String, path: req_id, path.clone(), ) + .instrument(ws_span.clone()) .await; } #[cfg(feature = "enterprise")] @@ -417,7 +431,10 @@ pub async fn handle_text_message(user_id: &str, req_id: &str, msg: String, path: // First handle the cancel event // send a cancel flag to the search task - if let Err(e) = handle_cancel_event(&trace_id).await { + if let Err(e) = handle_cancel_event(&trace_id) + .instrument(ws_span.clone()) + .await + { log::warn!("[WS_HANDLER]: Error in cancelling : {}", e); return; } @@ -596,6 +613,7 @@ async fn cleanup_and_close_session(req_id: &str, close_reason: Option Result<(), anyhow::Error> { let mut w = WS_SEARCH_REGISTRY.write().await; if let Some(state) = w.get_mut(trace_id) { @@ -756,6 +778,10 @@ async fn handle_cancel_event(trace_id: &str) -> Result<(), anyhow::Error> { Ok(()) } +#[tracing::instrument( + name = "service:websocket_events:search::handle_search_error", + skip_all +)] async fn handle_search_error(e: &Error, req_id: &str, trace_id: &str) -> Option { // if the error is due to search cancellation, return. // the cancel handler will close the session @@ -816,6 +842,10 @@ async fn cleanup_search_resources(trace_id: &str) { } // Main values handler +#[tracing::instrument( + name = "service:websocket_events:search::handle_values_event", + skip_all +)] async fn handle_values_event( values_req: &ValuesEventReq, org_id: &str, diff --git a/src/service/search/cache/mod.rs b/src/service/search/cache/mod.rs index 31ebffeede..fb1eaf2334 100644 --- a/src/service/search/cache/mod.rs +++ b/src/service/search/cache/mod.rs @@ -457,6 +457,7 @@ pub async fn search( // based on _timestamp of first record in config::meta::search::Response either add it in start // or end to cache response +#[tracing::instrument(name = "service:search:cache:merge_response", skip_all)] pub fn merge_response( trace_id: &str, cache_responses: &mut Vec, @@ -727,6 +728,7 @@ pub async fn _write_results( /// 6. **Cache to Disk**: /// - Saves the filtered response to a file named: /// `"___.json"`. +#[tracing::instrument(name = "service:search:cache:write_results_v2", skip_all)] #[allow(clippy::too_many_arguments)] pub async fn write_results_v2( trace_id: &str, diff --git a/src/service/websocket_events/search.rs b/src/service/websocket_events/search.rs index c1efd12633..ba27243789 100644 --- a/src/service/websocket_events/search.rs +++ b/src/service/websocket_events/search.rs @@ -49,11 +49,14 @@ use crate::{ self as SearchService, cache, datafusion::distributed_plan::streaming_aggs_exec, sql::Sql, }, - websocket_events::{TimeOffset, WsServerEvents, search_registry_utils}, + websocket_events::{ + TimeOffset, WsServerEvents, search_registry_utils, setup_tracing_with_trace_id, + }, }, }; #[cfg(feature = "enterprise")] +#[tracing::instrument(name = "service:websocket_events:search::handle_cancel", skip_all)] pub async fn handle_cancel(trace_id: &str, org_id: &str) -> WsServerEvents { match crate::service::search::cancel_query(org_id, trace_id).await { Ok(ret) => { @@ -81,6 +84,7 @@ pub async fn handle_cancel(trace_id: &str, org_id: &str) -> WsServerEvents { } } +#[tracing::instrument(name = "service:search:websocket::handle_search_request", skip_all)] pub async fn handle_search_request( req_id: &str, accumulated_results: &mut Vec, @@ -101,6 +105,13 @@ pub async fn handle_search_request( end_time ); + // Setup tracing + let ws_search_span = setup_tracing_with_trace_id( + &req.trace_id, + tracing::info_span!("src::service::websocket_events::search::handle_search_request"), + ) + .await; + // check and append search event type if req.payload.search_type.is_none() { req.payload.search_type = Some(req.search_type); @@ -185,6 +196,7 @@ pub async fn handle_search_request( if req.payload.query.from == 0 { let c_resp = cache::check_cache_v2(&trace_id, org_id, stream_type, &req.payload, req.use_cache) + .instrument(ws_search_span.clone()) .await?; let local_c_resp = c_resp.clone(); let cached_resp = local_c_resp.cached_response; @@ -242,6 +254,7 @@ pub async fn handle_search_request( remaining_query_range, &order_by, ) + .instrument(ws_search_span.clone()) .await?; } else { // Step 2: Search without cache @@ -262,12 +275,14 @@ pub async fn handle_search_request( accumulated_results, max_query_range, ) + .instrument(ws_search_span.clone()) .await?; } // Step 3: Write to results cache // cache only if from is 0 and is not an aggregate_query if req.payload.query.from == 0 { write_results_to_cache(c_resp, start_time, end_time, accumulated_results) + .instrument(ws_search_span.clone()) .await .map_err(|e| { log::error!( @@ -292,6 +307,7 @@ pub async fn handle_search_request( accumulated_results, max_query_range, ) + .instrument(ws_search_span) .await?; } @@ -305,17 +321,12 @@ pub async fn handle_search_request( Ok(()) } +#[tracing::instrument(name = "service:search:websocket::do_search", skip_all)] async fn do_search( req: &SearchEventReq, user_id: &str, use_cache: bool, ) -> Result { - let span = tracing::info_span!( - "src::handler::http::request::websocket::search::do_search", - trace_id = %req.trace_id, - org_id = %req.org_id, - ); - let mut req = req.clone(); req.payload.use_cache = Some(use_cache); let res = SearchService::cache::search( @@ -326,7 +337,6 @@ async fn do_search( &req.payload, "".to_string(), ) - .instrument(span) .await; res.map(handle_partial_response) @@ -345,6 +355,10 @@ fn handle_partial_response(mut res: Response) -> Response { res } +#[tracing::instrument( + name = "service:search:websocket::handle_cache_responses_and_deltas", + skip_all +)] #[allow(clippy::too_many_arguments)] pub async fn handle_cache_responses_and_deltas( req_id: &str, @@ -811,6 +825,7 @@ async fn send_cached_responses( } // Do partitioned search without cache +#[tracing::instrument(name = "service:search:websocket::do_partitioned_search", skip_all)] #[allow(clippy::too_many_arguments)] pub async fn do_partitioned_search( req_id: &str, @@ -925,10 +940,15 @@ pub async fn do_partitioned_search( } if req.search_type == SearchEventType::Values && req.values_event_context.is_some() { + let ws_search_span = tracing::info_span!( + "src::service::websocket_events::search::do_partitioned_search::get_top_k_values", + org_id = %req.org_id, + ); let instant = Instant::now(); let top_k_values = tokio::task::spawn_blocking(move || { get_top_k_values(&search_res.hits, &req.values_event_context.clone().unwrap()) }) + .instrument(ws_search_span.clone()) .await .unwrap(); search_res.hits = top_k_values?; @@ -1028,6 +1048,10 @@ async fn send_partial_search_resp( Ok(()) } +#[tracing::instrument( + name = "service:websocket_events:search::write_results_to_cache", + skip_all +)] pub async fn write_results_to_cache( c_resp: MultiCachedQueryResponse, start_time: i64, @@ -1115,7 +1139,7 @@ pub async fn write_results_to_cache( } /// This function will compute top k values for values request -/// This is a com +#[tracing::instrument(name = "service:websocket_events:search::get_top_k_values", skip_all)] pub fn get_top_k_values(hits: &Vec, ctx: &ValuesEventContext) -> Result, Error> { let mut top_k_values: Vec = Vec::new(); diff --git a/src/service/websocket_events/utils.rs b/src/service/websocket_events/utils.rs index 369f12ad0d..61e01fd104 100644 --- a/src/service/websocket_events/utils.rs +++ b/src/service/websocket_events/utils.rs @@ -14,10 +14,14 @@ // along with this program. If not, see . use actix_web::http::StatusCode; -use config::meta::websocket::{SearchEventReq, ValuesEventReq}; +use config::{ + ider, + meta::websocket::{SearchEventReq, ValuesEventReq}, +}; use infra::errors; use serde::{Deserialize, Serialize}; use tokio_tungstenite::tungstenite; +use tracing_opentelemetry::OpenTelemetrySpanExt; use crate::handler::http::request::search::error_utils::map_error_to_http_response; @@ -309,6 +313,21 @@ pub mod search_registry_utils { } } +/// Setup tracing with a trace ID +pub async fn setup_tracing_with_trace_id(trace_id: &str, span: tracing::Span) -> tracing::Span { + let mut headers: std::collections::HashMap = std::collections::HashMap::new(); + let traceparent = format!( + "00-{}-{}-01", /* 01 to indicate that the span is sampled i.e. needs to be + * recorded/exported */ + trace_id, + ider::generate_span_id() + ); + headers.insert("traceparent".to_string(), traceparent); + let parent_ctx = opentelemetry::global::get_text_map_propagator(|prop| prop.extract(&headers)); + span.set_parent(parent_ctx); + span +} + /// Represents the different types of WebSocket client messages that can be sent. /// /// The `WSClientMessage` enum is used to represent the different types of messages that can be sent diff --git a/src/service/websocket_events/values.rs b/src/service/websocket_events/values.rs index 2bf068a536..edbfe91cb9 100644 --- a/src/service/websocket_events/values.rs +++ b/src/service/websocket_events/values.rs @@ -21,6 +21,7 @@ use config::{ websocket::{SearchResultType, ValuesEventReq}, }, }; +use tracing::Instrument; #[cfg(feature = "enterprise")] use crate::service::websocket_events::enterprise_utils; @@ -36,10 +37,12 @@ use crate::{ search::{ do_partitioned_search, handle_cache_responses_and_deltas, write_results_to_cache, }, + setup_tracing_with_trace_id, }, }, }; +#[tracing::instrument(name = "service:search:websocket::handle_values_request", skip_all)] pub async fn handle_values_request( org_id: &str, user_id: &str, @@ -56,6 +59,13 @@ pub async fn handle_values_request( let stream_name = payload.stream_name.clone(); let no_count = payload.no_count; + // Setup tracing + let ws_values_span = setup_tracing_with_trace_id( + &trace_id, + tracing::info_span!("src::service::websocket_events::values::handle_values_request"), + ) + .await; + log::info!( "[WS_VALUES] trace_id: {} Received values request, start_time: {}, end_time: {}", trace_id, @@ -87,7 +97,9 @@ pub async fn handle_values_request( let top_k = payload.size.or(Some(cfg.limit.query_default_limit)); // get values req query - let reqs = build_search_request_per_field(&payload, org_id, stream_type, &stream_name).await?; + let reqs = build_search_request_per_field(&payload, org_id, stream_type, &stream_name) + .instrument(ws_values_span.clone()) + .await?; for (search_req, stream_type, field) in reqs { // Convert the request query to SearchQuery type if needed @@ -105,6 +117,7 @@ pub async fn handle_values_request( &search_req, search_req.use_cache.unwrap_or(false), ) + .instrument(ws_values_span.clone()) .await?; let local_c_resp = c_resp.clone(); let cached_resp = local_c_resp.cached_response; @@ -181,6 +194,7 @@ pub async fn handle_values_request( remaining_query_range, &order_by, ) + .instrument(ws_values_span.clone()) .await?; } else { // Step 2: Search without cache @@ -220,6 +234,7 @@ pub async fn handle_values_request( accumulated_results, max_query_range, ) + .instrument(ws_values_span.clone()) .await?; } // Step 3: Write to results cache @@ -230,6 +245,7 @@ pub async fn handle_values_request( let safe_end_time = end_time.unwrap_or(0); write_results_to_cache(c_resp, safe_start_time, safe_end_time, accumulated_results) + .instrument(ws_values_span.clone()) .await .map_err(|e| { log::error!( @@ -273,6 +289,7 @@ pub async fn handle_values_request( accumulated_results, max_query_range, ) + .instrument(ws_values_span.clone()) .await?; } }