diff --git a/rust/src/deferred.rs b/rust/src/deferred.rs index a22439f483..f0fa396d0a 100644 --- a/rust/src/deferred.rs +++ b/rust/src/deferred.rs @@ -228,11 +228,49 @@ pub(crate) async fn run_python_awaitable( ) -> PyResult> where F: for<'py> Fn(Python<'py>) -> PyResult> + Send + Sync + 'static, +{ + run_python_awaitable_observed( + runtime, + make_awaitable, + |value: &Bound<'_, PyAny>| Ok(value.clone().unbind()), + None, + ) + .await +} + +/// Told, on the reactor thread, when the awaitable of +/// [`run_python_awaitable_observed`] starts and when it finishes. +pub trait AwaitableObserver: Send + Sync + 'static { + /// The awaitable has been made and is running on the reactor. + fn started(&self); + /// The awaitable has a result, which the caller has yet to see. Also + /// called if it failed to start, without `started`. + fn finished(&self); +} + +/// As [`run_python_awaitable`], with two additions. +/// +/// - `convert` turns the result into a Rust value on the reactor thread, in +/// the GIL hold that receives it. The caller then needs no GIL hop of its +/// own to read the result. +/// - `observer` hears when the awaitable starts and finishes. A test can then +/// tell a task waiting on the reactor from one with work still to do. +pub(crate) async fn run_python_awaitable_observed( + runtime: &RustRuntime, + make_awaitable: F, + convert: C, + observer: Option>, +) -> PyResult +where + F: for<'py> Fn(Python<'py>) -> PyResult> + Send + Sync + 'static, + C: for<'py> Fn(&Bound<'py, PyAny>) -> PyResult + Send + Sync + 'static, + T: Send + 'static, { // Resolves when the awaitable completes; carries the resolved value or error. - let (tx, rx) = oneshot::channel::>>(); + let (tx, rx) = oneshot::channel::>(); // Shared between the success and error callbacks (only one ever fires). - let sender = Arc::new(Mutex::new(Some(tx))); + // The observer hears of the result before the channel carries it. + let sender = Arc::new(Mutex::new(Some((tx, observer)))); // Capture the logcontext of the calling tokio task (if any). We restore it // on the reactor thread before driving the awaitable, so Python code @@ -253,15 +291,18 @@ where None, None, move |args, _kwargs| -> PyResult> { - let value = args.get_item(0)?.unbind(); - if let Some(tx) = success_sender + let value = convert(&args.get_item(0)?); + if let Some((tx, observer)) = success_sender .lock() .map_err(|err| { anyhow::anyhow!("Failed to acquire lock on `success_sender`: {:#}", err) })? .take() { - let _ = tx.send(Ok(value)); + if let Some(observer) = observer { + observer.finished(); + } + let _ = tx.send(value); } Ok(args.py().None()) }, @@ -275,13 +316,16 @@ where None, move |args, _kwargs| -> PyResult> { let err = failure_to_pyerr(&args.get_item(0)?); - if let Some(tx) = error_sender + if let Some((tx, observer)) = error_sender .lock() .map_err(|err| { anyhow::anyhow!("Failed to acquire lock on `error_sender`: {:#}", err) })? .take() { + if let Some(observer) = observer { + observer.finished(); + } let _ = tx.send(Err(err)); } Ok(args.py().None()) @@ -317,6 +361,13 @@ where // logcontext and follows the logcontext rules from there. In // particular, it arranges for the reactor to be back at the // sentinel when the awaitable later completes. + let started_observer = sender + .lock() + .ok() + .and_then(|slot| slot.as_ref().and_then(|(_, observer)| observer.clone())); + if let Some(observer) = started_observer { + observer.started(); + } with_logcontext(py, context, || { let deferred = run_in_background(py, awaitable_factory.bind(py))?; deferred.call_method1( @@ -330,7 +381,10 @@ where if let Err(err) = started { // We failed before the awaitable's callbacks were attached, so // nothing else will ever resolve the channel. Do it ourselves. - if let Some(tx) = sender.lock().ok().and_then(|mut slot| slot.take()) { + if let Some((tx, observer)) = sender.lock().ok().and_then(|mut slot| slot.take()) { + if let Some(observer) = observer { + observer.finished(); + } let _ = tx.send(Err(err)); } } diff --git a/rust/src/federation/client/mod.rs b/rust/src/federation/client/mod.rs index d5a96bc7a9..c1747449d1 100644 --- a/rust/src/federation/client/mod.rs +++ b/rust/src/federation/client/mod.rs @@ -307,6 +307,21 @@ impl FederationClient { self.is_shutdown.store(true, Ordering::Relaxed); } + /// Whether we are backing off from `destination`: the check of + /// `get_retry_limiter`, which the transmission loop makes before it does + /// anything. + pub async fn backoff(&self, destination: &str) -> Result, ClientError> { + let checked = RetryLimiter::check( + &*self.retry_timings, + destination, + self.now_ms(), + LimiterOptions::default(), + ) + .await + .map_err(ClientError::Store)?; + Ok(checked.err()) + } + /// The destination may be back: end any sleep before a retry to it. pub fn wake_destination(&self, destination: &str) { self.retry_sleeper.wake(destination); diff --git a/rust/src/federation/mod.rs b/rust/src/federation/mod.rs index 00104243f9..909d8c4471 100644 --- a/rust/src/federation/mod.rs +++ b/rust/src/federation/mod.rs @@ -21,6 +21,7 @@ //! TLS policies, and the connection pool. //! - [`client`]: the request layer: signing, retries and the backoff. //! - [`gil`]: accounting for every time this code takes the GIL. +//! - [`queue`]: the per-destination queues and their transmission loops. //! - [`signing`]: the `X-Matrix` request signature. //! - [`transaction`]: builds and sends `/send` transactions. //! - [`transport`]: the seam between the request layer and the network, with @@ -31,6 +32,7 @@ use pyo3::prelude::*; pub mod agent; pub mod client; pub mod gil; +pub mod queue; pub mod signing; pub mod transaction; pub mod transport; @@ -91,6 +93,8 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> child_module.add_class::()?; child_module.add_class::()?; child_module.add_class::()?; + child_module.add_class::()?; + child_module.add_class::()?; m.add_submodule(&child_module)?; diff --git a/rust/src/federation/queue/mod.rs b/rust/src/federation/queue/mod.rs new file mode 100644 index 0000000000..ab60116f0b --- /dev/null +++ b/rust/src/federation/queue/mod.rs @@ -0,0 +1,1110 @@ +/* + * 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: + * . + * + */ + +//! Per-destination queues: Python's `PerDestinationQueue` and its +//! transmission loop. +//! +//! A [`DestinationQueue`] holds what is waiting to go to one destination. The +//! first thing queued starts its transmission loop as a tokio task, which +//! sends transactions until there is nothing left and then ends, as the +//! Python loop does. +//! +//! 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. + +use std::{ + any::Any, + sync::{ + atomic::{AtomicBool, AtomicIsize, AtomicU64, Ordering}, + Arc, Mutex, MutexGuard, + }, +}; + +use async_trait::async_trait; +use pyo3::prelude::*; +use serde_json::{json, Map, Value}; +use tokio::{runtime::Handle, task::AbortHandle}; + +use super::{ + client::{limiter::NotRetrying, ClientError, RequestStats}, + transaction::Transaction, +}; +use crate::{deferred::AwaitableObserver, events::Event}; + +mod ordered; +pub mod python; + +#[cfg(test)] +mod tests; + +pub use ordered::OrderedMap; + +/// If we back off from a destination for longer than this, we drop what we +/// cannot afford to keep and catch up later. In milliseconds. +pub const CATCHUP_RETRY_INTERVAL_MS: u64 = 60 * 60 * 1000; +/// Presence states per presence EDU. +pub const MAX_PRESENCE_STATES_PER_EDU: usize = 50; +/// PDUs per transaction: the spec's limit. +pub const MAX_PDUS_PER_TRANSACTION: usize = 50; +/// EDUs per transaction: `MAX_EDUS_PER_TRANSACTION`. +pub const MAX_EDUS_PER_TRANSACTION: usize = 100; +/// Room kept for device list updates when there are many to-device +/// messages: `NUMBER_OF_RESERVED_EDUS_PER_TRANSACTION`. +pub const NUMBER_OF_RESERVED_EDUS_PER_TRANSACTION: usize = 10; +/// Receipt EDUs per transaction. +const MAX_RECEIPT_EDUS_PER_TRANSACTION: usize = 5; + +/// The logger of Python's `PerDestinationQueue`. +const QUEUE_LOG: &str = "synapse.federation.sender.per_destination_queue"; + +/// An EDU: `synapse.federation.units.Edu`, less the origin and destination. +#[derive(Clone, Debug, PartialEq)] +pub struct Edu { + pub edu_type: String, + pub content: Value, +} + +impl Edu { + /// `Edu.get_dict`. + pub fn into_value(self) -> Value { + json!({"edu_type": self.edu_type, "content": self.content}) + } +} + +/// A user's presence, as much of `UserPresenceState` as goes over +/// federation. +#[derive(Clone, Debug, PartialEq)] +pub struct PresenceState { + pub user_id: String, + pub state: String, + pub last_active_ts: i64, + pub status_msg: Option, + pub currently_active: bool, +} + +impl PresenceState { + /// `format_user_presence_state(state, now)`. + pub fn format(&self, now_ms: i64) -> Value { + let mut content = Map::new(); + content.insert("presence".to_owned(), Value::from(self.state.as_str())); + content.insert("user_id".to_owned(), Value::from(self.user_id.as_str())); + if self.last_active_ts != 0 { + content.insert( + "last_active_ago".to_owned(), + Value::from(now_ms - self.last_active_ts), + ); + } + if let Some(status_msg) = self.status_msg.as_deref().filter(|msg| !msg.is_empty()) { + content.insert("status_msg".to_owned(), Value::from(status_msg)); + } + if self.state == "online" { + content.insert( + "currently_active".to_owned(), + Value::from(self.currently_active), + ); + } + Value::Object(content) + } +} + +/// A read receipt: `synapse.types.ReadReceipt`. +#[derive(Clone, Debug)] +pub struct Receipt { + pub room_id: String, + pub receipt_type: String, + pub user_id: String, + pub thread_id: Option, + pub event_ids: Vec, + pub data: Map, +} + +impl Receipt { + /// The receipt as the EDU carries it. + fn serialize(&self) -> Value { + let mut data = self.data.clone(); + if let Some(thread_id) = &self.thread_id { + data.insert("thread_id".to_owned(), Value::from(thread_id.as_str())); + } + json!({"event_ids": self.event_ids, "data": data}) + } +} + +/// Where a PDU's JSON comes from. +pub enum PduBody { + /// The event. Its JSON is read on tokio, without the GIL. + Event(Py), + /// The JSON itself, for tests. + #[cfg(test)] + Json(Value), +} + +/// A PDU waiting to be sent. +pub struct QueuedPdu { + pub event_id: String, + pub stream_ordering: Option, + pub origin_server_ts: i64, + pub body: PduBody, +} + +impl QueuedPdu { + /// The PDU as `get_pdu_json` gives it. + fn to_value(&self) -> serde_json::Result { + match &self.body { + PduBody::Event(event) => event.get().to_pdu_value(), + #[cfg(test)] + PduBody::Json(value) => Ok(value.clone()), + } + } +} + +/// Why the loop stopped: the exceptions Python's loop tells apart. +#[derive(Debug)] +pub enum LoopError { + /// `NotRetryingDestination`. + NotRetrying(NotRetrying), + /// `FederationDeniedError`, with its message. + Denied(String), + /// `HttpResponseException`. + Http { code: u16, message: String }, + /// `RequestSendFailed`, with its message. + SendFailed(String), + /// Anything else, with its message and, from Python, the traceback. + Other(String), +} + +impl LoopError { + pub fn from_client(err: ClientError, destination: &str) -> Self { + match err { + ClientError::Denied => { + LoopError::Denied(format!("Federation denied with {destination}.")) + } + ClientError::NotRetrying(not_retrying) => LoopError::NotRetrying(not_retrying), + ClientError::Http(http) => LoopError::Http { + code: http.code, + message: format!("{}: {}", http.code, http.phrase), + }, + err @ ClientError::SendFailed { .. } => { + LoopError::SendFailed(format!("Failed to send request: {err}")) + } + err => LoopError::Other(err.to_string()), + } + } +} + +/// What the device EDU query is asked for. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct DeviceEdusRequest { + pub last_device_stream_id: i64, + pub last_device_list_stream_id: i64, + /// How many to-device messages may go in. + pub to_device_limit: usize, + /// How many EDUs may go in, of both kinds. + pub edu_limit: usize, +} + +/// The device EDUs for a transaction. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct DeviceEdus { + /// The contents of the to-device EDUs. + pub to_device: Vec, + /// How far through the to-device stream these go. + pub to_device_stream_id: i64, + pub device_list: Vec, + /// How far through the device list stream these go. + pub device_list_stream_id: i64, +} + +/// A position in the sticky events stream. Python's, which Rust only hands +/// back. +pub type StickyStreamPosition = Arc; + +/// Backlogged sticky events from one room. +pub struct StickyBatch { + pub room_id: String, + pub position: StickyStreamPosition, + pub pdus: Vec>, +} + +/// A transaction that went out. +pub struct Sent { + pub pdus: Vec>, + pub edu_types: Vec, + /// The PDUs the remote accepted. + pub acknowledged: Vec, + pub stats: RequestStats, +} + +/// What a finished transaction moved forward. Only reported once the +/// transaction was sent, or had nothing to send. +pub struct Completion { + pub to_device_stream_id: Option, + pub device_list_stream_id: Option, + /// The stream ordering of the last PDU sent. + pub last_stream_ordering: Option, + pub sticky: Option<(String, StickyStreamPosition)>, + /// `None` for a sticky batch that visibility filtering emptied. + pub sent: Option, +} + +/// The parts of the loop that run in Python for now. +#[async_trait] +pub trait QueueHooks: Send + Sync + 'static { + /// `_catch_up_transmission_loop`. + async fn catch_up(&self) -> Result<(), LoopError>; + + /// The to-device messages and device list updates to send. + async fn device_edus(&self, request: DeviceEdusRequest) -> Result; + + /// `StickyEventBacklogTracker.prepare_transaction`: `None` once the + /// backlog is empty. + async fn sticky_backlog(&self) -> Result, LoopError>; + + /// Record a finished transaction: the stream positions, the module + /// callback and the metrics. + async fn complete(&self, completion: Completion) -> Result<(), LoopError>; +} + +/// The answer to a sent transaction. +pub struct SendOutcome { + pub acknowledged: Vec, + pub stats: RequestStats, +} + +/// Where transactions go. +#[async_trait] +pub trait TransactionSink: Send + Sync + 'static { + fn now_ms(&self) -> u64; + + /// `get_retry_limiter`'s check: whether we are backing off from the + /// destination. + async fn backoff(&self, destination: &str) -> Result, LoopError>; + + fn build(&self, destination: &str, pdus: Vec, edus: Vec) -> Transaction; + + async fn send( + &self, + transaction: Transaction, + observer: Arc, + ) -> Result; +} + +/// How many transmission loops are doing work, as opposed to waiting on the +/// reactor. Tests wait for this to reach zero. +#[derive(Default)] +pub struct Activity { + busy: AtomicIsize, +} + +impl Activity { + /// Whether every loop has finished or is waiting on the reactor. + pub fn idle(&self) -> bool { + self.busy.load(Ordering::SeqCst) <= 0 + } +} + +/// One queue's part in [`Activity`]. Python calls tell it when they start +/// and finish, on the reactor thread. +pub struct LoopActivity { + activity: Arc, + /// The loop is waiting on a Python call that has started. + parked: AtomicBool, +} + +impl LoopActivity { + pub fn new(activity: Arc) -> Arc { + Arc::new(LoopActivity { + activity, + parked: AtomicBool::new(false), + }) + } + + fn loop_started(&self) { + self.activity.busy.fetch_add(1, Ordering::SeqCst); + } + + fn loop_ended(&self) { + // A loop aborted while parked is no longer counted. + if !self.parked.swap(false, Ordering::SeqCst) { + self.activity.busy.fetch_sub(1, Ordering::SeqCst); + } + } +} + +impl AwaitableObserver for LoopActivity { + fn started(&self) { + self.parked.store(true, Ordering::SeqCst); + self.activity.busy.fetch_sub(1, Ordering::SeqCst); + } + + fn finished(&self) { + if self.parked.swap(false, Ordering::SeqCst) { + self.activity.busy.fetch_add(1, Ordering::SeqCst); + } + } +} + +/// The PDUs waiting, in the order they were queued. +type PendingPdus = OrderedMap>; + +/// A receipt EDU's content: room ID → receipt type → user ID → receipt. +type ReceiptContent = Map; + +/// Everything the queue holds. Python's attributes of the same names. +struct QueueState { + transmission_loop_running: bool, + /// Which loop is running, so an old loop ending cannot mark a newer one + /// as stopped. + loop_id: u64, + loop_task: Option, + transmission_loop_enabled: bool, + new_data_to_send: bool, + + catching_up: bool, + catchup_last_skipped: i64, + last_successful_stream_ordering: Option, + + pending_pdus: PendingPdus, + /// Never used in practice: see matrix-org/synapse#7549. + pending_edus: Vec, + /// EDUs that replace each other by key, such as typing. + pending_edus_keyed: OrderedMap<(String, String), Edu>, + pending_presence: OrderedMap, + pending_receipt_edus: Vec, + + last_device_stream_id: i64, + last_device_list_stream_id: i64, + + sticky_backlog_turn: bool, + /// Whether there may be backlogged sticky events. + sticky_backlogged: bool, +} + +/// A transaction ready to send: `_PreparedTransaction`. +struct Prepared { + pdus: Vec>, + edus: Vec, + to_device_stream_id: Option, + device_list_stream_id: Option, + last_stream_ordering: Option, + sticky: Option<(String, StickyStreamPosition)>, +} + +impl Prepared { + fn has_anything_to_send(&self) -> bool { + !self.pdus.is_empty() || !self.edus.is_empty() + } +} + +/// What the queue was built with. +pub struct QueueConfig { + pub destination: String, + /// Whether this instance sends to the destination. If not, nothing is + /// sent, as in Python. + pub should_send_on_this_instance: bool, + /// MSC4354: whether there is a sticky event backlog to send. + pub msc4354_enabled: bool, +} + +/// The queue for one destination: `PerDestinationQueue`. +pub struct DestinationQueue { + config: QueueConfig, + state: Mutex, + hooks: Arc, + sink: Arc, + activity: Arc, + tokio: Handle, +} + +/// Loop ids, unique across queues. +static NEXT_LOOP_ID: AtomicU64 = AtomicU64::new(1); + +impl DestinationQueue { + pub fn new( + config: QueueConfig, + hooks: Arc, + sink: Arc, + activity: Arc, + tokio: Handle, + ) -> Arc { + if !config.should_send_on_this_instance { + log::error!( + target: QUEUE_LOG, + "Create a per destination queue for {} on wrong worker", + config.destination + ); + } + Arc::new(DestinationQueue { + config, + state: Mutex::new(QueueState { + transmission_loop_running: false, + loop_id: 0, + loop_task: None, + transmission_loop_enabled: true, + new_data_to_send: false, + // We start out catching up, so we can catch up at startup. + catching_up: true, + catchup_last_skipped: 0, + last_successful_stream_ordering: None, + pending_pdus: OrderedMap::default(), + pending_edus: Vec::new(), + pending_edus_keyed: OrderedMap::default(), + pending_presence: OrderedMap::default(), + pending_receipt_edus: Vec::new(), + last_device_stream_id: 0, + last_device_list_stream_id: 0, + sticky_backlog_turn: false, + sticky_backlogged: true, + }), + hooks, + sink, + activity, + tokio, + }) + } + + pub fn destination(&self) -> &str { + &self.config.destination + } + + fn state(&self) -> MutexGuard<'_, QueueState> { + self.state.lock().expect("poisoned") + } + + /// Stop sending. The running loop, if any, is aborted where it stands. + pub fn shutdown(&self) { + let task = { + let mut state = self.state(); + state.transmission_loop_enabled = false; + state.loop_task.take() + }; + if let Some(task) = task { + task.abort(); + } + } + + pub fn pending_pdu_count(&self) -> usize { + self.state().pending_pdus.len() + } + + pub fn pending_edu_count(&self) -> usize { + let state = self.state(); + state.pending_edus.len() + state.pending_presence.len() + state.pending_edus_keyed.len() + } + + pub fn transmission_loop_running(&self) -> bool { + self.state().transmission_loop_running + } + + pub fn new_data_to_send(&self) -> bool { + self.state().new_data_to_send + } + + /// Queue a PDU and start the loop. While catching up, the PDU is left + /// for catch-up to find in the database, and only noted as skipped. + pub fn send_pdu(self: &Arc, pdu: Arc) { + { + let mut state = self.state(); + if !state.catching_up || state.last_successful_stream_ordering.is_none() { + state.pending_pdus.insert(pdu.event_id.clone(), pdu); + } else { + state.catchup_last_skipped = pdu.stream_ordering.unwrap_or(0); + } + } + self.attempt_new_transaction(); + } + + /// Queue presence updates, replacing any queued for the same users. + pub fn send_presence(self: &Arc, states: Vec, start_loop: bool) { + { + let mut state = self.state(); + for presence in states { + state + .pending_presence + .insert(presence.user_id.clone(), presence); + } + state.new_data_to_send = true; + } + if start_loop { + self.attempt_new_transaction(); + } + } + + /// Queue a read receipt, without starting the loop. + /// + /// Each EDU holds one receipt per (room, receipt type, user, thread). + /// The receipt replaces one in the first EDU that has the same tuple, or + /// goes into the first EDU without the user, or into a new EDU. + pub fn queue_read_receipt(&self, receipt: &Receipt) { + let serialized = receipt.serialize(); + let mut state = self.state(); + + let mut placed = false; + for edu in &mut state.pending_receipt_edus { + let room = edu + .entry(receipt.room_id.clone()) + .or_insert_with(|| Value::Object(Map::new())); + let Some(room) = room.as_object_mut() else { + continue; + }; + let by_user = room + .entry(receipt.receipt_type.clone()) + .or_insert_with(|| Value::Object(Map::new())); + let Some(by_user) = by_user.as_object_mut() else { + continue; + }; + // This reads `thread_id` from the top level of the queued + // receipt, as Python does. It is never there (it is under + // `data`), so a receipt with a thread never replaces one. + let fits = match by_user.get(&receipt.user_id) { + None => true, + Some(existing) => { + existing.get("thread_id").and_then(Value::as_str) + == receipt.thread_id.as_deref() + } + }; + if fits { + by_user.insert(receipt.user_id.clone(), serialized.clone()); + placed = true; + break; + } + } + if !placed { + let mut by_user = Map::new(); + by_user.insert(receipt.user_id.clone(), serialized); + let mut room = Map::new(); + room.insert(receipt.receipt_type.clone(), Value::Object(by_user)); + let mut edu = Map::new(); + edu.insert(receipt.room_id.clone(), Value::Object(room)); + state.pending_receipt_edus.push(edu); + } + + state.new_data_to_send = true; + } + + /// Queue an EDU that replaces any queued with the same type and key. + pub fn send_keyed_edu(self: &Arc, edu: Edu, key: String) { + self.state() + .pending_edus_keyed + .insert((edu.edu_type.clone(), key), edu); + self.attempt_new_transaction(); + } + + pub fn send_edu(self: &Arc, edu: Edu) { + self.state().pending_edus.push(edu); + self.attempt_new_transaction(); + } + + /// Note new data without starting the loop. A running loop picks it up. + pub fn mark_new_data(&self) { + self.state().new_data_to_send = true; + } + + /// Start the loop, unless it is running. + pub fn attempt_new_transaction(self: &Arc) { + let mut state = self.state(); + // A running loop rechecks for things to send before it ends. + state.new_data_to_send = true; + + if state.transmission_loop_running { + log::debug!( + target: QUEUE_LOG, + "TX [{}] Transaction already in progress", + self.destination() + ); + return; + } + if !self.config.should_send_on_this_instance { + log::error!( + target: QUEUE_LOG, + "Trying to start a transaction to {} on wrong worker", + self.destination() + ); + return; + } + if !state.transmission_loop_enabled { + log::warn!( + target: QUEUE_LOG, + "Shutdown has been requested. Not sending transaction" + ); + return; + } + + log::debug!( + target: QUEUE_LOG, + "TX [{}] Starting transaction loop", + self.destination() + ); + let loop_id = NEXT_LOOP_ID.fetch_add(1, Ordering::Relaxed); + state.transmission_loop_running = true; + state.loop_id = loop_id; + self.activity.loop_started(); + // The guard takes the lock when it drops, which it does at once if + // the runtime has shut down. + drop(state); + + // The guard is made here, so it drops even if the task never runs. + let guard = LoopGuard { + queue: Arc::clone(self), + loop_id, + }; + let this = Arc::clone(self); + let task = self.tokio.spawn(async move { + let _guard = guard; + this.transmission_loop().await; + }); + + let mut state = self.state(); + if state.loop_id == loop_id && state.transmission_loop_running { + state.loop_task = Some(task.abort_handle()); + } + } + + /// `_transaction_transmission_loop`. + async fn transmission_loop(&self) { + let mut in_flight = Vec::new(); + if let Err(err) = self.run_transmission_loop(&mut in_flight).await { + self.handle_error(err, &in_flight); + } + } + + async fn run_transmission_loop( + &self, + in_flight: &mut Vec>, + ) -> Result<(), LoopError> { + // Fail fast if we are backing off. The client checks again. + if let Some(not_retrying) = self.sink.backoff(self.destination()).await? { + return Err(LoopError::NotRetrying(not_retrying)); + } + + if self.state().catching_up { + self.hooks.catch_up().await?; + if self.state().catching_up { + // Not caught up yet. + return Ok(()); + } + } + + loop { + { + let mut state = self.state(); + if !state.transmission_loop_enabled { + return Ok(()); + } + state.new_data_to_send = false; + } + in_flight.clear(); + + let Some(prepared) = self.prepare_transaction().await? else { + log::debug!(target: QUEUE_LOG, "TX [{}] Nothing to send", self.destination()); + let mut state = self.state(); + if state.new_data_to_send { + // Something was queued while we looked. A queue call + // that saw the loop running did not start another. + continue; + } + state.transmission_loop_running = false; + state.loop_task = None; + return Ok(()); + }; + + if !prepared.has_anything_to_send() { + // Filtering emptied a sticky event batch. Preparing it still + // moved the backlog on, which must be recorded. + self.complete_transaction(prepared, None).await?; + continue; + } + + in_flight.extend(prepared.pdus.iter().cloned()); + let sent = self.send_transaction(&prepared).await?; + self.complete_transaction(prepared, Some(sent)).await?; + } + } + + /// Build and send the transaction. + async fn send_transaction(&self, prepared: &Prepared) -> Result { + let destination = self.destination(); + if !prepared.pdus.is_empty() { + log::debug!( + target: QUEUE_LOG, + "TX [{destination}] len(pending_pdus_by_dest[dest]) = {}", + prepared.pdus.len() + ); + } + + let pdus = prepared + .pdus + .iter() + .map(|pdu| pdu.to_value()) + .collect::, _>>() + .map_err(|err| LoopError::Other(format!("Failed to serialize event: {err}")))?; + let edus = prepared.edus.iter().cloned().map(Edu::into_value).collect(); + let transaction = self.sink.build(destination, pdus, edus); + + let observer: Arc = self.activity.clone(); + let outcome = self.sink.send(transaction, observer).await?; + Ok(Sent { + pdus: prepared.pdus.clone(), + edu_types: prepared + .edus + .iter() + .map(|edu| edu.edu_type.clone()) + .collect(), + acknowledged: outcome.acknowledged, + stats: outcome.stats, + }) + } + + /// The exception handlers of `_transaction_transmission_loop`. + fn handle_error(&self, err: LoopError, in_flight: &[Arc]) { + let destination = self.destination(); + let log_failed_pdus = || { + for pdu in in_flight { + log::info!( + target: QUEUE_LOG, + "Failed to send event {} to {destination}", + pdu.event_id + ); + } + }; + match err { + LoopError::NotRetrying(not_retrying) => { + log::debug!( + target: QUEUE_LOG, + "TX [{destination}] not ready for retry yet (next retry at {} ms) - dropping transaction for now", + not_retrying.retry_last_ts + not_retrying.retry_interval + ); + if not_retrying.retry_interval > CATCHUP_RETRY_INTERVAL_MS { + // We will not retry for over an hour. Drop the EDUs we + // can afford to lose (typing, receipts and presence), + // or they pile up. To-device messages and device lists + // are in the database. Catching up drops the PDUs. + let mut state = self.state(); + state.pending_edus.clear(); + state.pending_edus_keyed.clear(); + state.pending_presence.clear(); + state.pending_receipt_edus.clear(); + start_catching_up(&mut state); + } + } + LoopError::Denied(message) => log::info!(target: QUEUE_LOG, "{message}"), + LoopError::Http { code, message } => log::warn!( + target: QUEUE_LOG, + "TX [{destination}] Received {code} response to transaction: {message}" + ), + LoopError::SendFailed(message) => { + log::warn!( + target: QUEUE_LOG, + "TX [{destination}] Failed to send transaction: {message}" + ); + log_failed_pdus(); + } + LoopError::Other(message) => { + log::error!( + target: QUEUE_LOG, + "TX [{destination}] Failed to send transaction\n{message}" + ); + log_failed_pdus(); + } + } + } + + /// `_prepare_transaction`: take turns between the sticky event backlog + /// and the real-time queue, so neither starves the other. + async fn prepare_transaction(&self) -> Result, LoopError> { + let sticky_turn = std::mem::take(&mut self.state().sticky_backlog_turn); + if sticky_turn { + if let Some(prepared) = self.prepare_sticky_transaction().await? { + return Ok(Some(prepared)); + } + return self.prepare_main_queue_transaction().await; + } + + let prepared = self.prepare_main_queue_transaction().await?; + let backlogged = { + let mut state = self.state(); + state.sticky_backlog_turn = self.config.msc4354_enabled && state.sticky_backlogged; + state.sticky_backlog_turn + }; + if prepared.is_some() { + return Ok(prepared); + } + if backlogged { + return self.prepare_sticky_transaction().await; + } + Ok(None) + } + + /// `StickyEventBacklogTracker.prepare_transaction`. + async fn prepare_sticky_transaction(&self) -> Result, LoopError> { + if !self.config.msc4354_enabled || !self.state().sticky_backlogged { + return Ok(None); + } + + let Some(batch) = self.hooks.sticky_backlog().await? else { + log::info!( + target: QUEUE_LOG, + "Caught up on federation sticky event backlog for destination '{}'", + self.destination() + ); + self.state().sticky_backlogged = false; + return Ok(None); + }; + + Ok(Some(Prepared { + pdus: batch.pdus, + // No EDUs go with backlogged sticky events. + edus: Vec::new(), + to_device_stream_id: None, + device_list_stream_id: None, + last_stream_ordering: None, + sticky: Some((batch.room_id, batch.position)), + })) + } + + /// `_prepare_main_queue_transaction`. + /// + /// EDUs go in this order, up to 100: one presence EDU, up to five + /// receipt EDUs, to-device messages (leaving ten places for device list + /// updates), device list updates, then anything else. + async fn prepare_main_queue_transaction(&self) -> Result, LoopError> { + let now_ms = self.sink.now_ms() as i64; + let mut edus = Vec::new(); + let request = { + let mut state = self.state(); + + if !state.pending_presence.is_empty() { + let mut push = Vec::new(); + while push.len() < MAX_PRESENCE_STATES_PER_EDU { + let Some((_, presence)) = state.pending_presence.pop_front() else { + break; + }; + push.push(presence.format(now_ms)); + } + edus.push(Edu { + edu_type: "m.presence".to_owned(), + content: json!({ "push": push }), + }); + } + + let receipts = state + .pending_receipt_edus + .len() + .min(MAX_RECEIPT_EDUS_PER_TRANSACTION); + for content in state.pending_receipt_edus.drain(..receipts) { + edus.push(Edu { + edu_type: "m.receipt".to_owned(), + content: Value::Object(content), + }); + } + + let edu_limit = MAX_EDUS_PER_TRANSACTION - edus.len(); + DeviceEdusRequest { + last_device_stream_id: state.last_device_stream_id, + last_device_list_stream_id: state.last_device_list_stream_id, + to_device_limit: edu_limit - NUMBER_OF_RESERVED_EDUS_PER_TRANSACTION, + edu_limit, + } + }; + + let device = self.hooks.device_edus(request).await?; + + let mut state = self.state(); + let mut edu_limit = request.edu_limit; + + // With nothing to send in a stream, we can move on in it now. + // Otherwise only once the transaction is sent. + let to_device_stream_id = if device.to_device.is_empty() { + state.last_device_stream_id = device.to_device_stream_id; + None + } else { + Some(device.to_device_stream_id) + }; + edu_limit = edu_limit.saturating_sub(device.to_device.len()); + edus.extend(device.to_device.into_iter().map(|content| Edu { + edu_type: "m.direct_to_device".to_owned(), + content, + })); + + let device_list_stream_id = if device.device_list.is_empty() { + state.last_device_list_stream_id = device.device_list_stream_id; + None + } else { + Some(device.device_list_stream_id) + }; + edu_limit = edu_limit.saturating_sub(device.device_list.len()); + edus.extend(device.device_list); + + let others = state.pending_edus.len().min(edu_limit); + edus.extend(state.pending_edus.drain(..others)); + edu_limit -= others; + while edu_limit > 0 { + let Some((_, edu)) = state.pending_edus_keyed.pop_back() else { + break; + }; + edus.push(edu); + edu_limit -= 1; + } + + let pdus: Vec> = state + .pending_pdus + .values() + .take(MAX_PDUS_PER_TRANSACTION) + .cloned() + .collect(); + + if pdus.is_empty() && edus.is_empty() { + // Nothing to send, and nothing to record: the device stream + // positions have moved on above. + return Ok(None); + } + + let last_stream_ordering = pdus.last().and_then(|pdu| pdu.stream_ordering); + Ok(Some(Prepared { + pdus, + edus, + to_device_stream_id, + device_list_stream_id, + last_stream_ordering, + sticky: None, + })) + } + + /// `_complete_transaction`: the transaction went out, so drop its PDUs + /// and record how far we got. + async fn complete_transaction( + &self, + prepared: Prepared, + sent: Option, + ) -> Result<(), LoopError> { + { + let mut state = self.state(); + for pdu in &prepared.pdus { + state.pending_pdus.remove(pdu.event_id.as_str()); + } + } + + // Python skips a zero position, as it tests them for truth. + let to_device_stream_id = prepared.to_device_stream_id.filter(|id| *id != 0); + let device_list_stream_id = prepared.device_list_stream_id.filter(|id| *id != 0); + if let Some(device_list_stream_id) = device_list_stream_id { + log::info!( + target: QUEUE_LOG, + "Marking as sent '{}' {device_list_stream_id}", + self.destination() + ); + } + + self.hooks + .complete(Completion { + to_device_stream_id, + device_list_stream_id, + last_stream_ordering: prepared.last_stream_ordering.filter(|so| *so != 0), + sticky: prepared.sticky, + sent, + }) + .await?; + + let mut state = self.state(); + if let Some(id) = to_device_stream_id { + state.last_device_stream_id = id; + } + if let Some(id) = device_list_stream_id { + state.last_device_list_stream_id = id; + } + Ok(()) + } + + // The catch-up state, which the Python catch-up loop reads and writes. + + pub fn catching_up(&self) -> bool { + self.state().catching_up + } + + pub fn set_catching_up(&self, catching_up: bool) { + self.state().catching_up = catching_up; + } + + pub fn catchup_last_skipped(&self) -> i64 { + self.state().catchup_last_skipped + } + + pub fn set_catchup_last_skipped(&self, stream_ordering: i64) { + self.state().catchup_last_skipped = stream_ordering; + } + + pub fn last_successful_stream_ordering(&self) -> Option { + self.state().last_successful_stream_ordering + } + + pub fn set_last_successful_stream_ordering(&self, stream_ordering: Option) { + self.state().last_successful_stream_ordering = stream_ordering; + } + + pub fn transmission_loop_enabled(&self) -> bool { + self.state().transmission_loop_enabled + } + + /// `_start_catching_up`: enter catch-up and drop the queued PDUs. + pub fn start_catching_up(&self) { + start_catching_up(&mut self.state()); + } + + /// Whether there may be backlogged sticky events: `is_backlogged`. + pub fn sticky_backlogged(&self) -> bool { + self.config.msc4354_enabled && self.state().sticky_backlogged + } + + pub fn set_sticky_backlogged(&self, backlogged: bool) { + self.state().sticky_backlogged = backlogged; + } + + /// The event IDs of the queued PDUs, oldest first. For tests. + pub fn pending_pdu_ids(&self) -> Vec { + self.state() + .pending_pdus + .values() + .map(|pdu| pdu.event_id.clone()) + .collect() + } +} + +fn start_catching_up(state: &mut QueueState) { + state.catching_up = true; + state.pending_pdus.clear(); +} + +/// Marks the loop as stopped when its task ends, however it ends. +struct LoopGuard { + queue: Arc, + loop_id: u64, +} + +impl Drop for LoopGuard { + fn drop(&mut self) { + { + let mut state = self.queue.state(); + if state.loop_id == self.loop_id { + state.transmission_loop_running = false; + state.loop_task = None; + } + } + self.queue.activity.loop_ended(); + } +} diff --git a/rust/src/federation/queue/ordered.rs b/rust/src/federation/queue/ordered.rs new file mode 100644 index 0000000000..20585f66ef --- /dev/null +++ b/rust/src/federation/queue/ordered.rs @@ -0,0 +1,130 @@ +/* + * 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: + * . + * + */ + +//! [`OrderedMap`]: the parts of Python's `OrderedDict` the queue uses. + +use std::{ + borrow::Borrow, + collections::{BTreeMap, HashMap}, + hash::Hash, +}; + +/// A map that remembers insertion order, as a Python `dict` does. +/// +/// Inserting a key that is already there replaces its value and keeps its +/// place. Removing from either end, or by key, is `O(log n)`. An `IndexMap` +/// would make removing from the front `O(n)`, and the queue does that for +/// every transaction. +#[derive(Debug)] +pub struct OrderedMap { + /// Each key's place in `entries`. + places: HashMap, + entries: BTreeMap, + next_place: u64, +} + +impl Default for OrderedMap { + fn default() -> Self { + OrderedMap { + places: HashMap::new(), + entries: BTreeMap::new(), + next_place: 0, + } + } +} + +impl OrderedMap { + pub fn len(&self) -> usize { + self.places.len() + } + + pub fn is_empty(&self) -> bool { + self.places.is_empty() + } + + /// `map[key] = value`. + pub fn insert(&mut self, key: K, value: V) { + if let Some(place) = self.places.get(&key) { + if let Some(entry) = self.entries.get_mut(place) { + entry.1 = value; + } + return; + } + let place = self.next_place; + self.next_place += 1; + self.places.insert(key.clone(), place); + self.entries.insert(place, (key, value)); + } + + /// `map.pop(key, None)`. + pub fn remove(&mut self, key: &Q) -> Option + where + K: Borrow, + Q: Hash + Eq + ?Sized, + { + let place = self.places.remove(key)?; + self.entries.remove(&place).map(|(_, value)| value) + } + + /// `map.popitem(last=False)`. + pub fn pop_front(&mut self) -> Option<(K, V)> { + let (_, (key, value)) = self.entries.pop_first()?; + self.places.remove(&key); + Some((key, value)) + } + + /// `map.popitem()`. + pub fn pop_back(&mut self) -> Option<(K, V)> { + let (_, (key, value)) = self.entries.pop_last()?; + self.places.remove(&key); + Some((key, value)) + } + + /// The values, oldest first. + pub fn values(&self) -> impl Iterator { + self.entries.values().map(|(_, value)| value) + } + + pub fn clear(&mut self) { + self.places.clear(); + self.entries.clear(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn keeps_insertion_order() { + let mut map = OrderedMap::default(); + map.insert("a", 1); + map.insert("b", 2); + map.insert("c", 3); + // Replacing keeps the place, as a Python dict does. + map.insert("a", 10); + assert_eq!(map.values().copied().collect::>(), [10, 2, 3]); + assert_eq!(map.len(), 3); + + assert_eq!(map.remove("b"), Some(2)); + assert_eq!(map.remove("b"), None); + assert_eq!(map.pop_back(), Some(("c", 3))); + map.insert("d", 4); + assert_eq!(map.pop_front(), Some(("a", 10))); + assert_eq!(map.pop_front(), Some(("d", 4))); + assert!(map.is_empty()); + assert_eq!(map.pop_front(), None); + } +} diff --git a/rust/src/federation/queue/python.rs b/rust/src/federation/queue/python.rs new file mode 100644 index 0000000000..4a1b3b122a --- /dev/null +++ b/rust/src/federation/queue/python.rs @@ -0,0 +1,665 @@ +/* + * 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 queues as Python's `FederationSender` uses them. +//! +//! [`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`. + +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use pyo3::{ + intern, + prelude::*, + types::{PyDict, PyList, PyString, PyTuple}, +}; +use serde_json::Value; + +use super::{ + Activity, Completion, DestinationQueue, DeviceEdus, DeviceEdusRequest, Edu, LoopActivity, + LoopError, PduBody, PresenceState, QueueConfig, QueueHooks, QueuedPdu, Receipt, SendOutcome, + StickyBatch, StickyStreamPosition, TransactionSink, +}; +use crate::{ + deferred::{run_python_awaitable_observed, AwaitableObserver}, + events::Event, + federation::{ + client::{limiter::NotRetrying, RequestStats}, + gil::{account, with_gil}, + transaction::{ + acknowledged_pdus, + python::{python_acknowledged, RustTransaction, RustTransactionManager}, + Transaction, TransactionResponse, TransactionSender, + }, + }, + runtime::RustRuntime, +}; + +/// Makes the queues, and tells tests when their loops are idle. +#[pyclass(frozen, name = "RustDestinationQueues")] +pub struct RustDestinationQueues { + sink: Arc, + runtime: RustRuntime, + activity: Arc, + msc4354_enabled: bool, +} + +#[pymethods] +impl RustDestinationQueues { + #[new] + fn py_new(transactions: &Bound<'_, RustTransactionManager>, msc4354_enabled: bool) -> Self { + let transactions = transactions.get(); + RustDestinationQueues { + sink: Arc::new(ClientSink { + sender: Arc::clone(&transactions.sender), + runtime: transactions.runtime.clone(), + record_metrics: Arc::clone(&transactions.record_metrics), + send_for_test: Mutex::new(None), + }), + runtime: transactions.runtime.clone(), + activity: Arc::default(), + msc4354_enabled, + } + } + + /// The queue for `destination`. `facade` is the Python queue, which + /// provides the hooks. The Rust queue holds it weakly. + fn queue( + &self, + py: Python<'_>, + destination: String, + facade: &Bound<'_, PyAny>, + should_send_on_this_instance: bool, + ) -> PyResult { + let facade = py + .import("weakref")? + .getattr("ref")? + .call1((facade,))? + .unbind(); + let activity = LoopActivity::new(Arc::clone(&self.activity)); + let hooks = PyHooks { + runtime: self.runtime.clone(), + facade: Arc::new(facade), + activity: Arc::clone(&activity), + destination: destination.clone(), + }; + let sink: Arc = self.sink.clone(); + Ok(RustDestinationQueue { + queue: DestinationQueue::new( + QueueConfig { + destination, + should_send_on_this_instance, + msc4354_enabled: self.msc4354_enabled, + }, + Arc::new(hooks), + sink, + activity, + self.runtime.tokio_handle()?, + ), + }) + } + + /// Whether every loop has finished or is waiting on the reactor. For + /// tests, which pump the reactor until it is. + fn idle(&self) -> bool { + self.activity.idle() + } + + /// Send transactions through `send(transaction)`, a coroutine function + /// taking a `RustTransaction` and returning the response dict, in place + /// of the network. For tests. + fn set_send_for_test(&self, send: Option>) { + *self.sink.send_for_test.lock().expect("poisoned") = send.map(Arc::new); + } +} + +/// One destination's queue. Each method is the `PerDestinationQueue` method +/// of the same name. +#[pyclass(frozen, name = "RustDestinationQueue")] +pub struct RustDestinationQueue { + queue: Arc, +} + +#[pymethods] +impl RustDestinationQueue { + fn shutdown(&self) { + self.queue.shutdown(); + } + + fn pending_pdu_count(&self) -> usize { + self.queue.pending_pdu_count() + } + + fn pending_edu_count(&self) -> usize { + self.queue.pending_edu_count() + } + + #[getter] + fn transmission_loop_running(&self) -> bool { + self.queue.transmission_loop_running() + } + + #[getter] + fn new_data_to_send(&self) -> bool { + self.queue.new_data_to_send() + } + + fn send_pdu(&self, pdu: &Bound<'_, Event>) { + self.queue.send_pdu(Arc::new(queued_pdu(pdu))); + } + + #[pyo3(signature = (states, start_loop = true))] + fn send_presence(&self, states: Vec>, start_loop: bool) -> PyResult<()> { + let states = states + .iter() + .map(presence_state) + .collect::>>()?; + self.queue.send_presence(states, start_loop); + Ok(()) + } + + fn queue_read_receipt(&self, receipt: &Bound<'_, PyAny>) -> PyResult<()> { + let py = receipt.py(); + let data: Value = pythonize::depythonize(&receipt.getattr(intern!(py, "data"))?)?; + let Value::Object(data) = data else { + return Err(pyo3::exceptions::PyTypeError::new_err( + "receipt data must be a dict", + )); + }; + self.queue.queue_read_receipt(&Receipt { + room_id: receipt.getattr(intern!(py, "room_id"))?.extract()?, + receipt_type: receipt.getattr(intern!(py, "receipt_type"))?.extract()?, + user_id: receipt.getattr(intern!(py, "user_id"))?.extract()?, + thread_id: receipt.getattr(intern!(py, "thread_id"))?.extract()?, + event_ids: receipt.getattr(intern!(py, "event_ids"))?.extract()?, + data, + }); + Ok(()) + } + + /// `key` is the `repr` of Python's key, which is any hashable. + fn send_keyed_edu(&self, edu: &Bound<'_, PyAny>, key: String) -> PyResult<()> { + self.queue.send_keyed_edu(edu_from_py(edu)?, key); + Ok(()) + } + + fn send_edu(&self, edu: &Bound<'_, PyAny>) -> PyResult<()> { + self.queue.send_edu(edu_from_py(edu)?); + Ok(()) + } + + fn mark_new_data(&self) { + self.queue.mark_new_data(); + } + + fn attempt_new_transaction(&self) { + self.queue.attempt_new_transaction(); + } + + // The catch-up state, for the Python catch-up loop. + + #[getter] + fn catching_up(&self) -> bool { + self.queue.catching_up() + } + + #[setter] + fn set_catching_up(&self, catching_up: bool) { + self.queue.set_catching_up(catching_up); + } + + #[getter] + fn catchup_last_skipped(&self) -> i64 { + self.queue.catchup_last_skipped() + } + + #[setter] + fn set_catchup_last_skipped(&self, stream_ordering: i64) { + self.queue.set_catchup_last_skipped(stream_ordering); + } + + #[getter] + fn last_successful_stream_ordering(&self) -> Option { + self.queue.last_successful_stream_ordering() + } + + #[setter] + fn set_last_successful_stream_ordering(&self, stream_ordering: Option) { + self.queue + .set_last_successful_stream_ordering(stream_ordering); + } + + #[getter] + fn transmission_loop_enabled(&self) -> bool { + self.queue.transmission_loop_enabled() + } + + fn start_catching_up(&self) { + self.queue.start_catching_up(); + } + + #[getter] + fn sticky_backlogged(&self) -> bool { + self.queue.sticky_backlogged() + } + + #[setter] + fn set_sticky_backlogged(&self, backlogged: bool) { + self.queue.set_sticky_backlogged(backlogged); + } + + /// The event IDs of the queued PDUs, oldest first. For tests. + fn pending_pdu_ids(&self) -> Vec { + self.queue.pending_pdu_ids() + } +} + +/// A PDU for the queue. Its JSON is read later, on tokio. +fn queued_pdu(pdu: &Bound<'_, Event>) -> QueuedPdu { + let event = pdu.get(); + QueuedPdu { + event_id: event.event_id().to_owned(), + stream_ordering: event + .internal_metadata() + .stream_ordering() + .map(|so| so.get()), + origin_server_ts: event.origin_server_ts(), + body: PduBody::Event(pdu.clone().unbind()), + } +} + +/// The parts of a `UserPresenceState` that go over federation. +fn presence_state(state: &Bound<'_, PyAny>) -> PyResult { + let py = state.py(); + Ok(PresenceState { + user_id: state.getattr(intern!(py, "user_id"))?.extract()?, + state: state.getattr(intern!(py, "state"))?.extract()?, + last_active_ts: state.getattr(intern!(py, "last_active_ts"))?.extract()?, + status_msg: state.getattr(intern!(py, "status_msg"))?.extract()?, + currently_active: state + .getattr(intern!(py, "currently_active"))? + .is_truthy()?, + }) +} + +/// An `Edu`'s type and content. +fn edu_from_py(edu: &Bound<'_, PyAny>) -> PyResult { + let py = edu.py(); + Ok(Edu { + edu_type: edu.getattr(intern!(py, "edu_type"))?.extract()?, + content: pythonize::depythonize(&edu.getattr(intern!(py, "content"))?)?, + }) +} + +/// The hooks, as coroutine methods of the Python queue. +struct PyHooks { + runtime: RustRuntime, + /// A weak reference to the Python queue. + facade: Arc>, + activity: Arc, + destination: String, +} + +impl PyHooks { + /// Await `facade.(*args(py))` on the reactor, and read its result + /// with `convert` there. + async fn call( + &self, + site: &'static str, + method: &'static str, + args: A, + convert: C, + ) -> Result + where + A: for<'py> Fn(Python<'py>) -> PyResult> + Send + Sync + 'static, + C: for<'py> Fn(&Bound<'py, PyAny>) -> PyResult + Send + Sync + 'static, + T: Send + 'static, + { + let facade = Arc::clone(&self.facade); + let observer: Arc = self.activity.clone(); + let result = run_python_awaitable_observed( + &self.runtime, + move |py| { + account(site, py, |_| { + let facade = facade.bind(py).call0()?; + if facade.is_none() { + return Err(pyo3::exceptions::PyRuntimeError::new_err( + "the federation queue has gone", + )); + } + facade.call_method1(method, args(py)?) + }) + }, + convert, + Some(observer), + ) + .await; + result.map_err(|err| python_error(err, &self.destination)) + } +} + +#[async_trait] +impl QueueHooks for PyHooks { + async fn catch_up(&self) -> Result<(), LoopError> { + self.call( + "queue_catch_up", + "_rust_catch_up", + |py| Ok(PyTuple::empty(py)), + |_| Ok(()), + ) + .await + } + + async fn device_edus(&self, request: DeviceEdusRequest) -> Result { + self.call( + "queue_device_edus", + "_rust_device_edus", + move |py| { + PyTuple::new( + py, + [ + request.last_device_stream_id.into_pyobject(py)?, + request.last_device_list_stream_id.into_pyobject(py)?, + request.to_device_limit.into_pyobject(py)?, + request.edu_limit.into_pyobject(py)?, + ], + ) + }, + |result| { + let (to_device, to_device_stream_id, device_list, device_list_stream_id): ( + Vec>, + i64, + Vec<(String, Bound<'_, PyAny>)>, + i64, + ) = result.extract()?; + Ok(DeviceEdus { + to_device: to_device + .iter() + .map(|content| pythonize::depythonize(content)) + .collect::>()?, + to_device_stream_id, + device_list: device_list + .iter() + .map(|(edu_type, content)| -> PyResult { + Ok(Edu { + edu_type: edu_type.clone(), + content: pythonize::depythonize(content)?, + }) + }) + .collect::>()?, + device_list_stream_id, + }) + }, + ) + .await + } + + async fn sticky_backlog(&self) -> Result, LoopError> { + self.call( + "queue_sticky_backlog", + "_rust_sticky_backlog", + |py| Ok(PyTuple::empty(py)), + |result| { + if result.is_none() { + return Ok(None); + } + let (room_id, position, events): (String, Py, Vec>) = + result.extract()?; + let pdus = events + .iter() + .map(|event| Ok(Arc::new(queued_pdu(event.cast::()?)))) + .collect::>()?; + let position: StickyStreamPosition = Arc::new(position); + Ok(Some(StickyBatch { + room_id, + position, + pdus, + })) + }, + ) + .await + } + + async fn complete(&self, completion: Completion) -> Result<(), LoopError> { + let completion = Arc::new(completion); + self.call( + "queue_complete", + "_rust_complete", + move |py| completion_args(py, &completion), + |_| Ok(()), + ) + .await + } +} + +/// The arguments of `_rust_complete`. +fn completion_args<'py>(py: Python<'py>, completion: &Completion) -> PyResult> { + let sticky = match &completion.sticky { + Some((room_id, position)) => { + let position = position + .downcast_ref::>() + .map(|position| position.clone_ref(py)) + .unwrap_or_else(|| py.None()); + (room_id.as_str(), position).into_pyobject(py)?.into_any() + } + None => py.None().into_bound(py), + }; + let sent = match &completion.sent { + Some(sent) => { + let pdus = PyList::empty(py); + for pdu in &sent.pdus { + match &pdu.body { + PduBody::Event(event) => pdus.append(event.clone_ref(py))?, + #[cfg(test)] + PduBody::Json(_) => {} + } + } + ( + pdus, + &sent.edu_types, + &sent.acknowledged, + sent.stats.attempts, + &sent.stats.responses, + ) + .into_pyobject(py)? + .into_any() + } + None => py.None().into_bound(py), + }; + PyTuple::new( + py, + [ + completion.to_device_stream_id.into_pyobject(py)?, + completion.device_list_stream_id.into_pyobject(py)?, + completion.last_stream_ordering.into_pyobject(py)?, + sticky, + sent, + ], + ) +} + +/// What a Python exception means to the loop. Takes the GIL, which only +/// error paths need. +fn python_error(err: PyErr, destination: &str) -> LoopError { + with_gil("queue_error", |py| { + classify_python_error(py, &err).unwrap_or_else(|classify_err| { + LoopError::Other(format!( + "{} (and classifying it failed: {classify_err}) for {destination}", + format_exception(py, &err) + )) + }) + }) +} + +fn classify_python_error(py: Python<'_>, err: &PyErr) -> PyResult { + let value = err.value(py); + let errors = py.import("synapse.api.errors")?; + let not_retrying = py + .import("synapse.util.retryutils")? + .getattr("NotRetryingDestination")?; + + Ok(if value.is_instance(¬_retrying)? { + LoopError::NotRetrying(NotRetrying { + retry_last_ts: value.getattr("retry_last_ts")?.extract()?, + retry_interval: value.getattr("retry_interval")?.extract()?, + }) + } else if value.is_instance(&errors.getattr("FederationDeniedError")?)? { + LoopError::Denied(value.str()?.to_string()) + } else if value.is_instance(&errors.getattr("HttpResponseException")?)? { + LoopError::Http { + code: value.getattr("code")?.extract()?, + message: value.str()?.to_string(), + } + } else if value.is_instance(&errors.getattr("RequestSendFailed")?)? { + LoopError::SendFailed(value.str()?.to_string()) + } else { + LoopError::Other(format_exception(py, err)) + }) +} + +/// The exception with its traceback, as `logger.exception` logs it. +fn format_exception(py: Python<'_>, err: &PyErr) -> String { + let formatted = (|| -> PyResult { + let lines = py + .import("traceback")? + .call_method1("format_exception", (err.value(py),))?; + Ok(PyString::new(py, "") + .call_method1("join", (lines,))? + .extract::()? + .trim_end() + .to_owned()) + })(); + formatted.unwrap_or_else(|_| err.to_string()) +} + +/// Sends through the Rust transaction sender, or a Python stand-in in tests. +struct ClientSink { + sender: Arc, + runtime: RustRuntime, + record_metrics: Arc>, + send_for_test: Mutex>>>, +} + +impl ClientSink { + /// Send through the Python stand-in. + async fn send_for_test( + &self, + send: Arc>, + transaction: Transaction, + observer: Arc, + ) -> Result { + let destination = transaction.destination.clone(); + let log_prefix = transaction.log_prefix(); + let transaction = Arc::new(transaction); + let response = run_python_awaitable_observed( + &self.runtime, + move |py| { + let built = RustTransaction { + inner: (*transaction).clone(), + }; + send.bind(py).call1((built,)) + }, + |response| { + Ok(match response.cast::() { + Ok(response) => pythonize::depythonize::(response)?, + // A mock that returns nothing. + Err(_) => Value::Null, + }) + }, + Some(observer), + ) + .await + .map_err(|err| python_error(err, &destination))?; + + let acknowledged = match response { + Value::Object(response) => acknowledged_pdus(&log_prefix, &response), + _ => Vec::new(), + }; + Ok(SendOutcome { + acknowledged, + stats: RequestStats::default(), + }) + } +} + +#[async_trait] +impl TransactionSink for ClientSink { + fn now_ms(&self) -> u64 { + self.sender.client().now_ms() + } + + async fn backoff(&self, destination: &str) -> Result, LoopError> { + self.sender + .client() + .backoff(destination) + .await + .map_err(|err| LoopError::Other(err.to_string())) + } + + fn build(&self, destination: &str, pdus: Vec, edus: Vec) -> Transaction { + self.sender.build(destination, pdus, edus) + } + + async fn send( + &self, + transaction: Transaction, + observer: Arc, + ) -> Result { + let send_for_test = self.send_for_test.lock().expect("poisoned").clone(); + if let Some(send) = send_for_test { + return self.send_for_test(send, transaction, observer).await; + } + + let destination = transaction.destination.clone(); + let transaction_id = transaction.transaction_id.clone(); + let mut stats = RequestStats::default(); + let result = self.sender.send(transaction, Vec::new(), &mut stats).await; + match result { + Ok(TransactionResponse::Acknowledged(acknowledged)) => Ok(SendOutcome { + acknowledged, + stats, + }), + Ok(TransactionResponse::Unparsed(body)) => with_gil("queue_response_fallback", |py| { + python_acknowledged(py, &destination, &transaction_id, &body).map_err(|err| { + classify_python_error(py, &err) + .unwrap_or_else(|_| LoopError::Other(format_exception(py, &err))) + }) + }) + .map(|acknowledged| SendOutcome { + acknowledged, + stats, + }), + Err(err) => { + // The request counters are Python's. A failure is the only + // case without a later trip to Python to carry them. + let record_metrics = Arc::clone(&self.record_metrics); + with_gil("queue_send_failed", |py| { + let recorded = 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}"); + } + }); + Err(LoopError::from_client(err, &destination)) + } + } + } +} diff --git a/rust/src/federation/queue/tests.rs b/rust/src/federation/queue/tests.rs new file mode 100644 index 0000000000..7f05fd2994 --- /dev/null +++ b/rust/src/federation/queue/tests.rs @@ -0,0 +1,734 @@ +/* + * 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: + * . + * + */ + +use std::{ + collections::VecDeque, + sync::{OnceLock, Weak}, + time::Duration, +}; + +use serde_json::json; +use tokio::sync::Notify; + +use super::*; + +const DEST: &str = "remote.example"; +const NOW_MS: u64 = 1_000_000; + +/// What the fake sink does with the next transaction. +enum Answer { + /// Accept every PDU. + Accept, + Fail(fn() -> LoopError), + /// Wait for the notify, then accept. + WaitThenAccept(Arc), +} + +#[derive(Default)] +struct FakeSink { + answers: Mutex>, + sent: Mutex>, + backoff: Mutex>, + next_txn_id: AtomicU64, +} + +impl FakeSink { + fn sent(&self) -> Vec { + self.sent.lock().unwrap().clone() + } + + fn answer(&self, answer: Answer) { + self.answers.lock().unwrap().push_back(answer); + } +} + +#[async_trait] +impl TransactionSink for FakeSink { + fn now_ms(&self) -> u64 { + NOW_MS + } + + async fn backoff(&self, _destination: &str) -> Result, LoopError> { + Ok(*self.backoff.lock().unwrap()) + } + + fn build(&self, destination: &str, pdus: Vec, edus: Vec) -> Transaction { + Transaction { + transaction_id: self.next_txn_id.fetch_add(1, Ordering::SeqCst).to_string(), + origin: "origin.example".to_owned(), + destination: destination.to_owned(), + origin_server_ts: NOW_MS, + pdus, + edus, + } + } + + async fn send( + &self, + transaction: Transaction, + _observer: Arc, + ) -> Result { + let answer = self + .answers + .lock() + .unwrap() + .pop_front() + .unwrap_or(Answer::Accept); + self.sent.lock().unwrap().push(transaction.clone()); + match answer { + Answer::Accept => {} + Answer::Fail(err) => return Err(err()), + Answer::WaitThenAccept(notify) => notify.notified().await, + } + let acknowledged = transaction + .pdus + .iter() + .filter_map(|pdu| pdu["event_id"].as_str().map(str::to_owned)) + .collect(); + Ok(SendOutcome { + acknowledged, + stats: RequestStats::default(), + }) + } +} + +/// A completion, without the parts that are not comparable. +#[derive(Debug, PartialEq)] +struct Completed { + to_device_stream_id: Option, + device_list_stream_id: Option, + last_stream_ordering: Option, + sticky_room: Option, + sent_pdus: Option>, + edu_types: Option>, +} + +#[derive(Default)] +struct FakeHooks { + queue: OnceLock>, + device_edus: Mutex>, + device_requests: Mutex>, + sticky: Mutex>>, + completed: Mutex>, + catch_ups: AtomicU64, +} + +impl FakeHooks { + fn completed(&self) -> Vec { + std::mem::take(&mut *self.completed.lock().unwrap()) + } +} + +#[async_trait] +impl QueueHooks for FakeHooks { + async fn catch_up(&self) -> Result<(), LoopError> { + self.catch_ups.fetch_add(1, Ordering::SeqCst); + // As Python's, with nothing in the database to catch up on. + let queue = self.queue.get().and_then(Weak::upgrade).unwrap(); + if queue.last_successful_stream_ordering().is_none() { + queue.set_catching_up(false); + } + Ok(()) + } + + async fn device_edus(&self, request: DeviceEdusRequest) -> Result { + self.device_requests.lock().unwrap().push(request); + Ok(self + .device_edus + .lock() + .unwrap() + .pop_front() + .unwrap_or(DeviceEdus { + to_device_stream_id: request.last_device_stream_id, + device_list_stream_id: request.last_device_list_stream_id, + ..DeviceEdus::default() + })) + } + + async fn sticky_backlog(&self) -> Result, LoopError> { + Ok(self.sticky.lock().unwrap().pop_front().flatten()) + } + + async fn complete(&self, completion: Completion) -> Result<(), LoopError> { + self.completed.lock().unwrap().push(Completed { + to_device_stream_id: completion.to_device_stream_id, + device_list_stream_id: completion.device_list_stream_id, + last_stream_ordering: completion.last_stream_ordering, + sticky_room: completion.sticky.map(|(room_id, _)| room_id), + sent_pdus: completion + .sent + .as_ref() + .map(|sent| sent.acknowledged.clone()), + edu_types: completion.sent.map(|sent| sent.edu_types), + }); + Ok(()) + } +} + +struct Harness { + queue: Arc, + hooks: Arc, + sink: Arc, + activity: Arc, +} + +fn harness_with(msc4354_enabled: bool) -> Harness { + let hooks = Arc::new(FakeHooks::default()); + let sink = Arc::new(FakeSink::default()); + let activity = Arc::new(Activity::default()); + let queue = DestinationQueue::new( + QueueConfig { + destination: DEST.to_owned(), + should_send_on_this_instance: true, + msc4354_enabled, + }, + hooks.clone(), + sink.clone(), + LoopActivity::new(Arc::clone(&activity)), + Handle::current(), + ); + hooks.queue.set(Arc::downgrade(&queue)).unwrap(); + Harness { + queue, + hooks, + sink, + activity, + } +} + +fn harness() -> Harness { + harness_with(false) +} + +impl Harness { + /// Wait for the loop to end. + async fn settle(&self) { + for _ in 0..1000 { + if !self.queue.transmission_loop_running() { + return; + } + tokio::time::sleep(Duration::from_millis(1)).await; + } + panic!("the loop did not end"); + } +} + +fn pdu(event_id: &str, stream_ordering: i64) -> Arc { + Arc::new(QueuedPdu { + event_id: event_id.to_owned(), + stream_ordering: Some(stream_ordering), + origin_server_ts: 1, + body: PduBody::Json(json!({"event_id": event_id, "depth": 1})), + }) +} + +fn edu_types(transaction: &Transaction) -> Vec<&str> { + transaction + .edus + .iter() + .map(|edu| edu["edu_type"].as_str().unwrap()) + .collect() +} + +fn pdu_ids(transaction: &Transaction) -> Vec<&str> { + transaction + .pdus + .iter() + .map(|pdu| pdu["event_id"].as_str().unwrap()) + .collect() +} + +fn presence(user_id: &str) -> PresenceState { + PresenceState { + user_id: user_id.to_owned(), + state: "online".to_owned(), + last_active_ts: 0, + status_msg: None, + currently_active: true, + } +} + +fn receipt(user_id: &str, event_id: &str, thread_id: Option<&str>) -> Receipt { + Receipt { + room_id: "!room".to_owned(), + receipt_type: "m.read".to_owned(), + user_id: user_id.to_owned(), + thread_id: thread_id.map(str::to_owned), + event_ids: vec![event_id.to_owned()], + data: Map::new(), + } +} + +#[tokio::test] +async fn sends_a_queued_pdu() { + let h = harness(); + h.queue.send_pdu(pdu("$a", 5)); + h.settle().await; + + let sent = h.sink.sent(); + assert_eq!(sent.len(), 1); + assert_eq!(pdu_ids(&sent[0]), ["$a"]); + // Nothing to catch up on, so the first loop leaves catch-up. + assert_eq!(h.hooks.catch_ups.load(Ordering::SeqCst), 1); + assert!(!h.queue.catching_up()); + + assert_eq!( + h.hooks.completed(), + [Completed { + to_device_stream_id: None, + device_list_stream_id: None, + last_stream_ordering: Some(5), + sticky_room: None, + sent_pdus: Some(vec!["$a".to_owned()]), + edu_types: Some(vec![]), + }] + ); + assert_eq!(h.queue.pending_pdu_count(), 0); + assert!(h.activity.idle()); +} + +#[tokio::test] +async fn fifty_pdus_per_transaction() { + let h = harness(); + h.queue.set_catching_up(false); + { + // Queue them all before the loop runs. + let mut state = h.queue.state(); + for i in 0..120 { + let pdu = pdu(&format!("${i}"), i + 1); + state.pending_pdus.insert(pdu.event_id.clone(), pdu); + } + } + h.queue.attempt_new_transaction(); + h.settle().await; + + let sizes: Vec<_> = h.sink.sent().iter().map(|t| t.pdus.len()).collect(); + assert_eq!(sizes, [50, 50, 20]); + assert_eq!(pdu_ids(&h.sink.sent()[1])[0], "$50"); +} + +#[tokio::test] +async fn edus_in_order_and_within_limits() { + let h = harness(); + h.queue.set_catching_up(false); + h.hooks.device_edus.lock().unwrap().push_back(DeviceEdus { + to_device: vec![json!({"message_id": "m1"})], + to_device_stream_id: 7, + device_list: vec![Edu { + edu_type: "m.device_list_update".to_owned(), + content: json!({}), + }], + device_list_stream_id: 9, + }); + + let users: Vec<_> = (0..60).map(|i| presence(&format!("@u{i}:x"))).collect(); + h.queue.send_presence(users, false); + for i in 0..7 { + // Each receipt has a thread, so each needs its own EDU. + h.queue + .queue_read_receipt(&receipt("@u:x", &format!("$e{i}"), Some("t"))); + } + h.queue.send_edu(Edu { + edu_type: "m.other".to_owned(), + content: json!({}), + }); + h.settle().await; + + let sent = h.sink.sent(); + assert_eq!( + edu_types(&sent[0]), + [ + "m.presence", + "m.receipt", + "m.receipt", + "m.receipt", + "m.receipt", + "m.receipt", + "m.direct_to_device", + "m.device_list_update", + "m.other", + ] + ); + assert_eq!( + sent[0].edus[0]["content"]["push"].as_array().unwrap().len(), + 50 + ); + // The rest go in the next transaction. + assert_eq!( + edu_types(&sent[1]), + ["m.presence", "m.receipt", "m.receipt"] + ); + assert_eq!( + sent[1].edus[0]["content"]["push"].as_array().unwrap().len(), + 10 + ); + + // One presence and five receipt EDUs leave 94 places, ten of them kept + // for device list updates. + let requests = h.hooks.device_requests.lock().unwrap().clone(); + assert_eq!( + requests[0], + DeviceEdusRequest { + last_device_stream_id: 0, + last_device_list_stream_id: 0, + to_device_limit: 84, + edu_limit: 94, + } + ); + // The positions move on once the transaction is sent. + assert_eq!(requests[1].last_device_stream_id, 7); + assert_eq!(requests[1].last_device_list_stream_id, 9); + + let completed = h.hooks.completed(); + assert_eq!(completed[0].to_device_stream_id, Some(7)); + assert_eq!(completed[0].device_list_stream_id, Some(9)); +} + +#[tokio::test] +async fn positions_move_on_when_nothing_to_send() { + let h = harness(); + h.queue.set_catching_up(false); + h.hooks.device_edus.lock().unwrap().push_back(DeviceEdus { + to_device_stream_id: 4, + device_list_stream_id: 6, + ..DeviceEdus::default() + }); + h.queue.attempt_new_transaction(); + h.settle().await; + + assert!(h.sink.sent().is_empty()); + let state = h.queue.state(); + assert_eq!(state.last_device_stream_id, 4); + assert_eq!(state.last_device_list_stream_id, 6); +} + +#[test] +fn receipts_merge_as_python_does() { + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + let _guard = rt.enter(); + let h = harness(); + + h.queue.queue_read_receipt(&receipt("@a:x", "$1", None)); + // The same user without a thread replaces the receipt. + h.queue.queue_read_receipt(&receipt("@a:x", "$2", None)); + // Another user goes in the same EDU. + h.queue.queue_read_receipt(&receipt("@b:x", "$3", None)); + // A thread receipt for a user already there needs a new EDU. + h.queue + .queue_read_receipt(&receipt("@a:x", "$4", Some("t"))); + + let state = h.queue.state(); + assert_eq!(state.pending_receipt_edus.len(), 2); + assert_eq!( + Value::Object(state.pending_receipt_edus[0].clone()), + json!({"!room": {"m.read": { + "@a:x": {"event_ids": ["$2"], "data": {}}, + "@b:x": {"event_ids": ["$3"], "data": {}}, + }}}) + ); + assert_eq!( + Value::Object(state.pending_receipt_edus[1].clone()), + json!({"!room": {"m.read": { + "@a:x": {"event_ids": ["$4"], "data": {"thread_id": "t"}}, + }}}) + ); + assert!(state.new_data_to_send); +} + +#[tokio::test] +async fn keyed_edus_replace_and_go_newest_first() { + let h = harness(); + h.queue.set_catching_up(false); + let typing = |room: &str, typing: bool| Edu { + edu_type: "m.typing".to_owned(), + content: json!({"room_id": room, "typing": typing}), + }; + { + let mut state = h.queue.state(); + state + .pending_edus_keyed + .insert(("m.typing".to_owned(), "a".to_owned()), typing("!a", true)); + state + .pending_edus_keyed + .insert(("m.typing".to_owned(), "b".to_owned()), typing("!b", true)); + } + h.queue.send_keyed_edu(typing("!a", false), "a".to_owned()); + h.settle().await; + + let sent = h.sink.sent(); + assert_eq!( + sent[0].edus, + [ + typing("!b", true).into_value(), + typing("!a", false).into_value() + ] + ); +} + +#[tokio::test] +async fn a_failed_send_keeps_the_pdus() { + let h = harness(); + h.queue.set_catching_up(false); + h.sink + .answer(Answer::Fail(|| LoopError::SendFailed("boom".to_owned()))); + h.queue.send_pdu(pdu("$a", 1)); + h.settle().await; + + assert_eq!(h.sink.sent().len(), 1); + assert!(h.hooks.completed().is_empty()); + assert_eq!(h.queue.pending_pdu_ids(), ["$a"]); + + // The next attempt sends it. + h.queue.attempt_new_transaction(); + h.settle().await; + assert_eq!(h.sink.sent().len(), 2); + assert_eq!(h.queue.pending_pdu_count(), 0); +} + +#[tokio::test] +async fn long_backoff_drops_edus_and_catches_up() { + let h = harness(); + h.queue.set_catching_up(false); + h.queue.send_presence(vec![presence("@u:x")], false); + h.queue.queue_read_receipt(&receipt("@u:x", "$r", None)); + { + let mut state = h.queue.state(); + let pdu = pdu("$a", 1); + state.pending_pdus.insert(pdu.event_id.clone(), pdu); + } + *h.sink.backoff.lock().unwrap() = Some(NotRetrying { + retry_last_ts: NOW_MS, + retry_interval: CATCHUP_RETRY_INTERVAL_MS + 1, + }); + h.queue.attempt_new_transaction(); + h.settle().await; + + assert!(h.sink.sent().is_empty()); + assert!(h.queue.catching_up()); + assert_eq!(h.queue.pending_pdu_count(), 0); + assert_eq!(h.queue.pending_edu_count(), 0); + assert!(h.queue.state().pending_receipt_edus.is_empty()); +} + +#[tokio::test] +async fn short_backoff_keeps_everything() { + let h = harness(); + h.queue.set_catching_up(false); + h.queue.send_presence(vec![presence("@u:x")], false); + *h.sink.backoff.lock().unwrap() = Some(NotRetrying { + retry_last_ts: NOW_MS, + retry_interval: 60_000, + }); + h.queue.send_pdu(pdu("$a", 1)); + h.settle().await; + + assert!(h.sink.sent().is_empty()); + assert!(!h.queue.catching_up()); + assert_eq!(h.queue.pending_pdu_count(), 1); + assert_eq!(h.queue.pending_edu_count(), 1); +} + +#[tokio::test] +async fn pdus_skipped_while_catching_up() { + let h = harness(); + h.queue.set_last_successful_stream_ordering(Some(3)); + // Keep the loop from running, so the queue stays catching up. + h.queue.state().transmission_loop_enabled = false; + h.queue.send_pdu(pdu("$a", 8)); + + assert_eq!(h.queue.pending_pdu_count(), 0); + assert_eq!(h.queue.catchup_last_skipped(), 8); +} + +#[tokio::test] +async fn data_queued_mid_send_is_not_lost() { + let h = harness(); + h.queue.set_catching_up(false); + let release = Arc::new(Notify::new()); + h.sink.answer(Answer::WaitThenAccept(Arc::clone(&release))); + + h.queue.send_pdu(pdu("$a", 1)); + // Wait for the first send to start. + for _ in 0..1000 { + if !h.sink.sent().is_empty() { + break; + } + tokio::time::sleep(Duration::from_millis(1)).await; + } + // The loop is running, so this only marks new data. + h.queue.send_pdu(pdu("$b", 2)); + assert!(h.queue.transmission_loop_running()); + release.notify_one(); + h.settle().await; + + let sent = h.sink.sent(); + assert_eq!(sent.len(), 2); + assert_eq!(pdu_ids(&sent[1]), ["$b"]); +} + +#[tokio::test] +async fn sticky_backlog_takes_turns() { + let h = harness_with(true); + h.queue.set_catching_up(false); + let batch = |room: &str, event_id: &str| { + Some(StickyBatch { + room_id: room.to_owned(), + position: Arc::new(0i64), + pdus: vec![pdu(event_id, 100)], + }) + }; + h.hooks + .sticky + .lock() + .unwrap() + .extend([batch("!s1", "$s1"), batch("!s2", "$s2"), None]); + { + let mut state = h.queue.state(); + for i in 0..120 { + let pdu = pdu(&format!("${i}"), i + 1); + state.pending_pdus.insert(pdu.event_id.clone(), pdu); + } + } + h.queue.attempt_new_transaction(); + h.settle().await; + + let firsts: Vec<_> = h + .sink + .sent() + .iter() + .map(|t| pdu_ids(t)[0].to_owned()) + .collect(); + assert_eq!(firsts, ["$0", "$s1", "$50", "$s2", "$100"]); + assert!(!h.queue.sticky_backlogged()); + let sticky: Vec<_> = h + .hooks + .completed() + .into_iter() + .filter_map(|c| c.sticky_room) + .collect(); + assert_eq!(sticky, ["!s1", "!s2"]); +} + +#[tokio::test] +async fn an_emptied_sticky_batch_is_still_completed() { + let h = harness_with(true); + h.queue.set_catching_up(false); + h.hooks.sticky.lock().unwrap().extend([ + Some(StickyBatch { + room_id: "!s".to_owned(), + position: Arc::new(0i64), + pdus: Vec::new(), + }), + None, + ]); + h.queue.attempt_new_transaction(); + h.settle().await; + + assert!(h.sink.sent().is_empty()); + assert_eq!( + h.hooks.completed(), + [Completed { + to_device_stream_id: None, + device_list_stream_id: None, + last_stream_ordering: None, + sticky_room: Some("!s".to_owned()), + sent_pdus: None, + edu_types: None, + }] + ); +} + +#[tokio::test] +async fn shutdown_stops_the_loop() { + let h = harness(); + h.queue.set_catching_up(false); + let release = Arc::new(Notify::new()); + h.sink.answer(Answer::WaitThenAccept(release)); + h.queue.send_pdu(pdu("$a", 1)); + for _ in 0..1000 { + if !h.sink.sent().is_empty() { + break; + } + tokio::time::sleep(Duration::from_millis(1)).await; + } + + h.queue.shutdown(); + h.settle().await; + assert!(h.activity.idle()); + + // A shut down queue starts nothing. + h.queue.send_pdu(pdu("$b", 2)); + assert!(!h.queue.transmission_loop_running()); +} + +#[test] +fn presence_format_matches_python() { + let state = PresenceState { + user_id: "@u:x".to_owned(), + state: "online".to_owned(), + last_active_ts: 400, + status_msg: Some("hi".to_owned()), + currently_active: false, + }; + assert_eq!( + state.format(1000), + json!({ + "presence": "online", + "user_id": "@u:x", + "last_active_ago": 600, + "status_msg": "hi", + "currently_active": false, + }) + ); + + // Unset fields are left out, and only "online" says whether the user is + // active. + let state = PresenceState { + state: "unavailable".to_owned(), + last_active_ts: 0, + status_msg: Some(String::new()), + ..state + }; + assert_eq!( + state.format(1000), + json!({"presence": "unavailable", "user_id": "@u:x"}) + ); +} + +#[test] +fn receipts_with_threads_share_an_edu() { + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + let _guard = rt.enter(); + let h = harness(); + + h.queue.queue_read_receipt(&receipt("alice", "$e", None)); + h.queue + .queue_read_receipt(&receipt("alice", "$e", Some("thread"))); + h.queue.queue_read_receipt(&receipt("bob", "$e", None)); + h.queue + .queue_read_receipt(&receipt("bob", "$e", Some("diff-thread"))); + + let state = h.queue.state(); + let users: Vec> = state + .pending_receipt_edus + .iter() + .map(|edu| edu["!room"]["m.read"].as_object().unwrap().keys().collect()) + .collect(); + assert_eq!(users, [["alice", "bob"], ["alice", "bob"]]); +} diff --git a/rust/src/federation/transaction/mod.rs b/rust/src/federation/transaction/mod.rs index f06f19f73c..9d5293d078 100644 --- a/rust/src/federation/transaction/mod.rs +++ b/rust/src/federation/transaction/mod.rs @@ -106,7 +106,7 @@ impl Transaction { } /// The prefix of the transaction's log lines, `TX [dest] {txn}`. - fn log_prefix(&self) -> String { + pub(crate) fn log_prefix(&self) -> String { log_prefix(&self.destination, &self.transaction_id) } } diff --git a/rust/src/federation/transaction/python.rs b/rust/src/federation/transaction/python.rs index 37b7cb0584..1fc276fd16 100644 --- a/rust/src/federation/transaction/python.rs +++ b/rust/src/federation/transaction/python.rs @@ -45,7 +45,7 @@ use crate::{ /// A built transaction. Tests read it as they read a Python `Transaction`. #[pyclass(frozen, name = "RustTransaction")] pub struct RustTransaction { - inner: Transaction, + pub(crate) inner: Transaction, } #[pymethods] @@ -89,9 +89,9 @@ impl RustTransaction { /// Builds and sends transactions through a [`RustFederationClient`]. #[pyclass(frozen, name = "RustTransactionManager")] pub struct RustTransactionManager { - sender: Arc, - runtime: RustRuntime, - record_metrics: Arc>, + pub(crate) sender: Arc, + pub(crate) runtime: RustRuntime, + pub(crate) record_metrics: Arc>, } #[pymethods] @@ -231,7 +231,7 @@ impl RustTransactionManager { /// The response handling for a body serde_json refused: decode it with /// Python's decoder, as `JsonParser` would, and read the `pdus` map from the /// result. -fn python_acknowledged( +pub(crate) fn python_acknowledged( py: Python<'_>, destination: &str, transaction_id: &str, diff --git a/synapse/synapse_rust/federation.pyi b/synapse/synapse_rust/federation.pyi index a7c7e04f36..5d9019da16 100644 --- a/synapse/synapse_rust/federation.pyi +++ b/synapse/synapse_rust/federation.pyi @@ -10,15 +10,17 @@ # See the GNU Affero General Public License for more details: # . -from typing import Any, Callable, Mapping, Sequence +from typing import Any, Awaitable, Callable, Mapping, Sequence from signedjson.types import SigningKey from twisted.internet.defer import Deferred +from synapse.api.presence import UserPresenceState +from synapse.federation.units import Edu from synapse.server import HomeServer from synapse.synapse_rust.events import Event -from synapse.types import JsonDict +from synapse.types import JsonDict, ReadReceipt def build_auth_header_for_test( signing_key: SigningKey, @@ -199,3 +201,54 @@ class RustTransactionManager: ) -> list[str]: """The ids of the PDUs `response` accepts, logging the ones it rejects. For tests.""" + +class RustDestinationQueues: + """Makes the Rust per-destination queues, at + `rust_federation_sending: queues`.""" + + def __init__( + self, transactions: RustTransactionManager, msc4354_enabled: bool + ) -> None: ... + def queue( + self, destination: str, facade: Any, should_send_on_this_instance: bool + ) -> "RustDestinationQueue": + """The queue for `destination`. `facade` provides the `_rust_*` + hooks, and is held weakly.""" + + def idle(self) -> bool: + """Whether every transmission loop has finished or is waiting on the + reactor. For tests.""" + + def set_send_for_test( + self, send: Callable[[RustTransaction], Awaitable[Any]] | None + ) -> None: + """Send transactions through `send` in place of the network. For + tests.""" + +class RustDestinationQueue: + """One destination's queue. Each method is the `PerDestinationQueue` + method of the same name.""" + + transmission_loop_running: bool + new_data_to_send: bool + catching_up: bool + catchup_last_skipped: int + last_successful_stream_ordering: int | None + transmission_loop_enabled: bool + sticky_backlogged: bool + + def shutdown(self) -> None: ... + def pending_pdu_count(self) -> int: ... + def pending_edu_count(self) -> int: ... + def send_pdu(self, pdu: Event) -> None: ... + def send_presence( + self, states: Sequence[UserPresenceState], start_loop: bool = True + ) -> None: ... + def queue_read_receipt(self, receipt: ReadReceipt) -> None: ... + def send_keyed_edu(self, edu: Edu, key: str) -> None: ... + def send_edu(self, edu: Edu) -> None: ... + def mark_new_data(self) -> None: ... + def attempt_new_transaction(self) -> None: ... + def start_catching_up(self) -> None: ... + def pending_pdu_ids(self) -> list[str]: + """The event IDs of the queued PDUs, oldest first. For tests."""