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

Coverage Report

Created: 2024-12-19 04:04