matrix-authentication-service/crates/context/src/lib.rs

// Copyright 2025, 2026 Element Creations Ltd.
// Copyright 2025 New Vector Ltd.
//
// SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial
// Please see LICENSE files in the repository root for full details.

mod fmt;
mod future;
mod layer;
mod service;

use std::{
    borrow::Cow,
    sync::{
        Arc, OnceLock,
        atomic::{AtomicU64, Ordering},
    },
    time::Duration,
};

use quanta::Instant;
use tokio::task_local;
use ulid::Ulid;

pub use self::{
    fmt::EventFormatter,
    future::{LogContextFuture, PollRecordingFuture},
    layer::LogContextLayer,
    service::LogContextService,
};

/// An identified principal making a request, recorded on the [`LogContext`].
///
/// This is the entity that *authenticated the request* (the caller), not the
/// subject an operation might be about — e.g. on the introspection endpoint
/// it is the homeserver or client doing the introspecting, not the user whose
/// token is being introspected.
#[derive(Clone, Debug)]
pub enum Requester {
    /// A user, identified by their Ulid and (when known) their username.
    User {
        id: Ulid,
        /// Sometimes only the `user_id` is known at auth-resolution time (e.g.
        /// `OAuth2`/compat bearer flows that don't load the `User` row).
        username: Option<String>,
    },

    /// An `OAuth2` client acting on its own behalf (e.g. at the token endpoint,
    /// or a client-credentials session), identified by its internal Ulid.
    OAuth2Client { id: Ulid },

    /// The homeserver, e.g. calling the introspection endpoint with its shared
    /// secret.
    Homeserver,
}

impl Requester {
    /// Build a `User` requester with both an id and a known username.
    #[must_use]
    pub fn user(id: Ulid, username: impl Into<String>) -> Self {
        Self::User {
            id,
            username: Some(username.into()),
        }
    }

    /// Build a `User` requester from a `user_id` alone (username unknown).
    #[must_use]
    pub fn user_id_only(id: Ulid) -> Self {
        Self::User { id, username: None }
    }

    /// Build an `OAuth2Client` requester from a client's internal Ulid.
    #[must_use]
    pub fn oauth2_client(id: Ulid) -> Self {
        Self::OAuth2Client { id }
    }
}

impl std::fmt::Display for Requester {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::User {
                id,
                username: Some(name),
            } => write!(f, "user:{id}({name})"),
            Self::User { id, username: None } => write!(f, "user:{id}"),
            Self::OAuth2Client { id } => write!(f, "oauth2-client:{id}"),
            Self::Homeserver => write!(f, "homeserver"),
        }
    }
}

/// A counter which increments each time we create a new log context
/// It will wrap around if we create more than [`u64::MAX`] contexts
static LOG_CONTEXT_INDEX: AtomicU64 = AtomicU64::new(0);
task_local! {
    pub static CURRENT_LOG_CONTEXT: LogContext;
}

/// A log context saves informations about the current task, such as the
/// elapsed time, the number of polls, and the poll time.
#[derive(Clone)]
pub struct LogContext {
    inner: Arc<LogContextInner>,
}

struct LogContextInner {
    /// A user-defined tag for the log context
    tag: Cow<'static, str>,

    /// A unique index for the log context
    index: u64,

    /// The time when the context was created
    start: Instant,

    /// The number of [`Future::poll`] recorded
    polls: AtomicU64,

    /// An approximation of the total CPU time spent in the context, in
    /// nanoseconds
    cpu_time: AtomicU64,

    /// The number of database queries executed in the context
    db_queries: AtomicU64,

    /// The cumulative wall-clock time spent executing database queries in the
    /// context, in nanoseconds
    db_time: AtomicU64,

    /// The number of rows fetched from the database in the context
    db_rows_fetched: AtomicU64,

    /// The identified principal making the request, if it has been resolved.
    /// Set once per context, first writer wins.
    requester: OnceLock<Requester>,
}

impl LogContext {
    /// Create a new log context with the given tag
    pub fn new(tag: impl Into<Cow<'static, str>>) -> Self {
        let tag = tag.into();
        let inner = LogContextInner {
            tag,
            index: LOG_CONTEXT_INDEX.fetch_add(1, Ordering::Relaxed),
            start: Instant::now(),
            polls: AtomicU64::new(0),
            cpu_time: AtomicU64::new(0),
            db_queries: AtomicU64::new(0),
            db_rows_fetched: AtomicU64::new(0),
            db_time: AtomicU64::new(0),
            requester: OnceLock::new(),
        };

        Self {
            inner: Arc::new(inner),
        }
    }

    /// Run a closure with the current log context, if any
    pub fn maybe_with<F, R>(f: F) -> Option<R>
    where
        F: FnOnce(&Self) -> R,
    {
        CURRENT_LOG_CONTEXT.try_with(f).ok()
    }

    /// Run the async function `f` with the given log context. It will wrap the
    /// output future to record poll and CPU statistics.
    pub fn run<F: FnOnce() -> Fut, Fut: Future>(&self, f: F) -> LogContextFuture<Fut> {
        let future = self.run_sync(f);
        self.wrap_future(future)
    }

    /// Run the sync function `f` with the given log context, recording the CPU
    /// time spent.
    pub fn run_sync<F: FnOnce() -> R, R>(&self, f: F) -> R {
        let start = Instant::now();
        let result = CURRENT_LOG_CONTEXT.sync_scope(self.clone(), f);
        let elapsed = start.elapsed().as_nanos().try_into().unwrap_or(u64::MAX);
        self.inner.cpu_time.fetch_add(elapsed, Ordering::Relaxed);
        result
    }

    /// Record the stats of a query on the current log context.
    pub fn maybe_record_query_stats(fetched: usize, duration: Duration) {
        LogContext::maybe_with(|ctx| ctx.record_query_stats(fetched, duration));
    }

    /// Record the stats of a query
    pub fn record_query_stats(&self, fetched: usize, duration: Duration) {
        let nanos = duration.as_nanos().try_into().unwrap_or(u64::MAX);
        self.inner.db_time.fetch_add(nanos, Ordering::Relaxed);
        self.inner
            .db_rows_fetched
            .fetch_add(fetched as u64, Ordering::Relaxed);
        self.inner.db_queries.fetch_add(1, Ordering::Relaxed);
    }

    /// Associate a [`Requester`] with this log context. Silently does nothing
    /// if a requester has already been recorded (first writer wins).
    pub fn set_requester(&self, requester: Requester) {
        let _ = self.inner.requester.set(requester);
    }

    /// Get the [`Requester`] associated with this log context, if one was set.
    #[must_use]
    pub fn requester(&self) -> Option<&Requester> {
        self.inner.requester.get()
    }

    /// Convenience: set the [`Requester`] on the current task-local
    /// `LogContext`, if one exists. No-op outside of a `LogContext` scope.
    pub fn maybe_set_requester(requester: Requester) {
        Self::maybe_with(|ctx| ctx.set_requester(requester));
    }

    /// Create a snapshot of the log context statistics
    #[must_use]
    pub fn stats(&self) -> LogContextStats {
        let polls = self.inner.polls.load(Ordering::Relaxed);
        let cpu_time = self.inner.cpu_time.load(Ordering::Relaxed);
        let cpu_time = Duration::from_nanos(cpu_time);
        let elapsed = self.inner.start.elapsed();
        let db_queries = self.inner.db_queries.load(Ordering::Relaxed);
        let db_rows_fetched = self.inner.db_rows_fetched.load(Ordering::Relaxed);
        let db_time = Duration::from_nanos(self.inner.db_time.load(Ordering::Relaxed));
        LogContextStats {
            polls,
            cpu_time,
            elapsed,
            db_queries,
            db_rows_fetched,
            db_time,
        }
    }
}

impl std::fmt::Display for LogContext {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let tag = &self.inner.tag;
        let index = self.inner.index;
        write!(f, "{tag}-{index}")
    }
}

/// A snapshot of a log context statistics
#[derive(Debug, Clone, Copy)]
pub struct LogContextStats {
    /// How many times the context was polled
    pub polls: u64,

    /// The approximate CPU time spent in the context
    pub cpu_time: Duration,

    /// How much time elapsed since the context was created
    pub elapsed: Duration,

    /// How many database queries were executed in the context
    pub db_queries: u64,

    /// The number of rows fetched from the database in the context
    pub db_rows_fetched: u64,

    /// The cumulative wall-clock time spent executing database queries
    pub db_time: Duration,
}

impl std::fmt::Display for LogContextStats {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let polls = self.polls;
        #[expect(clippy::cast_precision_loss)]
        let cpu_time_ms = self.cpu_time.as_nanos() as f64 / 1_000_000.;
        #[expect(clippy::cast_precision_loss)]
        let elapsed_ms = self.elapsed.as_nanos() as f64 / 1_000_000.;
        let db_queries = self.db_queries;
        let db_rows_fetched = self.db_rows_fetched;
        #[expect(clippy::cast_precision_loss)]
        let db_time_ms = self.db_time.as_nanos() as f64 / 1_000_000.;
        write!(
            f,
            "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}",
        )
    }
}