Add Tokio runtimes metrics (#4984)

.cargo/config.toml

@@ -1,8 +1,12 @@
+[build]
+rustflags = ["--cfg", "tokio_unstable"]
+rustdocflags = ["--cfg", "tokio_unstable"]
+
 [target.x86_64-unknown-linux-gnu]
-# Targetting x86-64-v2 gives a ~2% performance boost while only
+# Targeting x86-64-v2 gives a ~2% performance boost while only
 # disallowing Intel CPUs older than 2008 and AMD CPUs older than 2011.
 # None of those very old CPUs are used in GCP
 # (https://cloud.google.com/compute/docs/cpu-platforms). Unfortunately,
 # AWS does not seem to disclose the exact CPUs they use.
-rustflags = ["-C", "target-cpu=x86-64-v2"]
+rustflags = ["-C", "target-cpu=x86-64-v2", "--cfg", "tokio_unstable"]
 

.github/workflows/cbench.yml

@@ -13,11 +13,13 @@ on:
   # pull request.
   pull_request_target:
 
+env:
+  RUSTFLAGS: --cfg tokio_unstable
+
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
   cancel-in-progress: true
 
-
 jobs:
   tests:
     name: Benchmark

.github/workflows/ci.yml

@@ -16,7 +16,7 @@ env:
   QW_TEST_DATABASE_URL: postgres://quickwit-dev:quickwit-dev@localhost:5432/quickwit-metastore-dev
   RUST_BACKTRACE: 1
   RUSTDOCFLAGS: -Dwarnings -Arustdoc::private_intra_doc_links
-  RUSTFLAGS: -Dwarnings
+  RUSTFLAGS: -Dwarnings --cfg tokio_unstable
 
 # Ensures that we cancel running jobs for the same PR / same workflow.
 concurrency:

.github/workflows/coverage.yml

@@ -21,6 +21,7 @@ env:
   QW_S3_ENDPOINT: "http://localhost:4566" # Services are exposed as localhost because we are not running coverage in a container.
   QW_S3_FORCE_PATH_STYLE_ACCESS: 1
   QW_TEST_DATABASE_URL: postgres://quickwit-dev:quickwit-dev@localhost:5432/quickwit-metastore-dev
+  RUSTFLAGS: -Dwarnings --cfg tokio_unstable
 
 jobs:
   test:

Dockerfile

@@ -39,7 +39,8 @@ COPY --from=ui-builder /quickwit/quickwit-ui/build /quickwit/quickwit-ui/build
 WORKDIR /quickwit
 
 RUN echo "Building workspace with feature(s) '$CARGO_FEATURES' and profile '$CARGO_PROFILE'" \
-    && cargo build \
+    && ENV RUSTFLAGS="--cfg tokio_unstable" \
+        cargo build \
         -p quickwit-cli \
         --features $CARGO_FEATURES \
         --bin quickwit \

quickwit/Cargo.lock

@@ -5706,6 +5706,7 @@ dependencies = [
  "tempfile",
  "thiserror",
  "tokio",
+ "tokio-metrics",
  "tokio-stream",
  "tonic",
  "tower",
@@ -8380,6 +8381,18 @@ dependencies = [
  "syn 2.0.59",
 ]
 
+[[package]]
+name = "tokio-metrics"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "eace09241d62c98b7eeb1107d4c5c64ca3bd7da92e8c218c153ab3a78f9be112"
+dependencies = [
+ "futures-util",
+ "pin-project-lite",
+ "tokio",
+ "tokio-stream",
+]
+
 [[package]]
 name = "tokio-native-tls"
 version = "0.3.1"

quickwit/Cargo.toml

@@ -231,6 +231,7 @@ tikv-jemalloc-ctl = "0.5"
 tikv-jemallocator = "0.5"
 time = { version = "0.3", features = ["std", "formatting", "macros"] }
 tokio = { version = "1.37", features = ["full"] }
+tokio-metrics = { version = "0.3.1", features = ["rt"] }
 tokio-stream = { version = "0.1", features = ["sync"] }
 tokio-util = { version = "0.7", features = ["full"] }
 toml = "0.7.6"

quickwit/quickwit-cli/src/main.rs

@@ -21,6 +21,7 @@
 
 use std::collections::BTreeMap;
 
+use anyhow::Context;
 use colored::Colorize;
 use opentelemetry::global;
 use quickwit_cli::busy_detector;
@@ -29,17 +30,21 @@ use quickwit_cli::cli::{build_cli, CliCommand};
 #[cfg(feature = "jemalloc")]
 use quickwit_cli::jemalloc::start_jemalloc_metrics_loop;
 use quickwit_cli::logger::setup_logging_and_tracing;
+use quickwit_common::runtimes::scrape_tokio_runtime_metrics;
 use quickwit_serve::BuildInfo;
 use tracing::error;
 
 fn main() -> anyhow::Result<()> {
-    tokio::runtime::Builder::new_multi_thread()
+    let rt = tokio::runtime::Builder::new_multi_thread()
         .enable_all()
         .on_thread_unpark(busy_detector::thread_unpark)
         .on_thread_park(busy_detector::thread_park)
         .build()
-        .unwrap()
-        .block_on(main_impl())
+        .context("failed to start main Tokio runtime")?;
+
+    scrape_tokio_runtime_metrics(rt.handle(), "main");
+
+    rt.block_on(main_impl())
 }
 
 fn register_build_info_metric() {

quickwit/quickwit-cli/src/split.rs

@@ -340,14 +340,6 @@ async fn mark_splits_for_deletion_cli(args: MarkForDeletionArgs) -> anyhow::Resu
     Ok(())
 }
 
-#[derive(Tabled)]
-struct FileRow {
-    #[tabled(rename = "File Name")]
-    file_name: String,
-    #[tabled(rename = "Size")]
-    size: String,
-}
-
 async fn describe_split_cli(args: DescribeSplitArgs) -> anyhow::Result<()> {
     debug!(args=?args, "describe-split");
     let qw_client = args.client_args.client();

quickwit/quickwit-cluster/src/metrics.rs

@@ -93,26 +93,31 @@ impl Default for ClusterMetrics {
                 "gossip_recv_messages_total",
                 "Total number of gossip messages received.",
                 "cluster",
+                &[],
             ),
             gossip_recv_bytes_total: new_counter(
                 "gossip_recv_bytes_total",
                 "Total amount of gossip data received in bytes.",
                 "cluster",
+                &[],
             ),
             gossip_sent_messages_total: new_counter(
                 "gossip_sent_messages_total",
                 "Total number of gossip messages sent.",
                 "cluster",
+                &[],
             ),
             gossip_sent_bytes_total: new_counter(
                 "gossip_sent_bytes_total",
                 "Total amount of gossip data sent in bytes.",
                 "cluster",
+                &[],
             ),
             grpc_gossip_rounds_total: new_counter(
                 "grpc_gossip_rounds_total",
                 "Total number of gRPC gossip rounds performed with peer nodes.",
                 "cluster",
+                &[],
             ),
         }
     }

quickwit/quickwit-common/Cargo.toml

@@ -37,6 +37,7 @@ siphasher = { workspace = true }
 tempfile = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true }
+tokio-metrics ={ workspace = true }
 tokio-stream = { workspace = true }
 tonic = { workspace = true }
 tower = { workspace = true }

quickwit/quickwit-common/src/metrics.rs

@@ -26,7 +26,7 @@ pub use prometheus::{
     IntCounter, IntCounterVec as PrometheusIntCounterVec, IntGauge,
     IntGaugeVec as PrometheusIntGaugeVec,
 };
-use prometheus::{Encoder, HistogramOpts, Opts, TextEncoder};
+use prometheus::{Encoder, Gauge, HistogramOpts, Opts, TextEncoder};
 
 #[derive(Clone)]
 pub struct HistogramVec<const N: usize> {
@@ -71,10 +71,20 @@ pub fn register_info(name: &'static str, help: &'static str, kvs: BTreeMap<&'sta
     prometheus::register(Box::new(counter)).expect("failed to register counter");
 }
 
-pub fn new_counter(name: &str, help: &str, subsystem: &str) -> IntCounter {
+pub fn new_counter(
+    name: &str,
+    help: &str,
+    subsystem: &str,
+    const_labels: &[(&str, &str)],
+) -> IntCounter {
+    let owned_const_labels: HashMap<String, String> = const_labels
+        .iter()
+        .map(|(label_name, label_value)| (label_name.to_string(), label_value.to_string()))
+        .collect();
     let counter_opts = Opts::new(name, help)
         .namespace("quickwit")
-        .subsystem(subsystem);
+        .subsystem(subsystem)
+        .const_labels(owned_const_labels);
     let counter = IntCounter::with_opts(counter_opts).expect("failed to create counter");
     prometheus::register(Box::new(counter.clone())).expect("failed to register counter");
     counter
@@ -104,6 +114,25 @@ pub fn new_counter_vec<const N: usize>(
     IntCounterVec { underlying }
 }
 
+pub fn new_float_gauge(
+    name: &str,
+    help: &str,
+    subsystem: &str,
+    const_labels: &[(&str, &str)],
+) -> Gauge {
+    let owned_const_labels: HashMap<String, String> = const_labels
+        .iter()
+        .map(|(label_name, label_value)| (label_name.to_string(), label_value.to_string()))
+        .collect();
+    let gauge_opts = Opts::new(name, help)
+        .namespace("quickwit")
+        .subsystem(subsystem)
+        .const_labels(owned_const_labels);
+    let gauge = Gauge::with_opts(gauge_opts).expect("failed to create float gauge");
+    prometheus::register(Box::new(gauge.clone())).expect("failed to register float gauge");
+    gauge
+}
+
 pub fn new_gauge(
     name: &str,
     help: &str,

quickwit/quickwit-common/src/runtimes.rs

@@ -19,9 +19,14 @@
 
 use std::collections::HashMap;
 use std::sync::atomic::{AtomicUsize, Ordering};
+use std::time::Duration;
 
 use once_cell::sync::OnceCell;
+use prometheus::{Gauge, IntCounter, IntGauge};
 use tokio::runtime::Runtime;
+use tokio_metrics::{RuntimeMetrics, RuntimeMonitor};
+
+use crate::metrics::{new_counter, new_float_gauge, new_gauge};
 
 static RUNTIMES: OnceCell<HashMap<RuntimeType, tokio::runtime::Runtime>> = OnceCell::new();
 
@@ -63,7 +68,7 @@ impl RuntimesConfig {
     }
 
     pub fn with_num_cpus(num_cpus: usize) -> Self {
-        // Non blocking task are supposed to be io intensive, and  not require many threads...
+        // Non blocking task are supposed to be io intensive, and not require many threads...
         let num_threads_non_blocking = if num_cpus > 6 { 2 } else { 1 };
         // On the other hand the blocking actors are cpu intensive. We allocate
         // almost all of the threads to them.
@@ -83,7 +88,8 @@ impl Default for RuntimesConfig {
 }
 
 fn start_runtimes(config: RuntimesConfig) -> HashMap<RuntimeType, Runtime> {
-    let mut runtimes = HashMap::default();
+    let mut runtimes = HashMap::with_capacity(2);
+
     let blocking_runtime = tokio::runtime::Builder::new_multi_thread()
         .worker_threads(config.num_threads_blocking)
         .thread_name_fn(|| {
@@ -94,7 +100,10 @@ fn start_runtimes(config: RuntimesConfig) -> HashMap<RuntimeType, Runtime> {
         .enable_all()
         .build()
         .unwrap();
+
+    scrape_tokio_runtime_metrics(blocking_runtime.handle(), "blocking");
     runtimes.insert(RuntimeType::Blocking, blocking_runtime);
+
     let non_blocking_runtime = tokio::runtime::Builder::new_multi_thread()
         .worker_threads(config.num_threads_non_blocking)
         .thread_name_fn(|| {
@@ -105,7 +114,10 @@ fn start_runtimes(config: RuntimesConfig) -> HashMap<RuntimeType, Runtime> {
         .enable_all()
         .build()
         .unwrap();
+
+    scrape_tokio_runtime_metrics(non_blocking_runtime.handle(), "non_blocking");
     runtimes.insert(RuntimeType::NonBlocking, non_blocking_runtime);
+
     runtimes
 }
 
@@ -135,6 +147,69 @@ impl RuntimeType {
     }
 }
 
+/// Spawns a background task
+pub fn scrape_tokio_runtime_metrics(handle: &tokio::runtime::Handle, label: &'static str) {
+    let runtime_monitor = RuntimeMonitor::new(handle);
+    handle.spawn(async move {
+        let mut interval = tokio::time::interval(Duration::from_secs(1));
+        let mut prometheus_runtime_metrics = PrometheusRuntimeMetrics::new(label);
+
+        for tokio_runtime_metrics in runtime_monitor.intervals() {
+            interval.tick().await;
+            prometheus_runtime_metrics.update(&tokio_runtime_metrics);
+        }
+    });
+}
+
+struct PrometheusRuntimeMetrics {
+    scheduled_tasks: IntGauge,
+    worker_busy_duration_milliseconds_total: IntCounter,
+    worker_busy_ratio: Gauge,
+    worker_threads: IntGauge,
+}
+
+impl PrometheusRuntimeMetrics {
+    pub fn new(label: &'static str) -> Self {
+        Self {
+            scheduled_tasks: new_gauge(
+                "tokio_scheduled_tasks",
+                "The total number of tasks currently scheduled in workers' local queues.",
+                "runtime",
+                &[("runtime_type", label)],
+            ),
+            worker_busy_duration_milliseconds_total: new_counter(
+                "tokio_worker_busy_duration_milliseconds_total",
+                " The total amount of time worker threads were busy.",
+                "runtime",
+                &[("runtime_type", label)],
+            ),
+            worker_busy_ratio: new_float_gauge(
+                "tokio_worker_busy_ratio",
+                "The ratio of time worker threads were busy since the last time runtime metrics \
+                 were collected.",
+                "runtime",
+                &[("runtime_type", label)],
+            ),
+            worker_threads: new_gauge(
+                "tokio_worker_threads",
+                "The number of worker threads used by the runtime.",
+                "runtime",
+                &[("runtime_type", label)],
+            ),
+        }
+    }
+
+    pub fn update(&mut self, runtime_metrics: &RuntimeMetrics) {
+        self.scheduled_tasks
+            .set(runtime_metrics.total_local_queue_depth as i64);
+        self.worker_busy_duration_milliseconds_total
+            .inc_by(runtime_metrics.total_busy_duration.as_millis() as u64);
+        self.worker_busy_ratio.set(runtime_metrics.busy_ratio());
+        self.worker_threads
+            .set(runtime_metrics.workers_count as i64);
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

quickwit/quickwit-common/src/thread_pool.rs

@@ -51,7 +51,7 @@ impl ThreadPool {
         }
         let thread_pool = rayon_pool_builder
             .build()
-            .expect("failed to spawn the spawning pool");
+            .expect("failed to spawn thread pool");
         let ongoing_tasks = THREAD_POOL_METRICS.ongoing_tasks.with_label_values([name]);
         let pending_tasks = THREAD_POOL_METRICS.pending_tasks.with_label_values([name]);
         ThreadPool {

quickwit/quickwit-control-plane/src/metrics.rs

@@ -65,23 +65,27 @@ impl Default for ControlPlaneMetrics {
                 "restart_total",
                 "Number of control plane restart.",
                 "control_plane",
+                &[],
             ),
             schedule_total: new_counter(
                 "schedule_total",
                 "Number of control plane `schedule` operations.",
                 "control_plane",
+                &[],
             ),
             metastore_error_aborted: new_counter(
                 "metastore_error_aborted",
                 "Number of aborted metastore transaction (= do not trigger a control plane \
                  restart)",
                 "control_plane",
+                &[],
             ),
             metastore_error_maybe_executed: new_counter(
                 "metastore_error_maybe_executed",
                 "Number of metastore transaction with an uncertain outcome (= do trigger a \
                  control plane restart)",
                 "control_plane",
+                &[],
             ),
             open_shards_total: new_gauge_vec(
                 "open_shards_total",

quickwit/quickwit-ingest/src/metrics.rs

@@ -35,21 +35,25 @@ impl Default for IngestMetrics {
                 "ingested_num_bytes",
                 "Total size of the docs ingested in bytes",
                 "ingest",
+                &[],
             ),
             ingested_num_docs: new_counter(
                 "ingested_num_docs",
                 "Number of docs received to be ingested",
                 "ingest",
+                &[],
             ),
             replicated_num_bytes_total: new_counter(
                 "replicated_num_bytes_total",
                 "Total size in bytes of the replicated docs.",
                 "ingest",
+                &[],
             ),
             replicated_num_docs_total: new_counter(
                 "replicated_num_docs_total",
                 "Total number of docs replicated.",
                 "ingest",
+                &[],
             ),
             queue_count: new_gauge(
                 "queue_count",

quickwit/quickwit-search/src/collector.rs

@@ -471,48 +471,6 @@ fn get_score_extractor(
     })
 }
 
-/// PartialHitHeapItem order is the inverse of the natural order
-/// so that we actually have a min-heap.
-#[derive(Clone, Copy, Debug)]
-struct PartialHitHeapItem {
-    sort_value_opt1: Option<u64>,
-    sort_value_opt2: Option<u64>,
-    doc_id: DocId,
-}
-
-impl PartialOrd for PartialHitHeapItem {
-    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
-        Some(self.cmp(other))
-    }
-}
-
-impl Ord for PartialHitHeapItem {
-    #[inline]
-    fn cmp(&self, other: &Self) -> Ordering {
-        let by_sorting_field1 = other.sort_value_opt1.cmp(&self.sort_value_opt1);
-        let by_sorting_field2 = other.sort_value_opt2.cmp(&self.sort_value_opt2);
-
-        let lazy_order_by_doc_id = || {
-            self.doc_id
-                .partial_cmp(&other.doc_id)
-                .unwrap_or(Ordering::Equal)
-        };
-
-        // In case of a tie on the feature, we sort by ascending `DocId`.
-        by_sorting_field1
-            .then_with(|| by_sorting_field2)
-            .then_with(lazy_order_by_doc_id)
-    }
-}
-
-impl PartialEq for PartialHitHeapItem {
-    fn eq(&self, other: &Self) -> bool {
-        self.cmp(other) == Ordering::Equal
-    }
-}
-
-impl Eq for PartialHitHeapItem {}
-
 enum AggregationSegmentCollectors {
     FindTraceIdsSegmentCollector(Box<FindTraceIdsSegmentCollector>),
     TantivyAggregationSegmentCollector(AggregationSegmentCollector),
@@ -1300,52 +1258,9 @@ mod tests {
     use tantivy::collector::Collector;
     use tantivy::TantivyDocument;
 
-    use super::{make_merge_collector, IncrementalCollector, PartialHitHeapItem};
+    use super::{make_merge_collector, IncrementalCollector};
     use crate::collector::top_k_partial_hits;
 
-    #[test]
-    fn test_partial_hit_ordered_by_sorting_field() {
-        let lesser_score = PartialHitHeapItem {
-            doc_id: 1u32,
-            sort_value_opt1: Some(1u64),
-            sort_value_opt2: None,
-        };
-        let higher_score = PartialHitHeapItem {
-            sort_value_opt1: Some(2u64),
-            sort_value_opt2: None,
-            doc_id: 1u32,
-        };
-        assert_eq!(lesser_score.cmp(&higher_score), Ordering::Greater);
-    }
-    #[test]
-    fn test_partial_hit_ordered_by_sorting_field_2() {
-        let get_el = |val1, val2, docid| PartialHitHeapItem {
-            doc_id: docid,
-            sort_value_opt1: val1,
-            sort_value_opt2: val2,
-        };
-        let mut data = vec![
-            get_el(Some(1u64), None, 1u32),
-            get_el(Some(2u64), Some(2u64), 1u32),
-            get_el(Some(2u64), Some(1u64), 1u32),
-            get_el(Some(2u64), None, 1u32),
-            get_el(None, Some(1u64), 1u32),
-            get_el(None, None, 1u32),
-        ];
-        data.sort();
-        assert_eq!(
-            data,
-            vec![
-                get_el(Some(2u64), Some(2u64), 1u32),
-                get_el(Some(2u64), Some(1u64), 1u32),
-                get_el(Some(2u64), None, 1u32),
-                get_el(Some(1u64), None, 1u32),
-                get_el(None, Some(1u64), 1u32),
-                get_el(None, None, 1u32),
-            ]
-        );
-    }
-
     #[test]
     fn test_merge_partial_hits_no_tie() {
         let make_doc = |sort_value: u64| PartialHit {

quickwit/quickwit-search/src/error.rs

@@ -156,9 +156,3 @@ impl From<std::convert::Infallible> for SearchError {
         match infallible {}
     }
 }
-
-#[derive(Debug)]
-pub struct NodeSearchError {
-    pub search_error: SearchError,
-    pub split_ids: Vec<String>,
-}

quickwit/quickwit-search/src/metrics.rs

@@ -36,6 +36,7 @@ impl Default for SearchMetrics {
                 "leaf_searches_splits_total",
                 "Number of leaf searches (count of splits) started.",
                 "search",
+                &[],
             ),
             leaf_search_split_duration_secs: new_histogram(
                 "leaf_search_split_duration_secs",

quickwit/quickwit-serve/src/ingest_api/rest_handler.rs

@@ -29,7 +29,6 @@ use quickwit_proto::ingest::router::{
 };
 use quickwit_proto::types::IndexId;
 use serde::Deserialize;
-use thiserror::Error;
 use warp::{Filter, Rejection};
 
 use crate::decompression::get_body_bytes;
@@ -50,12 +49,6 @@ pub struct IngestApi;
 )))]
 pub struct IngestApiSchemas;
 
-#[derive(Debug, Error)]
-#[error("request body contains invalid UTF-8 characters")]
-struct InvalidUtf8;
-
-impl warp::reject::Reject for InvalidUtf8 {}
-
 #[derive(Clone, Debug, Default, Deserialize, PartialEq)]
 struct IngestOptions {
     #[serde(alias = "commit")]

quickwit/quickwit-serve/src/otlp_api/rest_handler.rs

@@ -38,10 +38,6 @@ use warp::{Filter, Rejection};
 use crate::rest_api_response::into_rest_api_response;
 use crate::{require, with_arg, BodyFormat};
 
-#[derive(utoipa::OpenApi)]
-#[openapi(paths())]
-pub(crate) struct OtlpApi;
-
 /// Setup OpenTelemetry API handlers.
 pub(crate) fn otlp_ingest_api_handlers(
     otlp_logs_service: Option<OtlpGrpcLogsService>,

quickwit/quickwit-storage/src/metrics.rs

@@ -51,27 +51,32 @@ impl Default for StorageMetrics {
                 "object_storage_gets_total",
                 "Number of objects fetched.",
                 "storage",
+                &[],
             ),
             object_storage_put_total: new_counter(
                 "object_storage_puts_total",
                 "Number of objects uploaded. May differ from object_storage_requests_parts due to \
                  multipart upload.",
                 "storage",
+                &[],
             ),
             object_storage_put_parts: new_counter(
                 "object_storage_puts_parts",
                 "Number of object parts uploaded.",
                 "",
+                &[],
             ),
             object_storage_download_num_bytes: new_counter(
                 "object_storage_download_num_bytes",
                 "Amount of data downloaded from an object storage.",
                 "storage",
+                &[],
             ),
             object_storage_upload_num_bytes: new_counter(
                 "object_storage_upload_num_bytes",
                 "Amount of data uploaded to an object storage.",
                 "storage",
+                &[],
             ),
         }
     }
@@ -109,16 +114,19 @@ impl CacheMetrics {
                 "cache_hits_total",
                 "Number of {component_name} cache hits",
                 &namespace,
+                &[],
             ),
             hits_num_bytes: new_counter(
                 "cache_hits_bytes",
                 "Number of {component_name} cache hits in bytes",
                 &namespace,
+                &[],
             ),
             misses_num_items: new_counter(
                 "cache_misses_total",
                 "Number of {component_name} cache misses",
                 &namespace,
+                &[],
             ),
         }
     }