From 4d002f82eb62b30c2c2bc061846e89338cd44f0f Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Mon, 12 Oct 2026 14:34:32 +0000 Subject: [PATCH] Move `is_feature_enabled` onto a cache that lives in Rust The first Rust database function to get a cache, and the first Python `@cached` function to be replaced by one. The Rust `/versions` handler calls `Store::is_feature_enabled_for_user` on every authenticated request, and until now each call ran a query. The Python `is_feature_enabled` had a cache of the same rows under another name, which every writer would have had to invalidate alongside. Now there is one cache, keyed on `(user_id, feature)`, and both sides read through it. The main datastore builds the Rust `Store` in its `__init__`, over its own `db_pool`, and `hs.get_rust_store()` hands it to the Rust handlers: the datastore owns the store, the handlers borrow it. The store's `is_feature_enabled_for_user` reaches Python as a callable cache object, bound on the datastore under that name. Calling it answers a hit with an already-fired Deferred, without leaving the reactor thread, and only a miss goes through tokio. Being a `RustCache`, it is also what `_invalidate_cache_and_stream`, replication and `register_cache` take, so `set_features_for_user` invalidates it with the one call it already made, and the name is spelled once in Rust, where the cache is built. `is_feature_enabled` keeps its global-config short cut in Python and delegates the per-user lookup. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01ULdieJyQoHNzeHq42H2M1p --- changelog.d/20296.misc | 1 + rust/src/deferred.rs | 8 + rust/src/handlers/mod.rs | 16 +- rust/src/handlers/versions.rs | 12 +- rust/src/homeserver.rs | 16 ++ rust/src/lib.rs | 1 + rust/src/storage/mod.rs | 1 + rust/src/storage/python.rs | 163 ++++++++++++++++++ rust/src/storage/store.rs | 121 +++++++++---- synapse/server.py | 6 + synapse/storage/databases/main/cache.py | 15 -- .../databases/main/experimental_features.py | 33 ++-- synapse/synapse_rust/storage.pyi | 44 +++++ tests/rest/client/test_versions.py | 75 ++++++++ tests/storage/databases/main/test_cache.py | 45 +++++ .../main/test_experimental_features.py | 103 +++++++++++ 16 files changed, 588 insertions(+), 72 deletions(-) create mode 100644 changelog.d/20296.misc create mode 100644 rust/src/storage/python.rs create mode 100644 synapse/synapse_rust/storage.pyi create mode 100644 tests/storage/databases/main/test_experimental_features.py diff --git a/changelog.d/20296.misc b/changelog.d/20296.misc new file mode 100644 index 0000000000..7b06cebb07 --- /dev/null +++ b/changelog.d/20296.misc @@ -0,0 +1 @@ +Add caching for database functions implemented in Rust, and move the per-user experimental features lookup onto it. diff --git a/rust/src/deferred.rs b/rust/src/deferred.rs index 73721bfe89..7e3db57d68 100644 --- a/rust/src/deferred.rs +++ b/rust/src/deferred.rs @@ -69,6 +69,14 @@ fn logging_context_module(py: Python<'_>) -> PyResult<&Bound<'_, PyAny>> { .bind(py)) } +/// An already-fired Twisted deferred holding `value`. +pub(crate) fn succeed<'py, V: IntoPyObject<'py>>( + py: Python<'py>, + value: V, +) -> PyResult> { + defer(py)?.call_method1("succeed", (value,)) +} + /// Creates a twisted deferred from the given future, spawning the task on the /// tokio runtime. /// diff --git a/rust/src/handlers/mod.rs b/rust/src/handlers/mod.rs index ce92d0db5a..6e231f5f15 100644 --- a/rust/src/handlers/mod.rs +++ b/rust/src/handlers/mod.rs @@ -24,8 +24,6 @@ use pyo3::{ use crate::homeserver::HomeServer; use crate::state::StateResolutionHandler; -use crate::storage::db::python_db_pool::PythonDatabasePoolWrapper; -use crate::storage::store::Store; pub mod versions; @@ -47,13 +45,9 @@ impl RustHandlers { // and the Twisted reactor. let runtime = homeserver.get_rust_runtime(py)?; - let main_db_pool = homeserver.main_database_pool(py)?; - let db_pool = PythonDatabasePoolWrapper::new(&main_db_pool, runtime.clone())?; - - // Store is shared across all of the handlers so let's use an `Arc` - let store = Arc::new(Store { - db_pool: Box::new(db_pool), - }); + // The main datastore builds the store and keeps it; we share it. See + // `crate::storage::python`. + let store = homeserver.get_rust_store(py)?; let global_unstable_feature_map = Arc::new( versions::synapse_config_to_global_unstable_feature_map(&config), @@ -70,7 +64,9 @@ impl RustHandlers { // The Rust state resolution reads the database through the native // driver's pool, so it needs the main database to use that driver. - let connection_pool = main_db_pool.getattr(intern!(py, "_db_pool"))?; + let connection_pool = homeserver + .main_database_pool(py)? + .getattr(intern!(py, "_db_pool"))?; let rust_connection_pool_class = py .import("synapse.storage.rust_pool")? .getattr("RustConnectionPool")?; diff --git a/rust/src/handlers/versions.rs b/rust/src/handlers/versions.rs index e3fb2fd4d9..d3412ab101 100644 --- a/rust/src/handlers/versions.rs +++ b/rust/src/handlers/versions.rs @@ -22,7 +22,7 @@ use serde::Serialize; use crate::config::{types::RoomCreationPreset, SynapseHomeServerConfig}; use crate::deferred::create_deferred; use crate::runtime::RustRuntime; -use crate::storage::store::{PerUserExperimentalFeature, Store}; +use crate::storage::store::{IsFeatureEnabledKey, PerUserExperimentalFeature, Store}; /// `GET /_matrix/client/versions` response #[derive(Serialize, Clone, Debug)] @@ -95,7 +95,10 @@ async fn build_versions_response( } else { // Look up whether it's explicitly enabled/disabled for this user store - .is_feature_enabled_for_user(user_id, PerUserExperimentalFeature::MSC3881) + .is_feature_enabled_for_user(IsFeatureEnabledKey::new( + user_id, + PerUserExperimentalFeature::MSC3881, + )) .await? // Default to false if there is no entry for this user .unwrap_or(false) @@ -112,7 +115,10 @@ async fn build_versions_response( } else { // Look up whether it's explicitly enabled/disabled for this user store - .is_feature_enabled_for_user(user_id, PerUserExperimentalFeature::MSC3575) + .is_feature_enabled_for_user(IsFeatureEnabledKey::new( + user_id, + PerUserExperimentalFeature::MSC3575, + )) .await? // Default to false if there is no entry for this user .unwrap_or(false) diff --git a/rust/src/homeserver.rs b/rust/src/homeserver.rs index 44f8ba905e..956455e8d4 100644 --- a/rust/src/homeserver.rs +++ b/rust/src/homeserver.rs @@ -15,11 +15,14 @@ //! A typed wrapper around the Python `HomeServer`. +use std::sync::Arc; + use pyo3::{intern, prelude::*, types::PyDict}; use crate::config::SynapseHomeServerConfig; use crate::reactor::Reactor; use crate::runtime::RustRuntime; +use crate::storage::{python::RustStore, store::Store}; /// The Python `HomeServer`, as seen from Rust. /// @@ -99,4 +102,17 @@ impl HomeServer { .getattr(intern!(py, "main"))? .getattr(intern!(py, "db_pool")) } + + /// The Rust `Store` (`hs.get_rust_store()`), which the main datastore + /// builds and owns. See [`crate::storage::python`]. + pub fn get_rust_store(&self, py: Python<'_>) -> PyResult> { + Ok(self + .0 + .bind(py) + .call_method0(intern!(py, "get_rust_store"))? + .cast::()? + .get() + .store() + .clone()) + } } diff --git a/rust/src/lib.rs b/rust/src/lib.rs index add7f2f6d5..95f4073cf1 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -102,6 +102,7 @@ fn synapse_rust(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { http_client::register_module(py, m)?; rendezvous::register_module(py, m)?; runtime::register_module(py, m)?; + storage::python::register_module(py, m)?; msc4388_rendezvous::register_module(py, m)?; segmenter::register_module(py, m)?; state::register_module(py, m)?; diff --git a/rust/src/storage/mod.rs b/rust/src/storage/mod.rs index 735fcecb26..886578a080 100644 --- a/rust/src/storage/mod.rs +++ b/rust/src/storage/mod.rs @@ -14,4 +14,5 @@ */ pub mod db; +pub mod python; pub mod store; diff --git a/rust/src/storage/python.rs b/rust/src/storage/python.rs new file mode 100644 index 0000000000..f31006903f --- /dev/null +++ b/rust/src/storage/python.rs @@ -0,0 +1,163 @@ +/* + * 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: + * . + * + */ + +//! The Python view of the Rust [`Store`]. +//! +//! The main datastore builds a [`RustStore`] in its `__init__`, over its own +//! `db_pool`, and keeps it as `rust_store`. `hs.get_rust_store()` hands the +//! same store to the Rust handlers. So the datastore owns the store, as it +//! owns everything else that reads the database, and the handlers borrow it. +//! +//! A store function that Python calls as well as Rust is exposed as a +//! callable cache object, such as [`IsFeatureEnabledForUser`]. The datastore +//! binds it under the cache's name, so Python callers call it like the +//! `@cached` function it replaces, and Synapse's invalidation machinery finds +//! it by that name exactly as it would a `@cached` function. See +//! [`crate::caches::python`]. + +use std::sync::Arc; + +use pyo3::{ + exceptions::{PyRuntimeError, PyValueError}, + pyclass, pymethods, + types::{PyAnyMethods, PyModule, PyModuleMethods}, + Bound, Py, PyAny, PyClassInitializer, PyResult, Python, +}; + +use crate::{ + caches::python::RustCache, + deferred::{create_deferred, succeed}, + runtime::RustRuntime, + storage::{ + db::python_db_pool::PythonDatabasePoolWrapper, + store::{IsFeatureEnabledKey, PerUserExperimentalFeature, Store}, + }, +}; + +/// The Rust [`Store`], as the main datastore holds it. +#[pyclass(frozen, name = "RustStore")] +pub struct RustStore { + store: Arc, + is_feature_enabled_for_user: Py, +} + +impl RustStore { + pub fn store(&self) -> &Arc { + &self.store + } +} + +#[pymethods] +impl RustStore { + /// Build the store over `db_pool`, the datastore's `DatabasePool`, with + /// its caches in the runtime's registry. + #[new] + fn py_new(py: Python<'_>, runtime: &RustRuntime, db_pool: &Bound<'_, PyAny>) -> PyResult { + let db_pool = PythonDatabasePoolWrapper::new(db_pool, runtime.clone())?; + let store = Arc::new(Store::new(Box::new(db_pool), runtime.caches())); + + let is_feature_enabled_for_user = Py::new( + py, + PyClassInitializer::from(RustCache::new( + store.is_feature_enabled_for_user_cache.clone(), + )) + .add_subclass(IsFeatureEnabledForUser { + store: Arc::clone(&store), + runtime: runtime.clone(), + }), + )?; + + Ok(RustStore { + store, + is_feature_enabled_for_user, + }) + } + + /// `Store::is_feature_enabled_for_user` as a callable cache, for the + /// datastore to bind under that name. + #[getter] + fn is_feature_enabled_for_user(&self, py: Python<'_>) -> Py { + self.is_feature_enabled_for_user.clone_ref(py) + } +} + +/// `Store::is_feature_enabled_for_user`, callable from Python, and the cache +/// behind it. +/// +/// Calling it returns a Twisted deferred of the answer. A hit is answered +/// synchronously, with an already-fired deferred, so a cached lookup costs +/// Python no more than its own caches do. Only a miss goes through tokio. +/// +/// As a subclass of [`RustCache`], it is also the cache: `invalidate`, +/// `__name__` and the counters are all here, so `_invalidate_cache_and_stream` +/// and `register_cache` take it as they would a `@cached` function. +#[pyclass(frozen, extends = RustCache, name = "IsFeatureEnabledForUser")] +pub struct IsFeatureEnabledForUser { + store: Arc, + runtime: RustRuntime, +} + +#[pymethods] +impl IsFeatureEnabledForUser { + /// Whether `feature` is enabled for `user_id`. `False` when there is no + /// row, as for the Python function this replaces. + /// + /// Raises `ValueError` for a feature Rust doesn't know. Python's + /// `ExperimentalFeature` and [`PerUserExperimentalFeature`] list the same + /// features, so that means one side has gained a feature without the + /// other, which should fail loudly. + fn __call__<'py>( + &self, + py: Python<'py>, + user_id: String, + feature: String, + ) -> PyResult> { + let feature = PerUserExperimentalFeature::from_name(&feature).ok_or_else(|| { + PyValueError::new_err(format!("unknown per-user experimental feature: {feature}")) + })?; + let key = IsFeatureEnabledKey::new(&user_id, feature); + if let Some(enabled) = self.store.is_feature_enabled_for_user_cache.get(&key) { + return succeed(py, enabled.unwrap_or(false)); + } + + let store = Arc::clone(&self.store); + create_deferred(py, &self.runtime, async move { + store + .is_feature_enabled_for_user(key) + .await + .map(|enabled| enabled.unwrap_or(false)) + .map_err(|err| { + PyRuntimeError::new_err(format!("is_feature_enabled_for_user failed: {err:#}")) + }) + }) + } +} + +pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { + let child_module = PyModule::new(py, "storage")?; + + child_module.add_class::()?; + child_module.add_class::()?; + + m.add_submodule(&child_module)?; + + // We need to manually add the module to sys.modules to make `from + // synapse.synapse_rust import storage` work. + py.import("sys")? + .getattr("modules")? + .set_item("synapse.synapse_rust.storage", child_module)?; + + Ok(()) +} diff --git a/rust/src/storage/store.rs b/rust/src/storage/store.rs index c099d6f501..cc713b8c8c 100644 --- a/rust/src/storage/store.rs +++ b/rust/src/storage/store.rs @@ -16,61 +16,124 @@ use std::sync::Arc; use futures::FutureExt; -use serde::Serialize; +use pyo3::{ + types::{PyAnyMethods, PyTuple}, + Bound, +}; +use crate::caches::{python::PyCacheKey, registry::CacheRegistry, Cache}; use crate::storage::db::{DatabasePool, DatabasePoolExt, DbRowExt}; -/// Currently supported per-user features -#[derive(Serialize, Debug)] +/// Currently supported per-user features. +/// +/// Mirrors Python's `ExperimentalFeature`: the two must list the same +/// features, since Python reads the cache keyed by this through +/// [`crate::storage::python::IsFeatureEnabledForUser`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum PerUserExperimentalFeature { - #[serde(rename = "msc3881")] MSC3881, - #[serde(rename = "msc3575")] MSC3575, } -impl std::fmt::Display for PerUserExperimentalFeature { - fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - write!( - f, - "{}", - // Serialize so we can use the serde name of the variant as the source of truth - serde_json::to_string(self) - .unwrap_or_else(|err| format!( - "", - self, err - )) - // Remove the surrounding quotes from JSON serialization - .trim_matches('"') - ) +impl PerUserExperimentalFeature { + const ALL: [Self; 2] = [Self::MSC3881, Self::MSC3575]; + + /// The feature's name as stored in `per_user_experimental_features.feature`. + /// This is the one place the name is spelled. + pub fn name(&self) -> &'static str { + match self { + PerUserExperimentalFeature::MSC3881 => "msc3881", + PerUserExperimentalFeature::MSC3575 => "msc3575", + } + } + + /// The feature called `name`, if Rust knows it. + pub fn from_name(name: &str) -> Option { + Self::ALL.into_iter().find(|feature| feature.name() == name) + } +} + +/// The key of the `is_feature_enabled_for_user` cache. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct IsFeatureEnabledKey { + user_id: Arc, + feature: PerUserExperimentalFeature, +} + +impl IsFeatureEnabledKey { + pub fn new(user_id: &str, feature: PerUserExperimentalFeature) -> Self { + IsFeatureEnabledKey { + user_id: user_id.into(), + feature, + } + } +} + +impl PyCacheKey for IsFeatureEnabledKey { + fn from_py_key(key: &Bound<'_, PyTuple>) -> Option { + // `ExperimentalFeature` subclasses `str`, so the feature extracts as + // its value ("msc3881", ...) whether it arrives as the enum from this + // worker or as a plain string over replication. Anything else matches + // nothing: a key of a different shape, or a feature name this version + // doesn't know, which can't be cached here anyway. + let (user_id, feature): (String, String) = key.extract().ok()?; + let feature = PerUserExperimentalFeature::from_name(&feature)?; + + Some(IsFeatureEnabledKey::new(&user_id, feature)) } } pub struct Store { pub db_pool: Box, + + /// Caches [`Store::is_feature_enabled_for_user`]. Python reads through it + /// too, and invalidates it, as + /// [`crate::storage::python::IsFeatureEnabledForUser`]. Sized to match the + /// Python cache it replaced. + pub(crate) is_feature_enabled_for_user_cache: Arc>>, } impl Store { + /// A store over `db_pool`, with its caches built in `registry`, the + /// registry of the homeserver's Rust runtime. + pub fn new(db_pool: Box, registry: &CacheRegistry) -> Store { + Store { + db_pool, + is_feature_enabled_for_user_cache: registry + .new_cache("is_feature_enabled_for_user", 1000), + } + } + /// Checks whether a given feature is enabled/disabled for this user /// /// If there is no entry, returns None pub async fn is_feature_enabled_for_user( &self, - user_id: &str, + key: IsFeatureEnabledKey, + ) -> Result, anyhow::Error> { + // The database callback needs its own copy of the user id. `Arc` + // makes that a refcount bump. + let user_id = Arc::clone(&key.user_id); + let feature = key.feature; + + self.is_feature_enabled_for_user_cache + .get_or_load(key, || { + self.load_is_feature_enabled_for_user(Arc::clone(&user_id), feature) + }) + .await + .map_err(anyhow::Error::new) + } + + /// The uncached body of [`Store::is_feature_enabled_for_user`]. + async fn load_is_feature_enabled_for_user( + &self, + user_id: Arc, feature: PerUserExperimentalFeature, ) -> Result, anyhow::Error> { - // We need owned copies to move into the callback because it is `'static` (it - // may be moved to another thread). We use `Arc` rather than `String` so - // the per-call clone is just a cheap refcount bump rather than a fresh - // allocation. - let user_id: Arc = user_id.into(); - let feature: Arc = feature.to_string().into(); - let is_feature_enabled_for_user = self .db_pool .run_interaction("is_feature_enabled_for_user", move |txn| { let user_id = user_id.clone(); - let feature = feature.clone(); async move { let rows = txn .query( @@ -79,7 +142,7 @@ impl Store { FROM per_user_experimental_features WHERE user_id = ? AND feature = ? "#, - &[user_id.as_ref(), feature.as_ref()], + &[user_id.as_ref(), feature.name()], ) .await?; diff --git a/synapse/server.py b/synapse/server.py index 1438e997eb..ec1a57de2a 100644 --- a/synapse/server.py +++ b/synapse/server.py @@ -178,6 +178,7 @@ from synapse.synapse_rust.handlers import RustHandlers from synapse.synapse_rust.msc4388_rendezvous import MSC4388RendezvousHandler from synapse.synapse_rust.rendezvous import RendezvousHandler from synapse.synapse_rust.runtime import RustRuntime +from synapse.synapse_rust.storage import RustStore from synapse.types import DomainSpecificString, ISynapseReactor from synapse.util import SYNAPSE_VERSION from synapse.util.caches import CACHE_METRIC_REGISTRY @@ -973,6 +974,11 @@ class HomeServer(metaclass=abc.ABCMeta): def get_rust_handlers(self) -> RustHandlers: return RustHandlers(self) + def get_rust_store(self) -> RustStore: + """The Rust side of the main datastore, which builds and owns it. The + Rust handlers read the database through this.""" + return self.get_datastores().main.rust_store + @cache_in_self def get_rust_runtime(self) -> RustRuntime: """The per-homeserver state for the Rust side of Synapse: the tokio diff --git a/synapse/storage/databases/main/cache.py b/synapse/storage/databases/main/cache.py index b3f11fac3c..be580a4b22 100644 --- a/synapse/storage/databases/main/cache.py +++ b/synapse/storage/databases/main/cache.py @@ -670,21 +670,6 @@ class CacheInvalidationWorkerStore(SQLBaseStore): txn.call_after(cache_func.invalidate, keys) self._send_invalidation_to_replication(txn, cache_func.__name__, keys) - def _invalidate_rust_cache_and_stream( - self, txn: LoggingTransaction, cache_name: str, keys: tuple[Any, ...] - ) -> None: - """`_invalidate_cache_and_stream` for a cache that lives in Rust. - - Rust caches are addressed by name through the runtime rather than held - as attributes of the store, and a worker only has one once its Rust - `Store` has been built. The row goes on the stream either way, for the - workers that do have the cache. - """ - cache = self.hs.get_rust_runtime().get_cache(cache_name) - if cache is not None: - txn.call_after(cache.invalidate, keys) - self._send_invalidation_to_replication(txn, cache_name, keys) - def _invalidate_cache_and_stream_bulk( self, txn: LoggingTransaction, diff --git a/synapse/storage/databases/main/experimental_features.py b/synapse/storage/databases/main/experimental_features.py index 77b6c36884..7538e6a2a1 100644 --- a/synapse/storage/databases/main/experimental_features.py +++ b/synapse/storage/databases/main/experimental_features.py @@ -27,6 +27,8 @@ from synapse.storage.database import ( LoggingTransaction, ) from synapse.storage.databases.main import CacheInvalidationWorkerStore +from synapse.synapse_rust.storage import RustStore +from synapse.util.caches import register_rust_cache from synapse.util.caches.descriptors import cached if TYPE_CHECKING: @@ -43,6 +45,17 @@ class ExperimentalFeaturesStore(CacheInvalidationWorkerStore): ) -> None: super().__init__(database, db_conn, hs) + # The Rust side of this store, over the same database pool. The Rust + # handlers share it, via `hs.get_rust_store()`. + self.rust_store = RustStore(hs.get_rust_runtime(), self.db_pool) + + # `is_feature_enabled_for_user` is a cache that lives in Rust, and the + # function that reads through it. It is bound here under its own name, + # which is how Synapse's invalidation machinery finds a cache, and + # registered for metrics and cache factors like any other. + self.is_feature_enabled_for_user = self.rust_store.is_feature_enabled_for_user + register_rust_cache(self.is_feature_enabled_for_user, server_name=hs.hostname) + @cached() async def list_enabled_features(self, user_id: str) -> frozenset[str]: """ @@ -89,7 +102,7 @@ class ExperimentalFeaturesStore(CacheInvalidationWorkerStore): ) self._invalidate_cache_and_stream( - txn, self.is_feature_enabled, (user, feature) + txn, self.is_feature_enabled_for_user, (user, feature) ) self._invalidate_cache_and_stream(txn, self.list_enabled_features, (user,)) @@ -98,7 +111,6 @@ class ExperimentalFeaturesStore(CacheInvalidationWorkerStore): "set_features_for_user", set_features_for_user_txn ) - @cached() async def is_feature_enabled( self, user_id: str, feature: "ExperimentalFeature" ) -> bool: @@ -115,16 +127,7 @@ class ExperimentalFeaturesStore(CacheInvalidationWorkerStore): if feature.is_globally_enabled(self.hs.config): return True - # if it's not enabled globally, check if it is enabled per-user - res = await self.db_pool.simple_select_one_onecol( - table="per_user_experimental_features", - keyvalues={"user_id": user_id, "feature": feature}, - retcol="enabled", - allow_none=True, - desc="get_feature_enabled", - ) - - # None and false are treated the same - db_enabled = bool(res) - - return db_enabled + # If it's not enabled globally, check if it is enabled per-user. That + # lookup, and its cache, live in Rust, shared with the Rust `/versions` + # handler. + return await self.is_feature_enabled_for_user(user_id, feature) diff --git a/synapse/synapse_rust/storage.pyi b/synapse/synapse_rust/storage.pyi new file mode 100644 index 0000000000..95e0dbb802 --- /dev/null +++ b/synapse/synapse_rust/storage.pyi @@ -0,0 +1,44 @@ +# 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: +# . + +from twisted.internet.defer import Deferred + +from synapse.storage.database import DatabasePool +from synapse.synapse_rust.caches import RustCache +from synapse.synapse_rust.runtime import RustRuntime + +class IsFeatureEnabledForUser(RustCache): + """The Rust `Store::is_feature_enabled_for_user`, and the cache behind it. + + Calling it answers from the cache without leaving the reactor thread, with + an already-fired Deferred. Only a miss goes to the database, through tokio. + """ + + def __call__(self, user_id: str, feature: str) -> "Deferred[bool]": + """Whether `feature` is enabled for `user_id`. False when there is no + row for them. + + Raises ValueError for a feature name the Rust side doesn't know, which + means `ExperimentalFeature` and `PerUserExperimentalFeature` have + drifted apart.""" + +class RustStore: + """The Rust side of the main datastore. + + The datastore builds this in its `__init__`, over its own `db_pool`, and + binds each callable cache here under that cache's name. Get it from + `hs.get_rust_store()`. + """ + + def __init__(self, runtime: RustRuntime, db_pool: DatabasePool) -> None: ... + @property + def is_feature_enabled_for_user(self) -> IsFeatureEnabledForUser: ... diff --git a/tests/rest/client/test_versions.py b/tests/rest/client/test_versions.py index 6ecff9de6f..df7c266240 100644 --- a/tests/rest/client/test_versions.py +++ b/tests/rest/client/test_versions.py @@ -18,6 +18,7 @@ from synapse.rest import admin from synapse.rest.client import login, versions from synapse.server import HomeServer from synapse.types import JsonDict +from synapse.util.caches.lrucache import setup_expire_lru_cache_entries from synapse.util.clock import Clock from tests import unittest @@ -145,6 +146,70 @@ class VersionsTestCase(unittest.HomeserverTestCase): self.assertEqual(channel.code, 200, channel.result) self.assertTrue(channel.json_body["unstable_features"]["io.element.msc4502"]) + def test_per_user_feature_can_be_disabled_again(self) -> None: + """The Rust cache must be invalidated when a feature is turned back off. + + Enabling is the easier direction, because the value being replaced is + the "no row for this user" one. Turning a feature back off only works + if invalidation reaches an entry that was populated from a real row. + """ + user1_id = self.register_user("user1", "pass") + user1_tok = self.login(user1_id, "pass") + + self._enable_experimental_feature_for_user( + target_user_id=user1_id, features={"msc3881": True} + ) + self.assertTrue(self._unstable_feature(user1_tok, "org.matrix.msc3881")) + + self._enable_experimental_feature_for_user( + target_user_id=user1_id, features={"msc3881": False} + ) + self.assertFalse(self._unstable_feature(user1_tok, "org.matrix.msc3881")) + + # `msc3575_enabled` defaults to True, in which case the per-user lookup is + # skipped altogether. Turn it off so that this feature actually reaches + # the cache. + @unittest.override_config({"experimental_features": {"msc3575_enabled": False}}) + def test_per_user_msc3575(self) -> None: + """The other feature `/versions` looks up per user, which is otherwise + only covered when it is globally enabled.""" + user1_id = self.register_user("user1", "pass") + user1_tok = self.login(user1_id, "pass") + + feature = "org.matrix.simplified_msc3575" + self.assertFalse(self._unstable_feature(user1_tok, feature)) + + self._enable_experimental_feature_for_user( + target_user_id=user1_id, features={"msc3575": True} + ) + self.assertTrue(self._unstable_feature(user1_tok, feature)) + + self._enable_experimental_feature_for_user( + target_user_id=user1_id, features={"msc3575": False} + ) + self.assertFalse(self._unstable_feature(user1_tok, feature)) + + @unittest.override_config({"caches": {"expiry_time": "30m"}}) + def test_cache_entries_expire_when_unread(self) -> None: + """The Rust cache must take part in Synapse's age-based cache expiry, + off the same looping call and the same config as the Python ones.""" + setup_expire_lru_cache_entries(self.hs) + + user1_id = self.register_user("user1", "pass") + user1_tok = self.login(user1_id, "pass") + + cache = self.hs.get_datastores().main.is_feature_enabled_for_user + self._unstable_feature(user1_tok, "org.matrix.msc3881") + self.assertEqual(len(cache), 1) + + # Not stale yet. + self.reactor.advance(20 * 60) + self.assertEqual(len(cache), 1) + + # Now it has gone 30 minutes unread. + self.reactor.advance(20 * 60) + self.assertEqual(len(cache), 0) + def _sanity_check_versions_response(self, versions_response: JsonDict) -> None: """ Make sure this looks like a `/_matrix/client/versions` response @@ -175,3 +240,13 @@ class VersionsTestCase(unittest.HomeserverTestCase): access_token=self.admin_user_tok, ) self.assertEqual(channel.code, 200) + + def _unstable_feature(self, access_token: str, feature: str) -> bool: + """Ask `/versions` whether `feature` is on for whoever holds this token.""" + channel = self.make_request( + "GET", + "/_matrix/client/versions", + access_token=access_token, + ) + self.assertEqual(channel.code, 200, channel.result) + return channel.json_body["unstable_features"][feature] diff --git a/tests/storage/databases/main/test_cache.py b/tests/storage/databases/main/test_cache.py index 1e7ae7e5a3..0863715358 100644 --- a/tests/storage/databases/main/test_cache.py +++ b/tests/storage/databases/main/test_cache.py @@ -22,6 +22,7 @@ from unittest.mock import Mock, call from signedjson.key import generate_signing_key, get_verify_key +from synapse.rest.admin.experimental_features import ExperimentalFeature from synapse.storage.database import LoggingTransaction from synapse.storage.keys import FetchKeyResult @@ -173,3 +174,47 @@ class CacheInvalidationOverReplicationTestCase(BaseMultiWorkerStreamTestCase): expected_key = (("server1", key_id),) master_invalidate.assert_called_once_with(expected_key) worker_invalidate.assert_called_once_with(expected_key) + + def test_rust_cache_invalidation_replicates(self) -> None: + """A write on one worker empties the Rust cache on another. + + The Rust cache is bound on the datastore under its name, like a + `@cached` function, so a row on the caches stream naming it reaches it + through `getattr(store, name).invalidate(keys)`. The keys it sees are + strings, because the stream's column is `keys text[]`. + """ + worker = self.make_worker_hs("synapse.app.generic_worker") + worker_store = worker.get_datastores().main + + user_id = self.get_success( + self.hs.get_registration_handler().register_user(localpart="user") + ) + + # Read through the worker's cache, so that it holds this user. + self.assertFalse( + self.get_success( + worker_store.is_feature_enabled_for_user(user_id, "msc3881") + ) + ) + self.assertEqual(len(worker_store.is_feature_enabled_for_user), 1) + + assert self.store._cache_id_gen is not None + self.get_success( + self.store.set_features_for_user( + user_id, {ExperimentalFeature.MSC3881: True} + ) + ) + token = self.store._cache_id_gen.get_current_token() + + self.get_success( + worker.get_replication_data_handler().wait_for_stream_position( + "master", "caches", token + ) + ) + + self.assertEqual(len(worker_store.is_feature_enabled_for_user), 0) + self.assertTrue( + self.get_success( + worker_store.is_feature_enabled_for_user(user_id, "msc3881") + ) + ) diff --git a/tests/storage/databases/main/test_experimental_features.py b/tests/storage/databases/main/test_experimental_features.py new file mode 100644 index 0000000000..2c4101530b --- /dev/null +++ b/tests/storage/databases/main/test_experimental_features.py @@ -0,0 +1,103 @@ +# +# 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: +# . +# +# + +from twisted.internet.testing import MemoryReactor + +from synapse.rest.admin.experimental_features import ExperimentalFeature +from synapse.server import HomeServer +from synapse.util.clock import Clock + +from tests import unittest + + +class ExperimentalFeaturesStoreTestCase(unittest.HomeserverTestCase): + """`is_feature_enabled` reads through a cache that lives in Rust.""" + + def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None: + self.store = hs.get_datastores().main + self.user_id = self.get_success( + hs.get_registration_handler().register_user(localpart="user") + ) + + def _is_enabled(self, feature: ExperimentalFeature) -> bool: + return self.get_success(self.store.is_feature_enabled(self.user_id, feature)) + + def test_reads_through_the_rust_cache(self) -> None: + cache = self.store.is_feature_enabled_for_user + + # No row means disabled, and that answer is cached too. + self.assertFalse(self._is_enabled(ExperimentalFeature.MSC3881)) + self.assertEqual((cache.hits, cache.misses), (0, 1)) + self.assertFalse(self._is_enabled(ExperimentalFeature.MSC3881)) + self.assertEqual((cache.hits, cache.misses), (1, 1)) + + # Writing invalidates, so the next read sees the new row... + self.get_success( + self.store.set_features_for_user( + self.user_id, {ExperimentalFeature.MSC3881: True} + ) + ) + self.assertTrue(self._is_enabled(ExperimentalFeature.MSC3881)) + self.assertEqual((cache.hits, cache.misses), (1, 2)) + + # ...in both directions. + self.get_success( + self.store.set_features_for_user( + self.user_id, {ExperimentalFeature.MSC3881: False} + ) + ) + self.assertFalse(self._is_enabled(ExperimentalFeature.MSC3881)) + + def test_invalidation_with_plain_string_keys(self) -> None: + """Invalidations from other workers arrive as strings. + + The cache invalidation stream's column is `keys text[]`, so whatever a + key started life as, a worker receiving it sees a string, and the Rust + side has to decode that back into its own key type. + """ + cache = self.store.is_feature_enabled_for_user + self._is_enabled(ExperimentalFeature.MSC3881) + self.assertEqual(len(cache), 1) + + # A key that matches nothing is a no-op rather than an error: a worker + # may well see feature names this version doesn't know about. + cache.invalidate((self.user_id, "not_a_feature")) + cache.invalidate(("@nobody:test", "msc3881")) + self.assertEqual(len(cache), 1) + + # So is a key of the wrong shape, which a newer worker might stream. A + # Python cache shrugs those off, and raising here would fail the whole + # batch of invalidations the row arrived in. + cache.invalidate((self.user_id,)) + cache.invalidate((self.user_id, "msc3881", "extra")) + cache.invalidate((self.user_id, 42)) + self.assertEqual(len(cache), 1) + + cache.invalidate((self.user_id, "msc3881")) + self.assertEqual(len(cache), 0) + + def test_an_unknown_feature_is_an_error(self) -> None: + """The Rust side keys the cache by its own enum of features, which + mirrors `ExperimentalFeature`. A name it doesn't know means the two + have drifted apart, and should fail rather than quietly read as off.""" + with self.assertRaises(ValueError): + self.store.is_feature_enabled_for_user(self.user_id, "not_a_feature") + + @unittest.override_config({"experimental_features": {"msc3881_enabled": True}}) + def test_a_globally_enabled_feature_skips_the_lookup(self) -> None: + cache = self.store.is_feature_enabled_for_user + + self.assertTrue(self._is_enabled(ExperimentalFeature.MSC3881)) + self.assertEqual((cache.hits, cache.misses), (0, 0))