mas_tower/metrics/
make_attributes.rs

1// Copyright 2024 New Vector Ltd.
2// Copyright 2023, 2024 The Matrix.org Foundation C.I.C.
3//
4// SPDX-License-Identifier: AGPL-3.0-only
5// Please see LICENSE in the repository root for full details.
6
7use opentelemetry::{KeyValue, Value};
8
9use crate::{FnWrapper, utils::KV};
10
11/// Make metrics attributes from a type.
12pub trait MetricsAttributes<T> {
13    type Iter<'a>: Iterator<Item = KeyValue>
14    where
15        Self: 'a,
16        T: 'a;
17
18    fn attributes<'a>(&'a self, t: &'a T) -> Self::Iter<'a>;
19}
20
21pub fn metrics_attributes_fn<T, F>(f: F) -> FnWrapper<F>
22where
23    F: Fn(&T) -> Vec<KeyValue> + 'static,
24    T: 'static,
25{
26    FnWrapper(f)
27}
28
29impl<T, F> MetricsAttributes<T> for FnWrapper<F>
30where
31    F: Fn(&T) -> Vec<KeyValue> + 'static,
32    T: 'static,
33{
34    type Iter<'a> = std::vec::IntoIter<KeyValue>;
35
36    fn attributes<'a>(&'a self, t: &'a T) -> Self::Iter<'a> {
37        let values: Vec<KeyValue> = self.0(t);
38        values.into_iter()
39    }
40}
41
42impl<T> MetricsAttributes<T> for ()
43where
44    T: 'static,
45{
46    type Iter<'a> = std::iter::Empty<KeyValue>;
47
48    fn attributes(&self, _t: &T) -> Self::Iter<'_> {
49        std::iter::empty()
50    }
51}
52
53impl<V, T> MetricsAttributes<T> for Vec<V>
54where
55    V: MetricsAttributes<T> + 'static,
56    T: 'static,
57{
58    type Iter<'a> = Box<dyn Iterator<Item = KeyValue> + 'a>;
59    fn attributes<'a>(&'a self, t: &'a T) -> Self::Iter<'a> {
60        Box::new(self.iter().flat_map(|v| v.attributes(t)))
61    }
62}
63
64impl<V, T, const N: usize> MetricsAttributes<T> for [V; N]
65where
66    V: MetricsAttributes<T> + 'static,
67    T: 'static,
68{
69    type Iter<'a> = Box<dyn Iterator<Item = KeyValue> + 'a>;
70    fn attributes<'a>(&'a self, t: &'a T) -> Self::Iter<'a> {
71        Box::new(self.iter().flat_map(|v| v.attributes(t)))
72    }
73}
74
75impl<V, T> MetricsAttributes<T> for KV<V>
76where
77    V: Into<Value> + Clone + 'static,
78    T: 'static,
79{
80    type Iter<'a> = std::iter::Once<KeyValue>;
81    fn attributes(&self, _t: &T) -> Self::Iter<'_> {
82        std::iter::once(KeyValue::new(self.0, self.1.clone().into()))
83    }
84}
85
86impl<T> MetricsAttributes<T> for KeyValue
87where
88    T: 'static,
89{
90    type Iter<'a> = std::iter::Once<KeyValue>;
91    fn attributes(&self, _t: &T) -> Self::Iter<'_> {
92        std::iter::once(self.clone())
93    }
94}
95
96impl<V, T> MetricsAttributes<T> for Option<V>
97where
98    V: MetricsAttributes<T> + 'static,
99    T: 'static,
100{
101    type Iter<'a> = std::iter::Flatten<std::option::IntoIter<V::Iter<'a>>>;
102
103    fn attributes<'a>(&'a self, t: &'a T) -> Self::Iter<'a> {
104        self.as_ref().map(|v| v.attributes(t)).into_iter().flatten()
105    }
106}
107
108macro_rules! chain_for {
109    // Sub-macro for reversing the list of types.
110    (@reverse ($( $reversed:ident ,)*)) => {
111        chain_for!(@build_chain $($reversed),*)
112    };
113    (@reverse ($($reversed:ident,)*) $head:ident $(, $tail:ident)*) => {
114        chain_for!(@reverse ($head, $($reversed,)*) $($tail),*)
115    };
116
117    // Sub-macro for building the chain of iterators.
118    (@build_chain $last:ident) => {
119        $last::Iter<'a>
120    };
121    (@build_chain $head:ident, $($tail:ident),*) => {
122        std::iter::Chain<chain_for!(@build_chain $($tail),*), $head::Iter<'a>>
123    };
124
125    ($($idents:ident),+) => {
126        chain_for!(@reverse () $($idents),+)
127    };
128}
129
130macro_rules! impl_for_tuple {
131    ($first:ident $(,$rest:ident)*) => {
132        impl<T, $first, $($rest,)*> MetricsAttributes<T> for ($first, $($rest,)*)
133        where
134            T: 'static,
135            $first: MetricsAttributes<T> + 'static,
136            $($rest: MetricsAttributes<T> + 'static,)*
137        {
138            type Iter<'a> = chain_for!($first $(, $rest)*);
139            fn attributes<'a>(&'a self, t: &'a T) -> Self::Iter<'a> {
140                #[allow(non_snake_case)]
141                let (head, $($rest,)*) = self;
142                head.attributes(t)
143                    $(.chain($rest.attributes(t)))*
144            }
145        }
146    };
147}
148
149impl_for_tuple!(V1);
150impl_for_tuple!(V1, V2);
151impl_for_tuple!(V1, V2, V3);
152impl_for_tuple!(V1, V2, V3, V4);
153impl_for_tuple!(V1, V2, V3, V4, V5);
154impl_for_tuple!(V1, V2, V3, V4, V5, V6);
155impl_for_tuple!(V1, V2, V3, V4, V5, V6, V7);
156impl_for_tuple!(V1, V2, V3, V4, V5, V6, V7, V8);
157impl_for_tuple!(V1, V2, V3, V4, V5, V6, V7, V8, V9);