Merge pull request #31694 from bkirwi/tokio-metrics

Add a bunch of tokio metrics to prometheus

Author: bkirwi

src/clusterd/src/lib.rs

@@ -25,7 +25,7 @@ use mz_http_util::DynamicFilterTarget;
 use mz_orchestrator_tracing::{StaticTracingConfig, TracingCliArgs};
 use mz_ore::cli::{self, CliConfig};
 use mz_ore::error::ErrorExt;
-use mz_ore::metrics::MetricsRegistry;
+use mz_ore::metrics::{register_runtime_metrics, MetricsRegistry};
 use mz_ore::netio::{Listener, SocketAddr};
 use mz_ore::now::SYSTEM_TIME;
 use mz_persist_client::cache::PersistClientCache;
@@ -38,6 +38,7 @@ use mz_storage::storage_state::StorageInstanceContext;
 use mz_storage_client::client::proto_storage_server::ProtoStorageServer;
 use mz_storage_types::connections::ConnectionContext;
 use mz_txn_wal::operator::TxnsContext;
+use tokio::runtime::Handle;
 use tower::Service;
 use tracing::{error, info};
 
@@ -184,6 +185,7 @@ async fn run(args: Args) -> Result<(), anyhow::Error> {
         .await?;
 
     let tracing_handle = Arc::new(tracing_handle);
+    register_runtime_metrics("main", Handle::current().metrics(), &metrics_registry);
 
     // Keep this _after_ the mz_ore::tracing::configure call so that its panic
     // hook runs _before_ the one that sends things to sentry.

src/environmentd/src/environmentd/main.rs

@@ -54,7 +54,7 @@ use mz_orchestrator_tracing::{StaticTracingConfig, TracingCliArgs, TracingOrches
 use mz_ore::cli::{self, CliConfig, KeyValueArg};
 use mz_ore::error::ErrorExt;
 use mz_ore::metric;
-use mz_ore::metrics::MetricsRegistry;
+use mz_ore::metrics::{register_runtime_metrics, MetricsRegistry};
 use mz_ore::now::SYSTEM_TIME;
 use mz_ore::task::RuntimeExt;
 use mz_ore::url::SensitiveUrl;
@@ -724,6 +724,7 @@ fn run(mut args: Args) -> Result<(), anyhow::Error> {
         },
         metrics_registry.clone(),
     ))?;
+    register_runtime_metrics("main", runtime.metrics(), &metrics_registry);
 
     let span = tracing::info_span!("environmentd::run").entered();
 

src/ore/src/metrics.rs

@@ -182,10 +182,10 @@ pub type IntGaugeVec = DeleteOnDropWrapper<prometheus::IntGaugeVec>;
 /// Delete-on-drop shadow of Prometheus [raw::UIntGaugeVec].
 pub type UIntGaugeVec = DeleteOnDropWrapper<raw::UIntGaugeVec>;
 
-pub use prometheus::{Counter, Histogram, IntCounter, IntGauge};
-
 use crate::assert_none;
 
+pub use prometheus::{Counter, Histogram, IntCounter, IntGauge};
+
 /// Access to non-delete-on-drop vector types
 pub mod raw {
     use prometheus::core::{AtomicU64, GenericGaugeVec};
@@ -217,13 +217,12 @@ impl MetricsRegistry {
     }
 
     /// Registers a gauge whose value is computed when observed.
-    pub fn register_computed_gauge<F, P>(
+    pub fn register_computed_gauge<P>(
         &self,
         opts: MakeCollectorOpts,
-        f: F,
+        f: impl Fn() -> P::T + Send + Sync + 'static,
     ) -> ComputedGenericGauge<P>
     where
-        F: Fn() -> P::T + Send + Sync + 'static,
         P: Atomic + 'static,
     {
         let gauge = ComputedGenericGauge {
@@ -769,6 +768,65 @@ impl DurationMetric for &'_ mut f64 {
     }
 }
 
+/// Register the Tokio runtime's metrics in our metrics registry.
+#[cfg(feature = "async")]
+pub fn register_runtime_metrics(
+    name: &'static str,
+    runtime_metrics: tokio::runtime::RuntimeMetrics,
+    registry: &MetricsRegistry,
+) {
+    macro_rules! register {
+        ($method:ident, $doc:literal) => {
+            let metrics = runtime_metrics.clone();
+            registry.register_computed_gauge::<prometheus::core::AtomicU64>(
+                crate::metric!(
+                    name: concat!("mz_tokio_", stringify!($method)),
+                    help: $doc,
+                    const_labels: {"runtime" => name},
+                ),
+                move || <u64 as crate::cast::CastFrom<_>>::cast_from(metrics.$method()),
+            );
+        };
+    }
+
+    register!(
+        num_workers,
+        "The number of worker threads used by the runtime."
+    );
+    register!(
+        num_alive_tasks,
+        "The current number of alive tasks in the runtime."
+    );
+    register!(
+        global_queue_depth,
+        "The number of tasks currently scheduled in the runtime's global queue."
+    );
+    register!(
+        num_blocking_threads,
+        "The number of additional threads spawned by the runtime."
+    );
+    register!(
+        num_idle_blocking_threads,
+        "The number of idle threads which have spawned by the runtime for spawn_blocking calls."
+    );
+    register!(
+        spawned_tasks_count,
+        "The number of tasks spawned in this runtime since it was created."
+    );
+    register!(
+        remote_schedule_count,
+        "The number of tasks scheduled from outside of the runtime."
+    );
+    register!(
+        budget_forced_yield_count,
+        "The number of times that tasks have been forced to yield back to the scheduler after exhausting their task budgets."
+    );
+    register!(
+        blocking_queue_depth,
+        "The number of tasks currently scheduled in the blocking thread pool, spawned using spawn_blocking."
+    );
+}
+
 #[cfg(test)]
 mod tests {
     use std::time::Duration;

src/persist-client/src/async_runtime.rs

@@ -12,6 +12,7 @@
 use std::future::Future;
 use std::sync::atomic::{AtomicUsize, Ordering};
 
+use mz_ore::metrics::{register_runtime_metrics, MetricsRegistry};
 use mz_ore::task::{JoinHandle, RuntimeExt};
 use tokio::runtime::{Builder, Runtime};
 
@@ -41,9 +42,12 @@ pub struct IsolatedRuntime {
 
 impl IsolatedRuntime {
     /// Creates a new isolated runtime.
-    pub fn new(worker_threads: usize) -> IsolatedRuntime {
-        let runtime = Builder::new_multi_thread()
-            .worker_threads(worker_threads)
+    pub fn new(metrics: &MetricsRegistry, worker_threads: Option<usize>) -> IsolatedRuntime {
+        let mut runtime = Builder::new_multi_thread();
+        if let Some(worker_threads) = worker_threads {
+            runtime.worker_threads(worker_threads);
+        }
+        let runtime = runtime
             .thread_name_fn(|| {
                 static ATOMIC_ID: AtomicUsize = AtomicUsize::new(0);
                 let id = ATOMIC_ID.fetch_add(1, Ordering::SeqCst);
@@ -54,6 +58,7 @@ impl IsolatedRuntime {
             .enable_all()
             .build()
             .expect("known to be valid");
+        register_runtime_metrics("persist", runtime.metrics(), metrics);
         IsolatedRuntime {
             inner: Some(runtime),
         }
@@ -79,7 +84,7 @@ impl IsolatedRuntime {
 
 impl Default for IsolatedRuntime {
     fn default() -> Self {
-        IsolatedRuntime::new(num_cpus::get())
+        IsolatedRuntime::new(&MetricsRegistry::new(), None)
     }
 }
 

src/persist-client/src/cache.rs

@@ -86,7 +86,8 @@ impl PersistClientCache {
             Arc::clone(&state_cache),
             pubsub_client.receiver,
         );
-        let isolated_runtime = IsolatedRuntime::new(cfg.isolated_runtime_worker_threads);
+        let isolated_runtime =
+            IsolatedRuntime::new(registry, Some(cfg.isolated_runtime_worker_threads));
 
         PersistClientCache {
             cfg,

src/persist-client/src/cli/admin.rs

@@ -223,7 +223,7 @@ pub async fn run(command: AdminArgs) -> Result<(), anyhow::Error> {
             .await?;
 
             if force_downgrade_upper {
-                let isolated_runtime = Arc::new(IsolatedRuntime::default());
+                let isolated_runtime = Arc::new(IsolatedRuntime::new(&metrics_registry, None));
                 let pubsub_sender: Arc<dyn PubSubSender> = Arc::new(NoopPubSubSender);
                 let shared_states = Arc::new(StateCache::new(
                     &cfg,

src/persist-client/src/cli/inspect.rs

@@ -647,7 +647,7 @@ pub async fn blob_usage(args: &StateArgs) -> Result<(), anyhow::Error> {
     let consensus =
         make_consensus(&cfg, &args.consensus_uri, NO_COMMIT, Arc::clone(&metrics)).await?;
     let blob = make_blob(&cfg, &args.blob_uri, NO_COMMIT, Arc::clone(&metrics)).await?;
-    let isolated_runtime = Arc::new(IsolatedRuntime::default());
+    let isolated_runtime = Arc::new(IsolatedRuntime::new(&metrics_registry, None));
     let state_cache = Arc::new(StateCache::new(
         &cfg,
         Arc::clone(&metrics),