mas_tower/tracing/
layer.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 tower::Layer;
16use tracing::Span;
17
18use crate::{enrich_span_fn, make_span_fn, utils::FnWrapper};
19
20#[derive(Clone, Debug)]
21pub struct TraceLayer<MakeSpan, OnResponse = (), OnError = ()> {
22    make_span: MakeSpan,
23    on_response: OnResponse,
24    on_error: OnError,
25}
26
27impl<F> TraceLayer<FnWrapper<F>> {
28    #[must_use]
29    pub fn from_fn<T>(f: F) -> TraceLayer<FnWrapper<F>>
30    where
31        F: Fn(&T) -> Span,
32    {
33        TraceLayer::new(make_span_fn(f))
34    }
35}
36
37impl<MakeSpan> TraceLayer<MakeSpan> {
38    #[must_use]
39    pub fn new(make_span: MakeSpan) -> Self {
40        Self {
41            make_span,
42            on_response: (),
43            on_error: (),
44        }
45    }
46}
47
48impl<MakeSpan, OnResponse, OnError> TraceLayer<MakeSpan, OnResponse, OnError> {
49    #[must_use]
50    pub fn on_response<NewOnResponse>(
51        self,
52        on_response: NewOnResponse,
53    ) -> TraceLayer<MakeSpan, NewOnResponse, OnError> {
54        TraceLayer {
55            make_span: self.make_span,
56            on_response,
57            on_error: self.on_error,
58        }
59    }
60
61    #[must_use]
62    pub fn on_response_fn<F, T>(self, f: F) -> TraceLayer<MakeSpan, FnWrapper<F>, OnError>
63    where
64        F: Fn(&Span, &T),
65    {
66        self.on_response(enrich_span_fn(f))
67    }
68
69    #[must_use]
70    pub fn on_error<NewOnError>(
71        self,
72        on_error: NewOnError,
73    ) -> TraceLayer<MakeSpan, OnResponse, NewOnError> {
74        TraceLayer {
75            make_span: self.make_span,
76            on_response: self.on_response,
77            on_error,
78        }
79    }
80
81    pub fn on_error_fn<F, E>(self, f: F) -> TraceLayer<MakeSpan, OnResponse, FnWrapper<F>>
82    where
83        F: Fn(&Span, &E),
84    {
85        self.on_error(enrich_span_fn(f))
86    }
87}
88
89impl<S, MakeSpan, OnResponse, OnError> Layer<S> for TraceLayer<MakeSpan, OnResponse, OnError>
90where
91    MakeSpan: Clone,
92    OnResponse: Clone,
93    OnError: Clone,
94{
95    type Service = super::service::TraceService<S, MakeSpan, OnResponse, OnError>;
96
97    fn layer(&self, inner: S) -> Self::Service {
98        super::service::TraceService::new(
99            inner,
100            self.make_span.clone(),
101            self.on_response.clone(),
102            self.on_error.clone(),
103        )
104    }
105}
mas_tower/tracing/layer.rs

mas_tower/tracing/
layer.rs
