/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 std::mem::forget;
use std::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.
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-04-19 16:54

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