mas_tower/tracing/
service.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::Service;
16
17use super::future::TraceFuture;
18
19#[derive(Clone, Debug)]
20pub struct TraceService<S, MakeSpan, OnResponse = (), OnError = ()> {
21    inner: S,
22    make_span: MakeSpan,
23    on_response: OnResponse,
24    on_error: OnError,
25}
26
27impl<S, MakeSpan, OnResponse, OnError> TraceService<S, MakeSpan, OnResponse, OnError> {
28    /// Create a new [`TraceService`].
29    #[must_use]
30    pub fn new(inner: S, make_span: MakeSpan, on_response: OnResponse, on_error: OnError) -> Self {
31        Self {
32            inner,
33            make_span,
34            on_response,
35            on_error,
36        }
37    }
38}
39
40impl<R, S, MakeSpan, OnResponse, OnError> Service<R>
41    for TraceService<S, MakeSpan, OnResponse, OnError>
42where
43    S: Service<R>,
44    MakeSpan: super::make_span::MakeSpan<R>,
45    OnResponse: super::enrich_span::EnrichSpan<S::Response> + Clone,
46    OnError: super::enrich_span::EnrichSpan<S::Error> + Clone,
47{
48    type Response = S::Response;
49    type Error = S::Error;
50    type Future = TraceFuture<S::Future, OnResponse, OnError>;
51
52    fn poll_ready(
53        &mut self,
54        cx: &mut std::task::Context<'_>,
55    ) -> std::task::Poll<Result<(), Self::Error>> {
56        self.inner.poll_ready(cx)
57    }
58
59    fn call(&mut self, request: R) -> Self::Future {
60        let span = self.make_span.make_span(&request);
61        let guard = span.enter();
62        let inner = self.inner.call(request);
63        drop(guard);
64
65        TraceFuture::new(inner, span, self.on_response.clone(), self.on_error.clone())
66    }
67}
mas_tower/tracing/service.rs

mas_tower/tracing/
service.rs
