mas_tower/tracing/
future.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 std::{future::Future, task::ready};
16
17use pin_project_lite::pin_project;
18use tracing::Span;
19
20pin_project! {
21    pub struct TraceFuture<F, OnResponse, OnError> {
22        #[pin]
23        inner: F,
24        span: Span,
25        on_response: OnResponse,
26        on_error: OnError,
27    }
28}
29
30impl<F, OnResponse, OnError> TraceFuture<F, OnResponse, OnError> {
31    pub fn new(inner: F, span: Span, on_response: OnResponse, on_error: OnError) -> Self {
32        Self {
33            inner,
34            span,
35            on_response,
36            on_error,
37        }
38    }
39}
40
41impl<F, R, E, OnResponse, OnError> Future for TraceFuture<F, OnResponse, OnError>
42where
43    F: Future<Output = Result<R, E>>,
44    OnResponse: super::enrich_span::EnrichSpan<R>,
45    OnError: super::enrich_span::EnrichSpan<E>,
46{
47    type Output = Result<R, E>;
48
49    fn poll(
50        self: std::pin::Pin<&mut Self>,
51        cx: &mut std::task::Context<'_>,
52    ) -> std::task::Poll<Self::Output> {
53        let this = self.project();
54
55        // Poll the inner future, with the span entered. This is effectively what
56        // [`tracing::Instrumented`] does.
57        let _guard = this.span.enter();
58        let result = ready!(this.inner.poll(cx));
59
60        match &result {
61            Ok(response) => {
62                this.on_response.enrich_span(this.span, response);
63            }
64            Err(error) => {
65                this.on_error.enrich_span(this.span, error);
66            }
67        }
68
69        std::task::Poll::Ready(result)
70    }
71}
mas_tower/tracing/future.rs

mas_tower/tracing/
future.rs
