[Metrics] Create a thread local version to be used in hot paths

wqfish · wqfish · commit 3ffc93867eb1 · 2025-09-04T18:01:28.000-07:00
We have lots of counters in the codebase. Often, the way we use them is that we
call `with_label_values` and then record values.

`with_label_values` would need to acquire a read lock, which is usually not
expensive but can be nontrivial in hot paths if there are many threads.

This change creates an alternative which is a thread local version. The caller
would just record the value in the thread local counter/histogram, and flush
periodically. This should reduce the contention and speed things up a little.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/aptos-move/block-executor/src/counters.rs b/aptos-move/block-executor/src/counters.rs
@@ -241,7 +241,7 @@ pub static BLOCK_VIEW_BASE_VALUES_MEMORY_USAGE: Lazy<HistogramVec> = Lazy::new(|
     )
 });
 
-fn observe_gas(counter: &Lazy<HistogramVec>, mode_str: &str, fee_statement: &FeeStatement) {
+fn observe_gas(counter: &'static Lazy<HistogramVec>, mode_str: &str, fee_statement: &FeeStatement) {
     counter.observe_with(
         &[mode_str, GasType::TOTAL_GAS],
         fee_statement.gas_used() as f64,
diff --git a/crates/aptos-metrics-core/Cargo.toml b/crates/aptos-metrics-core/Cargo.toml
@@ -14,6 +14,8 @@ rust-version = { workspace = true }
 
 [dependencies]
 anyhow = { workspace = true }
+once_cell = { workspace = true }
+paste = { workspace = true }
 prometheus = { workspace = true }
 
 [dev-dependencies]
diff --git a/crates/aptos-metrics-core/src/lib.rs b/crates/aptos-metrics-core/src/lib.rs
@@ -2,6 +2,10 @@
 // Parts of the project are originally copyright © Meta Platforms, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
+pub use crate::{
+    avg_counter::{register_avg_counter, register_avg_counter_vec},
+    thread_local::{ThreadLocalHistogramVec, ThreadLocalIntCounter, ThreadLocalIntCounterVec},
+};
 // Re-export counter types from prometheus crate
 pub use prometheus::{
     exponential_buckets, gather, histogram_opts, register_counter, register_gauge,
@@ -12,23 +16,29 @@ pub use prometheus::{
 };
 
 mod avg_counter;
-pub use avg_counter::{register_avg_counter, register_avg_counter_vec};
 pub mod const_metric;
 pub mod op_counters;
+pub mod thread_local;
 
 pub trait TimerHelper {
-    fn timer_with(&self, labels: &[&str]) -> HistogramTimer;
+    type TimerType<'a>
+    where
+        Self: 'a;
+
+    fn timer_with<'a>(&'static self, labels: &'a [&str]) -> Self::TimerType<'a>;
 
-    fn observe_with(&self, labels: &[&str], val: f64);
+    fn observe_with(&'static self, labels: &[&str], val: f64);
 }
 
 impl TimerHelper for HistogramVec {
-    fn timer_with(&self, vals: &[&str]) -> HistogramTimer {
-        self.with_label_values(vals).start_timer()
+    type TimerType<'a> = HistogramTimer;
+
+    fn timer_with<'a>(&'static self, labels: &'a [&str]) -> Self::TimerType<'a> {
+        self.with_label_values(labels).start_timer()
     }
 
-    fn observe_with(&self, labels: &[&str], val: f64) {
-        self.with_label_values(labels).observe(val)
+    fn observe_with(&'static self, labels: &[&str], val: f64) {
+        self.with_label_values(labels).observe(val);
     }
 }
 
@@ -68,19 +78,27 @@ impl IntGaugeVecHelper for IntGaugeVec {
 pub trait IntCounterVecHelper {
     type IntType;
 
-    fn inc_with(&self, labels: &[&str]);
+    fn inc_with(&'static self, labels: &[&str]);
 
-    fn inc_with_by(&self, labels: &[&str], by: Self::IntType);
+    fn inc_with_by(&'static self, labels: &[&str], by: Self::IntType);
 }
 
 impl IntCounterVecHelper for IntCounterVec {
     type IntType = u64;
 
-    fn inc_with(&self, labels: &[&str]) {
+    fn inc_with(&'static self, labels: &[&str]) {
         self.with_label_values(labels).inc()
     }
 
-    fn inc_with_by(&self, labels: &[&str], v: Self::IntType) {
+    fn inc_with_by(&'static self, labels: &[&str], v: Self::IntType) {
         self.with_label_values(labels).inc_by(v)
     }
 }
+
+pub trait IntCounterHelper {
+    type IntType;
+
+    fn inc(&'static self);
+
+    fn inc_by(&'static self, v: Self::IntType);
+}
diff --git a/crates/aptos-metrics-core/src/thread_local.rs b/crates/aptos-metrics-core/src/thread_local.rs
@@ -0,0 +1,285 @@
+// Copyright © Aptos Foundation
+// SPDX-License-Identifier: Apache-2.0
+
+pub mod __private {
+    pub use once_cell::sync::Lazy;
+    pub use paste::paste;
+}
+
+use crate::{IntCounterHelper, IntCounterVecHelper, TimerHelper};
+use std::{
+    cell::RefCell,
+    thread::LocalKey,
+    time::{Duration, Instant},
+};
+
+const FLUSH_INTERVAL: Duration = Duration::from_secs(1);
+
+pub struct ThreadLocalIntCounter {
+    inner: prometheus::local::LocalIntCounter,
+    last_flush: Instant,
+}
+
+impl ThreadLocalIntCounter {
+    pub fn new(shared: &prometheus::IntCounter) -> Self {
+        Self {
+            inner: shared.local(),
+            last_flush: Instant::now(),
+        }
+    }
+
+    fn maybe_flush(&mut self) {
+        let now = Instant::now();
+        if now.duration_since(self.last_flush) > FLUSH_INTERVAL {
+            self.inner.flush();
+        }
+        self.last_flush = now;
+    }
+}
+
+impl IntCounterHelper for LocalKey<RefCell<ThreadLocalIntCounter>> {
+    type IntType = u64;
+
+    fn inc(&'static self) {
+        self.inc_by(1);
+    }
+
+    fn inc_by(&'static self, v: Self::IntType) {
+        self.with_borrow_mut(|x| {
+            x.inner.inc_by(v);
+            x.maybe_flush();
+        })
+    }
+}
+
+pub struct ThreadLocalIntCounterVec {
+    inner: prometheus::local::LocalIntCounterVec,
+    last_flush: Instant,
+}
+
+impl ThreadLocalIntCounterVec {
+    pub fn new(shared: &prometheus::IntCounterVec) -> Self {
+        Self {
+            inner: shared.local(),
+            last_flush: Instant::now(),
+        }
+    }
+
+    fn maybe_flush(&mut self) {
+        let now = Instant::now();
+        if now.duration_since(self.last_flush) > FLUSH_INTERVAL {
+            self.inner.flush();
+        }
+        self.last_flush = now;
+    }
+}
+
+impl IntCounterVecHelper for LocalKey<RefCell<ThreadLocalIntCounterVec>> {
+    type IntType = u64;
+
+    fn inc_with(&'static self, labels: &[&str]) {
+        self.inc_with_by(labels, 1);
+    }
+
+    fn inc_with_by(&'static self, labels: &[&str], v: Self::IntType) {
+        self.with_borrow_mut(|x| {
+            x.inner.with_label_values(labels).inc_by(v);
+            x.maybe_flush();
+        });
+    }
+}
+
+pub struct ThreadLocalHistogramTimer<'a> {
+    start: Instant,
+    labels: &'a [&'a str],
+    parent: &'static LocalKey<RefCell<ThreadLocalHistogramVec>>,
+}
+
+impl<'a> ThreadLocalHistogramTimer<'a> {
+    fn new(
+        labels: &'a [&'a str],
+        parent: &'static LocalKey<RefCell<ThreadLocalHistogramVec>>,
+    ) -> Self {
+        Self {
+            start: Instant::now(),
+            labels,
+            parent,
+        }
+    }
+}
+
+impl<'a> Drop for ThreadLocalHistogramTimer<'a> {
+    fn drop(&mut self) {
+        self.parent
+            .observe_with(self.labels, self.start.elapsed().as_secs_f64());
+    }
+}
+
+pub struct ThreadLocalHistogramVec {
+    inner: prometheus::local::LocalHistogramVec,
+    last_flush: Instant,
+}
+
+impl ThreadLocalHistogramVec {
+    pub fn new(shared: &prometheus::HistogramVec) -> Self {
+        Self {
+            inner: shared.local(),
+            last_flush: Instant::now(),
+        }
+    }
+
+    fn maybe_flush(&mut self) {
+        let now = Instant::now();
+        if now.duration_since(self.last_flush) > FLUSH_INTERVAL {
+            self.inner.flush();
+        }
+        self.last_flush = now;
+    }
+}
+
+impl TimerHelper for LocalKey<RefCell<ThreadLocalHistogramVec>> {
+    type TimerType<'a> = ThreadLocalHistogramTimer<'a>;
+
+    fn timer_with<'a>(&'static self, labels: &'a [&str]) -> Self::TimerType<'a> {
+        // We could use `self.with_borrow_mut(|x| x.inner.with_label_values(labels).start_timer())`.
+        // However, this creates a `LocalHistogramTimer`, which internally stores a copy of
+        // `LocalHistogram`:
+        // https://github.com/tikv/rust-prometheus/blob/1d3174bf5ddf056dcb0fe59e06cad4ef42ebec68/src/histogram.rs#L1077-L1080.
+        // When the timer is dropped, the copied `LocalHistogram` is also dropped, and this always
+        // causes a flush:
+        // https://github.com/tikv/rust-prometheus/blob/1d3174bf5ddf056dcb0fe59e06cad4ef42ebec68/src/histogram.rs#L1142-L1146
+        ThreadLocalHistogramTimer::new(labels, self)
+    }
+
+    fn observe_with(&'static self, labels: &[&str], val: f64) {
+        self.with_borrow_mut(|x| {
+            x.inner.with_label_values(labels).observe(val);
+            x.maybe_flush();
+        });
+    }
+}
+
+#[macro_export]
+macro_rules! make_thread_local_int_counter {
+    (
+        $(#[$attr:meta])*
+        $vis:vis,
+        $var_name:ident,
+        $name:expr,
+        $help:expr $(,)?
+    ) => {
+        $crate::thread_local::__private::paste! {
+            static [<__ $var_name>]: $crate::thread_local::__private::Lazy<$crate::IntCounter> =
+                $crate::thread_local::__private::Lazy::new(|| {
+                    $crate::register_int_counter!($name, $help)
+                        .expect("register_int_counter should succeed")
+                });
+            ::std::thread_local! {
+                $(#[$attr])*
+                $vis static $var_name: ::std::cell::RefCell<$crate::thread_local::ThreadLocalIntCounter> =
+                    ::std::cell::RefCell::new(
+                        $crate::thread_local::ThreadLocalIntCounter::new(&[<__ $var_name>]),
+                    );
+            }
+        }
+    }
+}
+
+#[macro_export]
+macro_rules! make_thread_local_int_counter_vec {
+    (
+        $(#[$attr:meta])*
+        $vis:vis,
+        $var_name:ident,
+        $name:expr,
+        $help:expr,
+        $labels_names:expr $(,)?
+    ) => {
+        $crate::thread_local::__private::paste! {
+            static [<__ $var_name>]: $crate::thread_local::__private::Lazy<$crate::IntCounterVec> =
+                $crate::thread_local::__private::Lazy::new(|| {
+                    $crate::register_int_counter_vec!($name, $help, $labels_names)
+                        .expect("register_int_counter_vec should succeed")
+                });
+            ::std::thread_local! {
+                $(#[$attr])*
+                $vis static $var_name: ::std::cell::RefCell<$crate::thread_local::ThreadLocalIntCounterVec> =
+                    ::std::cell::RefCell::new(
+                        $crate::thread_local::ThreadLocalIntCounterVec::new(&[<__ $var_name>]),
+                    );
+            }
+        }
+    }
+}
+
+#[macro_export]
+macro_rules! make_thread_local_histogram_vec {
+    (
+        $(#[$attr:meta])*
+        $vis:vis,
+        $var_name:ident,
+        $name:expr,
+        $help:expr,
+        $labels_names:expr
+        $(, $buckets:expr)? $(,)?
+    ) => {
+        $crate::thread_local::__private::paste! {
+            static [<__ $var_name>]: $crate::thread_local::__private::Lazy<$crate::HistogramVec> =
+                $crate::thread_local::__private::Lazy::new(|| {
+                    $crate::register_histogram_vec!($name, $help, $labels_names $(, $buckets)?)
+                        .expect("register_histogram_vec should succeed")
+                });
+            ::std::thread_local! {
+                $(#[$attr])*
+                $vis static $var_name: ::std::cell::RefCell<$crate::thread_local::ThreadLocalHistogramVec> =
+                    ::std::cell::RefCell::new(
+                        $crate::thread_local::ThreadLocalHistogramVec::new(&[<__ $var_name>]),
+                    );
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{IntCounterHelper, IntCounterVecHelper, TimerHelper};
+
+    make_thread_local_int_counter!(
+        pub(self),
+        TEST_INT_COUNTER,
+        "aptos_test_int_counter",
+        "this is a help message",
+    );
+    make_thread_local_int_counter_vec!(
+        pub(self),
+        TEST_INT_COUNTER_VEC,
+        "aptos_test_int_counter_vec",
+        "this is a help message",
+        &["label"],
+    );
+    make_thread_local_histogram_vec!(
+        pub(self),
+        TEST_HISTOGRAM_VEC,
+        "aptos_test_histogram_vec",
+        "this is a help message",
+        &["label"],
+    );
+
+    #[test]
+    fn test_thread_local_int_counter() {
+        TEST_INT_COUNTER.inc();
+        TEST_INT_COUNTER.inc_by(2);
+    }
+
+    #[test]
+    fn test_thread_local_int_counter_vec() {
+        TEST_INT_COUNTER_VEC.inc_with(&["foo"]);
+        TEST_INT_COUNTER_VEC.inc_with_by(&["foo"], 2);
+    }
+
+    #[test]
+    fn test_thread_local_histogram_vec() {
+        let _timer = TEST_HISTOGRAM_VEC.timer_with(&["foo"]);
+        TEST_HISTOGRAM_VEC.observe_with(&["bar"], 1.0);
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -241,7 +241,7 @@ pub static BLOCK_VIEW_BASE_VALUES_MEMORY_USAGE: Lazy<HistogramVec> = Lazy::new(\|`
`241`	`241`	`)`
`242`	`242`	`});`
`243`	`243`
`244`		`-fn observe_gas(counter: &Lazy<HistogramVec>, mode_str: &str, fee_statement: &FeeStatement) {`
	`244`	`+fn observe_gas(counter: &'static Lazy<HistogramVec>, mode_str: &str, fee_statement: &FeeStatement) {`
`245`	`245`	`counter.observe_with(`
`246`	`246`	`&[mode_str, GasType::TOTAL_GAS],`
`247`	`247`	`fee_statement.gas_used() as f64,`