From fa91dc5e3c4899b6194363ce904034492f75d8a9 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Wed, 15 Jul 2026 13:34:53 +0000 Subject: [PATCH] Port ContextResourceUsage to a Rust pyclass Replace the Python `ContextResourceUsage` class with a native `#[pyclass]` (exposed via `synapse.synapse_rust.logcontext`, re-exported under the same name from `synapse.logging.context`). The public surface is preserved drop-in: the six get/set attributes, `copy`/`reset`, `+`/`-`/`+=`/`-=`, the `ContextResourceUsage(copy_from=...)` constructor and the exact `repr`. This is the first step of moving the logcontext machinery to Rust: keeping the usage accounting in a native struct lets the switch path (a later commit) do its rusage bookkeeping without allocating a Python object per operation. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_011LnJSXQ86AAWNtihR4C5PH --- rust/src/logging/context.rs | 116 ++++++++++++++++++++++++++++ synapse/logging/context.py | 102 +----------------------- synapse/synapse_rust/logcontext.pyi | 20 +++++ 3 files changed, 137 insertions(+), 101 deletions(-) diff --git a/rust/src/logging/context.rs b/rust/src/logging/context.rs index 7af6093f69..3afe3ccf8f 100644 --- a/rust/src/logging/context.rs +++ b/rust/src/logging/context.rs @@ -120,6 +120,121 @@ impl LogContext { } } +/// Tracks the resources used by a log context. +/// +/// A native drop-in for the former Python `ContextResourceUsage` class; the +/// public attribute surface, operators and `repr` are preserved so callers +/// (Measure, request/background-process metrics, task scheduler, ...) are +/// unaffected. Keeping this native lets the switch machinery do its rusage +/// accounting without allocating a Python object per operation. +#[pyclass( + name = "ContextResourceUsage", + module = "synapse.logging.context", + from_py_object +)] +#[derive(Clone, Default)] +pub struct ContextResourceUsage { + /// System CPU time, in seconds. + #[pyo3(get, set)] + pub ru_stime: f64, + /// User CPU time, in seconds. + #[pyo3(get, set)] + pub ru_utime: f64, + /// Number of database transactions done. + #[pyo3(get, set)] + pub db_txn_count: i64, + /// Time spent doing database transactions (excluding scheduling), in seconds. + #[pyo3(get, set)] + pub db_txn_duration_sec: f64, + /// Time spent waiting for a database connection, in seconds. + #[pyo3(get, set)] + pub db_sched_duration_sec: f64, + /// Number of events requested from the database. + #[pyo3(get, set)] + pub evt_db_fetch_count: i64, +} + +impl ContextResourceUsage { + fn add_assign(&mut self, other: &ContextResourceUsage) { + self.ru_utime += other.ru_utime; + self.ru_stime += other.ru_stime; + self.db_txn_count += other.db_txn_count; + self.db_txn_duration_sec += other.db_txn_duration_sec; + self.db_sched_duration_sec += other.db_sched_duration_sec; + self.evt_db_fetch_count += other.evt_db_fetch_count; + } + + fn sub_assign(&mut self, other: &ContextResourceUsage) { + self.ru_utime -= other.ru_utime; + self.ru_stime -= other.ru_stime; + self.db_txn_count -= other.db_txn_count; + self.db_txn_duration_sec -= other.db_txn_duration_sec; + self.db_sched_duration_sec -= other.db_sched_duration_sec; + self.evt_db_fetch_count -= other.evt_db_fetch_count; + } +} + +#[pymethods] +impl ContextResourceUsage { + /// `ContextResourceUsage(copy_from=None)` — if `copy_from` is given, copy its + /// stats; otherwise start at zero. + #[new] + #[pyo3(signature = (copy_from=None))] + fn new(copy_from: Option) -> Self { + copy_from.unwrap_or_default() + } + + /// Return a copy of this object. + fn copy(&self) -> ContextResourceUsage { + self.clone() + } + + /// Reset all stats to zero. + fn reset(&mut self) { + *self = ContextResourceUsage::default(); + } + + fn __repr__(&self) -> String { + // Matches the historical Python `__repr__` (values were interpolated with + // `%r`, i.e. `repr()`, inside single quotes). + format!( + "", + self.ru_stime, + self.ru_utime, + self.db_txn_count, + self.db_txn_duration_sec, + self.db_sched_duration_sec, + self.evt_db_fetch_count, + ) + } + + /// `self += other`; mutate in place. pyo3 returns `self` for the in-place slot. + fn __iadd__(&mut self, other: ContextResourceUsage) { + self.add_assign(&other); + } + + /// `self -= other`; mutate in place. pyo3 returns `self` for the in-place slot. + fn __isub__(&mut self, other: ContextResourceUsage) { + self.sub_assign(&other); + } + + /// `self + other`, returning a new object. + fn __add__(&self, other: ContextResourceUsage) -> ContextResourceUsage { + let mut res = self.clone(); + res.add_assign(&other); + res + } + + /// `self - other`, returning a new object. + fn __sub__(&self, other: ContextResourceUsage) -> ContextResourceUsage { + let mut res = self.clone(); + res.sub_assign(&other); + res + } +} + /// Register the Python sentinel logcontext. /// /// Called once from `synapse.logging.context` at import time. Registering twice @@ -183,6 +298,7 @@ pub fn swap_current_context(py: Python<'_>, context: Py) -> Py { /// Called when registering modules with python. pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { let child_module: Bound<'_, PyModule> = PyModule::new(py, "logcontext")?; + child_module.add_class::()?; child_module.add_function(wrap_pyfunction!(current_context, &child_module)?)?; child_module.add_function(wrap_pyfunction!(swap_current_context, &child_module)?)?; child_module.add_function(wrap_pyfunction!(register_sentinel, &child_module)?)?; diff --git a/synapse/logging/context.py b/synapse/logging/context.py index e0729ec1bd..e25e9355b6 100644 --- a/synapse/logging/context.py +++ b/synapse/logging/context.py @@ -53,6 +53,7 @@ from twisted.python.threadpool import ThreadPool from synapse.logging.loggers import ExplicitlyConfiguredLogger from synapse.synapse_rust.logcontext import ( + ContextResourceUsage as ContextResourceUsage, current_context as current_context, register_sentinel, swap_current_context, @@ -118,107 +119,6 @@ def logcontext_error(msg: str) -> None: get_thread_id = threading.get_ident -class ContextResourceUsage: - """Object for tracking the resources used by a log context - - Attributes: - ru_utime (float): user CPU time (in seconds) - ru_stime (float): system CPU time (in seconds) - db_txn_count (int): number of database transactions done - db_sched_duration_sec (float): amount of time spent waiting for a - database connection - db_txn_duration_sec (float): amount of time spent doing database - transactions (excluding scheduling time) - evt_db_fetch_count (int): number of events requested from the database - """ - - __slots__ = [ - "ru_stime", - "ru_utime", - "db_txn_count", - "db_txn_duration_sec", - "db_sched_duration_sec", - "evt_db_fetch_count", - ] - - def __init__(self, copy_from: "ContextResourceUsage | None" = None) -> None: - """Create a new ContextResourceUsage - - Args: - copy_from: if not None, an object to copy stats from - """ - if copy_from is None: - self.reset() - else: - # FIXME: mypy can't infer the types set via reset() above, so specify explicitly for now - self.ru_utime: float = copy_from.ru_utime - self.ru_stime: float = copy_from.ru_stime - self.db_txn_count: int = copy_from.db_txn_count - - self.db_txn_duration_sec: float = copy_from.db_txn_duration_sec - self.db_sched_duration_sec: float = copy_from.db_sched_duration_sec - self.evt_db_fetch_count: int = copy_from.evt_db_fetch_count - - def copy(self) -> "ContextResourceUsage": - return ContextResourceUsage(copy_from=self) - - def reset(self) -> None: - self.ru_stime = 0.0 - self.ru_utime = 0.0 - self.db_txn_count = 0 - - self.db_txn_duration_sec = 0.0 - self.db_sched_duration_sec = 0.0 - self.evt_db_fetch_count = 0 - - def __repr__(self) -> str: - return ( - "" - ) % ( - self.ru_stime, - self.ru_utime, - self.db_txn_count, - self.db_txn_duration_sec, - self.db_sched_duration_sec, - self.evt_db_fetch_count, - ) - - def __iadd__(self, other: "ContextResourceUsage") -> "ContextResourceUsage": - """Add another ContextResourceUsage's stats to this one's. - - Args: - other: the other resource usage object - """ - self.ru_utime += other.ru_utime - self.ru_stime += other.ru_stime - self.db_txn_count += other.db_txn_count - self.db_txn_duration_sec += other.db_txn_duration_sec - self.db_sched_duration_sec += other.db_sched_duration_sec - self.evt_db_fetch_count += other.evt_db_fetch_count - return self - - def __isub__(self, other: "ContextResourceUsage") -> "ContextResourceUsage": - self.ru_utime -= other.ru_utime - self.ru_stime -= other.ru_stime - self.db_txn_count -= other.db_txn_count - self.db_txn_duration_sec -= other.db_txn_duration_sec - self.db_sched_duration_sec -= other.db_sched_duration_sec - self.evt_db_fetch_count -= other.evt_db_fetch_count - return self - - def __add__(self, other: "ContextResourceUsage") -> "ContextResourceUsage": - res = ContextResourceUsage(copy_from=self) - res += other - return res - - def __sub__(self, other: "ContextResourceUsage") -> "ContextResourceUsage": - res = ContextResourceUsage(copy_from=self) - res -= other - return res - - @attr.s(slots=True, auto_attribs=True) class ContextRequest: """ diff --git a/synapse/synapse_rust/logcontext.pyi b/synapse/synapse_rust/logcontext.pyi index 092874a196..c0c79c16a7 100644 --- a/synapse/synapse_rust/logcontext.pyi +++ b/synapse/synapse_rust/logcontext.pyi @@ -10,8 +10,28 @@ # See the GNU Affero General Public License for more details: # . +from typing import Optional + from synapse.logging.context import LoggingContextOrSentinel +class ContextResourceUsage: + """Tracks the resources used by a log context.""" + + ru_stime: float + ru_utime: float + db_txn_count: int + db_txn_duration_sec: float + db_sched_duration_sec: float + evt_db_fetch_count: int + + def __init__(self, copy_from: "Optional[ContextResourceUsage]" = None) -> None: ... + def copy(self) -> "ContextResourceUsage": ... + def reset(self) -> None: ... + def __iadd__(self, other: "ContextResourceUsage") -> "ContextResourceUsage": ... + def __isub__(self, other: "ContextResourceUsage") -> "ContextResourceUsage": ... + def __add__(self, other: "ContextResourceUsage") -> "ContextResourceUsage": ... + def __sub__(self, other: "ContextResourceUsage") -> "ContextResourceUsage": ... + def current_context() -> LoggingContextOrSentinel: """Get the current logging context.