/build/source/nativelink-metric-collector/src/metrics_visitors.rs

Source (jump to first uncovered line)
// 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::borrow::Cow;
use std::fmt::Debug;

use nativelink_metric::MetricKind;
use serde::Serialize;
use tracing::field::{Field, Visit};

use crate::metrics_collection::{CollectedMetricPrimitive, CollectedMetricPrimitiveValue};

/// The type of the collected primitive metric.
#[derive(Default, Debug, Serialize)]
pub enum CollectionKind {
    #[default]
    Counter = 0,
    String = 1,
}

impl From<MetricKind> for CollectionKind {
    fn from(kind: MetricKind) -> Self {
        match kind {
            MetricKind::Counter => CollectionKind::Counter,
            MetricKind::String => CollectionKind::String,
            _ => CollectionKind::String,
        }
    }
}

/// The final-metric primitive value and type that was collected.
#[derive(Debug)]
enum ValueWithPrimitiveType {
    String(String),
    U64(u64),
}

impl Default for ValueWithPrimitiveType {
    fn default() -> Self {
        ValueWithPrimitiveType::U64(0)
    }
}

/// An intermediate structed that will have it's contents populated
/// by the `tracing` layer for a given field.
/// This is done by implementing the `Visit` trait and asking the
/// `tracing` library to visit the fields of the captured event
/// and populate this struct.
#[derive(Default, Debug)]
pub struct MetricDataVisitor {
    pub name: String,
    value: ValueWithPrimitiveType,
    help: String,
    value_type: Option<CollectionKind>,
}

impl From<MetricDataVisitor> for CollectedMetricPrimitive {
    fn from(visitor: MetricDataVisitor) -> Self {
        let (value, derived_type) = match visitor.value {
            ValueWithPrimitiveType::String(s) => (
                CollectedMetricPrimitiveValue::String(Cow::Owned(s)),
                CollectionKind::String,
            ),
            ValueWithPrimitiveType::U64(u) => (
                CollectedMetricPrimitiveValue::Counter(u),
                CollectionKind::Counter,
            ),
        };
        CollectedMetricPrimitive {
            value: Some(value),
            help: visitor.help,
            value_type: visitor.value_type.unwrap_or(derived_type),
        }
    }
}

impl Visit for MetricDataVisitor {
    fn record_debug(&mut self, _field: &Field, _value: &dyn Debug) {}

    fn record_f64(&mut self, field: &Field, value: f64) {
        if field.name() == "__value" {
            self.value = ValueWithPrimitiveType::String(value.to_string());
        }
    }
    fn record_i64(&mut self, field: &Field, value: i64) {
        if field.name() == "__value" {
            match u64::try_from(value) {
                Ok(v) => self.value = ValueWithPrimitiveType::U64(v),
                Err(_) => self.value = ValueWithPrimitiveType::String(value.to_string()),
            }
        }
    }
    fn record_u64(&mut self, field: &Field, value: u64) {
        match field.name() {
            "__value" => self.value = ValueWithPrimitiveType::U64(value)6,
            "__type" => self.value_type = Some(MetricKind::from(value).into()),
            "__help" => self.help = value.to_string(),
            "__name" => self.name = value.to_string(),
            field => panic!("UNKNOWN FIELD {field}"),
        }
    }
    fn record_i128(&mut self, field: &Field, value: i128) {
        if field.name() == "__value" {
            match u64::try_from(value) {
                Ok(v) => self.value = ValueWithPrimitiveType::U64(v),
                Err(_) => self.value = ValueWithPrimitiveType::String(value.to_string()),
            }
        }
    }
    fn record_u128(&mut self, field: &Field, value: u128) {
        if field.name() == "__value" {
            match u64::try_from(value) {
                Ok(v) => self.value = ValueWithPrimitiveType::U64(v),
                Err(_) => self.value = ValueWithPrimitiveType::String(value.to_string()),
            }
        }
    }
    fn record_bool(&mut self, field: &Field, value: bool) {
        if field.name() == "__value" {
            self.value = ValueWithPrimitiveType::U64(u64::from(value));
        }
    }
    fn record_str(&mut self, field: &Field, value: &str) {
        match field.name() {
            "__value" => self.value = ValueWithPrimitiveType::String(value.to_string())8,
            "__help" => self.help = value.to_string()14,
            "__name" => self.name = value.to_string(),
            field => panic!("UNKNOWN FIELD {field}"),
        }
    }
    fn record_error(&mut self, _field: &Field, _value: &(dyn std::error::Error + 'static)) {}
}

/// An intermediate structed that will have it's contents populated
/// by the `tracing` layer for a given field.
/// This is the same as `MetricDataVisitor` but only captures info
/// about a given span on span creation.
pub struct SpanFields {
    pub name: Cow<'static, str>,
}

impl Visit for SpanFields {
    fn record_debug(&mut self, _field: &Field, _value: &dyn Debug) {}

    fn record_str(&mut self, field: &Field, value: &str) {
        if field.name() == "__name" {
            self.name = Cow::Owned(value.to_string());
        }0
    }
}

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::fmt::Debug;
17
18		use nativelink_metric::MetricKind;
19		use serde::Serialize;
20		use tracing::field::{Field, Visit};
21
22		use crate::metrics_collection::{CollectedMetricPrimitive, CollectedMetricPrimitiveValue};
23
24		/// The type of the collected primitive metric.
25		#[derive(Default, Debug, Serialize)]
26		pub enum CollectionKind {
27		#[default]
28		Counter = 0,
29		String = 1,
30		}
31
32		impl From<MetricKind> for CollectionKind {
33	14	fn from(kind: MetricKind) -> Self {
34	14	match kind {
35	6	MetricKind::Counter => CollectionKind::Counter,
36	8	MetricKind::String => CollectionKind::String,
37	0	_ => CollectionKind::String,
38		}
39	14	}
40		}
41
42		/// The final-metric primitive value and type that was collected.
43		#[derive(Debug)]
44		enum ValueWithPrimitiveType {
45		String(String),
46		U64(u64),
47		}
48
49		impl Default for ValueWithPrimitiveType {
50	14	fn default() -> Self {
51	14	ValueWithPrimitiveType::U64(0)
52	14	}
53		}
54
55		/// An intermediate structed that will have it's contents populated
56		/// by the `tracing` layer for a given field.
57		/// This is done by implementing the `Visit` trait and asking the
58		/// `tracing` library to visit the fields of the captured event
59		/// and populate this struct.
60		#[derive(Default, Debug)]
61		pub struct MetricDataVisitor {
62		pub name: String,
63		value: ValueWithPrimitiveType,
64		help: String,
65		value_type: Option<CollectionKind>,
66		}
67
68		impl From<MetricDataVisitor> for CollectedMetricPrimitive {
69	14	fn from(visitor: MetricDataVisitor) -> Self {
70	14	let (value, derived_type) = match visitor.value {
71	8	ValueWithPrimitiveType::String(s) => (
72	8	CollectedMetricPrimitiveValue::String(Cow::Owned(s)),
73	8	CollectionKind::String,
74	8	),
75	6	ValueWithPrimitiveType::U64(u) => (
76	6	CollectedMetricPrimitiveValue::Counter(u),
77	6	CollectionKind::Counter,
78	6	),
79		};
80	14	CollectedMetricPrimitive {
81	14	value: Some(value),
82	14	help: visitor.help,
83	14	value_type: visitor.value_type.unwrap_or(derived_type),
84	14	}
85	14	}
86		}
87
88		impl Visit for MetricDataVisitor {
89	0	fn record_debug(&mut self, _field: &Field, _value: &dyn Debug) {}
90
91	0	fn record_f64(&mut self, field: &Field, value: f64) {
92	0	if field.name() == "__value" { Branch (92:12): [True: 0, False: 0] Branch (92:12): [Folded - Ignored]
93	0	self.value = ValueWithPrimitiveType::String(value.to_string());
94	0	}
95	0	}
96	0	fn record_i64(&mut self, field: &Field, value: i64) {
97	0	if field.name() == "__value" { Branch (97:12): [True: 0, False: 0] Branch (97:12): [Folded - Ignored]
98	0	match u64::try_from(value) {
99	0	Ok(v) => self.value = ValueWithPrimitiveType::U64(v),
100	0	Err(_) => self.value = ValueWithPrimitiveType::String(value.to_string()),
101		}
102	0	}
103	0	}
104	20	fn record_u64(&mut self, field: &Field, value: u64) {
105	20	match field.name() {
106	20	"__value" => self.value = ValueWithPrimitiveType::U64(value)6 ,
107	14	"__type" => self.value_type = Some(MetricKind::from(value).into()),
108	0	"__help" => self.help = value.to_string(),
109	0	"__name" => self.name = value.to_string(),
110	0	field => panic!("UNKNOWN FIELD {field}"),
111		}
112	20	}
113	0	fn record_i128(&mut self, field: &Field, value: i128) {
114	0	if field.name() == "__value" { Branch (114:12): [True: 0, False: 0] Branch (114:12): [Folded - Ignored]
115	0	match u64::try_from(value) {
116	0	Ok(v) => self.value = ValueWithPrimitiveType::U64(v),
117	0	Err(_) => self.value = ValueWithPrimitiveType::String(value.to_string()),
118		}
119	0	}
120	0	}
121	0	fn record_u128(&mut self, field: &Field, value: u128) {
122	0	if field.name() == "__value" { Branch (122:12): [True: 0, False: 0] Branch (122:12): [Folded - Ignored]
123	0	match u64::try_from(value) {
124	0	Ok(v) => self.value = ValueWithPrimitiveType::U64(v),
125	0	Err(_) => self.value = ValueWithPrimitiveType::String(value.to_string()),
126		}
127	0	}
128	0	}
129	0	fn record_bool(&mut self, field: &Field, value: bool) {
130	0	if field.name() == "__value" { Branch (130:12): [True: 0, False: 0] Branch (130:12): [Folded - Ignored]
131	0	self.value = ValueWithPrimitiveType::U64(u64::from(value));
132	0	}
133	0	}
134	36	fn record_str(&mut self, field: &Field, value: &str) {
135	36	match field.name() {
136	36	"__value" => self.value = ValueWithPrimitiveType::String(value.to_string())8 ,
137	28	"__help" => self.help = value.to_string()14 ,
138	14	"__name" => self.name = value.to_string(),
139	0	field => panic!("UNKNOWN FIELD {field}"),
140		}
141	36	}
142	0	fn record_error(&mut self, _field: &Field, _value: &(dyn std::error::Error + 'static)) {}
143		}
144
145		/// An intermediate structed that will have it's contents populated
146		/// by the `tracing` layer for a given field.
147		/// This is the same as `MetricDataVisitor` but only captures info
148		/// about a given span on span creation.
149		pub struct SpanFields {
150		pub name: Cow<'static, str>,
151		}
152
153		impl Visit for SpanFields {
154	0	fn record_debug(&mut self, _field: &Field, _value: &dyn Debug) {}
155
156	4	fn record_str(&mut self, field: &Field, value: &str) {
157	4	if field.name() == "__name" { Branch (157:12): [True: 4, False: 0] Branch (157:12): [Folded - Ignored]
158	4	self.name = Cow::Owned(value.to_string());
159	4	} 0
160	4	}
161		}

Coverage Report

Created: 2024-10-22 12:33