mas_context/
future.rs

1// Copyright 2025, 2026 Element Creations Ltd.
2// Copyright 2025 New Vector Ltd.
3//
4// SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial
5// Please see LICENSE files in the repository root for full details.
6
7use std::{
8    pin::Pin,
9    sync::{Arc, atomic::Ordering},
10    task::{Context, Poll},
11};
12
13use quanta::Instant;
14use tokio::task::futures::TaskLocalFuture;
15
16use crate::LogContext;
17
18pub type LogContextFuture<F> = TaskLocalFuture<crate::LogContext, PollRecordingFuture<F>>;
19
20impl LogContext {
21    /// Wrap a future with the given log context
22    pub(crate) fn wrap_future<F: Future>(&self, future: F) -> LogContextFuture<F> {
23        let future = PollRecordingFuture::new(future, self.inner.hub.clone());
24        crate::CURRENT_LOG_CONTEXT.scope(self.clone(), future)
25    }
26}
27
28pin_project_lite::pin_project! {
29    /// A future which records the elapsed time and the number of polls in the
30    /// active log context, and polls the inner future with the Sentry hub of
31    /// the log context
32    pub struct PollRecordingFuture<F> {
33        #[pin]
34        inner: F,
35        hub: Arc<sentry::Hub>,
36    }
37}
38
39impl<F: Future> PollRecordingFuture<F> {
40    pub(crate) fn new(inner: F, hub: Arc<sentry::Hub>) -> Self {
41        Self { inner, hub }
42    }
43}
44
45impl<F: Future> Future for PollRecordingFuture<F> {
46    type Output = F::Output;
47
48    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
49        let start = Instant::now();
50        let this = self.project();
51        let result = sentry::Hub::run(this.hub.clone(), || this.inner.poll(cx));
52
53        // Record the number of polls and the time we spent polling the future
54        let elapsed = start.elapsed().as_nanos().try_into().unwrap_or(u64::MAX);
55        let _ = crate::CURRENT_LOG_CONTEXT.try_with(|c| {
56            c.inner.polls.fetch_add(1, Ordering::Relaxed);
57            c.inner.cpu_time.fetch_add(elapsed, Ordering::Relaxed);
58        });
59
60        result
61    }
62}
mas_context/future.rs

mas_context/
future.rs
