Coverage Report

Created: 2024-10-22 12:33

/build/source/nativelink-metric/src/lib.rs
Line
Count
Source (jump to first uncovered line)
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::borrow::Cow;
16
use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
17
use std::sync::atomic::{AtomicI64, AtomicU64, Ordering};
18
use std::sync::{Arc, Weak};
19
use std::time::{Duration, SystemTime, UNIX_EPOCH};
20
21
pub use nativelink_metric_macro_derive::MetricsComponent;
22
pub use tracing::{error as __metric_error, info as __metric_event, info_span as __metric_span};
23
24
/// Error type for the metrics library.
25
// Note: We do not use the nativelink-error struct because
26
// we'd end up in a circular dependency if we did, because
27
// nativelink-error uses the metrics library.
28
#[derive(Debug)]
29
pub struct Error(String);
30
31
impl std::fmt::Display for Error {
32
0
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33
0
        write!(f, "{}", self.0)
34
0
    }
35
}
36
37
impl std::error::Error for Error {}
38
39
/// Holds metadata about the field that is being published.
40
#[derive(Default, Clone)]
41
pub struct MetricFieldData<'a> {
42
    pub name: Cow<'a, str>,
43
    pub help: Cow<'a, str>,
44
    pub group: Cow<'a, str>,
45
}
46
47
/// The final primtive data that is being published with the kind.
48
#[derive(Debug)]
49
pub enum MetricPublishKnownKindData {
50
    Counter(u64),
51
    String(String),
52
    Component,
53
}
54
55
/// The kind of metric that is being published.
56
// Note: This enum will be translate in-and-out
57
// of a u64 when traversing the `tracing::event`
58
// boundary for efficiency reasons.
59
#[derive(Clone, Copy, Debug)]
60
#[repr(u8)]
61
pub enum MetricKind {
62
    Default = 0,
63
    Counter = 1,
64
    String = 2,
65
    Component = 3,
66
}
67
68
impl From<u64> for MetricKind {
69
0
    fn from(value: u64) -> Self {
70
0
        match value {
71
0
            0 => MetricKind::Default,
72
0
            1 => MetricKind::Counter,
73
0
            2 => MetricKind::String,
74
0
            3 => MetricKind::Component,
75
0
            _ => MetricKind::Default,
76
        }
77
0
    }
78
}
79
80
impl MetricKind {
81
14
    pub fn into_known_kind(&self, default_kind: MetricKind) -> MetricPublishKnownKindData {
82
14
        let mut this = *self;
83
14
        if 
matches!4
(self, MetricKind::Default) {
84
10
            this = default_kind;
85
10
        }
4
86
14
        match this {
87
6
            MetricKind::Counter => MetricPublishKnownKindData::Counter(0),
88
8
            MetricKind::String => MetricPublishKnownKindData::String(String::new()),
89
0
            MetricKind::Component => MetricPublishKnownKindData::Component,
90
0
            MetricKind::Default => unreachable!("Default should have been handled"),
91
        }
92
14
    }
93
}
94
95
/// The trait that all components that can be published must implement.
96
pub trait MetricsComponent {
97
    fn publish(
98
        &self,
99
        kind: MetricKind,
100
        _field_metadata: MetricFieldData,
101
    ) -> Result<MetricPublishKnownKindData, Error>;
102
}
103
104
pub trait RootMetricsComponent: MetricsComponent + Send + Sync {
105
0
    fn publish(
106
0
        &self,
107
0
        kind: MetricKind,
108
0
        field_metadata: MetricFieldData,
109
0
    ) -> Result<MetricPublishKnownKindData, Error> {
110
0
        MetricsComponent::publish(self, kind, field_metadata)
111
0
    }
112
}
113
114
impl<T: MetricsComponent> MetricsComponent for Option<T> {
115
1
    fn publish(
116
1
        &self,
117
1
        kind: MetricKind,
118
1
        field_metadata: MetricFieldData,
119
1
    ) -> Result<MetricPublishKnownKindData, Error> {
120
1
        match self {
121
0
            Some(value) => value.publish(kind, field_metadata),
122
1
            None => Ok(MetricPublishKnownKindData::Component),
123
        }
124
1
    }
125
}
126
127
impl<T: MetricsComponent> MetricsComponent for tokio::sync::watch::Sender<T> {
128
0
    fn publish(
129
0
        &self,
130
0
        kind: MetricKind,
131
0
        field_metadata: MetricFieldData,
132
0
    ) -> Result<MetricPublishKnownKindData, Error> {
133
0
        self.borrow().publish(kind, field_metadata)
134
0
    }
135
}
136
137
impl<T: MetricsComponent + ?Sized> MetricsComponent for Arc<T> {
138
2
    fn publish(
139
2
        &self,
140
2
        kind: MetricKind,
141
2
        field_metadata: MetricFieldData,
142
2
    ) -> Result<MetricPublishKnownKindData, Error> {
143
2
        self.as_ref().publish(kind, field_metadata)
144
2
    }
145
}
146
147
impl<T: MetricsComponent> MetricsComponent for HashSet<T> {
148
0
    fn publish(
149
0
        &self,
150
0
        kind: MetricKind,
151
0
        field_metadata: MetricFieldData,
152
0
    ) -> Result<MetricPublishKnownKindData, Error> {
153
0
        for (i, item) in self.iter().enumerate() {
154
0
            let guard = group!(i).entered();
155
0
            let publish_result = item.publish(kind, field_metadata.clone())?;
156
0
            drop(guard);
157
0
            match publish_result {
158
0
                MetricPublishKnownKindData::Counter(value) => {
159
0
                    publish!(
160
0
                        i,
161
0
                        &value,
162
0
                        MetricKind::Counter,
163
0
                        field_metadata.help.to_string()
164
                    );
165
                }
166
0
                MetricPublishKnownKindData::String(value) => {
167
0
                    publish!(
168
0
                        i,
169
0
                        &value,
170
0
                        MetricKind::String,
171
0
                        field_metadata.help.to_string()
172
                    );
173
                }
174
0
                MetricPublishKnownKindData::Component => {}
175
            }
176
        }
177
0
        Ok(MetricPublishKnownKindData::Component)
178
0
    }
179
}
180
181
impl<U: ToString, T: MetricsComponent> MetricsComponent for HashMap<U, T> {
182
1
    fn publish(
183
1
        &self,
184
1
        kind: MetricKind,
185
1
        field_metadata: MetricFieldData,
186
1
    ) -> Result<MetricPublishKnownKindData, Error> {
187
1
        for (key, item) in self.iter() {
188
1
            let guard = group!(key).entered();
189
1
            let publish_result = item.publish(kind, field_metadata.clone())
?0
;
190
1
            drop(guard);
191
1
            match publish_result {
192
0
                MetricPublishKnownKindData::Counter(value) => {
193
0
                    publish!(
194
0
                        key,
195
0
                        &value,
196
0
                        MetricKind::Counter,
197
0
                        field_metadata.help.to_string()
198
                    );
199
                }
200
0
                MetricPublishKnownKindData::String(value) => {
201
0
                    publish!(
202
0
                        key,
203
0
                        &value,
204
0
                        MetricKind::String,
205
0
                        field_metadata.help.to_string()
206
                    );
207
                }
208
1
                MetricPublishKnownKindData::Component => {}
209
            }
210
        }
211
1
        Ok(MetricPublishKnownKindData::Component)
212
1
    }
213
}
214
215
impl<U: ToString, T: MetricsComponent> MetricsComponent for BTreeMap<U, T> {
216
0
    fn publish(
217
0
        &self,
218
0
        kind: MetricKind,
219
0
        field_metadata: MetricFieldData,
220
0
    ) -> Result<MetricPublishKnownKindData, Error> {
221
0
        for (key, item) in self.iter() {
222
0
            group!(key).in_scope(|| item.publish(kind, field_metadata.clone()))?;
223
        }
224
0
        Ok(MetricPublishKnownKindData::Component)
225
0
    }
226
}
227
228
impl<T: MetricsComponent> MetricsComponent for BTreeSet<T> {
229
0
    fn publish(
230
0
        &self,
231
0
        kind: MetricKind,
232
0
        field_metadata: MetricFieldData,
233
0
    ) -> Result<MetricPublishKnownKindData, Error> {
234
0
        for (i, item) in self.iter().enumerate() {
235
0
            group!(i).in_scope(|| item.publish(kind, field_metadata.clone()))?;
236
        }
237
0
        Ok(MetricPublishKnownKindData::Component)
238
0
    }
239
}
240
241
impl<T: MetricsComponent> MetricsComponent for Vec<T> {
242
0
    fn publish(
243
0
        &self,
244
0
        kind: MetricKind,
245
0
        field_metadata: MetricFieldData,
246
0
    ) -> Result<MetricPublishKnownKindData, Error> {
247
0
        for (i, item) in self.iter().enumerate() {
248
0
            group!(i).in_scope(|| item.publish(kind, field_metadata.clone()))?;
249
        }
250
0
        Ok(MetricPublishKnownKindData::Component)
251
0
    }
252
}
253
254
impl<T: MetricsComponent + ?Sized> MetricsComponent for Weak<T> {
255
0
    fn publish(
256
0
        &self,
257
0
        kind: MetricKind,
258
0
        field_metadata: MetricFieldData,
259
0
    ) -> Result<MetricPublishKnownKindData, Error> {
260
0
        let Some(this) = self.upgrade() else {
  Branch (260:13): [Folded - Ignored]
  Branch (260:13): [True: 0, False: 0]
  Branch (260:13): [Folded - Ignored]
261
0
            return Ok(MetricPublishKnownKindData::Component);
262
        };
263
0
        this.as_ref().publish(kind, field_metadata)
264
0
    }
265
}
266
267
impl<T, E> MetricsComponent for Result<T, E>
268
where
269
    T: MetricsComponent,
270
    E: MetricsComponent,
271
{
272
0
    fn publish(
273
0
        &self,
274
0
        kind: MetricKind,
275
0
        field_metadata: MetricFieldData,
276
0
    ) -> Result<MetricPublishKnownKindData, Error> {
277
0
        match self {
278
0
            Ok(value) => value.publish(kind, field_metadata),
279
0
            Err(value) => value.publish(kind, field_metadata),
280
        }
281
0
    }
282
}
283
284
impl MetricsComponent for Duration {
285
0
    fn publish(
286
0
        &self,
287
0
        kind: MetricKind,
288
0
        field_metadata: MetricFieldData,
289
0
    ) -> Result<MetricPublishKnownKindData, Error> {
290
0
        self.as_secs_f64().publish(kind, field_metadata)
291
0
    }
292
}
293
294
impl MetricsComponent for SystemTime {
295
0
    fn publish(
296
0
        &self,
297
0
        kind: MetricKind,
298
0
        field_metadata: MetricFieldData,
299
0
    ) -> Result<MetricPublishKnownKindData, Error> {
300
0
        match SystemTime::now().duration_since(UNIX_EPOCH) {
301
0
            Ok(n) => n.as_secs().publish(kind, field_metadata),
302
0
            Err(_) => Err(Error("SystemTime before UNIX EPOCH!".to_string())),
303
        }
304
0
    }
305
}
306
307
impl MetricsComponent for f64 {
308
0
    fn publish(
309
0
        &self,
310
0
        _kind: MetricKind,
311
0
        _field_metadata: MetricFieldData,
312
0
    ) -> Result<MetricPublishKnownKindData, Error> {
313
0
        Ok(MetricPublishKnownKindData::String(self.to_string()))
314
0
    }
315
}
316
317
impl MetricsComponent for bool {
318
0
    fn publish(
319
0
        &self,
320
0
        kind: MetricKind,
321
0
        field_metadata: MetricFieldData,
322
0
    ) -> Result<MetricPublishKnownKindData, Error> {
323
0
        let value = u64::from(*self);
324
0
        value.publish(kind, field_metadata)
325
0
    }
326
}
327
328
impl MetricsComponent for i32 {
329
0
    fn publish(
330
0
        &self,
331
0
        kind: MetricKind,
332
0
        field_metadata: MetricFieldData,
333
0
    ) -> Result<MetricPublishKnownKindData, Error> {
334
0
        let value = u64::try_from(*self)
335
0
            .map_err(|_| Error(format!("Could not convert {self} to u64 in metrics lib")))?;
336
0
        value.publish(kind, field_metadata)
337
0
    }
338
}
339
340
impl MetricsComponent for u64 {
341
2
    fn publish(
342
2
        &self,
343
2
        kind: MetricKind,
344
2
        _field_metadata: MetricFieldData,
345
2
    ) -> Result<MetricPublishKnownKindData, Error> {
346
2
        let mut known_kind_data = kind.into_known_kind(MetricKind::Counter);
347
2
        match &mut known_kind_data {
348
2
            MetricPublishKnownKindData::Counter(data) => {
349
2
                *data = *self;
350
2
            }
351
0
            MetricPublishKnownKindData::String(data) => {
352
0
                *data = self.to_string();
353
0
            }
354
0
            MetricPublishKnownKindData::Component => {}
355
        }
356
2
        Ok(known_kind_data)
357
2
    }
358
}
359
360
impl MetricsComponent for i64 {
361
0
    fn publish(
362
0
        &self,
363
0
        kind: MetricKind,
364
0
        field_metadata: MetricFieldData,
365
0
    ) -> Result<MetricPublishKnownKindData, Error> {
366
0
        let value = u64::try_from(*self)
367
0
            .map_err(|_| Error(format!("Could not convert {self} to u64 in metrics lib")))?;
368
0
        value.publish(kind, field_metadata)
369
0
    }
370
}
371
372
impl MetricsComponent for u32 {
373
0
    fn publish(
374
0
        &self,
375
0
        kind: MetricKind,
376
0
        field_metadata: MetricFieldData,
377
0
    ) -> Result<MetricPublishKnownKindData, Error> {
378
0
        u64::from(*self).publish(kind, field_metadata)
379
0
    }
380
}
381
382
impl MetricsComponent for usize {
383
0
    fn publish(
384
0
        &self,
385
0
        kind: MetricKind,
386
0
        field_metadata: MetricFieldData,
387
0
    ) -> Result<MetricPublishKnownKindData, Error> {
388
0
        let value = u64::try_from(*self)
389
0
            .map_err(|_| Error(format!("Could not convert {self} to u64 in metrics lib")))?;
390
0
        value.publish(kind, field_metadata)
391
0
    }
392
}
393
394
impl MetricsComponent for AtomicU64 {
395
0
    fn publish(
396
0
        &self,
397
0
        kind: MetricKind,
398
0
        field_metadata: MetricFieldData,
399
0
    ) -> Result<MetricPublishKnownKindData, Error> {
400
0
        self.load(Ordering::Acquire).publish(kind, field_metadata)
401
0
    }
402
}
403
404
impl MetricsComponent for AtomicI64 {
405
0
    fn publish(
406
0
        &self,
407
0
        kind: MetricKind,
408
0
        field_metadata: MetricFieldData,
409
0
    ) -> Result<MetricPublishKnownKindData, Error> {
410
0
        self.load(Ordering::Acquire).publish(kind, field_metadata)
411
0
    }
412
}
413
414
impl MetricsComponent for String {
415
12
    fn publish(
416
12
        &self,
417
12
        kind: MetricKind,
418
12
        _field_metadata: MetricFieldData,
419
12
    ) -> Result<MetricPublishKnownKindData, Error> {
420
12
        let mut known_kind_data = kind.into_known_kind(MetricKind::String);
421
12
        match &mut known_kind_data {
422
4
            MetricPublishKnownKindData::Counter(data) => {
423
4
                *data = self.parse::<u64>().map_err(|_| {
424
0
                    Error(format!(
425
0
                        "Could not convert String '{self}' to u64 in metrics lib"
426
0
                    ))
427
4
                })
?0
;
428
            }
429
8
            MetricPublishKnownKindData::String(data) => {
430
8
                data.clone_from(self);
431
8
            }
432
0
            MetricPublishKnownKindData::Component => {}
433
        }
434
12
        Ok(known_kind_data)
435
12
    }
436
}
437
438
impl<T: MetricsComponent> MetricsComponent for async_lock::Mutex<T> {
439
0
    fn publish(
440
0
        &self,
441
0
        kind: MetricKind,
442
0
        field_metadata: MetricFieldData,
443
0
    ) -> Result<MetricPublishKnownKindData, Error> {
444
0
        // It is safe to block in the publishing thread.
445
0
        let lock = self.lock_blocking();
446
0
        lock.publish(kind, field_metadata)
447
0
    }
448
}
449
450
impl<T: MetricsComponent> MetricsComponent for parking_lot::Mutex<T> {
451
0
    fn publish(
452
0
        &self,
453
0
        kind: MetricKind,
454
0
        field_metadata: MetricFieldData,
455
0
    ) -> Result<MetricPublishKnownKindData, Error> {
456
0
        // It is safe to block in the publishing thread.
457
0
        let lock = self.lock();
458
0
        lock.publish(kind, field_metadata)
459
0
    }
460
}
461
462
impl<T: MetricsComponent> MetricsComponent for parking_lot::RwLock<T> {
463
1
    fn publish(
464
1
        &self,
465
1
        kind: MetricKind,
466
1
        field_metadata: MetricFieldData,
467
1
    ) -> Result<MetricPublishKnownKindData, Error> {
468
1
        // It is safe to block in the publishing thread.
469
1
        let lock = self.read();
470
1
        lock.publish(kind, field_metadata)
471
1
    }
472
}
473
474
#[macro_export]
475
macro_rules! group {
476
    ($name:expr) => {
477
        $crate::__metric_span!(target: "nativelink_metric", "", __name = $name.to_string())
478
    };
479
}
480
481
#[macro_export]
482
macro_rules! publish {
483
    ($name:expr, $value:expr, $metric_kind:expr, $help:expr) => {
484
        $crate::publish!($name, $value, $metric_kind, $help, "")
485
    };
486
    ($name:expr, $value:expr, $metric_kind:expr, $help:expr, $group:expr) => {
487
        {
488
            let _maybe_entered = if !$group.is_empty() {
489
                Some($crate::group!($group).entered())
490
            } else {
491
                None
492
            };
493
            let name = $name.to_string();
494
            let field_metadata = $crate::MetricFieldData {
495
                name: ::std::borrow::Cow::Borrowed(&name),
496
                help: $help.into(),
497
                group: $group.into(),
498
            };
499
            match $crate::MetricsComponent::publish($value, $metric_kind, field_metadata)? {
500
                $crate::MetricPublishKnownKindData::Counter(value) => {
501
                    $crate::__metric_event!(
502
                        target: "nativelink_metric",
503
                        __value = value,
504
                        __type = $crate::MetricKind::Counter as u8,
505
                        __help = $help.to_string(),
506
                        __name = name
507
                    );
508
                }
509
                $crate::MetricPublishKnownKindData::String(value) => {
510
                    $crate::__metric_event!(
511
                        target: "nativelink_metric",
512
                        __value = value,
513
                        __type = $crate::MetricKind::String as u8,
514
                        __help = $help.to_string(),
515
                        __name = name
516
                    );
517
                }
518
                $crate::MetricPublishKnownKindData::Component => {
519
                    // Do nothing, data already published.
520
                }
521
            }
522
        }
523
    };
524
}