mas_tower/tracing/
enrich_span.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 tracing::{Span, Value};
16
17use crate::utils::{FnWrapper, KV};
18
19/// A trait for enriching a span with information a structure.
20pub trait EnrichSpan<T> {
21    fn enrich_span(&self, span: &Span, t: &T);
22}
23
24impl<T, F> EnrichSpan<T> for FnWrapper<F>
25where
26    F: Fn(&Span, &T),
27{
28    fn enrich_span(&self, span: &Span, t: &T) {
29        (self.0)(span, t);
30    }
31}
32
33/// Enrich span from a function.
34#[must_use]
35pub fn enrich_span_fn<T, F>(f: F) -> FnWrapper<F>
36where
37    F: Fn(&Span, &T),
38{
39    FnWrapper(f)
40}
41
42impl<T> EnrichSpan<T> for () {
43    fn enrich_span(&self, _span: &Span, _t: &T) {}
44}
45
46impl<V, T> EnrichSpan<T> for KV<V>
47where
48    V: Value,
49{
50    fn enrich_span(&self, span: &Span, _t: &T) {
51        span.record(self.0, &self.1);
52    }
53}
54
55/// A macro to implement [`EnrichSpan`] for a tuple of types that implement
56/// [`EnrichSpan`].
57macro_rules! impl_for_tuple {
58    ($($T:ident),+) => {
59        impl<T, $($T),+> EnrichSpan<T> for ($($T,)+)
60        where
61            $($T: EnrichSpan<T>),+
62        {
63            fn enrich_span(&self, span: &Span, t: &T) {
64                #[allow(non_snake_case)]
65                let ($(ref $T,)+) = *self;
66                $(
67                    $T.enrich_span(span, t);
68                )+
69            }
70        }
71    };
72}
73
74impl_for_tuple!(T1);
75impl_for_tuple!(T1, T2);
76impl_for_tuple!(T1, T2, T3);
77impl_for_tuple!(T1, T2, T3, T4);
78impl_for_tuple!(T1, T2, T3, T4, T5);
79impl_for_tuple!(T1, T2, T3, T4, T5, T6);
80impl_for_tuple!(T1, T2, T3, T4, T5, T6, T7);
81impl_for_tuple!(T1, T2, T3, T4, T5, T6, T7, T8);
82
83impl<T, R> EnrichSpan<R> for Option<T>
84where
85    T: EnrichSpan<R>,
86{
87    fn enrich_span(&self, span: &Span, request: &R) {
88        if let Some(ref t) = *self {
89            t.enrich_span(span, request);
90        }
91    }
92}
93
94impl<T, R, const N: usize> EnrichSpan<R> for [T; N]
95where
96    T: EnrichSpan<R>,
97{
98    fn enrich_span(&self, span: &Span, request: &R) {
99        for t in self {
100            t.enrich_span(span, request);
101        }
102    }
103}
104
105impl<T, R> EnrichSpan<R> for Vec<T>
106where
107    T: EnrichSpan<R>,
108{
109    fn enrich_span(&self, span: &Span, request: &R) {
110        for t in self {
111            t.enrich_span(span, request);
112        }
113    }
114}
mas_tower/tracing/enrich_span.rs

mas_tower/tracing/
enrich_span.rs
