feat: tracing for ws v2 (#6602)

- [x] add tracing for `websocket_events` service: `search`, `cancel`,
`values`

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<Mutex<SnowflakeIdGenerator>> = 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

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,

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::<WsClientEvents>(&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<CloseReaso
 }
 
 // Main search handler
+#[tracing::instrument(name = "service:search:websocket::handle_search_event", skip_all)]
 async fn handle_search_event(
     search_req: &SearchEventReq,
     org_id: &str,
@@ -633,7 +651,7 @@ async fn handle_search_event(
     // Spawn the search task
     tokio::spawn(async move {
         // Handle the search request
-        // If search is cancelled, the task will exit
+        // If search is cancelled, the taek will exit
         // Otherwise, the task will complete and the results will be sent to the client
         // The task will also update the search state to completed
         // The task will also close the session
@@ -728,6 +746,10 @@ async fn handle_search_event(
 
 // Cancel handler
 #[cfg(feature = "enterprise")]
+#[tracing::instrument(
+    name = "service:websocket_events:search::handle_cancel_event",
+    skip_all
+)]
 async fn handle_cancel_event(trace_id: &str) -> 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<CloseReason> {
     // 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,

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<config::meta::search::Response>,
@@ -727,6 +728,7 @@ pub async fn _write_results(
 /// 6. **Cache to Disk**:
 ///    - Saves the filtered response to a file named:
 ///      `"<start_time>_<end_time>_<is_aggregate>_<is_descending>.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,

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<SearchResultType>,
@@ -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<Response, Error> {
-    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<Value>, ctx: &ValuesEventContext) -> Result<Vec<Value>, Error> {
     let mut top_k_values: Vec<Value> = Vec::new();
 

src/service/websocket_events/utils.rs

@@ -14,10 +14,14 @@
 // along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 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<String, String> = 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

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?;
         }
     }