mas_tower/metrics/
make_attributes.rs

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