search: fix trace fragmentation and restructure spans

quickwit/quickwit-common/src/tracing_utils.rs

@@ -80,16 +80,6 @@ pub fn extract_context(metadata: &MetadataMap) -> Context {
     global::get_text_map_propagator(|propagator| propagator.extract(&extractor))
 }
 
-/// Extracts a W3C trace context from incoming gRPC request metadata and
-/// installs it as the parent of the currently-active tracing span. Use this
-/// at the entry of a gRPC handler that is itself wrapped in a
-/// `#[tracing::instrument]` so the handler's span is stitched into the
-/// caller's trace.
-pub fn set_current_span_parent_from_metadata(metadata: &MetadataMap) {
-    let parent_context = extract_context(metadata);
-    let _ = Span::current().set_parent(parent_context);
-}
-
 /// Tonic interceptor that injects the active span's W3C trace context into
 /// the outgoing gRPC metadata. Drop-in replacement for the legacy
 /// `quickwit_proto::SpanContextInterceptor`.

quickwit/quickwit-search/src/fetch_docs.rs

@@ -28,7 +28,7 @@ use tantivy::schema::document::CompactDocValue;
 use tantivy::schema::{Document as DocumentTrait, Field, TantivyDocument, Value};
 use tantivy::snippet::SnippetGenerator;
 use tantivy::{ReloadPolicy, Score, Searcher, Term};
-use tracing::{Instrument, error};
+use tracing::{Instrument, error, instrument};
 
 use crate::leaf::open_index_with_caches;
 use crate::service::SearcherContext;
@@ -153,6 +153,7 @@ struct Document {
 }
 
 /// Fetching docs from a specific split.
+#[instrument(skip_all, fields(split_id = split.split_id, num_docs = global_doc_addrs.len()))]
 async fn fetch_docs_in_split(
     searcher_context: Arc<SearcherContext>,
     mut global_doc_addrs: Vec<GlobalDocAddress>,

quickwit/quickwit-search/src/leaf.rs

@@ -841,7 +841,10 @@ async fn leaf_search_single_split(
     }
 
     let split_num_docs = split.num_docs;
-    let span = info_span!("tantivy_search");
+    // Spans the CPU-pool queue wait: created here (right after warmup) and closed when the
+    // closure starts executing on a pool thread. `tantivy_search` is created inside the
+    // closure so it covers only the CPU execution, not this wait.
+    let cpu_wait_span = info_span!("waiting_for_cpu_pool");
 
     let split_clone = split.clone();
 
@@ -852,9 +855,12 @@ async fn leaf_search_single_split(
     let search_request_and_result: Option<(SearchRequest, LeafSearchResponse)> =
         crate::search_thread_pool()
             .run_cpu_intensive(move || {
+                // The CPU-pool queue wait ends as this closure starts executing.
+                drop(cpu_wait_span);
                 leaf_search_state_guard.set_state(SplitSearchState::Cpu);
                 let cpu_start = Instant::now();
                 let cpu_thread_pool_wait_microsecs = cpu_start.duration_since(warmup_end);
+                let span = info_span!("tantivy_search");
                 let _span_guard = span.enter();
                 // Our search execution has been scheduled, let's check if we can improve the
                 // request based on the results of the preceding searches
@@ -1545,7 +1551,7 @@ pub async fn multi_index_leaf_search(
         let searcher_context = searcher_context.clone();
         let search_request = search_request.clone();
 
-        leaf_request_futures.spawn({
+        leaf_request_futures.spawn(
             async move {
                 let storage = storage_resolver.resolve(&index_uri).await?;
                 single_doc_mapping_leaf_search(
@@ -1555,10 +1561,10 @@ pub async fn multi_index_leaf_search(
                     leaf_search_request_ref.split_offsets,
                     doc_mapper,
                 )
-                .in_current_span()
                 .await
             }
-        });
+            .instrument(Span::current()),
+        );
     }
 
     // Creates a collector which merges responses into one
@@ -1723,12 +1729,15 @@ async fn run_offloaded_search_tasks(
             }],
         };
         let invoker = lambda_invoker.clone();
-        lambda_tasks_joinset.spawn(async move {
-            (
-                batch_split_ids,
-                invoker.invoke_leaf_search(leaf_request).await,
-            )
-        });
+        lambda_tasks_joinset.spawn(
+            async move {
+                (
+                    batch_split_ids,
+                    invoker.invoke_leaf_search(leaf_request).await,
+                )
+            }
+            .in_current_span(),
+        );
     }
 
     while let Some(join_res) = lambda_tasks_joinset.join_next().await {
@@ -2021,9 +2030,15 @@ async fn run_local_search_tasks(
         search_permit_future,
     } in local_search_tasks
     {
-        let leaf_split_search_permit = search_permit_future
-            .instrument(info_span!("waiting_for_leaf_search_split_semaphore"))
-            .await;
+        // Per-split span covering both the permit wait and the search, so each split is a
+        // single subtree (wait + warmup + tantivy) rather than flat siblings.
+        let split_span = info_span!(
+            "leaf_search_single_split_wrapper",
+            split_id = split.split_id,
+            num_docs = split.num_docs
+        );
+        let wait_span = info_span!(parent: &split_span, "waiting_for_leaf_search_split_semaphore");
+        let leaf_split_search_permit = search_permit_future.instrument(wait_span).await;
 
         // We run simplify search request again: as we push split into the merge collector,
         // we may have discovered that we won't find any better candidates for top hits in this
@@ -2045,7 +2060,7 @@ async fn run_local_search_tasks(
                 split.clone(),
                 leaf_split_search_permit,
             )
-            .in_current_span(),
+            .instrument(split_span),
         );
         task_id_to_split_id_map.insert(handle.id(), split_id);
     }
@@ -2174,7 +2189,6 @@ struct LeafSearchContext {
     split_filter: Arc<RwLock<CanSplitDoBetter>>,
 }
 
-#[instrument(skip_all, fields(split_id = split.split_id, num_docs = split.num_docs))]
 async fn leaf_search_single_split_wrapper(
     request: SearchRequest,
     ctx: Arc<LeafSearchContext>,

quickwit/quickwit-search/src/lib.rs

@@ -188,6 +188,7 @@ pub async fn list_all_splits(
 }
 
 /// Extract the list of relevant splits for a given request.
+#[tracing::instrument(skip_all, fields(num_indexes = index_uids.len()))]
 pub async fn list_relevant_splits(
     index_uids: Vec<IndexUid>,
     start_timestamp: Option<i64>,

quickwit/quickwit-search/src/root.rs

@@ -789,7 +789,7 @@ fn compute_root_resource_stats(
 
 /// If this method fails for some splits, a partial search response is returned, with the list of
 /// faulty splits in the failed_splits field.
-#[instrument(level = "debug", skip_all)]
+#[instrument(skip_all)]
 pub(crate) async fn search_partial_hits_phase(
     searcher_context: &SearcherContext,
     indexes_metas_for_leaf_search: &IndexesMetasForLeafSearch,
@@ -1249,6 +1249,7 @@ async fn refine_and_list_matches(
 }
 
 /// Fetches the list of splits and their metadata from the metastore
+#[instrument(skip_all)]
 async fn plan_splits_for_root_search(
     search_request: &mut SearchRequest,
     metastore: &mut MetastoreServiceClient,

quickwit/quickwit-serve/src/search_api/grpc_adapter.rs

@@ -15,7 +15,6 @@
 use std::sync::Arc;
 
 use async_trait::async_trait;
-use quickwit_common::tracing_utils::set_current_span_parent_from_metadata;
 use quickwit_proto::error::convert_to_grpc_result;
 use quickwit_proto::search::{
     GetKvRequest, GetKvResponse, GetLoadRequest, GetLoadResponse, LeafListFieldsRequest,
@@ -24,7 +23,6 @@ use quickwit_proto::search::{
 };
 use quickwit_proto::tonic;
 use quickwit_search::SearchService;
-use tracing::instrument;
 
 #[derive(Clone)]
 pub struct GrpcSearchAdapter(Arc<dyn SearchService>);
@@ -35,58 +33,52 @@ impl From<Arc<dyn SearchService>> for GrpcSearchAdapter {
     }
 }
 
+// The `grpc-request` span installed by the tonic server's `trace_fn` already
+// extracts the caller's W3C trace context from the request headers and is the
+// active span for each handler, so these methods deliberately add no span of
+// their own to avoid collinear duplicates.
 #[async_trait]
 impl grpc::SearchService for GrpcSearchAdapter {
-    #[instrument(skip(self, request))]
     async fn root_search(
         &self,
         request: tonic::Request<quickwit_proto::search::SearchRequest>,
     ) -> Result<tonic::Response<quickwit_proto::search::SearchResponse>, tonic::Status> {
-        set_current_span_parent_from_metadata(request.metadata());
         let search_request = request.into_inner();
         let search_result = self.0.root_search(search_request).await;
         convert_to_grpc_result(search_result)
     }
 
-    #[instrument(skip(self, request))]
     async fn leaf_search(
         &self,
         request: tonic::Request<quickwit_proto::search::LeafSearchRequest>,
     ) -> Result<tonic::Response<quickwit_proto::search::LeafSearchResponse>, tonic::Status> {
-        set_current_span_parent_from_metadata(request.metadata());
         let leaf_search_request = request.into_inner();
         let leaf_search_result = self.0.leaf_search(leaf_search_request).await;
         convert_to_grpc_result(leaf_search_result)
     }
 
-    #[instrument(skip(self, request))]
     async fn fetch_docs(
         &self,
         request: tonic::Request<quickwit_proto::search::FetchDocsRequest>,
     ) -> Result<tonic::Response<quickwit_proto::search::FetchDocsResponse>, tonic::Status> {
-        set_current_span_parent_from_metadata(request.metadata());
         let fetch_docs_request = request.into_inner();
         let fetch_docs_result = self.0.fetch_docs(fetch_docs_request).await;
         convert_to_grpc_result(fetch_docs_result)
     }
 
-    #[instrument(skip(self, request))]
     async fn root_list_terms(
         &self,
         request: tonic::Request<quickwit_proto::search::ListTermsRequest>,
     ) -> Result<tonic::Response<quickwit_proto::search::ListTermsResponse>, tonic::Status> {
-        set_current_span_parent_from_metadata(request.metadata());
         let search_request = request.into_inner();
         let search_result = self.0.root_list_terms(search_request).await;
         convert_to_grpc_result(search_result)
     }
 
-    #[instrument(skip(self, request))]
     async fn leaf_list_terms(
         &self,
         request: tonic::Request<quickwit_proto::search::LeafListTermsRequest>,
     ) -> Result<tonic::Response<quickwit_proto::search::LeafListTermsResponse>, tonic::Status> {
-        set_current_span_parent_from_metadata(request.metadata());
         let leaf_search_request = request.into_inner();
         let leaf_search_result = self.0.leaf_list_terms(leaf_search_request).await;
         convert_to_grpc_result(leaf_search_result)
@@ -101,67 +93,56 @@ impl grpc::SearchService for GrpcSearchAdapter {
         convert_to_grpc_result(scroll_result)
     }
 
-    #[instrument(skip(self, request))]
     async fn put_kv(
         &self,
         request: tonic::Request<quickwit_proto::search::PutKvRequest>,
     ) -> Result<tonic::Response<quickwit_proto::search::PutKvResponse>, tonic::Status> {
-        set_current_span_parent_from_metadata(request.metadata());
         let put_request = request.into_inner();
         self.0.put_kv(put_request).await;
         Ok(tonic::Response::new(
             quickwit_proto::search::PutKvResponse {},
         ))
     }
 
-    #[instrument(skip(self, request))]
     async fn get_kv(
         &self,
         request: tonic::Request<GetKvRequest>,
     ) -> Result<tonic::Response<GetKvResponse>, tonic::Status> {
-        set_current_span_parent_from_metadata(request.metadata());
         let get_search_after_context_request = request.into_inner();
         let payload = self.0.get_kv(get_search_after_context_request).await;
         let get_response = GetKvResponse { payload };
         Ok(tonic::Response::new(get_response))
     }
 
-    #[instrument(skip(self, request))]
     async fn report_splits(
         &self,
         request: tonic::Request<ReportSplitsRequest>,
     ) -> Result<tonic::Response<ReportSplitsResponse>, tonic::Status> {
-        set_current_span_parent_from_metadata(request.metadata());
         let get_search_after_context_request = request.into_inner();
         self.0.report_splits(get_search_after_context_request).await;
         Ok(tonic::Response::new(ReportSplitsResponse {}))
     }
 
-    #[instrument(skip(self, request))]
     async fn list_fields(
         &self,
         request: tonic::Request<ListFieldsRequest>,
     ) -> Result<tonic::Response<ListFieldsResponse>, tonic::Status> {
-        set_current_span_parent_from_metadata(request.metadata());
         let resp = self.0.root_list_fields(request.into_inner()).await;
         convert_to_grpc_result(resp)
     }
-    #[instrument(skip(self, request))]
+
     async fn leaf_list_fields(
         &self,
         request: tonic::Request<LeafListFieldsRequest>,
     ) -> Result<tonic::Response<ListFieldsResponse>, tonic::Status> {
-        set_current_span_parent_from_metadata(request.metadata());
         let resp = self.0.leaf_list_fields(request.into_inner()).await;
         convert_to_grpc_result(resp)
     }
 
-    #[instrument(skip(self, request))]
     async fn search_plan(
         &self,
         request: tonic::Request<quickwit_proto::search::SearchRequest>,
     ) -> Result<tonic::Response<quickwit_proto::search::SearchPlanResponse>, tonic::Status> {
-        set_current_span_parent_from_metadata(request.metadata());
         let search_request = request.into_inner();
         let search_result = self.0.search_plan(search_request).await;
         convert_to_grpc_result(search_result)