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
<img width="1770" height="506" alt="image"
src="https://github.com/user-attachments/assets/55e30700-122c-4dec-bff7-8eed0f823bce"
/>

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<ActionType, Status>>;
     type DoExchangeStream = BoxStream<'static, Result<FlightData, Status>>;
 
-    #[tracing::instrument(name = "grpc:search:flight:do_get", skip_all)]
     async fn do_get(
         &self,
         request: Request<Ticket>,
@@ -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));

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<AsyncDefer>,
     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

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!(

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<dyn NodeInfo>,
     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<dyn NodeInfo>) -> Self {
+    pub fn new(node: Arc<dyn NodeInfo>) -> Self {
         Self {
-            parent_cx,
             node,
             is_super: false,
             is_querier: false,

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 <http://www.gnu.org/licenses/>.
 
-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<opentelemetry::trace::SpanContext, infra::errors::Error> {
-        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(),
+            )
+        )
     }
 }

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)

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<SendableRecordBatchStream> {
-    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,

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<super::QueryParams>,
     file_list: &mut Vec<FileKey>,
     index_condition: Option<IndexCondition>,