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

mas_tower/tracing/
future.rs
