/build/source/nativelink-util/src/metrics_utils.rs

Source
// Copyright 2024 The NativeLink Authors. All rights reserved.
//
// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    See LICENSE file for details
//
// 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.

use core::mem::forget;
use core::sync::atomic::{AtomicU64, Ordering};
use std::time::{Instant, SystemTime, UNIX_EPOCH};

use futures::Future;
use nativelink_metric::{
    MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent, group, publish,
};

#[derive(Debug, Default)]
pub struct FuncCounterWrapper {
    pub successes: AtomicU64,
    pub failures: AtomicU64,
}

impl FuncCounterWrapper {
    #[inline]
    pub fn wrap<T, E>(&self, func: impl FnOnce() -> Result<T, E>) -> Result<T, E> {
        let result = (func)();
        if result.is_ok() {
            self.successes.fetch_add(1, Ordering::Acquire);
        } else {
            self.failures.fetch_add(1, Ordering::Acquire);
        }
        result
    }
}

// Derive-macros have no way to tell the collector that the parent
// is now a group with the name of the group as the field so we
// can attach multiple values on the same group, so we need to
// manually implement the `MetricsComponent` trait to do so.
impl MetricsComponent for FuncCounterWrapper {
    fn publish(
        &self,
        _kind: MetricKind,
        field_metadata: MetricFieldData,
    ) -> Result<MetricPublishKnownKindData, nativelink_metric::Error> {
        let _enter = group!(field_metadata.name).entered();

        publish!(
            "successes",
            &self.successes,
            MetricKind::Counter,
            format!(
                "The number of times {} was successful.",
                field_metadata.name
            )
        );
        publish!(
            "failures",
            &self.failures,
            MetricKind::Counter,
            format!("The number of times {} failed.", field_metadata.name)
        );

        Ok(MetricPublishKnownKindData::Component)
    }
}

/// This is a utility that will only increment the referenced counter when it is dropped.
/// This struct is zero cost and has a runtime cost only when it is dropped.
/// This struct is very useful for tracking when futures are dropped.
#[derive(Debug)]
struct DropCounter<'a> {
    counter: &'a AtomicU64,
}

impl<'a> DropCounter<'a> {
    #[inline]
    pub(crate) const fn new(counter: &'a AtomicU64) -> Self {
        Self { counter }
    }
}

impl Drop for DropCounter<'_> {
    #[inline]
    fn drop(&mut self) {
        self.counter.fetch_add(1, Ordering::Acquire);
    }
}

#[derive(Debug)]
pub struct AsyncTimer<'a> {
    start: Instant,
    drop_counter: DropCounter<'a>,
    counter: &'a AsyncCounterWrapper,
}

impl AsyncTimer<'_> {
    #[inline]
    pub fn measure(self) {
        self.counter.sum_func_duration_ns.fetch_add(
            u64::try_from(self.start.elapsed().as_nanos()).unwrap_or(u64::MAX),
            Ordering::Acquire,
        );
        self.counter.calls.fetch_add(1, Ordering::Acquire);
        self.counter.successes.fetch_add(1, Ordering::Acquire);
        // This causes DropCounter's drop to never be called.
        forget(self.drop_counter);
    }
}

/// Tracks the number of calls, successes, failures, and drops of an async function.
/// call `.wrap(future)` to wrap a future and stats about the future are automatically
/// tracked and can be published to a `CollectorState`.
#[derive(Debug, Default)]
pub struct AsyncCounterWrapper {
    pub calls: AtomicU64,
    pub successes: AtomicU64,
    pub failures: AtomicU64,
    pub drops: AtomicU64,
    // Time spent in nanoseconds in the future.
    // 64 bit address space gives ~584 years of nanoseconds.
    pub sum_func_duration_ns: AtomicU64,
}

// Derive-macros have no way to tell the collector that the parent
// is now a group with the name of the group as the field so we
// can attach multiple values on the same group, so we need to
// manually implement the `MetricsComponent` trait to do so.
#[expect(
    clippy::cognitive_complexity,
    reason = "complexity arises from macro expansion"
)]
impl MetricsComponent for AsyncCounterWrapper {
    fn publish(
        &self,
        _kind: MetricKind,
        field_metadata: MetricFieldData,
    ) -> Result<MetricPublishKnownKindData, nativelink_metric::Error> {
        let _enter = group!(field_metadata.name).entered();

        publish!(
            "calls",
            &self.calls,
            MetricKind::Counter,
            format!("The number of times {} was called.", field_metadata.name)
        );
        publish!(
            "successes",
            &self.successes,
            MetricKind::Counter,
            format!(
                "The number of times {} was successful.",
                field_metadata.name
            )
        );
        publish!(
            "failures",
            &self.failures,
            MetricKind::Counter,
            format!("The number of times {} failed.", field_metadata.name)
        );
        publish!(
            "drops",
            &self.drops,
            MetricKind::Counter,
            format!("The number of times {} was dropped.", field_metadata.name)
        );
        publish!(
            "sum_func_duration_ns",
            &self.sum_func_duration_ns,
            MetricKind::Counter,
            format!(
                "The sum of the time spent in nanoseconds in {}.",
                field_metadata.name
            )
        );

        Ok(MetricPublishKnownKindData::Component)
    }
}

impl AsyncCounterWrapper {
    #[inline]
    pub fn wrap_fn<'a, T: 'a, E>(
        &'a self,
        func: impl FnOnce() -> Result<T, E> + 'a,
    ) -> Result<T, E> {
        self.calls.fetch_add(1, Ordering::Acquire);
        let result = (func)();
        if result.is_ok() {
            self.successes.fetch_add(1, Ordering::Acquire);
        } else {
            self.failures.fetch_add(1, Ordering::Acquire);
        }
        result
    }

    #[inline]
    pub async fn wrap<'a, T, E, F: Future<Output = Result<T, E>> + 'a>(
        &'a self,
        future: F,
    ) -> Result<T, E> {
        let result = self.wrap_no_capture_result(future).await;
        if result.is_ok() {
            self.successes.fetch_add(1, Ordering::Acquire);
        } else {
            self.failures.fetch_add(1, Ordering::Acquire);
        }
        result
    }

    #[inline]
    pub async fn wrap_no_capture_result<'a, T, F: Future<Output = T> + 'a>(
        &'a self,
        future: F,
    ) -> T {
        self.calls.fetch_add(1, Ordering::Acquire);
        let drop_counter = DropCounter::new(&self.drops);
        let instant = Instant::now();
        let result = future.await;
        // By default `drop_counter` will increment the drop counter when it goes out of scope.
        // This will ensure we don't increment the counter if we make it here with a zero cost.
        forget(drop_counter);
        self.sum_func_duration_ns.fetch_add(
            u64::try_from(instant.elapsed().as_nanos()).unwrap_or(u64::MAX),
            Ordering::Acquire,
        );
        result
    }

    #[inline]
    pub fn begin_timer(&self) -> AsyncTimer<'_> {
        AsyncTimer {
            start: Instant::now(),
            drop_counter: DropCounter::new(&self.drops),
            counter: self,
        }
    }
}

/// Tracks a number.
#[derive(Debug, Default)]
pub struct Counter(AtomicU64);

impl Counter {
    #[inline]
    pub fn inc(&self) {
        self.add(1);
    }

    #[inline]
    pub fn add(&self, value: u64) {
        self.0.fetch_add(value, Ordering::Acquire);
    }

    #[inline]
    pub fn sub(&self, value: u64) {
        self.0.fetch_sub(value, Ordering::Acquire);
    }
}

impl MetricsComponent for Counter {
    fn publish(
        &self,
        kind: MetricKind,
        field_metadata: MetricFieldData,
    ) -> Result<MetricPublishKnownKindData, nativelink_metric::Error> {
        self.0.publish(kind, field_metadata)
    }
}

/// Tracks an counter through time and the last time the counter was changed.
#[derive(Debug, Default)]
pub struct CounterWithTime {
    pub counter: AtomicU64,
    pub last_time: AtomicU64,
}

impl CounterWithTime {
    #[inline]
    pub fn inc(&self) {
        self.counter.fetch_add(1, Ordering::Acquire);
        self.last_time.store(
            SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_secs(),
            Ordering::Release,
        );
    }
}

// Derive-macros have no way to tell the collector that the parent
// is now a group with the name of the group as the field so we
// can attach multiple values on the same group, so we need to
// manually implement the `MetricsComponent` trait to do so.
impl MetricsComponent for CounterWithTime {
    fn publish(
        &self,
        _kind: MetricKind,
        field_metadata: MetricFieldData,
    ) -> Result<MetricPublishKnownKindData, nativelink_metric::Error> {
        let _enter = group!(field_metadata.name).entered();

        publish!(
            "counter",
            &self.counter,
            MetricKind::Counter,
            format!("Current count of {}.", field_metadata.name)
        );
        publish!(
            "last_time",
            &self.last_time,
            MetricKind::Counter,
            format!("Last timestamp {} was published.", field_metadata.name)
        );

        Ok(MetricPublishKnownKindData::Component)
    }
}

Coverage Report

Created: 2025-10-20 10:50

Line	Count	Source
1		// Copyright 2024 The NativeLink Authors. All rights reserved.
2		//
3		// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
4		// you may not use this file except in compliance with the License.
5		// You may obtain a copy of the License at
6		//
7		// See LICENSE file for details
8		//
9		// Unless required by applicable law or agreed to in writing, software
10		// distributed under the License is distributed on an "AS IS" BASIS,
11		// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12		// See the License for the specific language governing permissions and
13		// limitations under the License.
14
15		use core::mem::forget;
16		use core::sync::atomic::{AtomicU64, Ordering};
17		use std::time::{Instant, SystemTime, UNIX_EPOCH};
18
19		use futures::Future;
20		use nativelink_metric::{
21		MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent, group, publish,
22		};
23
24		#[derive(Debug, Default)]
25		pub struct FuncCounterWrapper {
26		pub successes: AtomicU64,
27		pub failures: AtomicU64,
28		}
29
30		impl FuncCounterWrapper {
31		#[inline]
32	1	pub fn wrap<T, E>(&self, func: impl FnOnce() -> Result<T, E>) -> Result<T, E> {
33	1	let result = (func)();
34	1	if result.is_ok() { Branch (34:12): [True: 1, False: 0] Branch (34:12): [Folded - Ignored] Branch (34:12): [Folded - Ignored]
35	1	self.successes.fetch_add(1, Ordering::Acquire);
36	1	} else {
37	0	self.failures.fetch_add(1, Ordering::Acquire);
38	0	}
39	1	result
40	1	}
41		}
42
43		// Derive-macros have no way to tell the collector that the parent
44		// is now a group with the name of the group as the field so we
45		// can attach multiple values on the same group, so we need to
46		// manually implement the `MetricsComponent` trait to do so.
47		impl MetricsComponent for FuncCounterWrapper {
48	0	fn publish(
49	0	&self,
50	0	_kind: MetricKind,
51	0	field_metadata: MetricFieldData,
52	0	) -> Result<MetricPublishKnownKindData, nativelink_metric::Error> {
53	0	let _enter = group!(field_metadata.name).entered();
54
55	0	publish!(
56	0	"successes",
57	0	&self.successes,
58	0	MetricKind::Counter,
59	0	format!(
60	0	"The number of times {} was successful.",
61		field_metadata.name
62		)
63		);
64	0	publish!(
65	0	"failures",
66	0	&self.failures,
67	0	MetricKind::Counter,
68	0	format!("The number of times {} failed.", field_metadata.name)
69		);
70
71	0	Ok(MetricPublishKnownKindData::Component)
72	0	}
73		}
74
75		/// This is a utility that will only increment the referenced counter when it is dropped.
76		/// This struct is zero cost and has a runtime cost only when it is dropped.
77		/// This struct is very useful for tracking when futures are dropped.
78		#[derive(Debug)]
79		struct DropCounter<'a> {
80		counter: &'a AtomicU64,
81		}
82
83		impl<'a> DropCounter<'a> {
84		#[inline]
85	263	pub(crate) const fn new(counter: &'a AtomicU64) -> Self {
86	263	Self { counter }
87	263	}
88		}
89
90		impl Drop for DropCounter<'_> {
91		#[inline]
92	4	fn drop(&mut self) {
93	4	self.counter.fetch_add(1, Ordering::Acquire);
94	4	}
95		}
96
97		#[derive(Debug)]
98		pub struct AsyncTimer<'a> {
99		start: Instant,
100		drop_counter: DropCounter<'a>,
101		counter: &'a AsyncCounterWrapper,
102		}
103
104		impl AsyncTimer<'_> {
105		#[inline]
106	0	pub fn measure(self) {
107	0	self.counter.sum_func_duration_ns.fetch_add(
108	0	u64::try_from(self.start.elapsed().as_nanos()).unwrap_or(u64::MAX),
109	0	Ordering::Acquire,
110		);
111	0	self.counter.calls.fetch_add(1, Ordering::Acquire);
112	0	self.counter.successes.fetch_add(1, Ordering::Acquire);
113		// This causes DropCounter's drop to never be called.
114	0	forget(self.drop_counter);
115	0	}
116		}
117
118		/// Tracks the number of calls, successes, failures, and drops of an async function.
119		/// call `.wrap(future)` to wrap a future and stats about the future are automatically
120		/// tracked and can be published to a `CollectorState`.
121		#[derive(Debug, Default)]
122		pub struct AsyncCounterWrapper {
123		pub calls: AtomicU64,
124		pub successes: AtomicU64,
125		pub failures: AtomicU64,
126		pub drops: AtomicU64,
127		// Time spent in nanoseconds in the future.
128		// 64 bit address space gives ~584 years of nanoseconds.
129		pub sum_func_duration_ns: AtomicU64,
130		}
131
132		// Derive-macros have no way to tell the collector that the parent
133		// is now a group with the name of the group as the field so we
134		// can attach multiple values on the same group, so we need to
135		// manually implement the `MetricsComponent` trait to do so.
136		#[expect(
137		clippy::cognitive_complexity,
138		reason = "complexity arises from macro expansion"
139		)]
140		impl MetricsComponent for AsyncCounterWrapper {
141	0	fn publish(
142	0	&self,
143	0	_kind: MetricKind,
144	0	field_metadata: MetricFieldData,
145	0	) -> Result<MetricPublishKnownKindData, nativelink_metric::Error> {
146	0	let _enter = group!(field_metadata.name).entered();
147
148	0	publish!(
149	0	"calls",
150	0	&self.calls,
151	0	MetricKind::Counter,
152	0	format!("The number of times {} was called.", field_metadata.name)
153		);
154	0	publish!(
155	0	"successes",
156	0	&self.successes,
157	0	MetricKind::Counter,
158	0	format!(
159	0	"The number of times {} was successful.",
160		field_metadata.name
161		)
162		);
163	0	publish!(
164	0	"failures",
165	0	&self.failures,
166	0	MetricKind::Counter,
167	0	format!("The number of times {} failed.", field_metadata.name)
168		);
169	0	publish!(
170	0	"drops",
171	0	&self.drops,
172	0	MetricKind::Counter,
173	0	format!("The number of times {} was dropped.", field_metadata.name)
174		);
175	0	publish!(
176	0	"sum_func_duration_ns",
177	0	&self.sum_func_duration_ns,
178	0	MetricKind::Counter,
179	0	format!(
180	0	"The sum of the time spent in nanoseconds in {}.",
181		field_metadata.name
182		)
183		);
184
185	0	Ok(MetricPublishKnownKindData::Component)
186	0	}
187		}
188
189		impl AsyncCounterWrapper {
190		#[inline]
191	0	pub fn wrap_fn<'a, T: 'a, E>(
192	0	&'a self,
193	0	func: impl FnOnce() -> Result<T, E> + 'a,
194	0	) -> Result<T, E> {
195	0	self.calls.fetch_add(1, Ordering::Acquire);
196	0	let result = (func)();
197	0	if result.is_ok() { Branch (197:12): [Folded - Ignored] Branch (197:12): [Folded - Ignored]
198	0	self.successes.fetch_add(1, Ordering::Acquire);
199	0	} else {
200	0	self.failures.fetch_add(1, Ordering::Acquire);
201	0	}
202	0	result
203	0	}
204
205		#[inline]
206	248	pub async fn wrap<'a, T, E, F: Future<Output = Result<T, E>> + 'a>(
207	248	&'a self,
208	248	future: F,
209	248	) -> Result<T, E> {
210	248	let result = self.wrap_no_capture_result(future).await;
211	248	if result.is_ok() { Branch (211:12): [True: 0, False: 0] Branch (211:12): [Folded - Ignored] Branch (211:12): [True: 3, False: 1] Branch (211:12): [Folded - Ignored] Branch (211:12): [True: 4, False: 0] Branch (211:12): [True: 17, False: 0] Branch (211:12): [True: 6, False: 0] Branch (211:12): [True: 15, False: 0] Branch (211:12): [True: 15, False: 0] Branch (211:12): [True: 30, False: 0] Branch (211:12): [True: 18, False: 1] Branch (211:12): [True: 15, False: 0] Branch (211:12): [True: 19, False: 0] Branch (211:12): [True: 19, False: 0] Branch (211:12): [True: 11, False: 0] Branch (211:12): [True: 11, False: 0] Branch (211:12): [True: 13, False: 0] Branch (211:12): [True: 11, False: 0] Branch (211:12): [True: 11, False: 0] Branch (211:12): [True: 26, False: 1] Branch (211:12): [True: 1, False: 0]
212	245	self.successes.fetch_add(1, Ordering::Acquire);
213	245	} else {
214	3	self.failures.fetch_add(1, Ordering::Acquire);
215	3	}
216	248	result
217	248	}
218
219		#[inline]
220	250	pub async fn wrap_no_capture_result<'a, T, F: Future<Output = T> + 'a>(
221	250	&'a self,
222	250	future: F,
223	250	) -> T {
224	250	self.calls.fetch_add(1, Ordering::Acquire);
225	250	let drop_counter = DropCounter::new(&self.drops);
226	250	let instant = Instant::now();
227	250	let result = future.await;
228		// By default `drop_counter` will increment the drop counter when it goes out of scope.
229		// This will ensure we don't increment the counter if we make it here with a zero cost.
230	250	forget(drop_counter);
231	250	self.sum_func_duration_ns.fetch_add(
232	250	u64::try_from(instant.elapsed().as_nanos()).unwrap_or(u64::MAX),
233	250	Ordering::Acquire,
234		);
235	250	result
236	250	}
237
238		#[inline]
239	0	pub fn begin_timer(&self) -> AsyncTimer<'_> {
240	0	AsyncTimer {
241	0	start: Instant::now(),
242	0	drop_counter: DropCounter::new(&self.drops),
243	0	counter: self,
244	0	}
245	0	}
246		}
247
248		/// Tracks a number.
249		#[derive(Debug, Default)]
250		pub struct Counter(AtomicU64);
251
252		impl Counter {
253		#[inline]
254	0	pub fn inc(&self) {
255	0	self.add(1);
256	0	}
257
258		#[inline]
259	6.01k	pub fn add(&self, value: u64) {
260	6.01k	self.0.fetch_add(value, Ordering::Acquire);
261	6.01k	}
262
263		#[inline]
264	0	pub fn sub(&self, value: u64) {
265	0	self.0.fetch_sub(value, Ordering::Acquire);
266	0	}
267		}
268
269		impl MetricsComponent for Counter {
270	0	fn publish(
271	0	&self,
272	0	kind: MetricKind,
273	0	field_metadata: MetricFieldData,
274	0	) -> Result<MetricPublishKnownKindData, nativelink_metric::Error> {
275	0	self.0.publish(kind, field_metadata)
276	0	}
277		}
278
279		/// Tracks an counter through time and the last time the counter was changed.
280		#[derive(Debug, Default)]
281		pub struct CounterWithTime {
282		pub counter: AtomicU64,
283		pub last_time: AtomicU64,
284		}
285
286		impl CounterWithTime {
287		#[inline]
288	126	pub fn inc(&self) {
289	126	self.counter.fetch_add(1, Ordering::Acquire);
290	126	self.last_time.store(
291	126	SystemTime::now()
292	126	.duration_since(UNIX_EPOCH)
293	126	.unwrap()
294	126	.as_secs(),
295	126	Ordering::Release,
296		);
297	126	}
298		}
299
300		// Derive-macros have no way to tell the collector that the parent
301		// is now a group with the name of the group as the field so we
302		// can attach multiple values on the same group, so we need to
303		// manually implement the `MetricsComponent` trait to do so.
304		impl MetricsComponent for CounterWithTime {
305	0	fn publish(
306	0	&self,
307	0	_kind: MetricKind,
308	0	field_metadata: MetricFieldData,
309	0	) -> Result<MetricPublishKnownKindData, nativelink_metric::Error> {
310	0	let _enter = group!(field_metadata.name).entered();
311
312	0	publish!(
313	0	"counter",
314	0	&self.counter,
315	0	MetricKind::Counter,
316	0	format!("Current count of {}.", field_metadata.name)
317		);
318	0	publish!(
319	0	"last_time",
320	0	&self.last_time,
321	0	MetricKind::Counter,
322	0	format!("Last timestamp {} was published.", field_metadata.name)
323		);
324
325	0	Ok(MetricPublishKnownKindData::Component)
326	0	}
327		}