mas_context/
lib.rs

1// Copyright 2025 New Vector Ltd.
2//
3// SPDX-License-Identifier: AGPL-3.0-only
4// Please see LICENSE in the repository root for full details.
5
6mod fmt;
7mod future;
8mod layer;
9mod service;
10
11use std::{
12    borrow::Cow,
13    sync::{
14        Arc,
15        atomic::{AtomicU64, Ordering},
16    },
17    time::Duration,
18};
19
20use quanta::Instant;
21use tokio::task_local;
22
23pub use self::{
24    fmt::EventFormatter,
25    future::{LogContextFuture, PollRecordingFuture},
26    layer::LogContextLayer,
27    service::LogContextService,
28};
29
30/// A counter which increments each time we create a new log context
31/// It will wrap around if we create more than [`u64::MAX`] contexts
32static LOG_CONTEXT_INDEX: AtomicU64 = AtomicU64::new(0);
33task_local! {
34    pub static CURRENT_LOG_CONTEXT: LogContext;
35}
36
37/// A log context saves informations about the current task, such as the
38/// elapsed time, the number of polls, and the poll time.
39#[derive(Clone)]
40pub struct LogContext {
41    inner: Arc<LogContextInner>,
42}
43
44struct LogContextInner {
45    /// A user-defined tag for the log context
46    tag: Cow<'static, str>,
47
48    /// A unique index for the log context
49    index: u64,
50
51    /// The time when the context was created
52    start: Instant,
53
54    /// The number of [`Future::poll`] recorded
55    polls: AtomicU64,
56
57    /// An approximation of the total CPU time spent in the context, in
58    /// nanoseconds
59    cpu_time: AtomicU64,
60}
61
62impl LogContext {
63    /// Create a new log context with the given tag
64    pub fn new(tag: impl Into<Cow<'static, str>>) -> Self {
65        let tag = tag.into();
66        let inner = LogContextInner {
67            tag,
68            index: LOG_CONTEXT_INDEX.fetch_add(1, Ordering::Relaxed),
69            start: Instant::now(),
70            polls: AtomicU64::new(0),
71            cpu_time: AtomicU64::new(0),
72        };
73
74        Self {
75            inner: Arc::new(inner),
76        }
77    }
78
79    /// Run a closure with the current log context, if any
80    pub fn maybe_with<F, R>(f: F) -> Option<R>
81    where
82        F: FnOnce(&Self) -> R,
83    {
84        CURRENT_LOG_CONTEXT.try_with(f).ok()
85    }
86
87    /// Run the async function `f` with the given log context. It will wrap the
88    /// output future to record poll and CPU statistics.
89    pub fn run<F: FnOnce() -> Fut, Fut: Future>(&self, f: F) -> LogContextFuture<Fut> {
90        let future = self.run_sync(f);
91        self.wrap_future(future)
92    }
93
94    /// Run the sync function `f` with the given log context, recording the CPU
95    /// time spent.
96    pub fn run_sync<F: FnOnce() -> R, R>(&self, f: F) -> R {
97        let start = Instant::now();
98        let result = CURRENT_LOG_CONTEXT.sync_scope(self.clone(), f);
99        let elapsed = start.elapsed().as_nanos().try_into().unwrap_or(u64::MAX);
100        self.inner.cpu_time.fetch_add(elapsed, Ordering::Relaxed);
101        result
102    }
103
104    /// Create a snapshot of the log context statistics
105    #[must_use]
106    pub fn stats(&self) -> LogContextStats {
107        let polls = self.inner.polls.load(Ordering::Relaxed);
108        let cpu_time = self.inner.cpu_time.load(Ordering::Relaxed);
109        let cpu_time = Duration::from_nanos(cpu_time);
110        let elapsed = self.inner.start.elapsed();
111        LogContextStats {
112            polls,
113            cpu_time,
114            elapsed,
115        }
116    }
117}
118
119impl std::fmt::Display for LogContext {
120    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
121        let tag = &self.inner.tag;
122        let index = self.inner.index;
123        write!(f, "{tag}-{index}")
124    }
125}
126
127/// A snapshot of a log context statistics
128#[derive(Debug, Clone, Copy)]
129pub struct LogContextStats {
130    /// How many times the context was polled
131    pub polls: u64,
132
133    /// The approximate CPU time spent in the context
134    pub cpu_time: Duration,
135
136    /// How much time elapsed since the context was created
137    pub elapsed: Duration,
138}
139
140impl std::fmt::Display for LogContextStats {
141    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
142        let polls = self.polls;
143        #[expect(clippy::cast_precision_loss)]
144        let cpu_time_ms = self.cpu_time.as_nanos() as f64 / 1_000_000.;
145        #[expect(clippy::cast_precision_loss)]
146        let elapsed_ms = self.elapsed.as_nanos() as f64 / 1_000_000.;
147        write!(
148            f,
149            "polls: {polls}, cpu: {cpu_time_ms:.1}ms, elapsed: {elapsed_ms:.1}ms",
150        )
151    }
152}
mas_context/lib.rs

mas_context/
lib.rs
