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

Source
// Copyright 2024 The NativeLink Authors. All rights reserved.
//
// 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.

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(self.start.elapsed().as_nanos() as u64, 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(instant.elapsed().as_nanos() as u64, 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-06-24 08:57

Line	Count	Source
1		// Copyright 2024 The NativeLink Authors. All rights reserved.
2		//
3		// Licensed under the Apache License, Version 2.0 (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		// http://www.apache.org/licenses/LICENSE-2.0
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	248	pub(crate) const fn new(counter: &'a AtomicU64) -> Self {
86	248	Self { counter }
87	248	}
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
108	0	.sum_func_duration_ns
109	0	.fetch_add(self.start.elapsed().as_nanos() as u64, Ordering::Acquire);
110	0	self.counter.calls.fetch_add(1, Ordering::Acquire);
111	0	self.counter.successes.fetch_add(1, Ordering::Acquire);
112		// This causes DropCounter's drop to never be called.
113	0	forget(self.drop_counter);
114	0	}
115		}
116
117		/// Tracks the number of calls, successes, failures, and drops of an async function.
118		/// call `.wrap(future)` to wrap a future and stats about the future are automatically
119		/// tracked and can be published to a `CollectorState`.
120		#[derive(Debug, Default)]
121		pub struct AsyncCounterWrapper {
122		pub calls: AtomicU64,
123		pub successes: AtomicU64,
124		pub failures: AtomicU64,
125		pub drops: AtomicU64,
126		// Time spent in nanoseconds in the future.
127		// 64 bit address space gives ~584 years of nanoseconds.
128		pub sum_func_duration_ns: AtomicU64,
129		}
130
131		// Derive-macros have no way to tell the collector that the parent
132		// is now a group with the name of the group as the field so we
133		// can attach multiple values on the same group, so we need to
134		// manually implement the `MetricsComponent` trait to do so.
135		#[expect(
136		clippy::cognitive_complexity,
137		reason = "complexity arises from macro expansion"
138		)]
139		impl MetricsComponent for AsyncCounterWrapper {
140	0	fn publish(
141	0	&self,
142	0	_kind: MetricKind,
143	0	field_metadata: MetricFieldData,
144	0	) -> Result<MetricPublishKnownKindData, nativelink_metric::Error> {
145	0	let _enter = group!(field_metadata.name).entered();
146
147	0	publish!(
148	0	"calls",
149	0	&self.calls,
150	0	MetricKind::Counter,
151	0	format!("The number of times {} was called.", field_metadata.name)
152		);
153	0	publish!(
154	0	"successes",
155	0	&self.successes,
156	0	MetricKind::Counter,
157	0	format!(
158	0	"The number of times {} was successful.",
159		field_metadata.name
160		)
161		);
162	0	publish!(
163	0	"failures",
164	0	&self.failures,
165	0	MetricKind::Counter,
166	0	format!("The number of times {} failed.", field_metadata.name)
167		);
168	0	publish!(
169	0	"drops",
170	0	&self.drops,
171	0	MetricKind::Counter,
172	0	format!("The number of times {} was dropped.", field_metadata.name)
173		);
174	0	publish!(
175	0	"sum_func_duration_ns",
176	0	&self.sum_func_duration_ns,
177	0	MetricKind::Counter,
178	0	format!(
179	0	"The sum of the time spent in nanoseconds in {}.",
180		field_metadata.name
181		)
182		);
183
184	0	Ok(MetricPublishKnownKindData::Component)
185	0	}
186		}
187
188		impl AsyncCounterWrapper {
189		#[inline]
190	0	pub fn wrap_fn<'a, T: 'a, E>(
191	0	&'a self,
192	0	func: impl FnOnce() -> Result<T, E> + 'a,
193	0	) -> Result<T, E> {
194	0	self.calls.fetch_add(1, Ordering::Acquire);
195	0	let result = (func)();
196	0	if result.is_ok() { Branch (196:12): [Folded - Ignored] Branch (196:12): [Folded - Ignored]
197	0	self.successes.fetch_add(1, Ordering::Acquire);
198	0	} else {
199	0	self.failures.fetch_add(1, Ordering::Acquire);
200	0	}
201	0	result
202	0	}
203
204		#[inline]
205	233	pub async fn wrap<'a, T, E, F: Future<Output = Result<T, E>> + 'a>(
206	233	&'a self,
207	233	future: F,
208	233	) -> Result<T, E> {
209	233	let result = self.wrap_no_capture_result(future).await;
210	233	if result.is_ok() { Branch (210:12): [True: 0, False: 0] Branch (210:12): [Folded - Ignored] Branch (210:12): [True: 3, False: 1] Branch (210:12): [Folded - Ignored] Branch (210:12): [True: 6, False: 0] Branch (210:12): [True: 30, False: 0] Branch (210:12): [True: 15, False: 1] Branch (210:12): [True: 15, False: 0] Branch (210:12): [True: 16, False: 0] Branch (210:12): [True: 16, False: 0] Branch (210:12): [True: 15, False: 0] Branch (210:12): [True: 13, False: 0] Branch (210:12): [True: 11, False: 0] Branch (210:12): [True: 11, False: 0] Branch (210:12): [True: 15, False: 0] Branch (210:12): [True: 15, False: 0] Branch (210:12): [True: 11, False: 0] Branch (210:12): [True: 11, False: 0] Branch (210:12): [True: 26, False: 1] Branch (210:12): [True: 1, False: 0]
211	230	self.successes.fetch_add(1, Ordering::Acquire);
212	230	} else {
213	3	self.failures.fetch_add(1, Ordering::Acquire);
214	3	}
215	233	result
216	233	}
217
218		#[inline]
219	235	pub async fn wrap_no_capture_result<'a, T, F: Future<Output = T> + 'a>(
220	235	&'a self,
221	235	future: F,
222	235	) -> T {
223	235	self.calls.fetch_add(1, Ordering::Acquire);
224	235	let drop_counter = DropCounter::new(&self.drops);
225	235	let instant = Instant::now();
226	235	let result = future.await;
227		// By default `drop_counter` will increment the drop counter when it goes out of scope.
228		// This will ensure we don't increment the counter if we make it here with a zero cost.
229	235	forget(drop_counter);
230	235	self.sum_func_duration_ns
231	235	.fetch_add(instant.elapsed().as_nanos() as u64, Ordering::Acquire);
232	235	result
233	235	}
234
235		#[inline]
236	0	pub fn begin_timer(&self) -> AsyncTimer<'_> {
237	0	AsyncTimer {
238	0	start: Instant::now(),
239	0	drop_counter: DropCounter::new(&self.drops),
240	0	counter: self,
241	0	}
242	0	}
243		}
244
245		/// Tracks a number.
246		#[derive(Debug, Default)]
247		pub struct Counter(AtomicU64);
248
249		impl Counter {
250		#[inline]
251	0	pub fn inc(&self) {
252	0	self.add(1);
253	0	}
254
255		#[inline]
256	6.00k	pub fn add(&self, value: u64) {
257	6.00k	self.0.fetch_add(value, Ordering::Acquire);
258	6.00k	}
259
260		#[inline]
261	0	pub fn sub(&self, value: u64) {
262	0	self.0.fetch_sub(value, Ordering::Acquire);
263	0	}
264		}
265
266		impl MetricsComponent for Counter {
267	0	fn publish(
268	0	&self,
269	0	kind: MetricKind,
270	0	field_metadata: MetricFieldData,
271	0	) -> Result<MetricPublishKnownKindData, nativelink_metric::Error> {
272	0	self.0.publish(kind, field_metadata)
273	0	}
274		}
275
276		/// Tracks an counter through time and the last time the counter was changed.
277		#[derive(Debug, Default)]
278		pub struct CounterWithTime {
279		pub counter: AtomicU64,
280		pub last_time: AtomicU64,
281		}
282
283		impl CounterWithTime {
284		#[inline]
285	117	pub fn inc(&self) {
286	117	self.counter.fetch_add(1, Ordering::Acquire);
287	117	self.last_time.store(
288	117	SystemTime::now()
289	117	.duration_since(UNIX_EPOCH)
290	117	.unwrap()
291	117	.as_secs(),
292	117	Ordering::Release,
293		);
294	117	}
295		}
296
297		// Derive-macros have no way to tell the collector that the parent
298		// is now a group with the name of the group as the field so we
299		// can attach multiple values on the same group, so we need to
300		// manually implement the `MetricsComponent` trait to do so.
301		impl MetricsComponent for CounterWithTime {
302	0	fn publish(
303	0	&self,
304	0	_kind: MetricKind,
305	0	field_metadata: MetricFieldData,
306	0	) -> Result<MetricPublishKnownKindData, nativelink_metric::Error> {
307	0	let _enter = group!(field_metadata.name).entered();
308
309	0	publish!(
310	0	"counter",
311	0	&self.counter,
312	0	MetricKind::Counter,
313	0	format!("Current count of {}.", field_metadata.name)
314		);
315	0	publish!(
316	0	"last_time",
317	0	&self.last_time,
318	0	MetricKind::Counter,
319	0	format!("Last timestamp {} was published.", field_metadata.name)
320		);
321
322	0	Ok(MetricPublishKnownKindData::Component)
323	0	}
324		}