From 20a5fce67cc3ed24893bfe01bf15bc0843c7b4bc Mon Sep 17 00:00:00 2001 From: Huaijin Date: Fri, 22 Aug 2025 19:28:05 +0800 Subject: [PATCH] feat: improve tracing for flight (#8097) - [x] add span for tantivy_search - [x] improve tracing for flight after this pr the tracing for arrow flight will like below image --- src/handler/grpc/flight/mod.rs | 45 ++++-- src/handler/grpc/flight/stream.rs | 19 ++- src/service/promql/search/grpc/storage.rs | 4 +- .../datafusion/distributed_plan/common.rs | 4 +- .../distributed_plan/decoder_stream.rs | 129 +++++------------- .../distributed_plan/remote_scan.rs | 2 +- .../datafusion/plan/tantivy_optimize_exec.rs | 4 +- src/service/search/grpc/storage.rs | 5 +- 8 files changed, 94 insertions(+), 118 deletions(-) diff --git a/src/handler/grpc/flight/mod.rs b/src/handler/grpc/flight/mod.rs index e5f61e84ab..c76e50d247 100644 --- a/src/handler/grpc/flight/mod.rs +++ b/src/handler/grpc/flight/mod.rs @@ -31,6 +31,7 @@ use futures::{StreamExt, stream::BoxStream}; use futures_util::pin_mut; use prost::Message; use tonic::{Request, Response, Status, Streaming}; +use tracing::Instrument; use tracing_opentelemetry::OpenTelemetrySpanExt; #[cfg(feature = "enterprise")] use { @@ -70,7 +71,6 @@ impl FlightService for FlightServiceImpl { type ListActionsStream = BoxStream<'static, Result>; type DoExchangeStream = BoxStream<'static, Result>; - #[tracing::instrument(name = "grpc:search:flight:do_get", skip_all)] async fn do_get( &self, request: Request, @@ -81,7 +81,8 @@ impl FlightService for FlightServiceImpl { let parent_cx = opentelemetry::global::get_text_map_propagator(|prop| { prop.extract(&MetadataMap(request.metadata())) }); - tracing::Span::current().set_parent(parent_cx.clone()); + let span = tracing::info_span!("grpc:search:flight:do_get"); + span.set_parent(parent_cx); // decode ticket to RemoteExecNode let ticket = request.into_inner(); @@ -99,14 +100,35 @@ impl FlightService for FlightServiceImpl { let timeout = req.search_info.timeout as u64; log::info!("[trace_id {trace_id}] flight->search: do_get, timeout: {timeout}s",); - #[cfg(feature = "enterprise")] - if is_super_cluster && !SEARCH_SERVER.contain_key(&trace_id).await { - SEARCH_SERVER - .insert(trace_id.clone(), TaskStatus::new_follower(vec![], false)) - .await; + // Note: all async should in this place, otherwise it will break tracing + // https://docs.rs/tracing/latest/tracing/span/struct.Span.html#in-asynchronous-code + let req_move = req.clone(); + let trace_id_move = trace_id.clone(); + let result = async move { + #[cfg(feature = "enterprise")] + if is_super_cluster && !SEARCH_SERVER.contain_key(&trace_id_move).await { + // this is for work_group check in super cluster follower leader + SEARCH_SERVER + .insert( + trace_id_move.clone(), + TaskStatus::new_follower(vec![], false), + ) + .await; + } + + let result = get_ctx_and_physical_plan(&trace_id_move, &req_move).await; + + #[cfg(feature = "enterprise")] + if is_super_cluster && !SEARCH_SERVER.is_leader(&trace_id_move).await { + // this is for work_group check in super cluster follower leader + SEARCH_SERVER.remove(&trace_id_move, false).await; + } + + result } + .instrument(span.clone()) + .await; - let result = get_ctx_and_physical_plan(&trace_id, &req).await; log::info!( "{}", search_inspector_fields( @@ -123,11 +145,6 @@ impl FlightService for FlightServiceImpl { ) ); - #[cfg(feature = "enterprise")] - if is_super_cluster && !SEARCH_SERVER.is_leader(&trace_id).await { - SEARCH_SERVER.remove(&trace_id, false).await; - } - // prepare dataufion context let (ctx, physical_plan, defer, scan_stats) = match result { Ok(v) => v, @@ -197,7 +214,7 @@ impl FlightService for FlightServiceImpl { .with_custom_message(PreCustomMessage::ScanStatsRef(scan_stats_ref)) .with_custom_message(PreCustomMessage::Metrics(metrics)) .with_custom_message(PreCustomMessage::MetricsRef(metrics_ref)) - .build(stream); + .build(stream, span); let stream = async_stream::stream! { let timeout = tokio::time::sleep(tokio::time::Duration::from_secs(timeout)); diff --git a/src/handler/grpc/flight/stream.rs b/src/handler/grpc/flight/stream.rs index 31476ff580..c4f50bb704 100644 --- a/src/handler/grpc/flight/stream.rs +++ b/src/handler/grpc/flight/stream.rs @@ -22,6 +22,8 @@ use datafusion::execution::SendableRecordBatchStream; use flight::{common::PreCustomMessage, encoder::FlightDataEncoder}; use futures::{Stream, StreamExt}; use futures_core::ready; +use tracing::info_span; +use tracing_opentelemetry::OpenTelemetrySpanExt; use crate::{handler::grpc::flight::clear_session_data, service::search::utils::AsyncDefer}; @@ -67,7 +69,13 @@ impl FlightEncoderStreamBuilder { self } - pub fn build(self, inner: SendableRecordBatchStream) -> FlightEncoderStream { + pub fn build( + self, + inner: SendableRecordBatchStream, + span: tracing::Span, + ) -> FlightEncoderStream { + let child_span = info_span!("grpc:search:flight:execute_physical_plan"); + child_span.set_parent(span.context()); FlightEncoderStream { inner, encoder: self.encoder, @@ -78,6 +86,8 @@ impl FlightEncoderStreamBuilder { defer: self.defer, start: self.start, first_batch: true, + span, + child_span, } } } @@ -93,6 +103,8 @@ pub struct FlightEncoderStream { trace_id: String, defer: Option, start: std::time::Instant, + span: tracing::Span, + child_span: tracing::Span, } impl FlightEncoderStream { @@ -202,6 +214,11 @@ impl Drop for FlightEncoderStream { let took = self.start.elapsed().as_millis(); log::info!("[trace_id {trace_id}] flight->search: stream end, took: {took} ms",); + let _child_enter = self.child_span.enter(); + let _enter = self.span.enter(); + drop(_child_enter); + drop(_enter); + // metrics let time = self.start.elapsed().as_secs_f64(); metrics::GRPC_RESPONSE_TIME diff --git a/src/service/promql/search/grpc/storage.rs b/src/service/promql/search/grpc/storage.rs index c65c98cbb5..15f7f0ef9a 100644 --- a/src/service/promql/search/grpc/storage.rs +++ b/src/service/promql/search/grpc/storage.rs @@ -46,7 +46,7 @@ use crate::service::{ datafusion::exec::register_table, grpc::{ QueryParams, - storage::{cache_files, filter_file_list_by_tantivy_index}, + storage::{cache_files, tantivy_search}, }, index::{Condition, IndexCondition}, match_source, @@ -227,7 +227,7 @@ pub(crate) async fn create_context( let index_condition = convert_matchers_to_index_condition(&matchers, &schema, &index_fields)?; if !index_condition.conditions.is_empty() && cfg.common.inverted_index_enabled { let (idx_took, ..) = - filter_file_list_by_tantivy_index(query.clone(), &mut files, Some(index_condition), None) + tantivy_search(query.clone(), &mut files, Some(index_condition), None) .await .map_err(|e| { log::error!( diff --git a/src/service/search/datafusion/distributed_plan/common.rs b/src/service/search/datafusion/distributed_plan/common.rs index 23db338dfa..368cfbd888 100644 --- a/src/service/search/datafusion/distributed_plan/common.rs +++ b/src/service/search/datafusion/distributed_plan/common.rs @@ -28,7 +28,6 @@ use parking_lot::Mutex; #[derive(Debug)] pub struct QueryContext { pub trace_id: String, - pub parent_cx: opentelemetry::Context, pub node: Arc, pub is_super: bool, pub is_querier: bool, @@ -40,9 +39,8 @@ pub struct QueryContext { } impl QueryContext { - pub fn new(parent_cx: opentelemetry::Context, node: Arc) -> Self { + pub fn new(node: Arc) -> Self { Self { - parent_cx, node, is_super: false, is_querier: false, diff --git a/src/service/search/datafusion/distributed_plan/decoder_stream.rs b/src/service/search/datafusion/distributed_plan/decoder_stream.rs index f276df3875..07f826fa98 100644 --- a/src/service/search/datafusion/distributed_plan/decoder_stream.rs +++ b/src/service/search/datafusion/distributed_plan/decoder_stream.rs @@ -13,7 +13,7 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::{fmt::Debug, ops::Sub, pin::Pin, task::Poll}; +use std::{fmt::Debug, pin::Pin, task::Poll}; use arrow::array::RecordBatch; use arrow_flight::{FlightData, error::Result}; @@ -24,7 +24,6 @@ use flight::{ decoder::FlightDataDecoder, }; use futures::{Stream, StreamExt, ready}; -use opentelemetry::trace::{Span, TraceId, Tracer}; use tonic::Streaming; use crate::service::search::{ @@ -111,37 +110,8 @@ impl Stream for FlightDecoderStream { impl Drop for FlightDecoderStream { fn drop(&mut self) { - let cfg = config::get_config(); - if (cfg.common.tracing_enabled || cfg.common.tracing_search_enabled) - && let Err(e) = self.create_stream_end_span() - { - log::error!("error creating stream span: {e}"); - } - let scan_stats = { *self.query_context.scan_stats.lock() }; - log::info!( - "[trace_id {}] flight->search: response node: {}, is_super: {}, is_querier: {}, files: {}, scan_size: {} mb, num_rows: {}, took: {} ms", - self.query_context.trace_id, - self.query_context.node.get_grpc_addr(), - self.query_context.is_super, - self.query_context.is_querier, - scan_stats.files, - scan_stats.original_size / 1024 / 1024, - self.query_context.num_rows, - self.query_context.start.elapsed().as_millis() - ); - } -} - -/// Create a span for the stream end -impl FlightDecoderStream { - fn create_stream_end_span( - &self, - ) -> std::result::Result { - let tracer = opentelemetry::global::tracer("FlightDecoderStream"); - let QueryContext { trace_id, - parent_cx, node, is_super, is_querier, @@ -151,69 +121,42 @@ impl FlightDecoderStream { .. } = &self.query_context; - let now = std::time::SystemTime::now(); - let duration = start.elapsed(); - let start_time = now.sub(duration); + let search_role = if *is_super { + "leader".to_string() + } else { + "follower".to_string() + }; + let scan_stats = { *scan_stats.lock() }; - let trace_id = trace_id - .split('-') - .next() - .and_then(|id| TraceId::from_hex(id).ok()); - match trace_id { - Some(trace_id) => { - let mut span = tracer - .span_builder("service:search:flight::do_get_stream") - .with_trace_id(trace_id) - .with_start_time(start_time) - .with_attributes(vec![opentelemetry::KeyValue::new( - "duration", - duration.as_nanos() as i64, - )]) - .start_with_context(&tracer, parent_cx); - - let span_context = span.span_context().clone(); - let search_role = if *is_super { - "leader".to_string() - } else { - "follower".to_string() - }; - let scan_stats = { *scan_stats.lock() }; - let event = search_inspector_fields( - format!( - "[trace_id {}] flight->search: response node: {}, is_super: {}, is_querier: {}, files: {}, scan_size: {} mb, num_rows: {}, took: {} ms", - trace_id, - node.get_grpc_addr(), - is_super, - is_querier, + log::info!( + "{}", + search_inspector_fields( + format!( + "[trace_id {}] flight->search: response node: {}, is_super: {}, is_querier: {}, files: {}, scan_size: {} mb, num_rows: {}, took: {} ms", + trace_id, + node.get_grpc_addr(), + is_super, + is_querier, + scan_stats.files, + scan_stats.original_size / 1024 / 1024, + num_rows, + start.elapsed().as_millis(), + ), + SearchInspectorFieldsBuilder::new() + .node_name(node.get_name()) + .region(node.get_region()) + .cluster(node.get_cluster()) + .component("remote scan streaming".to_string()) + .search_role(search_role) + .duration(start.elapsed().as_millis() as usize) + .desc(format!( + "remote scan search files: {}, scan_size: {}, num_rows: {}", scan_stats.files, - scan_stats.original_size / 1024 / 1024, - num_rows, - start.elapsed().as_millis(), - ), - SearchInspectorFieldsBuilder::new() - .node_name(node.get_name()) - .region(node.get_region()) - .cluster(node.get_cluster()) - .component("remote scan streaming".to_string()) - .search_role(search_role) - .duration(start.elapsed().as_millis() as usize) - .desc(format!( - "remote scan search files: {}, scan_size: {}, num_rows: {}", - scan_stats.files, - bytes_to_human_readable(scan_stats.original_size as f64), - num_rows - )) - .build(), - ); - - span.add_event_with_timestamp(event, now, vec![]); - span.end_with_timestamp(now); - Ok(span_context) - } - None => Err(infra::errors::Error::Message(format!( - "Invalid trace id: {}", - self.query_context.trace_id - ))), - } + bytes_to_human_readable(scan_stats.original_size as f64), + num_rows + )) + .build(), + ) + ) } } diff --git a/src/service/search/datafusion/distributed_plan/remote_scan.rs b/src/service/search/datafusion/distributed_plan/remote_scan.rs index 0b47befb65..451c4bb0d0 100644 --- a/src/service/search/datafusion/distributed_plan/remote_scan.rs +++ b/src/service/search/datafusion/distributed_plan/remote_scan.rs @@ -315,7 +315,7 @@ async fn get_remote_batch( "[trace_id {trace_id}] flight->search: prepare to response node: {grpc_addr}, is_super: {is_super}, is_querier: {is_querier}", ); - let query_context = QueryContext::new(context.clone(), node) + let query_context = QueryContext::new(node) .with_trace_id(&trace_id) .with_is_super(is_super) .with_is_querier(is_querier) diff --git a/src/service/search/datafusion/plan/tantivy_optimize_exec.rs b/src/service/search/datafusion/plan/tantivy_optimize_exec.rs index b03dbc7654..76d6f382c0 100644 --- a/src/service/search/datafusion/plan/tantivy_optimize_exec.rs +++ b/src/service/search/datafusion/plan/tantivy_optimize_exec.rs @@ -38,7 +38,7 @@ use datafusion::{ use futures::TryStreamExt; use crate::service::search::{ - grpc::{QueryParams, storage::filter_file_list_by_tantivy_index}, + grpc::{QueryParams, storage::tantivy_search}, index::IndexCondition, }; @@ -164,7 +164,7 @@ async fn adapt_tantivy_result( schema: SchemaRef, idx_optimize_mode: IndexOptimizeMode, ) -> Result { - let (idx_took, error, result) = filter_file_list_by_tantivy_index( + let (idx_took, error, result) = tantivy_search( query.clone(), &mut file_list, index_condition, diff --git a/src/service/search/grpc/storage.rs b/src/service/search/grpc/storage.rs index 58c79d7d67..00784e2552 100644 --- a/src/service/search/grpc/storage.rs +++ b/src/service/search/grpc/storage.rs @@ -142,7 +142,7 @@ pub async fn search( let mut idx_took = 0; let mut is_add_filter_back = false; if use_inverted_index { - (idx_took, is_add_filter_back, ..) = filter_file_list_by_tantivy_index( + (idx_took, is_add_filter_back, ..) = tantivy_search( query.clone(), &mut files, tantivy_condition.clone(), @@ -622,7 +622,8 @@ pub async fn cache_files( /// If the query does not match any FST in the index file, the file will be filtered out. /// If the query does match then the segment IDs for the file will be updated. /// If the query not find corresponding index file, the file will *not* be filtered out. -pub async fn filter_file_list_by_tantivy_index( +#[tracing::instrument(name = "service:search:grpc:storage:tantivy_search", skip_all)] +pub async fn tantivy_search( query: Arc, file_list: &mut Vec, index_condition: Option,