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

mas_tower/tracing/
enrich_span.rs
