Refactor metrics trackers to their own file

Author: rdettai

quickwit/quickwit-search/src/leaf.rs

@@ -1174,7 +1174,7 @@ impl CanSplitDoBetter {
     }
 }
 
-/// Searches multiple splits, potentially in multiple indices, sitting on different storages and
+/// 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_index_leaf_search(
@@ -1299,7 +1299,7 @@ fn disable_search_request_hits(search_request: &mut SearchRequest) {
     search_request.sort_fields.clear();
 }
 
-/// Searches multiple splits from a specific index and a single doc mapping
+/// 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

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?))
+    }
+}

quickwit/quickwit-search/src/root.rs

@@ -13,17 +13,13 @@
 // limitations under the License.
 
 use std::collections::{HashMap, HashSet};
-use std::future::Future;
-use std::pin::Pin;
 use std::sync::OnceLock;
 use std::sync::atomic::{AtomicU64, Ordering};
-use std::task::{Context as TaskContext, Poll, ready};
-use std::time::Duration;
+use std::time::{Duration, Instant};
 
 use anyhow::Context;
 use futures::future::try_join_all;
 use itertools::Itertools;
-use pin_project::{pin_project, pinned_drop};
 use quickwit_common::pretty::PrettySample;
 use quickwit_common::shared_consts;
 use quickwit_common::uri::Uri;
@@ -49,13 +45,12 @@ use tantivy::aggregation::agg_result::AggregationResults;
 use tantivy::aggregation::intermediate_agg_result::IntermediateAggregationResults;
 use tantivy::collector::Collector;
 use tantivy::schema::{Field, FieldEntry, FieldType, Schema};
-use tokio::time::Instant;
 use tracing::{debug, info_span, instrument};
 
 use crate::cluster_client::ClusterClient;
 use crate::collector::{QuickwitAggregations, make_merge_collector};
 use crate::find_trace_ids_collector::Span;
-use crate::metrics::SEARCH_METRICS;
+use crate::metrics_trackers::{RootSearchMetricsFuture, RootSearchMetricsStep};
 use crate::scroll_context::{ScrollContext, ScrollKeyAndStartOffset};
 use crate::search_job_placer::{Job, group_by, group_jobs_by_index_id};
 use crate::search_response_rest::StorageRequestCount;
@@ -960,7 +955,7 @@ fn get_sort_field_datetime_format(
 }
 
 /// Performs a distributed search.
-/// 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. Sends fetch docs requests to multiple leaf nodes.
 /// 4. Builds the response with docs and returns.
@@ -1189,7 +1184,7 @@ async fn plan_splits_for_root_search(
 }
 
 /// Performs a distributed search.
-/// 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. Sends fetch docs requests to multiple leaf nodes.
 /// 4. Builds the response with docs and returns.
@@ -1200,7 +1195,7 @@ pub async fn root_search(
     mut metastore: MetastoreServiceClient,
     cluster_client: &ClusterClient,
 ) -> crate::Result<SearchResponse> {
-    let start_instant = tokio::time::Instant::now();
+    let start_instant = Instant::now();
 
     let (split_metadatas, indexes_meta_for_leaf_search) = RootSearchMetricsFuture {
         start: start_instant,
@@ -1227,7 +1222,7 @@ pub async fn root_search(
         ),
         is_success: None,
         step: RootSearchMetricsStep::Exec {
-            targeted_splits: num_splits,
+            num_targeted_splits: num_splits,
         },
     }
     .await;
@@ -1765,69 +1760,6 @@ pub fn jobs_to_fetch_docs_requests(
     Ok(fetch_docs_requests)
 }
 
-enum RootSearchMetricsStep {
-    Plan,
-    Exec { targeted_splits: usize },
-}
-
-/// Wrapper around the plan and search futures to track metrics.
-#[pin_project(PinnedDrop)]
-struct RootSearchMetricsFuture<F> {
-    #[pin]
-    tracked: F,
-    start: Instant,
-    step: RootSearchMetricsStep,
-    is_success: Option<bool>,
-}
-
-#[pinned_drop]
-impl<F> PinnedDrop for RootSearchMetricsFuture<F> {
-    fn drop(self: Pin<&mut Self>) {
-        let (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 { targeted_splits }, Some(true)) => {
-                (*targeted_splits, "success")
-            }
-            (RootSearchMetricsStep::Exec { targeted_splits }, Some(false)) => {
-                (*targeted_splits, "error")
-            }
-            (RootSearchMetricsStep::Exec { targeted_splits }, None) => {
-                (*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(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 TaskContext<'_>) -> 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?))
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use std::ops::Range;