From 93287194f8393c8a7f964b3a1804aa572c184c3a Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Thu, 15 Oct 2026 15:43:17 +0000 Subject: [PATCH] Do the Rust transmission loop's database work natively On the native Postgres driver, the Rust queues' `NativeHooks` read the device EDUs and record finished transactions through `FederationStore`, on tokio. A transaction then takes no trip to the reactor, unless it finishes a sticky event batch or a module wants to hear about delivered events. Catch-up and the sticky event backlog stay Python. SQLite and the Python Postgres driver keep the Python hooks. Python skips the device queries for destinations with nothing new by asking two stream change caches. Rust keeps a mirror of each, per destination, fed by a new `StreamChangeCache.add_listener`, so it hears of every change the caches do. The current stream positions live in Python's ID generators, and are read under the GIL (without a reactor trip) only when a destination has changed. The loops now count their transactions, EDUs and requests in Rust. `synapse.metrics.pre_collect_hooks` adds the counts to the Python metrics when they are collected. This removes a GIL crossing on every failed send, and from `_rust_complete`. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01R5b7Ruy4cgt27Mjq6Ptxbb --- changelog.d/rust-fed-store.misc | 1 + rust/src/federation/queue/changes.rs | 114 +++++++++ rust/src/federation/queue/metrics.rs | 99 ++++++++ rust/src/federation/queue/mod.rs | 16 +- rust/src/federation/queue/native.rs | 273 +++++++++++++++++++++ rust/src/federation/queue/python.rs | 171 ++++++++++--- scripts-dev/mypy_synapse_plugin.py | 1 + synapse/federation/sender/__init__.py | 24 +- synapse/federation/sender/rust_queue.py | 123 +++++++--- synapse/metrics/__init__.py | 46 ++++ synapse/server.py | 4 + synapse/synapse_rust/federation.pyi | 45 +++- synapse/util/caches/stream_change_cache.py | 24 +- tests/federation/test_rust_queues.py | 188 ++++++++++++-- tests/util/test_stream_change_cache.py | 22 ++ 15 files changed, 1053 insertions(+), 98 deletions(-) create mode 100644 changelog.d/rust-fed-store.misc create mode 100644 rust/src/federation/queue/changes.rs create mode 100644 rust/src/federation/queue/metrics.rs create mode 100644 rust/src/federation/queue/native.rs diff --git a/changelog.d/rust-fed-store.misc b/changelog.d/rust-fed-store.misc new file mode 100644 index 0000000000..652a1438c3 --- /dev/null +++ b/changelog.d/rust-fed-store.misc @@ -0,0 +1 @@ +Do the database work of the Rust federation transmission loops through the native Postgres driver, when the main database uses it (`rust_federation_sending: queues`). diff --git a/rust/src/federation/queue/changes.rs b/rust/src/federation/queue/changes.rs new file mode 100644 index 0000000000..30357d7afb --- /dev/null +++ b/rust/src/federation/queue/changes.rs @@ -0,0 +1,114 @@ +/* + * 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: + * . + * + */ + +//! Which destinations have new to-device messages or device list updates. +//! +//! Python skips the two device queries for a destination with nothing new, +//! by asking a stream change cache. [`StreamChanges`] answers the same +//! question for the Rust loop. Python feeds it every change it feeds its +//! cache, through a listener on the cache, so the two agree. +//! +//! It keeps one position per destination and never evicts, as there is one +//! entry per destination we have sent device updates to. + +use std::{collections::HashMap, sync::Mutex}; + +/// A mirror of a Python `StreamChangeCache` keyed by destination. +pub struct StreamChanges { + inner: Mutex, +} + +struct Inner { + /// Nothing is known about changes at or before this position. + earliest: i64, + latest: HashMap, +} + +impl StreamChanges { + /// A mirror that knows about the changes after `earliest`. + pub fn new(earliest: i64) -> Self { + StreamChanges { + inner: Mutex::new(Inner { + earliest, + latest: HashMap::new(), + }), + } + } + + /// `entity_has_changed`. + pub fn entity_has_changed(&self, destination: &str, stream_pos: i64) { + let mut inner = self.inner.lock().expect("poisoned"); + if stream_pos <= inner.earliest { + return; + } + match inner.latest.get_mut(destination) { + Some(latest) => *latest = (*latest).max(stream_pos), + None => { + inner.latest.insert(destination.to_owned(), stream_pos); + } + } + } + + /// `all_entities_changed`. + pub fn all_entities_changed(&self, stream_pos: i64) { + let mut inner = self.inner.lock().expect("poisoned"); + inner.latest.clear(); + inner.earliest = stream_pos; + } + + /// `has_entity_changed`: whether the destination may have changed after + /// `stream_pos`. + pub fn has_entity_changed(&self, destination: &str, stream_pos: i64) -> bool { + let inner = self.inner.lock().expect("poisoned"); + if stream_pos < inner.earliest { + return true; + } + inner + .latest + .get(destination) + .is_some_and(|latest| stream_pos < *latest) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn as_the_python_cache() { + let changes = StreamChanges::new(10); + + // Nothing is known before the start. + assert!(changes.has_entity_changed("a", 9)); + assert!(!changes.has_entity_changed("a", 10)); + + // A change at or before the start is already accounted for. + changes.entity_has_changed("a", 10); + assert!(!changes.has_entity_changed("a", 10)); + + changes.entity_has_changed("a", 12); + assert!(changes.has_entity_changed("a", 11)); + assert!(!changes.has_entity_changed("a", 12)); + assert!(!changes.has_entity_changed("b", 11)); + + // Positions only move forward. + changes.entity_has_changed("a", 11); + assert!(changes.has_entity_changed("a", 11)); + + changes.all_entities_changed(20); + assert!(changes.has_entity_changed("b", 19)); + assert!(!changes.has_entity_changed("a", 20)); + } +} diff --git a/rust/src/federation/queue/metrics.rs b/rust/src/federation/queue/metrics.rs new file mode 100644 index 0000000000..86a9e2de3f --- /dev/null +++ b/rust/src/federation/queue/metrics.rs @@ -0,0 +1,99 @@ +/* + * 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 counters the transmission loops bump. +//! +//! They are Python's Prometheus metrics. Bumping those from tokio would take +//! the GIL for every transaction, so the loops count here, and Python adds +//! the counts to its metrics when they are collected. + +use std::{ + collections::{HashMap, HashSet}, + sync::Mutex, +}; + +use super::Sent; +use crate::federation::client::RequestStats; + +/// What happened since the counts were last taken. +#[derive(Debug, Default, PartialEq, Eq)] +pub struct Counts { + /// `synapse_federation_client_sent_transactions`. + pub sent_transactions: u64, + /// `synapse_federation_client_sent_edus_by_type`. Their sum is + /// `synapse_federation_client_sent_edus`. + pub sent_edus_by_type: HashMap, + /// `synapse_http_matrixfederationclient_requests` for `PUT`. + pub requests: u64, + /// `synapse_http_matrixfederationclient_responses` for `PUT`, by code. + pub responses: HashMap, + /// `synapse_federation_last_sent_pdu_time`, in milliseconds, for the + /// destinations in `federation_metrics_domains`. + pub last_sent_pdu_ts: HashMap, +} + +/// The counts, shared by the queues of a sender. +pub struct QueueMetrics { + counts: Mutex, + /// `federation_metrics_domains`. + metrics_domains: HashSet, +} + +impl QueueMetrics { + pub fn new(metrics_domains: HashSet) -> Self { + QueueMetrics { + counts: Mutex::new(Counts::default()), + metrics_domains, + } + } + + /// Count the requests a transaction took, sent or not. + pub fn record_requests(&self, stats: &RequestStats) { + let mut counts = self.counts.lock().expect("poisoned"); + add_requests(&mut counts, stats); + } + + /// Count a sent transaction and its requests. + pub fn record_sent(&self, destination: &str, sent: &Sent) { + let mut counts = self.counts.lock().expect("poisoned"); + add_requests(&mut counts, &sent.stats); + counts.sent_transactions += 1; + for edu_type in &sent.edu_types { + *counts + .sent_edus_by_type + .entry(edu_type.clone()) + .or_default() += 1; + } + if let Some(last) = sent.pdus.last() { + if self.metrics_domains.contains(destination) { + counts + .last_sent_pdu_ts + .insert(destination.to_owned(), last.origin_server_ts); + } + } + } + + /// The counts since the last call. + pub fn take(&self) -> Counts { + std::mem::take(&mut *self.counts.lock().expect("poisoned")) + } +} + +fn add_requests(counts: &mut Counts, stats: &RequestStats) { + counts.requests += u64::from(stats.attempts); + for code in &stats.responses { + *counts.responses.entry(*code).or_default() += 1; + } +} diff --git a/rust/src/federation/queue/mod.rs b/rust/src/federation/queue/mod.rs index a76bd377da..c57688f9c1 100644 --- a/rust/src/federation/queue/mod.rs +++ b/rust/src/federation/queue/mod.rs @@ -23,11 +23,14 @@ //! //! The loop needs the database for to-device messages, device list updates //! and the stream positions, and runs catch-up and the sticky event backlog. -//! Those stay in Python for now, behind [`QueueHooks`]. A transaction takes -//! two trips to the reactor: one for its device EDUs and one to record that -//! it was sent. A loop that runs out of things to send has asked for device -//! EDUs once more to find that out, as Python's does. Everything else, -//! building and sending the transaction included, runs on tokio. +//! It reaches those through [`QueueHooks`]. Catch-up and the sticky event +//! backlog are Python. The rest is Python too, unless the main database uses +//! the native driver: then [`native::NativeHooks`] does it on tokio, and a +//! transaction needs no trip to the reactor. Otherwise a transaction takes +//! two: one for its device EDUs and one to record that it was sent, and a +//! loop that runs out of things to send has asked for device EDUs once more +//! to find that out, as Python's does. Everything else, building and sending +//! the transaction included, runs on tokio. use std::{ any::Any, @@ -48,6 +51,9 @@ use super::{ }; use crate::{deferred::AwaitableObserver, events::Event}; +pub mod changes; +pub mod metrics; +pub mod native; mod ordered; pub mod python; diff --git a/rust/src/federation/queue/native.rs b/rust/src/federation/queue/native.rs new file mode 100644 index 0000000000..08414d1e58 --- /dev/null +++ b/rust/src/federation/queue/native.rs @@ -0,0 +1,273 @@ +/* + * 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 loop's database work through the native Postgres driver. +//! +//! When the main database uses the native driver, [`NativeHooks`] reads +//! the device EDUs and records finished transactions with the +//! [`FederationStore`], on tokio. A transaction then needs no trip to the +//! reactor, unless it finishes a sticky event batch or a module wants to +//! hear about delivered events. Catch-up and the sticky event backlog stay +//! Python, through [`PyHooks`]. + +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use pyo3::{exceptions::PyRuntimeError, intern, prelude::*}; + +use super::{ + changes::StreamChanges, metrics::QueueMetrics, python::PyHooks, Completion, DeviceEdus, + DeviceEdusRequest, LoopError, Observer, QueueHooks, StickyBatch, +}; +use crate::{ + database::postgres::pool::{ConnectionPool, PyConnectionPool}, + federation::gil::with_gil, + storage::federation_store::FederationStore, +}; + +/// The logger of `synapse.logging.issue9533_logger`. +const ISSUE_9533_LOG: &str = "synapse.9533_debug"; + +/// What the queues of a sender share when they use the native store. +pub struct Native { + pub store: FederationStore, + pub pool: NativePool, + /// `stream_tokens()`: the current positions of the to-device and device + /// list streams, which only Python's ID generators know. + pub stream_tokens: Py, + /// Which destinations have new to-device messages. + pub to_device: StreamChanges, + /// Which destinations have new device list updates. + pub device_lists: StreamChanges, +} + +impl Native { + /// The positions of the to-device and device list streams. Takes the + /// GIL, but not a trip to the reactor: the ID generators are read where + /// they stand. + fn stream_tokens(&self) -> Result<(i64, i64), LoopError> { + with_gil("queue_stream_tokens", |py| { + self.stream_tokens.bind(py).call0()?.extract() + }) + .map_err(|err: PyErr| LoopError::Other(format!("Failed to read the stream tokens: {err}"))) + } +} + +/// The native pool, as the Python `RustConnectionPool` holds it. That +/// replaces its pool when it is closed and started again, so a closed pool +/// is read again from Python. +pub struct NativePool { + /// The Python `RustConnectionPool`. + python: Py, + cached: Mutex>, +} + +impl NativePool { + pub fn new(python: Py) -> Self { + NativePool { + python, + cached: Mutex::new(None), + } + } + + pub fn get(&self) -> Result { + let mut cached = self.cached.lock().expect("poisoned"); + if let Some(pool) = cached.as_ref().filter(|pool| !pool.is_closed()) { + return Ok(pool.clone()); + } + let pool = with_gil("queue_db_pool", |py| -> PyResult { + let pool = self.python.bind(py).getattr(intern!(py, "_pool"))?; + if pool.is_none() { + return Err(PyRuntimeError::new_err("the database pool is closed")); + } + Ok(pool.cast::()?.get().pool.clone()) + }) + .map_err(|err| LoopError::Other(err.to_string()))?; + *cached = Some(pool.clone()); + Ok(pool) + } +} + +/// The hooks with the database work done natively. +pub(super) struct NativeHooks { + pub(super) native: Arc, + /// For what stays Python. + pub(super) python: PyHooks, + pub(super) metrics: Arc, + pub(super) destination: String, + /// Whether a module wants to hear about delivered events, when the + /// queue was made. + pub(super) notify_delivered: bool, +} + +fn db_error(err: anyhow::Error) -> LoopError { + LoopError::Other(format!("{err:#}")) +} + +#[async_trait] +impl QueueHooks for NativeHooks { + async fn catch_up(&self, observer: &Observer) -> Result<(), LoopError> { + self.python.catch_up(observer).await + } + + /// `get_new_device_msgs_for_remote` and `get_device_updates_by_remote`, + /// with the change trackers in place of their stream change caches. + async fn device_edus( + &self, + _observer: &Observer, + request: DeviceEdusRequest, + ) -> Result { + let destination = self.destination.as_str(); + let to_device_changed = self + .native + .to_device + .has_entity_changed(destination, request.last_device_stream_id); + let device_lists_changed = self + .native + .device_lists + .has_entity_changed(destination, request.last_device_list_stream_id); + if !to_device_changed && !device_lists_changed { + // Python moves both positions on to the current tokens here. We + // would need the GIL to read them, and stay put instead. + return Ok(DeviceEdus { + to_device: Vec::new(), + to_device_stream_id: request.last_device_stream_id, + device_list: Vec::new(), + device_list_stream_id: request.last_device_list_stream_id, + }); + } + + let (to_device_token, device_list_token) = self.native.stream_tokens()?; + let pool = self.native.pool.get()?; + let store = &self.native.store; + + let (to_device, to_device_stream_id) = if to_device_changed { + store + .get_new_device_msgs_for_remote( + &pool, + destination, + request.last_device_stream_id, + to_device_token, + request.to_device_limit, + ) + .await + .map_err(db_error)? + } else { + (Vec::new(), to_device_token) + }; + if !to_device.is_empty() { + log::debug!( + target: ISSUE_9533_LOG, + "Sending {} to-device messages to {destination}, up to stream id {to_device_stream_id}", + to_device.len() + ); + } + + let device_list_limit = request.edu_limit.saturating_sub(to_device.len()); + let (device_list_stream_id, device_list) = if device_lists_changed { + store + .get_device_updates_by_remote( + &pool, + destination, + request.last_device_list_stream_id, + device_list_token, + device_list_limit, + ) + .await + .map_err(db_error)? + } else { + (device_list_token, Vec::new()) + }; + if device_list.len() > device_list_limit { + return Err(LoopError::Other( + "get_device_updates_by_remote returned too many EDUs".to_owned(), + )); + } + + Ok(DeviceEdus { + to_device, + to_device_stream_id, + device_list, + device_list_stream_id, + }) + } + + async fn sticky_backlog(&self, observer: &Observer) -> Result, LoopError> { + self.python.sticky_backlog(observer).await + } + + /// The database writes of `_rust_complete`, natively. Python is only + /// called for a sticky event batch or a module that wants to hear about + /// the delivered events. + async fn complete(&self, observer: &Observer, completion: Completion) -> Result<(), LoopError> { + let destination = self.destination.as_str(); + if let Some(sent) = &completion.sent { + self.metrics.record_sent(destination, sent); + } + + let writes = completion.to_device_stream_id.is_some() + || completion.device_list_stream_id.is_some() + || completion.last_stream_ordering.is_some(); + if writes { + let pool = self.native.pool.get()?; + let store = &self.native.store; + if let Some(stream_id) = completion.to_device_stream_id { + store + .delete_device_msgs_for_remote(&pool, destination, stream_id) + .await + .map_err(db_error)?; + } + if let Some(stream_id) = completion.device_list_stream_id { + store + .mark_as_sent_devices_by_remote(&pool, destination, stream_id) + .await + .map_err(db_error)?; + } + if let Some(stream_ordering) = completion.last_stream_ordering { + store + .set_destination_last_successful_stream_ordering( + &pool, + destination, + stream_ordering, + ) + .await + .map_err(db_error)?; + } + } + + // The module hears only of acknowledged PDUs, and nothing if none. + let notify = self.notify_delivered + && completion.sent.as_ref().is_some_and(|sent| { + sent.pdus + .iter() + .any(|pdu| sent.acknowledged.contains(&pdu.event_id)) + }); + if completion.sticky.is_none() && !notify { + return Ok(()); + } + self.python + .call_complete( + observer, + Completion { + to_device_stream_id: None, + device_list_stream_id: None, + last_stream_ordering: None, + sticky: completion.sticky, + sent: if notify { completion.sent } else { None }, + }, + ) + .await + } +} diff --git a/rust/src/federation/queue/python.rs b/rust/src/federation/queue/python.rs index c0408a0163..43961b9535 100644 --- a/rust/src/federation/queue/python.rs +++ b/rust/src/federation/queue/python.rs @@ -18,10 +18,12 @@ //! [`RustDestinationQueues`] makes a [`RustDestinationQueue`] for each //! destination. The Python `RustPerDestinationQueue` wraps one and provides //! the [`QueueHooks`] as coroutines: `_rust_catch_up`, `_rust_device_edus`, -//! `_rust_sticky_backlog` and `_rust_complete`. +//! `_rust_sticky_backlog` and `_rust_complete`. With the native database +//! driver, [`NativeHooks`] does the database work of the last two instead. use std::{ - sync::{Arc, Mutex}, + collections::HashSet, + sync::{Arc, Mutex, OnceLock}, time::Instant, }; @@ -34,6 +36,9 @@ use pyo3::{ use serde_json::Value; use super::{ + changes::StreamChanges, + metrics::QueueMetrics, + native::{Native, NativeHooks, NativePool}, Activity, Completion, DestinationQueue, DeviceEdus, DeviceEdusRequest, Edu, LoopError, Observer, PduBody, PresenceState, QueueConfig, QueueHooks, QueuedPdu, Receipt, SendOutcome, StickyBatch, StickyStreamPosition, TransactionSink, @@ -51,6 +56,7 @@ use crate::{ }, }, runtime::RustRuntime, + storage::federation_store::{FederationStore, FederationStoreConfig}, }; /// Makes the queues, and tells tests when their loops are idle. @@ -59,45 +65,140 @@ pub struct RustDestinationQueues { sink: Arc, runtime: RustRuntime, activity: Arc, + metrics: Arc, msc4354_enabled: bool, + /// Set when the main database uses the native driver. + native: OnceLock>, } #[pymethods] impl RustDestinationQueues { + /// `metrics_domains` is `federation_metrics_domains`. #[new] - fn py_new(transactions: &Bound<'_, RustTransactionManager>, msc4354_enabled: bool) -> Self { + fn py_new( + transactions: &Bound<'_, RustTransactionManager>, + msc4354_enabled: bool, + metrics_domains: HashSet, + ) -> Self { let transactions = transactions.get(); + let metrics = Arc::new(QueueMetrics::new(metrics_domains)); RustDestinationQueues { sink: Arc::new(ClientSink { sender: Arc::clone(&transactions.sender), runtime: transactions.runtime.clone(), - record_metrics: Arc::clone(&transactions.record_metrics), + metrics: Arc::clone(&metrics), send_for_test: Mutex::new(None), }), runtime: transactions.runtime.clone(), activity: Arc::default(), + metrics, msc4354_enabled, + native: OnceLock::new(), } } + /// Do the loops' database work through the native driver. `db_pool` is + /// the `RustConnectionPool`. `stream_tokens()` returns the current + /// to-device and device list stream positions. The change trackers know + /// about the changes after `to_device_since` and `device_lists_since`, + /// and Python must feed them every later change through + /// `to_device_changed` and `device_list_changed`. + /// + /// Only the queues made after this use the native driver. + fn use_native_store( + &self, + db_pool: Py, + stream_tokens: Py, + allow_device_name_lookup_over_federation: bool, + to_device_since: i64, + device_lists_since: i64, + ) -> PyResult<()> { + let native = Arc::new(Native { + store: FederationStore::new(FederationStoreConfig { + allow_device_name_lookup_over_federation, + }), + pool: NativePool::new(db_pool), + stream_tokens, + to_device: StreamChanges::new(to_device_since), + device_lists: StreamChanges::new(device_lists_since), + }); + self.native.set(native).map_err(|_| { + pyo3::exceptions::PyRuntimeError::new_err("the native store is already in use") + }) + } + + /// A destination has new to-device messages at `stream_pos`, or with no + /// destination, any may have. A `StreamChangeCache` listener. + #[pyo3(signature = (destination, stream_pos))] + fn to_device_changed(&self, destination: Option<&str>, stream_pos: i64) { + if let Some(native) = self.native.get() { + stream_changed(&native.to_device, destination, stream_pos); + } + } + + /// As `to_device_changed`, for device list updates. + #[pyo3(signature = (destination, stream_pos))] + fn device_list_changed(&self, destination: Option<&str>, stream_pos: i64) { + if let Some(native) = self.native.get() { + stream_changed(&native.device_lists, destination, stream_pos); + } + } + + /// The counts since the last call, for Python's metrics: + /// `(sent_transactions, sent_edus_by_type, put_requests, + /// put_responses_by_code, last_sent_pdu_ts_by_destination)`. + #[allow(clippy::type_complexity)] + fn take_metrics( + &self, + ) -> ( + u64, + std::collections::HashMap, + u64, + std::collections::HashMap, + std::collections::HashMap, + ) { + let counts = self.metrics.take(); + ( + counts.sent_transactions, + counts.sent_edus_by_type, + counts.requests, + counts.responses, + counts.last_sent_pdu_ts, + ) + } + /// The queue for `destination`. `facade` is the Python queue, which /// provides the hooks. The Rust queue holds it weakly. + /// `notify_delivered` is whether a module wants to hear about delivered + /// events. fn queue( &self, py: Python<'_>, destination: String, facade: &Bound<'_, PyAny>, should_send_on_this_instance: bool, + notify_delivered: bool, ) -> PyResult { let facade = py .import("weakref")? .getattr("ref")? .call1((facade,))? .unbind(); - let hooks = PyHooks { + let python = PyHooks { runtime: self.runtime.clone(), facade: Arc::new(facade), destination: destination.clone(), + metrics: Arc::clone(&self.metrics), + }; + let hooks: Arc = match self.native.get() { + Some(native) => Arc::new(NativeHooks { + native: Arc::clone(native), + python, + metrics: Arc::clone(&self.metrics), + destination: destination.clone(), + notify_delivered, + }), + None => Arc::new(python), }; let sink: Arc = self.sink.clone(); Ok(RustDestinationQueue { @@ -107,7 +208,7 @@ impl RustDestinationQueues { should_send_on_this_instance, msc4354_enabled: self.msc4354_enabled, }, - Arc::new(hooks), + hooks, sink, Arc::clone(&self.activity), self.runtime.tokio_handle()?, @@ -307,12 +408,20 @@ fn edu_from_py(edu: &Bound<'_, PyAny>) -> PyResult { }) } +fn stream_changed(changes: &StreamChanges, destination: Option<&str>, stream_pos: i64) { + match destination { + Some(destination) => changes.entity_has_changed(destination, stream_pos), + None => changes.all_entities_changed(stream_pos), + } +} + /// The hooks, as coroutine methods of the Python queue. -struct PyHooks { +pub(super) struct PyHooks { runtime: RustRuntime, /// A weak reference to the Python queue. facade: Arc>, destination: String, + metrics: Arc, } impl PyHooks { @@ -448,6 +557,20 @@ impl QueueHooks for PyHooks { } async fn complete(&self, observer: &Observer, completion: Completion) -> Result<(), LoopError> { + if let Some(sent) = &completion.sent { + self.metrics.record_sent(&self.destination, sent); + } + self.call_complete(observer, completion).await + } +} + +impl PyHooks { + /// `_rust_complete`, without counting the transaction. + pub(super) async fn call_complete( + &self, + observer: &Observer, + completion: Completion, + ) -> Result<(), LoopError> { let completion = Arc::new(completion); self.call( observer, @@ -482,15 +605,7 @@ fn completion_args<'py>(py: Python<'py>, completion: &Completion) -> PyResult {} } } - ( - pdus, - &sent.edu_types, - &sent.acknowledged, - sent.stats.attempts, - &sent.stats.responses, - ) - .into_pyobject(py)? - .into_any() + (pdus, &sent.acknowledged).into_pyobject(py)?.into_any() } None => py.None().into_bound(py), }; @@ -564,23 +679,13 @@ fn format_exception(py: Python<'_>, err: &PyErr) -> String { struct ClientSink { sender: Arc, runtime: RustRuntime, - record_metrics: Arc>, + /// Counts the requests of a failed transaction. A sent one's are counted + /// with its completion. + metrics: Arc, send_for_test: Mutex>>>, } impl ClientSink { - /// Record the request counters, which are Python's. A sent transaction's - /// go with its completion. A failed one needs this. - fn record_metrics(&self, py: Python<'_>, stats: &RequestStats) { - let recorded = - self.record_metrics - .bind(py) - .call1(("PUT", stats.attempts, stats.responses.clone())); - if let Err(err) = recorded { - log::warn!("Failed to record federation request metrics: {err}"); - } - } - /// Send through the Python stand-in. async fn send_for_test( &self, @@ -661,8 +766,8 @@ impl TransactionSink for ClientSink { }), Ok(TransactionResponse::Unparsed(body)) => with_gil("queue_response_fallback", |py| { python_acknowledged(py, &destination, &transaction_id, &body).map_err(|err| { - // No later trip to Python will carry the counters. - self.record_metrics(py, &stats); + // There will be no completion to count the requests. + self.metrics.record_requests(&stats); classify_python_error(py, &err) .unwrap_or_else(|_| LoopError::Other(format_exception(py, &err))) }) @@ -672,8 +777,8 @@ impl TransactionSink for ClientSink { stats, }), Err(err) => { - // As above, there is no later trip to Python. - with_gil("queue_send_failed", |py| self.record_metrics(py, &stats)); + // As above, there will be no completion. + self.metrics.record_requests(&stats); Err(LoopError::from_client(err, &destination)) } } diff --git a/scripts-dev/mypy_synapse_plugin.py b/scripts-dev/mypy_synapse_plugin.py index ac1fee7f67..a57f07af56 100644 --- a/scripts-dev/mypy_synapse_plugin.py +++ b/scripts-dev/mypy_synapse_plugin.py @@ -147,6 +147,7 @@ prometheus_metric_fullname_to_label_arg_map: Mapping[str, ArgLocation | None] = "synapse.metrics._gc.PyPyGCStats": None, "synapse.metrics._reactor_metrics.ReactorLastSeenMetric": None, "synapse.metrics.CPUMetrics": None, + "synapse.metrics._PreCollectHooks": None, "synapse.metrics.jemalloc.JemallocCollector": None, "synapse.util.metrics.DynamicCollectorRegistry": None, "synapse.metrics.background_process_metrics._Collector": None, diff --git a/synapse/federation/sender/__init__.py b/synapse/federation/sender/__init__.py index caba8575b9..4d0afbf7ad 100644 --- a/synapse/federation/sender/__init__.py +++ b/synapse/federation/sender/__init__.py @@ -154,7 +154,11 @@ from synapse.federation.sender.per_destination_queue import ( CATCHUP_RETRY_INTERVAL, PerDestinationQueue, ) -from synapse.federation.sender.rust_queue import RustPerDestinationQueue +from synapse.federation.sender.rust_queue import ( + RustPerDestinationQueue, + flush_rust_queue_metrics, + use_native_store, +) from synapse.federation.sender.transaction_manager import TransactionManager from synapse.federation.units import Edu from synapse.logging.context import make_deferred_yieldable, run_in_background @@ -164,6 +168,7 @@ from synapse.metrics import ( event_processing_loop_counter, event_processing_loop_room_count, events_processed_counter, + pre_collect_hooks, ) from synapse.metrics.background_process_metrics import ( wrap_as_background_process, @@ -433,9 +438,17 @@ class FederationSender(AbstractFederationSender): hs.config.experimental.rust_federation_sending_at_least("queues") and rust_transactions is not None ): - self._rust_queues = RustDestinationQueues( - rust_transactions, hs.config.experimental.msc4354_enabled + rust_queues = RustDestinationQueues( + rust_transactions, + hs.config.experimental.msc4354_enabled, + set(hs.config.federation.federation_metrics_domains), ) + use_native_store(hs, rust_queues) + pre_collect_hooks.register_hook( + homeserver_instance_id=hs.get_instance_id(), + hook=self._flush_rust_queue_metrics, + ) + self._rust_queues = rust_queues transaction_queue_pending_destinations_gauge.register_hook( homeserver_instance_id=hs.get_instance_id(), @@ -497,6 +510,11 @@ class FederationSender(AbstractFederationSender): for queue in self._per_destination_queues.values(): queue.shutdown() + def _flush_rust_queue_metrics(self) -> None: + """Add what the Rust loops counted to the metrics.""" + if self._rust_queues is not None: + flush_rust_queue_metrics(self._rust_queues, self.server_name) + def _get_per_destination_queue( self, destination: str ) -> PerDestinationQueue | RustPerDestinationQueue | None: diff --git a/synapse/federation/sender/rust_queue.py b/synapse/federation/sender/rust_queue.py index 03b0f39a8a..cd1d3773fc 100644 --- a/synapse/federation/sender/rust_queue.py +++ b/synapse/federation/sender/rust_queue.py @@ -17,7 +17,13 @@ The queue and its transmission loop are Rust's. `RustPerDestinationQueue` has the methods of `PerDestinationQueue` that `FederationSender` calls, and forwards them. The loop calls back into it for the parts still in Python: catch-up, the sticky event backlog, and the database work around each -transaction. Those are the `_rust_*` coroutines. +transaction. Those are the `_rust_*` coroutines. When the main database uses +the native Rust driver, Rust does that database work itself, and calls +`_rust_complete` only to finish a sticky event batch or to tell modules about +delivered events. + +The loops count what they send in Rust. `flush_rust_queue_metrics` adds the +counts to the Python metrics when those are collected. """ import logging @@ -34,11 +40,16 @@ from synapse.federation.sender.per_destination_queue import ( ) from synapse.federation.sender.transaction_manager import last_pdu_ts_metric from synapse.federation.units import Edu +from synapse.http.matrixfederationclient import ( + incoming_responses_counter, + outgoing_requests_counter, +) from synapse.logging import issue9533_logger from synapse.logging.opentracing import whitelisted_homeserver from synapse.metrics import SERVER_NAME_LABEL, sent_transactions_counter from synapse.metrics.background_process_metrics import BackgroundProcessLoggingContext from synapse.module_api.callbacks.federation import FederatedEventDeliveryMethod +from synapse.storage.rust_pool import RustConnectionPool from synapse.synapse_rust.federation import ( RustDestinationQueue, RustDestinationQueues, @@ -52,6 +63,68 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +def use_native_store( + hs: "synapse.server.HomeServer", rust_queues: RustDestinationQueues +) -> bool: + """Have the Rust loops do their database work through the native driver, + if the main database uses it. Returns whether they do.""" + store = hs.get_datastores().main + db_pool = store.db_pool._db_pool + device_lists_cache = store._device_list_federation_stream_cache + if not isinstance(db_pool, RustConnectionPool) or device_lists_cache is None: + return False + + # Rust skips the device queries for destinations with nothing new, as + # these caches let Python do. So it hears of every change they do. + to_device_since = store._device_federation_outbox_stream_cache.add_listener( + rust_queues.to_device_changed + ) + device_lists_since = device_lists_cache.add_listener( + rust_queues.device_list_changed + ) + rust_queues.use_native_store( + db_pool, + lambda: ( + store.get_to_device_stream_token(), + store.get_device_stream_token().stream, + ), + hs.config.federation.allow_device_name_lookup_over_federation, + to_device_since, + device_lists_since, + ) + return True + + +def flush_rust_queue_metrics( + rust_queues: RustDestinationQueues, server_name: str +) -> None: + """Add what the Rust loops counted since the last call to the metrics.""" + ( + sent_transactions, + sent_edus_by_type_counts, + put_requests, + put_responses, + last_sent_pdu_ts, + ) = rust_queues.take_metrics() + labels = {SERVER_NAME_LABEL: server_name} + + if sent_transactions: + sent_transactions_counter.labels(**labels).inc(sent_transactions) + sent_edus = sum(sent_edus_by_type_counts.values()) + if sent_edus: + sent_edus_counter.labels(**labels).inc(sent_edus) + for edu_type, count in sent_edus_by_type_counts.items(): + sent_edus_by_type.labels(type=edu_type, **labels).inc(count) + if put_requests: + outgoing_requests_counter.labels(method="PUT", **labels).inc(put_requests) + for code, count in put_responses.items(): + incoming_responses_counter.labels(method="PUT", code=code, **labels).inc(count) + for destination, ts in last_sent_pdu_ts.items(): + last_pdu_ts_metric.labels(destination_server_name=destination, **labels).set( + ts / 1000 + ) + + class _RustStickyBacklog: """`StickyEventBacklogTracker`'s flag, which the Rust queue holds.""" @@ -91,12 +164,6 @@ class RustPerDestinationQueue(_CatchUpLoop): self._transaction_manager = transaction_manager self._destination = destination self._federation_callbacks = hs.get_module_api_callbacks().federation - self._federation_metrics_domains = ( - hs.config.federation.federation_metrics_domains - ) - self._record_request_metrics = ( - hs.get_federation_http_client()._record_request_metrics - ) # Its database half. The Rust queue holds its flag. self._sticky_backlog_store = StickyEventBacklogTracker( destination, hs, hs.config.experimental.msc4354_enabled @@ -107,7 +174,12 @@ class RustPerDestinationQueue(_CatchUpLoop): hs.get_instance_name(), destination ) ) - self._rust = rust_queues.queue(destination, self, should_send_on_this_instance) + self._rust = rust_queues.queue( + destination, + self, + should_send_on_this_instance, + self._federation_callbacks.interested_in_events_delivered_over_federation(), + ) self._sticky_event_backlog_tracker = _RustStickyBacklog(self._rust) def __str__(self) -> str: @@ -260,13 +332,16 @@ class RustPerDestinationQueue(_CatchUpLoop): device_list_stream_id: int | None, last_stream_ordering: int | None, sticky: tuple[str, _StickyEventsTransactionInfo] | None, - sent: tuple[list[EventBase], list[str], list[str], int, list[int]] | None, + sent: tuple[list[EventBase], list[str]] | None, ) -> None: """Record a finished transaction: what `send_new_transaction` does - after the response, the counters, and `_complete_transaction`.""" + after the response, less the counters, and `_complete_transaction`. + + `sent` is the PDUs sent and the IDs of those the destination accepted. + """ with self._logcontext("federation_queue_complete"): if sent is not None: - await self._record_sent(*sent) + await self._notify_delivered(*sent) if to_device_stream_id: await self._store.delete_device_msgs_for_remote( @@ -283,16 +358,9 @@ class RustPerDestinationQueue(_CatchUpLoop): if sticky is not None: await self._sticky_backlog_store.complete_transaction(sticky[1]) - async def _record_sent( - self, - pdus: list[EventBase], - edu_types: list[str], - acknowledged: list[str], - attempts: int, - responses: list[int], + async def _notify_delivered( + self, pdus: list[EventBase], acknowledged: list[str] ) -> None: - self._record_request_metrics("PUT", attempts, responses) - if self._federation_callbacks.interested_in_events_delivered_over_federation(): acknowledged_ids = set(acknowledged) await self._federation_callbacks.notify_on_event_delivered_over_federation( @@ -300,18 +368,3 @@ class RustPerDestinationQueue(_CatchUpLoop): [p for p in pdus if p.event_id in acknowledged_ids], FederatedEventDeliveryMethod.SEND, ) - - if pdus and self._destination in self._federation_metrics_domains: - last_pdu_ts_metric.labels( - destination_server_name=self._destination, - **{SERVER_NAME_LABEL: self.server_name}, - ).set(pdus[-1].origin_server_ts / 1000) - - sent_transactions_counter.labels(**{SERVER_NAME_LABEL: self.server_name}).inc() - sent_edus_counter.labels(**{SERVER_NAME_LABEL: self.server_name}).inc( - len(edu_types) - ) - for edu_type in edu_types: - sent_edus_by_type.labels( - type=edu_type, **{SERVER_NAME_LABEL: self.server_name} - ).inc() diff --git a/synapse/metrics/__init__.py b/synapse/metrics/__init__.py index cf9aaa7f2d..803fb58dac 100644 --- a/synapse/metrics/__init__.py +++ b/synapse/metrics/__init__.py @@ -615,6 +615,52 @@ class CPUMetrics(Collector): REGISTRY.register(CPUMetrics()) # type: ignore[missing-server-name-label] +class _PreCollectHooks(Collector): + """Runs hooks when the metrics are collected, before collecting the + metrics registered after it. + + For counts kept outside Python, such as in Rust, which a hook adds to the + Python metrics only when they are read. It is registered here, so it runs + before every metric defined after this module. + """ + + def __init__(self) -> None: + self._instance_id_to_hook_map: dict[str, Callable[[], None]] = {} + + def describe(self) -> Iterable[Metric]: + return [] + + def collect(self) -> Iterable[Metric]: + for homeserver_instance_id, hook in list(self._instance_id_to_hook_map.items()): + try: + hook() + except Exception: + logger.exception( + "Exception running a pre-collect hook for homeserver_instance_id=%s", + homeserver_instance_id, + ) + return [] + + def register_hook( + self, *, homeserver_instance_id: str, hook: Callable[[], None] + ) -> None: + """Run `hook` at each collection, until the homeserver shuts down.""" + assert homeserver_instance_id not in self._instance_id_to_hook_map, ( + f"A pre-collect hook is already registered for homeserver_instance_id={homeserver_instance_id}" + ) + self._instance_id_to_hook_map[homeserver_instance_id] = hook + + def unregister_hooks_for_homeserver_instance_id( + self, homeserver_instance_id: str + ) -> None: + self._instance_id_to_hook_map.pop(homeserver_instance_id, None) + + +# It adds to metrics with the server name label, and has none of its own. +pre_collect_hooks = _PreCollectHooks() # type: ignore[missing-server-name-label] +REGISTRY.register(pre_collect_hooks) + + # # Federation Metrics # diff --git a/synapse/server.py b/synapse/server.py index ec1a57de2a..fffc144aed 100644 --- a/synapse/server.py +++ b/synapse/server.py @@ -149,6 +149,7 @@ from synapse.media.media_repository import MediaRepository from synapse.metrics import ( SERVER_NAME_LABEL, all_later_gauges_to_clean_up_on_shutdown, + pre_collect_hooks, register_threadpool, synapse_server_name_info, ) @@ -494,6 +495,9 @@ class HomeServer(metaclass=abc.ABCMeta): later_gauge.unregister_hooks_for_homeserver_instance_id( self.get_instance_id() ) + pre_collect_hooks.unregister_hooks_for_homeserver_instance_id( + self.get_instance_id() + ) CACHE_METRIC_REGISTRY.unregister_hooks_for_homeserver( self.config.server.server_name diff --git a/synapse/synapse_rust/federation.pyi b/synapse/synapse_rust/federation.pyi index 59a238dafa..2136b0da7d 100644 --- a/synapse/synapse_rust/federation.pyi +++ b/synapse/synapse_rust/federation.pyi @@ -207,13 +207,50 @@ class RustDestinationQueues: `rust_federation_sending: queues`.""" def __init__( - self, transactions: RustTransactionManager, msc4354_enabled: bool - ) -> None: ... + self, + transactions: RustTransactionManager, + msc4354_enabled: bool, + metrics_domains: set[str], + ) -> None: + """`metrics_domains` is `federation_metrics_domains`.""" + + def use_native_store( + self, + db_pool: Any, + stream_tokens: Callable[[], tuple[int, int]], + allow_device_name_lookup_over_federation: bool, + to_device_since: int, + device_lists_since: int, + ) -> None: + """Do the loops' database work through the native driver. `db_pool` + is the `RustConnectionPool`, and `stream_tokens()` returns the current + to-device and device list stream positions. The change trackers know + about the changes after the two positions given, and must be told of + every later one.""" + + def to_device_changed(self, destination: str | None, stream_pos: int) -> None: + """A `StreamChangeCache` listener for to-device messages.""" + + def device_list_changed(self, destination: str | None, stream_pos: int) -> None: + """A `StreamChangeCache` listener for device list updates.""" + + def take_metrics( + self, + ) -> tuple[int, dict[str, int], int, dict[int, int], dict[str, int]]: + """The counts since the last call: `(sent_transactions, + sent_edus_by_type, put_requests, put_responses_by_code, + last_sent_pdu_ts_by_destination)`.""" + def queue( - self, destination: str, facade: Any, should_send_on_this_instance: bool + self, + destination: str, + facade: Any, + should_send_on_this_instance: bool, + notify_delivered: bool, ) -> "RustDestinationQueue": """The queue for `destination`. `facade` provides the `_rust_*` - hooks, and is held weakly.""" + hooks, and is held weakly. `notify_delivered` is whether a module + wants to hear about delivered events.""" def idle(self) -> bool: """Whether every transmission loop has finished or is waiting on the diff --git a/synapse/util/caches/stream_change_cache.py b/synapse/util/caches/stream_change_cache.py index 7c6c9bc572..0540af11b0 100644 --- a/synapse/util/caches/stream_change_cache.py +++ b/synapse/util/caches/stream_change_cache.py @@ -21,7 +21,7 @@ import logging import math -from typing import Collection, Mapping +from typing import Callable, Collection, Mapping import attr from sortedcontainers import SortedDict @@ -106,6 +106,9 @@ class StreamChangeCache: # self._earliest_known_stream_pos = current_stream_pos + # Called with every change: see `add_listener`. + self._listeners: list[Callable[[EntityType | None, int], None]] = [] + self.name = name self.metrics = caches.register_cache( cache_type="cache", @@ -307,6 +310,11 @@ class StreamChangeCache: """ assert isinstance(stream_pos, int) + # A listener may know of positions we have evicted, so it hears of + # this change even if we do nothing with it. + for listener in self._listeners: + listener(entity, stream_pos) + # For a change before _cache is valid (e.g. at or before the earliest known # stream position) there's nothing to do. if stream_pos <= self._earliest_known_stream_pos: @@ -338,6 +346,20 @@ class StreamChangeCache: self._cache.clear() self._entity_to_key.clear() self._earliest_known_stream_pos = stream_pos + for listener in self._listeners: + listener(None, stream_pos) + + def add_listener(self, listener: Callable[[EntityType | None, int], None]) -> int: + """Call `listener(entity, stream_pos)` with every later change, and + `listener(None, stream_pos)` when all entities change. + + Returns: + The position after which the listener hears of every change. Changes + at or before it are unknown to the listener. + """ + self._listeners.append(listener) + latest = self._cache.peekitem(-1)[0] if self._cache else 0 + return max(self._earliest_known_stream_pos, latest) def _evict(self) -> None: """ diff --git a/tests/federation/test_rust_queues.py b/tests/federation/test_rust_queues.py index ea1047f09e..f3fbe8b64a 100644 --- a/tests/federation/test_rust_queues.py +++ b/tests/federation/test_rust_queues.py @@ -21,6 +21,7 @@ network, as `tests.http.test_rust_federation_client` does. import json from typing import Any +from unittest import SkipTest from unittest.mock import AsyncMock from twisted.internet.testing import MemoryReactor @@ -28,7 +29,11 @@ from twisted.internet.testing import MemoryReactor from synapse.api.presence import UserPresenceState from synapse.federation.sender import FederationSender from synapse.federation.sender.per_destination_queue import PerDestinationQueue -from synapse.federation.sender.rust_queue import RustPerDestinationQueue +from synapse.federation.sender.rust_queue import ( + RustPerDestinationQueue, + use_native_store, +) +from synapse.federation.sender.transaction_manager import last_pdu_ts_metric from synapse.federation.units import Edu from synapse.metrics import SERVER_NAME_LABEL, sent_transactions_counter from synapse.module_api.callbacks.federation import ( @@ -38,6 +43,7 @@ from synapse.module_api.callbacks.federation import ( from synapse.rest import admin from synapse.rest.client import login, room from synapse.server import HomeServer +from synapse.storage.rust_pool import RustConnectionPool from synapse.synapse_rust.federation import ( RustDestinationQueues, RustFederationClient, @@ -83,10 +89,17 @@ class RustQueuesTestCase(unittest.HomeserverTestCase): self.transport, ) + # Modules register their callbacks before the queues are made, which + # read whether there are any. + self.delivered = AsyncMock(return_value=None) + hs.get_module_api_callbacks().federation.register_callbacks( + on_event_delivered_over_federation=self.delivered + ) + sender = hs.get_federation_sender() assert isinstance(sender, FederationSender) self.sender = sender - self.queues = RustDestinationQueues(RustTransactionManager(rust_client), False) + self.queues = self._make_queues(RustTransactionManager(rust_client)) sender._rust_queues = self.queues user = self.register_user("user", "pass") @@ -102,20 +115,32 @@ class RustQueuesTestCase(unittest.HomeserverTestCase): assert isinstance(queue, RustPerDestinationQueue) self.queue = queue + def _make_queues( + self, transactions: RustTransactionManager + ) -> RustDestinationQueues: + return RustDestinationQueues(transactions, False, {DEST}) + def _settle(self) -> None: wait_for_rust_queues(self.reactor, self.queues) + def _queue_calls(self, before: dict[str, tuple[int, float]]) -> dict[str, int]: + """The GIL crossings of the queues since `before`.""" + after = federation_gil_stats() + return { + site: after[site][0] - before.get(site, (0, 0.0))[0] + for site in after + if site.startswith("queue_") + and after[site][0] != before.get(site, (0, 0.0))[0] + } + def _sent_transactions(self) -> float: + self.sender._flush_rust_queue_metrics() return sent_transactions_counter.labels( **{SERVER_NAME_LABEL: self.hs.hostname} )._value.get() def test_send_pdu(self) -> None: """The queue sends the PDU through the Rust client and records it.""" - callback = AsyncMock(return_value=None) - self.hs.get_module_api_callbacks().federation.register_callbacks( - on_event_delivered_over_federation=callback - ) self.transport.responses.append( json_response(200, {"pdus": {self.event.event_id: {}}}) ) @@ -139,7 +164,7 @@ class RustQueuesTestCase(unittest.HomeserverTestCase): ), self.event.internal_metadata.stream_ordering, ) - callback.assert_awaited_once_with( + self.delivered.assert_awaited_once_with( FederationEventDeliveryEvent( server_name=DEST, events=[self.event], @@ -220,17 +245,10 @@ class RustQueuesTestCase(unittest.HomeserverTestCase): before = federation_gil_stats() self.queue.send_pdu(self.event) self._settle() - after = federation_gil_stats() - calls = { - site: after[site][0] - before.get(site, (0, 0.0))[0] - for site in after - if after[site][0] != before.get(site, (0, 0.0))[0] - } - queue_calls = { - site: count for site, count in calls.items() if site.startswith("queue_") - } - self.assertEqual(queue_calls, {"queue_device_edus": 2, "queue_complete": 1}) + self.assertEqual( + self._queue_calls(before), {"queue_device_edus": 2, "queue_complete": 1} + ) def test_shutdown(self) -> None: self.sender.shutdown() @@ -241,6 +259,142 @@ class RustQueuesTestCase(unittest.HomeserverTestCase): self.assertFalse(self.queue.transmission_loop_running) +class RustQueuesNativeTestCase(RustQueuesTestCase): + """The same, with the database work done through the native driver.""" + + def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer: + hs = self.setup_test_homeserver() + if not isinstance( + hs.get_datastores().main.db_pool._db_pool, RustConnectionPool + ): + raise SkipTest("needs the native Rust Postgres driver") + return hs + + def _make_queues( + self, transactions: RustTransactionManager + ) -> RustDestinationQueues: + queues = super()._make_queues(transactions) + self.assertTrue(use_native_store(self.hs, queues)) + return queues + + def _catch_up(self) -> None: + """Send a first transaction, which leaves catch-up and reads the + device streams for the first time.""" + self.transport.responses.append(json_response(200, {})) + self.queue.send_pdu(self.event) + self._settle() + self.transport.requests.clear() + + def test_gil_crossings(self) -> None: + """Once caught up, a transaction with no device updates takes the GIL + only to tell the module about the delivered event.""" + self._catch_up() + self.transport.responses.append( + json_response(200, {"pdus": {self.event.event_id: {}}}) + ) + + before = federation_gil_stats() + self.queue.send_pdu(self.event) + self._settle() + + self.assertEqual(len(self.transport.requests), 1) + self.assertEqual(self._queue_calls(before), {"queue_complete": 1}) + + def test_gil_crossings_without_modules(self) -> None: + """With no module to tell, it takes the GIL nowhere in the loop.""" + self.hs.get_module_api_callbacks().federation._on_event_delivered_over_federation_callbacks.clear() + queue = self.sender._get_per_destination_queue("other.example.com") + assert isinstance(queue, RustPerDestinationQueue) + self.transport.responses.extend([json_response(200, {})] * 2) + queue.send_pdu(self.event) + self._settle() + self.transport.requests.clear() + + before = federation_gil_stats() + queue.send_pdu(self.event) + self._settle() + + self.assertEqual(len(self.transport.requests), 1) + self.assertEqual(self._queue_calls(before), {}) + + def test_to_device(self) -> None: + """A to-device message is sent, and deleted once sent. Reading the + stream position is the one GIL crossing.""" + self._catch_up() + self.transport.responses.append(json_response(200, {})) + message = { + "sender": "@user:test", + "type": "m.test", + "message_id": "m1", + "messages": {"@u:remote": {"D": {"body": "hi"}}}, + } + + before = federation_gil_stats() + self.get_success( + self.store.add_remote_messages_from_client_to_device_inbox({DEST: message}) + ) + self.queue.attempt_new_transaction() + self._settle() + + [request] = self.transport.requests + body = json.loads(request.body) + self.assertEqual( + body["edus"], [{"edu_type": "m.direct_to_device", "content": message}] + ) + self.assertEqual(self._queue_calls(before), {"queue_stream_tokens": 1}) + current = self.store.get_to_device_stream_token() + self.assertEqual( + self.get_success( + self.store.get_new_device_msgs_for_remote(DEST, 0, current, 10) + ), + ([], current), + ) + + def test_device_list_update(self) -> None: + """A device list update is sent, and marked as sent.""" + self._catch_up() + self.transport.responses.append(json_response(200, {})) + user = "@user:test" + self.get_success(self.store.store_device(user, "NEW", None)) + self.get_success(self.store.add_device_change_to_streams(user, ["NEW"], [])) + self.get_success( + self.store.add_device_list_outbound_pokes( + user_id=user, device_id="NEW", room_id="!r", hosts=[DEST], context={} + ) + ) + + self.queue.attempt_new_transaction() + self._settle() + + [request] = self.transport.requests + [edu] = json.loads(request.body)["edus"] + self.assertEqual(edu["edu_type"], "m.device_list_update") + self.assertEqual(edu["content"]["device_id"], "NEW") + pokes = self.get_success( + self.store.db_pool.simple_select_list( + "device_lists_outbound_pokes", {"destination": DEST}, ["stream_id"] + ) + ) + self.assertEqual(pokes, []) + + def test_metrics(self) -> None: + """The loops' counts reach the Python metrics when collected.""" + self._catch_up() + self.transport.responses.append(json_response(200, {})) + sent_before = self._sent_transactions() + + self.queue.send_pdu(self.event) + self._settle() + + self.assertEqual(self._sent_transactions(), sent_before + 1) + self.assertEqual( + last_pdu_ts_metric.labels( + destination_server_name=DEST, **{SERVER_NAME_LABEL: self.hs.hostname} + )._value.get(), + self.event.origin_server_ts / 1000, + ) + + class RustQueuesLevelTestCase(unittest.HomeserverTestCase): """Which queue each level gets.""" diff --git a/tests/util/test_stream_change_cache.py b/tests/util/test_stream_change_cache.py index 69a072cd36..d4e5a8b899 100644 --- a/tests/util/test_stream_change_cache.py +++ b/tests/util/test_stream_change_cache.py @@ -302,3 +302,25 @@ class StreamChangeCacheTests(unittest.HomeserverTestCase): self.assertFalse(cache.has_entity_changed("user@foo.com", 5)) self.assertFalse(cache.has_entity_changed("bar@baz.net", 5)) self.assertFalse(cache.has_entity_changed("user@elsewhere.org", 5)) + + def test_listener(self) -> None: + """A listener hears of every change after the position it is given, + including the ones the cache has no room for.""" + cache = StreamChangeCache( + name="#test", + server_name=self.hs.hostname, + current_stream_pos=1, + max_size=1, + prefilled_cache={"a": 3}, + ) + heard: list[tuple[str | None, int]] = [] + self.assertEqual(cache.add_listener(lambda e, pos: heard.append((e, pos))), 3) + + cache.entity_has_changed("b", 4) + # The cache only holds one position, so this evicts "b". + cache.entity_has_changed("c", 5) + # Before the cache's earliest position, which it ignores. + cache.entity_has_changed("b", 4) + cache.all_entities_changed(6) + + self.assertEqual(heard, [("b", 4), ("c", 5), ("b", 4), (None, 6)])