mas_context/
lib.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
7mod fmt;
8mod future;
9mod layer;
10mod service;
11
12use std::{
13    borrow::Cow,
14    sync::{
15        Arc, OnceLock,
16        atomic::{AtomicU64, Ordering},
17    },
18    time::Duration,
19};
20
21use quanta::Instant;
22use tokio::task_local;
23use ulid::Ulid;
24
25pub use self::{
26    fmt::EventFormatter,
27    future::{LogContextFuture, PollRecordingFuture},
28    layer::LogContextLayer,
29    service::LogContextService,
30};
31
32/// An identified principal making a request, recorded on the [`LogContext`].
33///
34/// This is the entity that *authenticated the request* (the caller), not the
35/// subject an operation might be about — e.g. on the introspection endpoint
36/// it is the homeserver or client doing the introspecting, not the user whose
37/// token is being introspected.
38#[derive(Clone, Debug)]
39pub enum Requester {
40    /// A user, identified by their Ulid and (when known) their username.
41    User {
42        id: Ulid,
43        /// Sometimes only the `user_id` is known at auth-resolution time (e.g.
44        /// `OAuth2`/compat bearer flows that don't load the `User` row).
45        username: Option<String>,
46    },
47
48    /// An `OAuth2` client acting on its own behalf (e.g. at the token endpoint,
49    /// or a client-credentials session), identified by its internal Ulid.
50    OAuth2Client { id: Ulid },
51
52    /// The homeserver, e.g. calling the introspection endpoint with its shared
53    /// secret.
54    Homeserver,
55}
56
57impl Requester {
58    /// Build a `User` requester with both an id and a known username.
59    #[must_use]
60    pub fn user(id: Ulid, username: impl Into<String>) -> Self {
61        Self::User {
62            id,
63            username: Some(username.into()),
64        }
65    }
66
67    /// Build a `User` requester from a `user_id` alone (username unknown).
68    #[must_use]
69    pub fn user_id_only(id: Ulid) -> Self {
70        Self::User { id, username: None }
71    }
72
73    /// Build an `OAuth2Client` requester from a client's internal Ulid.
74    #[must_use]
75    pub fn oauth2_client(id: Ulid) -> Self {
76        Self::OAuth2Client { id }
77    }
78}
79
80impl std::fmt::Display for Requester {
81    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
82        match self {
83            Self::User {
84                id,
85                username: Some(name),
86            } => write!(f, "user:{id}({name})"),
87            Self::User { id, username: None } => write!(f, "user:{id}"),
88            Self::OAuth2Client { id } => write!(f, "oauth2-client:{id}"),
89            Self::Homeserver => write!(f, "homeserver"),
90        }
91    }
92}
93
94/// A counter which increments each time we create a new log context
95/// It will wrap around if we create more than [`u64::MAX`] contexts
96static LOG_CONTEXT_INDEX: AtomicU64 = AtomicU64::new(0);
97task_local! {
98    pub static CURRENT_LOG_CONTEXT: LogContext;
99}
100
101/// A log context saves informations about the current task, such as the
102/// elapsed time, the number of polls, and the poll time.
103#[derive(Clone)]
104pub struct LogContext {
105    inner: Arc<LogContextInner>,
106}
107
108struct LogContextInner {
109    /// A user-defined tag for the log context
110    tag: Cow<'static, str>,
111
112    /// A unique index for the log context
113    index: u64,
114
115    /// The time when the context was created
116    start: Instant,
117
118    /// The number of [`Future::poll`] recorded
119    polls: AtomicU64,
120
121    /// An approximation of the total CPU time spent in the context, in
122    /// nanoseconds
123    cpu_time: AtomicU64,
124
125    /// The number of database queries executed in the context
126    db_queries: AtomicU64,
127
128    /// The cumulative wall-clock time spent executing database queries in the
129    /// context, in nanoseconds
130    db_time: AtomicU64,
131
132    /// The number of rows fetched from the database in the context
133    db_rows_fetched: AtomicU64,
134
135    /// The identified principal making the request, if it has been resolved.
136    /// Set once per context, first writer wins.
137    requester: OnceLock<Requester>,
138}
139
140impl LogContext {
141    /// Create a new log context with the given tag
142    pub fn new(tag: impl Into<Cow<'static, str>>) -> Self {
143        let tag = tag.into();
144        let inner = LogContextInner {
145            tag,
146            index: LOG_CONTEXT_INDEX.fetch_add(1, Ordering::Relaxed),
147            start: Instant::now(),
148            polls: AtomicU64::new(0),
149            cpu_time: AtomicU64::new(0),
150            db_queries: AtomicU64::new(0),
151            db_rows_fetched: AtomicU64::new(0),
152            db_time: AtomicU64::new(0),
153            requester: OnceLock::new(),
154        };
155
156        Self {
157            inner: Arc::new(inner),
158        }
159    }
160
161    /// Run a closure with the current log context, if any
162    pub fn maybe_with<F, R>(f: F) -> Option<R>
163    where
164        F: FnOnce(&Self) -> R,
165    {
166        CURRENT_LOG_CONTEXT.try_with(f).ok()
167    }
168
169    /// Run the async function `f` with the given log context. It will wrap the
170    /// output future to record poll and CPU statistics.
171    pub fn run<F: FnOnce() -> Fut, Fut: Future>(&self, f: F) -> LogContextFuture<Fut> {
172        let future = self.run_sync(f);
173        self.wrap_future(future)
174    }
175
176    /// Run the sync function `f` with the given log context, recording the CPU
177    /// time spent.
178    pub fn run_sync<F: FnOnce() -> R, R>(&self, f: F) -> R {
179        let start = Instant::now();
180        let result = CURRENT_LOG_CONTEXT.sync_scope(self.clone(), f);
181        let elapsed = start.elapsed().as_nanos().try_into().unwrap_or(u64::MAX);
182        self.inner.cpu_time.fetch_add(elapsed, Ordering::Relaxed);
183        result
184    }
185
186    /// Record the stats of a query on the current log context.
187    pub fn maybe_record_query_stats(fetched: usize, duration: Duration) {
188        LogContext::maybe_with(|ctx| ctx.record_query_stats(fetched, duration));
189    }
190
191    /// Record the stats of a query
192    pub fn record_query_stats(&self, fetched: usize, duration: Duration) {
193        let nanos = duration.as_nanos().try_into().unwrap_or(u64::MAX);
194        self.inner.db_time.fetch_add(nanos, Ordering::Relaxed);
195        self.inner
196            .db_rows_fetched
197            .fetch_add(fetched as u64, Ordering::Relaxed);
198        self.inner.db_queries.fetch_add(1, Ordering::Relaxed);
199    }
200
201    /// Associate a [`Requester`] with this log context. Silently does nothing
202    /// if a requester has already been recorded (first writer wins).
203    pub fn set_requester(&self, requester: Requester) {
204        let _ = self.inner.requester.set(requester);
205    }
206
207    /// Get the [`Requester`] associated with this log context, if one was set.
208    #[must_use]
209    pub fn requester(&self) -> Option<&Requester> {
210        self.inner.requester.get()
211    }
212
213    /// Convenience: set the [`Requester`] on the current task-local
214    /// `LogContext`, if one exists. No-op outside of a `LogContext` scope.
215    pub fn maybe_set_requester(requester: Requester) {
216        Self::maybe_with(|ctx| ctx.set_requester(requester));
217    }
218
219    /// Create a snapshot of the log context statistics
220    #[must_use]
221    pub fn stats(&self) -> LogContextStats {
222        let polls = self.inner.polls.load(Ordering::Relaxed);
223        let cpu_time = self.inner.cpu_time.load(Ordering::Relaxed);
224        let cpu_time = Duration::from_nanos(cpu_time);
225        let elapsed = self.inner.start.elapsed();
226        let db_queries = self.inner.db_queries.load(Ordering::Relaxed);
227        let db_rows_fetched = self.inner.db_rows_fetched.load(Ordering::Relaxed);
228        let db_time = Duration::from_nanos(self.inner.db_time.load(Ordering::Relaxed));
229        LogContextStats {
230            polls,
231            cpu_time,
232            elapsed,
233            db_queries,
234            db_rows_fetched,
235            db_time,
236        }
237    }
238}
239
240impl std::fmt::Display for LogContext {
241    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
242        let tag = &self.inner.tag;
243        let index = self.inner.index;
244        write!(f, "{tag}-{index}")
245    }
246}
247
248/// A snapshot of a log context statistics
249#[derive(Debug, Clone, Copy)]
250pub struct LogContextStats {
251    /// How many times the context was polled
252    pub polls: u64,
253
254    /// The approximate CPU time spent in the context
255    pub cpu_time: Duration,
256
257    /// How much time elapsed since the context was created
258    pub elapsed: Duration,
259
260    /// How many database queries were executed in the context
261    pub db_queries: u64,
262
263    /// The number of rows fetched from the database in the context
264    pub db_rows_fetched: u64,
265
266    /// The cumulative wall-clock time spent executing database queries
267    pub db_time: Duration,
268}
269
270impl std::fmt::Display for LogContextStats {
271    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
272        let polls = self.polls;
273        #[expect(clippy::cast_precision_loss)]
274        let cpu_time_ms = self.cpu_time.as_nanos() as f64 / 1_000_000.;
275        #[expect(clippy::cast_precision_loss)]
276        let elapsed_ms = self.elapsed.as_nanos() as f64 / 1_000_000.;
277        let db_queries = self.db_queries;
278        let db_rows_fetched = self.db_rows_fetched;
279        #[expect(clippy::cast_precision_loss)]
280        let db_time_ms = self.db_time.as_nanos() as f64 / 1_000_000.;
281        write!(
282            f,
283            "polls: {polls}, cpu: {cpu_time_ms:.1}ms, db: {db_time_ms:.1}ms, elapsed: {elapsed_ms:.1}ms, queries: {db_queries}, fetched: {db_rows_fetched}",
284        )
285    }
286}
mas_context/lib.rs

mas_context/
lib.rs
