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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011LnJSXQ86AAWNtihR4C5PH
This commit is contained in:
Erik Johnston
2026-07-16 13:22:04 +00:00
co-authored by Claude Fable 5
parent ec2894b4a7
commit fa91dc5e3c
3 changed files with 137 additions and 101 deletions
+116
View File
@@ -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<ContextResourceUsage>) -> 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!(
"<ContextResourceUsage ru_stime='{:?}', ru_utime='{:?}', \
db_txn_count='{}', db_txn_duration_sec='{:?}', \
db_sched_duration_sec='{:?}', evt_db_fetch_count='{}'>",
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<PyAny>) -> Py<PyAny> {
/// 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::<ContextResourceUsage>()?;
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)?)?;
+1 -101
View File
@@ -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 (
"<ContextResourceUsage ru_stime='%r', ru_utime='%r', "
"db_txn_count='%r', db_txn_duration_sec='%r', "
"db_sched_duration_sec='%r', evt_db_fetch_count='%r'>"
) % (
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:
"""
+20
View File
@@ -10,8 +10,28 @@
# See the GNU Affero General Public License for more details:
# <https://www.gnu.org/licenses/agpl-3.0.html>.
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.