Add a per-homeserver clock to the Rust side (#20294)

This is a prerequisite to being able to add caching on the Rust side.

changelog.d/20294.misc

@@ -0,0 +1 @@
+Add a clock to the Rust side of Synapse, and use it in the rendezvous handlers.

rust/src/clock.rs

@@ -0,0 +1,136 @@
+/*
+ * This file is licensed under the Affero General Public License (AGPL) version 3.
+ *
+ * Copyright (C) 2026 Element Creations Ltd
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as
+ * published by the Free Software Foundation, either version 3 of the
+ * License, or (at your option) any later version.
+ *
+ * See the GNU Affero General Public License for more details:
+ * <https://www.gnu.org/licenses/agpl-3.0.html>.
+ *
+ */
+
+//! A clock for the Rust side of Synapse.
+//!
+//! Rust code often needs the current time somewhere it cannot cheaply ask
+//! Python for it, such as on a tokio worker without the GIL. Each
+//! homeserver's [`RustRuntime`](crate::runtime::RustRuntime) owns a [`Clock`]
+//! that can be read without the GIL.
+//!
+//! Synapse's unit tests run against a virtual reactor clock, where time only
+//! moves when a test says so. The test reactor keeps this clock in step via
+//! [`set_virtual_time_msec`]. See `ThreadedMemoryReactorClock` in
+//! `tests/server.py`.
+
+use std::{
+    sync::atomic::{AtomicU64, Ordering},
+    time::{Duration, SystemTime, UNIX_EPOCH},
+};
+
+use pyo3::pyfunction;
+
+/// A sentinel value stored in [`VIRTUAL_MILLIS`] while the clock is
+/// reading the real time. This value indicates that the clock should read the
+/// real system time rather than a virtual time.
+///
+/// We choose a value that will never be used in tests, namely the maximum `u64`
+/// value. This is 584 million years after the epoch.
+const REAL_TIME_SENTINEL: u64 = u64::MAX;
+
+/// The virtual time in milliseconds since the Unix epoch, or
+/// [`REAL_TIME_SENTINEL`]. Only ever set by tests, via [`set_virtual_time_msec`].
+///
+/// We use an AtomicU64 (rather than e.g. `Option<..>`) so that reads are fast
+/// and do not require locking.
+static VIRTUAL_MILLIS: AtomicU64 = AtomicU64::new(REAL_TIME_SENTINEL);
+
+/// Pin the clock to the given time, in milliseconds since the Unix epoch.
+///
+/// Only tests call this. Nothing unpins the clock, since a test reactor
+/// and the runtimes attached to it live for a single test.
+#[pyfunction]
+pub fn set_virtual_time_msec(millis: u64) {
+    if millis == REAL_TIME_SENTINEL {
+        panic!("Cannot set virtual time to the sentinel value REAL_TIME_SENTINEL");
+    }
+
+    VIRTUAL_MILLIS.store(millis, Ordering::Relaxed);
+}
+
+/// The current time as the Rust side of a homeserver sees it.
+///
+/// Like `synapse.util.clock.Clock` on the Python side, this is a wall clock.
+/// While these are two separate, independent implementations, they both read
+/// the same system clock (on Linux both are
+/// `clock_gettime(CLOCK_REALTIME)`).Code that needs a monotonic clock should
+/// use [`std::time::Instant`] directly.
+///
+/// In future, this may be extended to support per-homeserver looping calls
+/// (which is why this is not a global clock).
+///
+/// See [`crate::clock`] docs for more details.
+pub struct Clock {}
+
+impl Default for Clock {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Clock {
+    /// A clock reading the real time.
+    pub fn new() -> Self {
+        Clock {}
+    }
+
+    /// The current time.
+    ///
+    /// Reads the same system clock as `synapse.util.clock.Clock.time_msec()`
+    /// on the Python side. See the module docs.
+    pub fn now(&self) -> SystemTime {
+        // Check if virtual time is set, otherwise fall back to real time.
+        match VIRTUAL_MILLIS.load(Ordering::Relaxed) {
+            // On Linux this is a vDSO call rather than a syscall, so it is
+            // cheap enough to make on every read.
+            REAL_TIME_SENTINEL => SystemTime::now(),
+            virtual_millis => UNIX_EPOCH + Duration::from_millis(virtual_millis),
+        }
+    }
+
+    /// [`Clock::now`] in milliseconds since the Unix epoch.
+    ///
+    /// Panics if the system clock is set before the Unix epoch, 1970. This was
+    /// checked at startup, so this panic should only ever occur if the system
+    /// clock is manually set before the Unix epoch during runtime.
+    pub fn now_millis(&self) -> u64 {
+        self.now()
+            .duration_since(UNIX_EPOCH)
+            .expect("system clock was set before the Unix epoch, 1970")
+            .as_millis() as u64
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn virtual_time_overrides_the_real_clock() {
+        // Any real wall clock reads later than this (2020-09-13).
+        const PLAUSIBLE_WALL_CLOCK_MILLIS: u64 = 1_600_000_000_000;
+
+        let clock = Clock::new();
+        assert!(clock.now_millis() > PLAUSIBLE_WALL_CLOCK_MILLIS);
+
+        set_virtual_time_msec(12345);
+        assert_eq!(clock.now_millis(), 12345);
+        assert_eq!(clock.now(), UNIX_EPOCH + Duration::from_millis(12345));
+
+        // Zero is a legitimate virtual time. The memory reactor starts there.
+        set_virtual_time_msec(0);
+        assert_eq!(clock.now_millis(), 0);
+    }
+}

rust/src/lib.rs

@@ -1,11 +1,15 @@
-use std::convert::Infallible;
+use std::{
+    convert::Infallible,
+    time::{SystemTime, UNIX_EPOCH},
+};
 
 use lazy_static::lazy_static;
-use pyo3::prelude::*;
+use pyo3::{exceptions::PyRuntimeError, prelude::*};
 use pyo3_log::ResetHandle;
 
 pub mod acl;
 pub mod canonical_json;
+pub mod clock;
 pub mod config;
 pub mod deferred;
 pub mod duration;
@@ -68,6 +72,14 @@ fn reset_logging_config() {
 /// The entry point for defining the Python module.
 #[pymodule]
 fn synapse_rust(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
+    // We expect that system clock is set for after the Unix epoch, 1970. We'll
+    // panic later if it is not, so we check here and fail early.
+    if SystemTime::now().duration_since(UNIX_EPOCH).is_err() {
+        return Err(PyRuntimeError::new_err(
+            "system clock was set before the Unix epoch, 1970",
+        ));
+    }
+
     m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
     m.add_function(wrap_pyfunction!(get_rust_file_digest, m)?)?;
     m.add_function(wrap_pyfunction!(get_rustc_version, m)?)?;

rust/src/msc4388_rendezvous/mod.rs

@@ -33,13 +33,14 @@ use crate::{
     homeserver::HomeServer,
     http::http_request_from_twisted,
     msc4388_rendezvous::session::{GetResponse, PostResponse, PutResponse},
+    runtime::RustRuntime,
 };
 
 mod session;
 
 #[pyclass]
 struct MSC4388RendezvousHandler {
-    clock: Py<PyAny>,
+    runtime: RustRuntime,
     sessions: BTreeMap<Ulid, Session>,
     soft_limit: usize,
     hard_limit: usize,
@@ -99,6 +100,7 @@ impl MSC4388RendezvousHandler {
         eviction_interval: u64,
         ttl: u64,
     ) -> PyResult<Py<Self>> {
+        let runtime = homeserver.get_rust_runtime(py)?;
         let clock = homeserver.get_clock(py)?;
 
         // Construct a Python object so that we can get a reference to the
@@ -106,7 +108,7 @@ impl MSC4388RendezvousHandler {
         let self_ = Py::new(
             py,
             Self {
-                clock: clock.clone_ref(py),
+                runtime,
                 sessions: BTreeMap::new(),
                 soft_limit,
                 hard_limit,
@@ -125,23 +127,15 @@ impl MSC4388RendezvousHandler {
         Ok(self_)
     }
 
-    fn _evict(&mut self, py: Python<'_>) -> PyResult<()> {
-        let clock = self.clock.bind(py);
-        let now: u64 = clock.call_method0("time_msec")?.extract()?;
-        let now = SystemTime::UNIX_EPOCH + Duration::from_millis(now);
+    fn _evict(&mut self) -> PyResult<()> {
+        let now = self.runtime.clock().now();
         self.evict(now);
 
         Ok(())
     }
 
-    fn handle_post(
-        &mut self,
-        py: Python<'_>,
-        twisted_request: &Bound<'_, PyAny>,
-    ) -> PyResult<(u8, PostResponse)> {
-        let clock = self.clock.bind(py);
-        let now: u64 = clock.call_method0("time_msec")?.extract()?;
-        let now = SystemTime::UNIX_EPOCH + Duration::from_millis(now);
+    fn handle_post(&mut self, twisted_request: &Bound<'_, PyAny>) -> PyResult<(u8, PostResponse)> {
+        let now = self.runtime.clock().now();
 
         // We trigger an immediate eviction if we're at the hard limit
         if self.sessions.len() >= self.hard_limit {
@@ -176,7 +170,6 @@ impl MSC4388RendezvousHandler {
 
     fn handle_get(
         &mut self,
-        py: Python<'_>,
         id: &str,
         twisted_request: &Bound<'_, PyAny>,
     ) -> PyResult<(u8, GetResponse)> {
@@ -274,9 +267,7 @@ impl MSC4388RendezvousHandler {
             ));
         }
 
-        let clock = self.clock.bind(py);
-        let now: u64 = clock.call_method0("time_msec")?.extract()?;
-        let now = SystemTime::UNIX_EPOCH + Duration::from_millis(now);
+        let now = self.runtime.clock().now();
 
         let id: Ulid = id.parse().map_err(|_| NotFoundError::new())?;
         let session = self
@@ -290,7 +281,6 @@ impl MSC4388RendezvousHandler {
 
     fn handle_put(
         &mut self,
-        py: Python<'_>,
         id: &str,
         twisted_request: &Bound<'_, PyAny>,
     ) -> PyResult<(u8, PutResponse)> {
@@ -313,9 +303,7 @@ impl MSC4388RendezvousHandler {
 
         self.check_data_length(&data)?;
 
-        let clock = self.clock.bind(py);
-        let now: u64 = clock.call_method0("time_msec")?.extract()?;
-        let now = SystemTime::UNIX_EPOCH + Duration::from_millis(now);
+        let now = self.runtime.clock().now();
 
         let id: Ulid = id.parse().map_err(|_| NotFoundError::new())?;
         let session = self

rust/src/rendezvous/mod.rs

@@ -39,6 +39,7 @@ use crate::{
     errors::{NotFoundError, SynapseError},
     homeserver::HomeServer,
     http::{http_request_from_twisted, http_response_to_twisted, HeaderMapPyExt},
+    runtime::RustRuntime,
 };
 
 mod session;
@@ -57,7 +58,7 @@ fn prepare_headers(headers: &mut HeaderMap, session: &Session) {
 #[pyclass]
 struct RendezvousHandler {
     base: Uri,
-    clock: Py<PyAny>,
+    runtime: RustRuntime,
     sessions: BTreeMap<Ulid, Session>,
     capacity: usize,
     max_content_length: u64,
@@ -123,6 +124,7 @@ impl RendezvousHandler {
         let base = Uri::try_from(format!("{base}_synapse/client/rendezvous"))
             .map_err(|_| PyValueError::new_err("Invalid base URI"))?;
 
+        let runtime = homeserver.get_rust_runtime(py)?;
         let clock = homeserver.get_clock(py)?;
 
         let eviction_duration = SynapseDuration::from_milliseconds(eviction_interval);
@@ -133,7 +135,7 @@ impl RendezvousHandler {
             py,
             Self {
                 base,
-                clock: clock.clone_ref(py),
+                runtime,
                 sessions: BTreeMap::new(),
                 capacity,
                 max_content_length,
@@ -149,23 +151,19 @@ impl RendezvousHandler {
         Ok(self_)
     }
 
-    fn _evict(&mut self, py: Python<'_>) -> PyResult<()> {
-        let clock = self.clock.bind(py);
-        let now: u64 = clock.call_method0("time_msec")?.extract()?;
-        let now = SystemTime::UNIX_EPOCH + Duration::from_millis(now);
+    fn _evict(&mut self) -> PyResult<()> {
+        let now = self.runtime.clock().now();
         self.evict(now);
 
         Ok(())
     }
 
-    fn handle_post(&mut self, py: Python<'_>, twisted_request: &Bound<'_, PyAny>) -> PyResult<()> {
+    fn handle_post(&mut self, twisted_request: &Bound<'_, PyAny>) -> PyResult<()> {
         let request = http_request_from_twisted(twisted_request)?;
 
         let content_type = self.check_input_headers(request.headers())?;
 
-        let clock = self.clock.bind(py);
-        let now: u64 = clock.call_method0("time_msec")?.extract()?;
-        let now = SystemTime::UNIX_EPOCH + Duration::from_millis(now);
+        let now = self.runtime.clock().now();
 
         // We trigger an immediate eviction if we're at 2x the capacity
         if self.sessions.len() >= self.capacity * 2 {
@@ -199,18 +197,12 @@ impl RendezvousHandler {
         Ok(())
     }
 
-    fn handle_get(
-        &mut self,
-        py: Python<'_>,
-        twisted_request: &Bound<'_, PyAny>,
-        id: &str,
-    ) -> PyResult<()> {
+    fn handle_get(&mut self, twisted_request: &Bound<'_, PyAny>, id: &str) -> PyResult<()> {
         let request = http_request_from_twisted(twisted_request)?;
 
         let if_none_match: Option<IfNoneMatch> = request.headers().typed_get_optional()?;
 
-        let now: u64 = self.clock.call_method0(py, "time_msec")?.extract(py)?;
-        let now = SystemTime::UNIX_EPOCH + Duration::from_millis(now);
+        let now = self.runtime.clock().now();
 
         let id: Ulid = id.parse().map_err(|_| NotFoundError::new())?;
         let session = self
@@ -240,12 +232,7 @@ impl RendezvousHandler {
         Ok(())
     }
 
-    fn handle_put(
-        &mut self,
-        py: Python<'_>,
-        twisted_request: &Bound<'_, PyAny>,
-        id: &str,
-    ) -> PyResult<()> {
+    fn handle_put(&mut self, twisted_request: &Bound<'_, PyAny>, id: &str) -> PyResult<()> {
         let request = http_request_from_twisted(twisted_request)?;
 
         let content_type = self.check_input_headers(request.headers())?;
@@ -254,8 +241,7 @@ impl RendezvousHandler {
 
         let data = request.into_body();
 
-        let now: u64 = self.clock.call_method0(py, "time_msec")?.extract(py)?;
-        let now = SystemTime::UNIX_EPOCH + Duration::from_millis(now);
+        let now = self.runtime.clock().now();
 
         let id: Ulid = id.parse().map_err(|_| NotFoundError::new())?;
         let session = self

rust/src/runtime.rs

@@ -32,6 +32,7 @@ use anyhow::Context;
 use pyo3::{exceptions::PyRuntimeError, prelude::*};
 use tokio::runtime::{Handle, Runtime};
 
+use crate::clock::{set_virtual_time_msec, Clock};
 use crate::homeserver::HomeServer;
 use crate::reactor::Reactor;
 use crate::twisted_dispatch::{self, TwistedDispatchReader, TwistedDispatcher};
@@ -60,6 +61,7 @@ pub struct RustRuntimeInner {
     reactor: Reactor,
     tokio: Mutex<TokioState>,
     worker_threads: usize,
+    clock: Clock,
 
     /// Runs closures on the Twisted reactor thread without taking the GIL on
     /// the calling thread. See [`crate::twisted_dispatch`].
@@ -70,6 +72,11 @@ pub struct RustRuntimeInner {
 }
 
 impl RustRuntimeInner {
+    /// This homeserver's clock. See [`crate::clock`].
+    pub fn clock(&self) -> &Clock {
+        &self.clock
+    }
+
     /// Queue `f` to run on the Twisted reactor thread with the GIL held, and
     /// wake the reactor. Never takes the GIL itself, so a tokio task can call
     /// it to hand a result back to Twisted. See [`crate::twisted_dispatch`].
@@ -170,7 +177,7 @@ impl Drop for RustRuntimeInner {
 /// `HomeServer.get_rust_runtime()`. Rust classes that need it take it as a
 /// constructor argument and store their own clone, which is just an `Arc`
 /// refcount bump. Derefs to [`RustRuntimeInner`].
-#[pyclass(frozen, skip_from_py_object)]
+#[pyclass(frozen, weakref, skip_from_py_object)]
 #[derive(Clone)]
 pub struct RustRuntime {
     inner: Arc<RustRuntimeInner>,
@@ -202,6 +209,7 @@ impl RustRuntime {
             reactor,
             tokio: Mutex::new(TokioState::NotStarted),
             worker_threads,
+            clock: Clock::new(),
             dispatcher,
             dispatch_reader,
         });
@@ -220,6 +228,11 @@ impl RustRuntime {
 
         Ok(RustRuntime { inner })
     }
+
+    /// The current time, in milliseconds since the Unix epoch.
+    fn time_msec(&self) -> u64 {
+        self.inner.clock.now_millis()
+    }
 }
 
 /// The callable registered with `HomeServer.register_sync_shutdown_handler`,
@@ -245,6 +258,7 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()>
     let child_module = PyModule::new(py, "runtime")?;
 
     child_module.add_class::<RustRuntime>()?;
+    child_module.add_function(wrap_pyfunction!(set_virtual_time_msec, m)?)?;
 
     m.add_submodule(&child_module)?;
 

synapse/synapse_rust/runtime.pyi

@@ -13,3 +13,15 @@ class RustRuntime:
         hs: HomeServer,
         worker_threads: int = 4,
     ) -> None: ...
+    def time_msec(self) -> int:
+        """Test utility only. Normal code should use
+        `hs.get_clock().time_msec()`.
+
+        The current time in milliseconds since the Unix epoch, as the Rust side
+        sees it. This and `Clock.time_msec()` agree because both read the system
+        clock, `clock_gettime(CLOCK_REALTIME)` on Linux.
+        """
+
+def set_virtual_time_msec(millis: int) -> None:
+    """Pin the Rust clock to the given time. Only for tests, which run
+    against a virtual reactor clock."""

tests/app/test_homeserver_shutdown.py

@@ -103,6 +103,39 @@ class HomeserverCleanShutdownTestCase(HomeserverTestCase):
                 f"{get_memory_debug_info_for_object(hs_after_shutdown)}",
             )
 
+    def test_rust_runtime_is_garbage_collected_with_the_homeserver(self) -> None:
+        """The `RustRuntime` should be garbage collected with its homeserver.
+
+        This catches reference cycles through Rust fields, which Python's
+        garbage collector cannot see into.
+        """
+        self.reactor, self.clock = get_clock()
+        self.hs = setup_test_homeserver(
+            cleanup_func=self.addCleanup,
+            reactor=self.reactor,
+            homeserver_to_use=SynapseHomeServer,
+            clock=self.clock,
+        )
+
+        # Trigger the homeserver to create the Rust handlers, otherwise nothing
+        # in these tests would create them (as we don't register any servlets).
+        self.hs.get_rust_handlers()
+
+        # Get a weakref to the Rust runtime so that we can check it got garbage
+        # collected.
+        runtime_ref = weakref.ref(self.hs.get_rust_runtime())
+
+        # The same shutdown sequence as `test_clean_homeserver_shutdown`.
+        self.reactor.run()
+        cleanup_test_reactor_system_event_triggers(self.reactor)
+        self.get_success(self.hs.shutdown())
+        self.reactor.advance(0)
+        del self.hs
+        gc.collect()
+
+        # Check the RustRuntime has been garbage collected.
+        self.assertIsNone(runtime_ref(), "RustRuntime outlived its HomeServer")
+
     @logcontext_clean
     def test_clean_homeserver_shutdown_mid_background_updates(self) -> None:
         """Ensure the `SynapseHomeServer` can be fully shutdown and garbage collected

tests/server.py

@@ -48,7 +48,7 @@ from unittest.mock import Mock, patch
 
 import attr
 from incremental import Version
-from typing_extensions import ParamSpec
+from typing_extensions import ParamSpec, override
 from zope.interface import implementer
 
 import twisted
@@ -101,6 +101,7 @@ from synapse.storage import DataStore
 from synapse.storage.database import LoggingDatabaseConnection, make_pool
 from synapse.storage.engines import BaseDatabaseEngine, create_engine
 from synapse.storage.prepare_database import prepare_database
+from synapse.synapse_rust.runtime import set_virtual_time_msec
 from synapse.types import ISynapseReactor, JsonDict
 from synapse.util.clock import Clock
 from synapse.util.duration import Duration
@@ -594,6 +595,9 @@ class ThreadedMemoryReactorClock(MemoryReactorClock):
         self.lookups: dict[str, str] = {}
         self._thread_callbacks: deque[Callable[..., R]] = deque()
 
+        # Pin the Rust clock to our virtual time. `advance()` keeps it in step.
+        set_virtual_time_msec(int(self.seconds() * 1000))
+
         lookups = self.lookups
 
         @implementer(IResolverSimple)
@@ -735,7 +739,12 @@ class ThreadedMemoryReactorClock(MemoryReactorClock):
 
         return conn
 
+    @override
     def advance(self, amount: float) -> None:
+        # Move the Rust clock before `super().advance()` fires any callbacks,
+        # since those may read it.
+        set_virtual_time_msec(int((self.seconds() + amount) * 1000))
+
         # first advance our reactor's time, and run any "callLater" callbacks that
         # makes ready
         super().advance(amount)

tests/util/test_clock.py

@@ -15,6 +15,7 @@
 
 import weakref
 
+from synapse.synapse_rust.runtime import RustRuntime
 from synapse.util.duration import Duration
 
 from tests.unittest import HomeserverTestCase
@@ -75,3 +76,45 @@ class ClockTestCase(HomeserverTestCase):
         was_called = False
         self.reactor.advance(2)
         self.assertFalse(was_called)
+
+
+class RustClockTestCase(HomeserverTestCase):
+    """Tests for the Rust side's view of the time (`RustRuntime.time_msec`)."""
+
+    def test_follows_the_reactor(self) -> None:
+        runtime = self.hs.get_rust_runtime()
+        before = runtime.time_msec()
+
+        self.reactor.advance(37)
+
+        self.assertEqual(runtime.time_msec(), before + 37 * 1000)
+        self.assertEqual(runtime.time_msec(), self.hs.get_clock().time_msec())
+
+    def test_is_up_to_date_inside_a_looping_call(self) -> None:
+        """Rust must see the new time from callbacks fired *during* an advance.
+
+        This is why the test reactor moves the Rust clock *before* it calls
+        `super().advance()`, which is what fires those callbacks.
+        """
+        runtime = self.hs.get_rust_runtime()
+        seen: list[int] = []
+        self.hs.get_clock().looping_call(
+            lambda: seen.append(runtime.time_msec()), Duration(seconds=1)
+        )
+
+        base_time = self.hs.get_clock().time_msec()
+
+        self.reactor.advance(1)
+        self.assertEqual(seen, [base_time + 1000])
+
+        self.reactor.advance(1)
+        self.assertEqual(seen, [base_time + 1000, base_time + 2000])
+
+    def test_every_runtime_on_the_reactor_follows_it(self) -> None:
+        """Worker tests run several homeservers, each with its own runtime, on
+        one reactor. All of their clocks must move together."""
+        other_runtime = RustRuntime(hs=self.hs)
+
+        self.reactor.advance(5)
+
+        self.assertEqual(other_runtime.time_msec(), self.hs.get_clock().time_msec())