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    /// The Sentry hub bound to the context, which starts with an empty scope
140    /// so that unrelated tasks don't share breadcrumbs or spans. Scope
141    /// changes made on the main hub are not inherited: global data goes on
142    /// the client options or in `before_send`.
143    hub: Arc<sentry::Hub>,
144}
145
146impl LogContext {
147    /// Create a new log context with the given tag
148    pub fn new(tag: impl Into<Cow<'static, str>>) -> Self {
149        let tag = tag.into();
150        let inner = LogContextInner {
151            tag,
152            index: LOG_CONTEXT_INDEX.fetch_add(1, Ordering::Relaxed),
153            start: Instant::now(),
154            polls: AtomicU64::new(0),
155            cpu_time: AtomicU64::new(0),
156            db_queries: AtomicU64::new(0),
157            db_rows_fetched: AtomicU64::new(0),
158            db_time: AtomicU64::new(0),
159            requester: OnceLock::new(),
160            hub: Arc::new(sentry::Hub::new(
161                sentry::Hub::main().client(),
162                Arc::default(),
163            )),
164        };
165
166        Self {
167            inner: Arc::new(inner),
168        }
169    }
170
171    /// Get the tag of the log context
172    #[must_use]
173    pub fn tag(&self) -> &Cow<'static, str> {
174        &self.inner.tag
175    }
176
177    /// Run a closure with the current log context, if any
178    pub fn maybe_with<F, R>(f: F) -> Option<R>
179    where
180        F: FnOnce(&Self) -> R,
181    {
182        CURRENT_LOG_CONTEXT.try_with(f).ok()
183    }
184
185    /// Run the async function `f` with the given log context. It will wrap the
186    /// output future to record poll and CPU statistics.
187    pub fn run<F: FnOnce() -> Fut, Fut: Future>(&self, f: F) -> LogContextFuture<Fut> {
188        let future = self.run_sync(f);
189        self.wrap_future(future)
190    }
191
192    /// Run the sync function `f` with the given log context, recording the CPU
193    /// time spent.
194    pub fn run_sync<F: FnOnce() -> R, R>(&self, f: F) -> R {
195        let start = Instant::now();
196        let result = sentry::Hub::run(self.inner.hub.clone(), || {
197            CURRENT_LOG_CONTEXT.sync_scope(self.clone(), f)
198        });
199        let elapsed = start.elapsed().as_nanos().try_into().unwrap_or(u64::MAX);
200        self.inner.cpu_time.fetch_add(elapsed, Ordering::Relaxed);
201        result
202    }
203
204    /// Record the stats of a query on the current log context.
205    pub fn maybe_record_query_stats(fetched: usize, duration: Duration) {
206        LogContext::maybe_with(|ctx| ctx.record_query_stats(fetched, duration));
207    }
208
209    /// Record the stats of a query
210    pub fn record_query_stats(&self, fetched: usize, duration: Duration) {
211        let nanos = duration.as_nanos().try_into().unwrap_or(u64::MAX);
212        self.inner.db_time.fetch_add(nanos, Ordering::Relaxed);
213        self.inner
214            .db_rows_fetched
215            .fetch_add(fetched as u64, Ordering::Relaxed);
216        self.inner.db_queries.fetch_add(1, Ordering::Relaxed);
217    }
218
219    /// Associate a [`Requester`] with this log context. Silently does nothing
220    /// if a requester has already been recorded (first writer wins).
221    pub fn set_requester(&self, requester: Requester) {
222        let _ = self.inner.requester.set(requester);
223    }
224
225    /// Get the [`Requester`] associated with this log context, if one was set.
226    #[must_use]
227    pub fn requester(&self) -> Option<&Requester> {
228        self.inner.requester.get()
229    }
230
231    /// Convenience: set the [`Requester`] on the current task-local
232    /// `LogContext`, if one exists. No-op outside of a `LogContext` scope.
233    pub fn maybe_set_requester(requester: Requester) {
234        Self::maybe_with(|ctx| ctx.set_requester(requester));
235    }
236
237    /// Create a snapshot of the log context statistics
238    #[must_use]
239    pub fn stats(&self) -> LogContextStats {
240        let polls = self.inner.polls.load(Ordering::Relaxed);
241        let cpu_time = self.inner.cpu_time.load(Ordering::Relaxed);
242        let cpu_time = Duration::from_nanos(cpu_time);
243        let elapsed = self.inner.start.elapsed();
244        let db_queries = self.inner.db_queries.load(Ordering::Relaxed);
245        let db_rows_fetched = self.inner.db_rows_fetched.load(Ordering::Relaxed);
246        let db_time = Duration::from_nanos(self.inner.db_time.load(Ordering::Relaxed));
247        LogContextStats {
248            polls,
249            cpu_time,
250            elapsed,
251            db_queries,
252            db_rows_fetched,
253            db_time,
254        }
255    }
256}
257
258impl std::fmt::Display for LogContext {
259    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
260        let tag = &self.inner.tag;
261        let index = self.inner.index;
262        write!(f, "{tag}-{index}")
263    }
264}
265
266/// A snapshot of a log context statistics
267#[derive(Debug, Clone, Copy)]
268pub struct LogContextStats {
269    /// How many times the context was polled
270    pub polls: u64,
271
272    /// The approximate CPU time spent in the context
273    pub cpu_time: Duration,
274
275    /// How much time elapsed since the context was created
276    pub elapsed: Duration,
277
278    /// How many database queries were executed in the context
279    pub db_queries: u64,
280
281    /// The number of rows fetched from the database in the context
282    pub db_rows_fetched: u64,
283
284    /// The cumulative wall-clock time spent executing database queries
285    pub db_time: Duration,
286}
287
288impl std::fmt::Display for LogContextStats {
289    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
290        let polls = self.polls;
291        #[expect(clippy::cast_precision_loss)]
292        let cpu_time_ms = self.cpu_time.as_nanos() as f64 / 1_000_000.;
293        #[expect(clippy::cast_precision_loss)]
294        let elapsed_ms = self.elapsed.as_nanos() as f64 / 1_000_000.;
295        let db_queries = self.db_queries;
296        let db_rows_fetched = self.db_rows_fetched;
297        #[expect(clippy::cast_precision_loss)]
298        let db_time_ms = self.db_time.as_nanos() as f64 / 1_000_000.;
299        write!(
300            f,
301            "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}",
302        )
303    }
304}
mas_context/lib.rs

mas_context/
lib.rs
