Record search metrics on cancelation

quickwit/quickwit-search/Cargo.toml

@@ -22,6 +22,7 @@ http = { workspace = true }
 itertools = { workspace = true }
 mockall = { workspace = true }
 once_cell = { workspace = true }
+pin-project = { workspace = true }
 postcard = { workspace = true }
 prost = { workspace = true }
 rayon = { workspace = true }

quickwit/quickwit-search/src/leaf.rs

@@ -45,7 +45,6 @@ use tokio::task::JoinError;
 use tracing::*;
 
 use crate::collector::{IncrementalCollector, make_collector_for_split, make_merge_collector};
-use crate::metrics::SEARCH_METRICS;
 use crate::root::is_metadata_count_request_with_ast;
 use crate::search_permit_provider::{SearchPermit, compute_initial_memory_allocation};
 use crate::service::{SearcherContext, deserialize_doc_mapper};
@@ -1175,9 +1174,10 @@ impl CanSplitDoBetter {
     }
 }
 
-/// `multi_leaf_search` searches multiple indices and multiple splits.
+/// Searches multiple splits, potentially in multiple indexes, sitting on different storages and
+/// having different doc mappings.
 #[instrument(skip_all, fields(index = ?leaf_search_request.search_request.as_ref().unwrap().index_id_patterns))]
-pub async fn multi_leaf_search(
+pub async fn multi_index_leaf_search(
     searcher_context: Arc<SearcherContext>,
     leaf_search_request: LeafSearchRequest,
     storage_resolver: &StorageResolver,
@@ -1225,18 +1225,25 @@ pub async fn multi_leaf_search(
             })?
             .clone();
 
-        let leaf_request_future = tokio::spawn(
-            resolve_storage_and_leaf_search(
-                searcher_context.clone(),
-                search_request.clone(),
-                index_uri,
-                storage_resolver.clone(),
-                leaf_search_request_ref.split_offsets,
-                doc_mapper,
-                aggregation_limits.clone(),
-            )
-            .in_current_span(),
-        );
+        let leaf_request_future = tokio::spawn({
+            let storage_resolver = storage_resolver.clone();
+            let searcher_context = searcher_context.clone();
+            let search_request = search_request.clone();
+            let aggregation_limits = aggregation_limits.clone();
+            async move {
+                let storage = storage_resolver.resolve(&index_uri).await?;
+                single_doc_mapping_leaf_search(
+                    searcher_context,
+                    search_request,
+                    storage,
+                    leaf_search_request_ref.split_offsets,
+                    doc_mapper,
+                    aggregation_limits,
+                )
+                .await
+            }
+            .in_current_span()
+        });
         leaf_request_tasks.push(leaf_request_future);
     }
 
@@ -1269,29 +1276,6 @@ pub async fn multi_leaf_search(
         .context("failed to merge split search responses")?
 }
 
-/// Resolves storage and calls leaf_search
-#[allow(clippy::too_many_arguments)]
-async fn resolve_storage_and_leaf_search(
-    searcher_context: Arc<SearcherContext>,
-    search_request: Arc<SearchRequest>,
-    index_uri: quickwit_common::uri::Uri,
-    storage_resolver: StorageResolver,
-    splits: Vec<SplitIdAndFooterOffsets>,
-    doc_mapper: Arc<DocMapper>,
-    aggregations_limits: AggregationLimitsGuard,
-) -> crate::Result<LeafSearchResponse> {
-    let storage = storage_resolver.resolve(&index_uri).await?;
-    leaf_search(
-        searcher_context.clone(),
-        search_request.clone(),
-        storage.clone(),
-        splits,
-        doc_mapper,
-        aggregations_limits,
-    )
-    .await
-}
-
 /// Optimizes the search_request based on CanSplitDoBetter
 /// Returns true if the split can return better results
 fn check_optimize_search_request(
@@ -1315,14 +1299,14 @@ fn disable_search_request_hits(search_request: &mut SearchRequest) {
     search_request.sort_fields.clear();
 }
 
-/// `leaf` step of search.
+/// Searches multiple splits for a specific index and a single doc mapping
 ///
 /// The leaf search collects all kind of information, and returns a set of
 /// [PartialHit](quickwit_proto::search::PartialHit) candidates. The root will be in
 /// charge to consolidate, identify the actual final top hits to display, and
 /// fetch the actual documents to convert the partial hits into actual Hits.
 #[instrument(skip_all, fields(index = ?request.index_id_patterns))]
-pub async fn leaf_search(
+pub async fn single_doc_mapping_leaf_search(
     searcher_context: Arc<SearcherContext>,
     request: Arc<SearchRequest>,
     index_storage: Arc<dyn Storage>,
@@ -1444,15 +1428,6 @@ pub async fn leaf_search(
             .await
             .context("failed to merge split search responses");
 
-    let label_values = match leaf_search_response_reresult {
-        Ok(Ok(_)) => ["success"],
-        _ => ["error"],
-    };
-    SEARCH_METRICS
-        .leaf_search_targeted_splits
-        .with_label_values(label_values)
-        .observe(num_splits as f64);
-
     Ok(leaf_search_response_reresult??)
 }
 

quickwit/quickwit-search/src/lib.rs

@@ -29,6 +29,7 @@ mod leaf_cache;
 mod list_fields;
 mod list_fields_cache;
 mod list_terms;
+mod metrics_trackers;
 mod retry;
 mod root;
 mod scroll_context;

quickwit/quickwit-search/src/list_fields.rs

@@ -279,7 +279,7 @@ pub struct IndexMetasForLeafSearch {
 }
 
 /// Performs a distributed list fields request.
-/// 1. Sends leaf request over gRPC to multiple leaf nodes.
+/// 1. Sends leaf requests over gRPC to multiple leaf nodes.
 /// 2. Merges the search results.
 /// 3. Builds the response and returns.
 pub async fn root_list_fields(

quickwit/quickwit-search/src/list_terms.rs

@@ -39,7 +39,7 @@ use crate::search_permit_provider::compute_initial_memory_allocation;
 use crate::{ClusterClient, SearchError, SearchJob, SearcherContext, resolve_index_patterns};
 
 /// Performs a distributed list terms.
-/// 1. Sends leaf request over gRPC to multiple leaf nodes.
+/// 1. Sends leaf requests over gRPC to multiple leaf nodes.
 /// 2. Merges the search results.
 /// 3. Builds the response and returns.
 /// this is much simpler than `root_search` as it doesn't need to get actual docs.

quickwit/quickwit-search/src/metrics_trackers.rs

@@ -0,0 +1,151 @@
+// Copyright 2021-Present Datadog, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// See https://prometheus.io/docs/practices/naming/
+
+use std::pin::Pin;
+use std::task::{Context, Poll, ready};
+use std::time::Instant;
+
+use pin_project::{pin_project, pinned_drop};
+use quickwit_proto::search::LeafSearchResponse;
+
+use crate::SearchError;
+use crate::metrics::SEARCH_METRICS;
+
+// root
+
+pub enum RootSearchMetricsStep {
+    Plan,
+    Exec { num_targeted_splits: usize },
+}
+
+/// Wrapper around the plan and search futures to track metrics.
+#[pin_project(PinnedDrop)]
+pub struct RootSearchMetricsFuture<F> {
+    #[pin]
+    pub tracked: F,
+    pub start: Instant,
+    pub step: RootSearchMetricsStep,
+    pub is_success: Option<bool>,
+}
+
+#[pinned_drop]
+impl<F> PinnedDrop for RootSearchMetricsFuture<F> {
+    fn drop(self: Pin<&mut Self>) {
+        let (num_targeted_splits, status) = match (&self.step, self.is_success) {
+            // is is a partial success, actual success is recorded during the search step
+            (RootSearchMetricsStep::Plan, Some(true)) => return,
+            (RootSearchMetricsStep::Plan, Some(false)) => (0, "plan-error"),
+            (RootSearchMetricsStep::Plan, None) => (0, "plan-cancelled"),
+            (
+                RootSearchMetricsStep::Exec {
+                    num_targeted_splits,
+                },
+                Some(true),
+            ) => (*num_targeted_splits, "success"),
+            (
+                RootSearchMetricsStep::Exec {
+                    num_targeted_splits,
+                },
+                Some(false),
+            ) => (*num_targeted_splits, "error"),
+            (
+                RootSearchMetricsStep::Exec {
+                    num_targeted_splits,
+                },
+                None,
+            ) => (*num_targeted_splits, "cancelled"),
+        };
+
+        let label_values = [status];
+        SEARCH_METRICS
+            .root_search_requests_total
+            .with_label_values(label_values)
+            .inc();
+        SEARCH_METRICS
+            .root_search_request_duration_seconds
+            .with_label_values(label_values)
+            .observe(self.start.elapsed().as_secs_f64());
+        SEARCH_METRICS
+            .root_search_targeted_splits
+            .with_label_values(label_values)
+            .observe(num_targeted_splits as f64);
+    }
+}
+
+impl<F, R, E> Future for RootSearchMetricsFuture<F>
+where F: Future<Output = Result<R, E>>
+{
+    type Output = Result<R, E>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.project();
+        let response = ready!(this.tracked.poll(cx));
+        *this.is_success = Some(response.is_ok());
+        Poll::Ready(Ok(response?))
+    }
+}
+
+// leaf
+
+/// Wrapper around the search future to track metrics.
+#[pin_project(PinnedDrop)]
+pub struct LeafSearchMetricsFuture<F>
+where F: Future<Output = Result<LeafSearchResponse, SearchError>>
+{
+    #[pin]
+    pub tracked: F,
+    pub start: Instant,
+    pub targeted_splits: usize,
+    pub status: Option<&'static str>,
+}
+
+#[pinned_drop]
+impl<F> PinnedDrop for LeafSearchMetricsFuture<F>
+where F: Future<Output = Result<LeafSearchResponse, SearchError>>
+{
+    fn drop(self: Pin<&mut Self>) {
+        let label_values = [self.status.unwrap_or("cancelled")];
+        SEARCH_METRICS
+            .leaf_search_requests_total
+            .with_label_values(label_values)
+            .inc();
+        SEARCH_METRICS
+            .leaf_search_request_duration_seconds
+            .with_label_values(label_values)
+            .observe(self.start.elapsed().as_secs_f64());
+        SEARCH_METRICS
+            .leaf_search_targeted_splits
+            .with_label_values(label_values)
+            .observe(self.targeted_splits as f64);
+    }
+}
+
+impl<F> Future for LeafSearchMetricsFuture<F>
+where F: Future<Output = Result<LeafSearchResponse, SearchError>>
+{
+    type Output = Result<LeafSearchResponse, SearchError>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.project();
+        let response = ready!(this.tracked.poll(cx));
+        *this.status = if response.is_ok() {
+            Some("success")
+        } else {
+            Some("error")
+        };
+        Poll::Ready(Ok(response?))
+    }
+}