diff --git a/changelog.d/rust-stream-id-tracker.misc b/changelog.d/rust-stream-id-tracker.misc new file mode 100644 index 0000000000..3e06c67e1a --- /dev/null +++ b/changelog.d/rust-stream-id-tracker.misc @@ -0,0 +1 @@ +Keep the stream ID bookkeeping of multi-writer streams in Rust, so Rust code can allocate stream IDs. diff --git a/rust/src/lib.rs b/rust/src/lib.rs index e514f32822..415fa3f5bb 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -108,6 +108,7 @@ fn synapse_rust(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { rendezvous::register_module(py, m)?; runtime::register_module(py, m)?; storage::python::register_module(py, m)?; + storage::stream_ids::register_module(py, m)?; msc4388_rendezvous::register_module(py, m)?; segmenter::register_module(py, m)?; state::register_module(py, m)?; diff --git a/rust/src/storage/mod.rs b/rust/src/storage/mod.rs index 13e2a5c0e5..8c61cb7331 100644 --- a/rust/src/storage/mod.rs +++ b/rust/src/storage/mod.rs @@ -20,4 +20,5 @@ pub mod federation_store; pub mod pool_querier; pub mod python; pub mod store; +pub mod stream_ids; pub mod sync_store; diff --git a/rust/src/storage/stream_ids.rs b/rust/src/storage/stream_ids.rs new file mode 100644 index 0000000000..3bc08a238e --- /dev/null +++ b/rust/src/storage/stream_ids.rs @@ -0,0 +1,663 @@ +/* + * 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 in-memory bookkeeping of a multi-writer stream: which stream IDs this +//! instance has allocated and not yet finished, how far each writer has got, +//! and from those the current token. +//! +//! `MultiWriterIdGenerator` in `synapse/storage/util/id_generators.py` keeps +//! the parts that touch the database or the reactor: loading the start +//! positions, the sequence, the `stream_positions` table and notifying +//! replication. It hands everything else to a [`StreamIdTracker`], so that +//! Rust code can allocate and finish stream IDs without taking the GIL. +//! +//! The current token must never pass an ID that this instance has allocated +//! but not finished. Otherwise a reader could see a later event and skip an +//! earlier one that commits after it. Allocation has two halves, around the +//! `nextval`: [`StreamIdTracker::begin_fetch`] records the highest ID seen so +//! far as "in flight", since the IDs that come back will be above it, and +//! [`StreamIdTracker::end_fetch`] swaps that for the allocated IDs. +//! +//! Units: the tracker stores IDs as positive numbers, also for streams that +//! count down such as backfill. Methods that take or return positions of +//! the stream as other code sees them (`advance`, the `get_*` methods) use +//! signed IDs. The bookkeeping methods (`begin_fetch`, `end_fetch`, +//! `mark_finished`, `load`) take positive IDs, as the sequence returns them. + +use std::{ + cmp::Reverse, + collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet}, + sync::{Mutex, MutexGuard}, +}; + +use log::debug; +use pyo3::prelude::*; + +/// The logger of the debug line for issue 19795. +const ISSUE_19795_LOG_TARGET: &str = "synapse.9533_debug"; + +/// See the module docs. +pub struct StreamIdTracker { + stream_name: String, + instance_name: String, + /// `1` for a stream that counts up, `-1` for one that counts down. + return_factor: i64, + state: Mutex, +} + +#[derive(Debug)] +struct State { + /// Each writer's position. Includes our own once we have written. + current_positions: HashMap, + + /// IDs this instance has allocated and not yet finished. The current + /// position must stay below the smallest. + unfinished_ids: BTreeSet, + + /// For each allocation in progress, the highest ID seen when it started, + /// counted per value, since two allocations can start at the same one. + /// The IDs they will get are above it, so the current position must not + /// pass it. + in_flight_fetches: BTreeMap, + + /// IDs this instance has finished that are above the current position, + /// because a lower ID is still unfinished. + finished_ids: HashSet, + + /// The position up to which every writer has persisted everything. + persisted_upto_position: i64, + + /// Persisted positions above `persisted_upto_position`. When the lowest + /// is one above it, the position moves up. Gaps in the sequence (from + /// rolled back transactions) stop that until the writers' positions pass + /// them. + known_persisted_positions: BinaryHeap>, + + /// The highest ID allocated by any writer that we know of. + max_seen_allocated_stream_id: i64, + + /// Our own position. It can be above our entry in `current_positions` + /// when we have nothing in flight, since our next write will be above + /// every ID allocated so far. + max_position_of_local_instance: i64, +} + +impl StreamIdTracker { + /// A tracker before [`StreamIdTracker::load`]. Like a stream with no + /// writes: ID 1 counts as persisted, so the first ID is 2. + pub fn new(stream_name: String, instance_name: String, positive: bool) -> Self { + StreamIdTracker { + stream_name, + instance_name, + return_factor: if positive { 1 } else { -1 }, + state: Mutex::new(State { + current_positions: HashMap::new(), + unfinished_ids: BTreeSet::new(), + in_flight_fetches: BTreeMap::new(), + finished_ids: HashSet::new(), + persisted_upto_position: 1, + known_persisted_positions: BinaryHeap::new(), + max_seen_allocated_stream_id: 1, + max_position_of_local_instance: 1, + }), + } + } + + fn lock(&self) -> MutexGuard<'_, State> { + // The state is a set of integers that every method leaves + // consistent before it can panic, so a poisoned lock is still usable. + self.state.lock().unwrap_or_else(|e| e.into_inner()) + } + + /// Set the start state, from what the generator read from the database. + /// + /// - `current_positions`: each writer's position. + /// - `persisted_upto_position`: the lowest of those, or if there are + /// none, the highest ID in the stream's tables. + /// - `persisted_rows`: the `(instance, stream ID)` rows above that + /// position, sorted by stream ID. Our own rows also move our position. + /// Old rows can have no instance. + /// - `writers`: the configured writers. Each gets a position. With none, + /// any instance may write, so we start ourselves at the persisted + /// position. + /// + /// All IDs are positive (see the module docs). + pub fn load( + &self, + current_positions: HashMap, + persisted_upto_position: i64, + persisted_rows: &[(Option, i64)], + writers: &[String], + ) { + let mut state = self.lock(); + state.current_positions = current_positions; + state.persisted_upto_position = persisted_upto_position; + + for (instance, stream_id) in persisted_rows { + self.add_persisted_position(&mut state, *stream_id, false); + if instance.as_deref() == Some(self.instance_name.as_str()) { + state + .current_positions + .insert(self.instance_name.clone(), *stream_id); + } + } + + let persisted_upto_position = state.persisted_upto_position; + for writer in writers { + state + .current_positions + .entry(writer.clone()) + .or_insert(persisted_upto_position); + } + + // `stream_positions` may lag the tables, so the persisted position + // can be the higher one. + let max_seen = state + .current_positions + .values() + .copied() + .max() + .unwrap_or(1) + .max(state.persisted_upto_position); + state.max_seen_allocated_stream_id = max_seen; + state.max_position_of_local_instance = max_seen; + + if writers.is_empty() { + state + .current_positions + .insert(self.instance_name.clone(), persisted_upto_position); + } + } + + /// Start allocating IDs. Until [`StreamIdTracker::end_fetch`] the + /// current position stays at or below the returned value, the highest + /// ID seen so far, since the new IDs will be above it. + pub fn begin_fetch(&self) -> i64 { + let mut state = self.lock(); + let current_max = state.max_seen_allocated_stream_id; + *state.in_flight_fetches.entry(current_max).or_insert(0) += 1; + current_max + } + + /// Finish allocating IDs: `token` is what + /// [`StreamIdTracker::begin_fetch`] returned and `allocated` the IDs, or + /// `None` if the allocation failed. The IDs are unfinished until + /// [`StreamIdTracker::mark_finished`]. + pub fn end_fetch(&self, token: i64, allocated: Option<&[i64]>) { + let mut state = self.lock(); + if let Some(allocated) = allocated { + state.unfinished_ids.extend(allocated.iter().copied()); + if let Some(&max_unfinished) = state.unfinished_ids.last() { + state.max_seen_allocated_stream_id = + state.max_seen_allocated_stream_id.max(max_unfinished); + } + } + match state.in_flight_fetches.get_mut(&token) { + Some(count) if *count > 1 => *count -= 1, + Some(_) => { + state.in_flight_fetches.remove(&token); + } + None => panic!("end_fetch({token}) without a matching begin_fetch"), + } + } + + /// The transactions that wrote these IDs have finished, committed or + /// not. Moves the current position up if it can. + pub fn mark_finished(&self, ids: &[i64]) { + let log_issue_19795 = self.wants_issue_19795_log(); + let mut state = self.lock(); + for id in ids { + state.unfinished_ids.remove(id); + } + state.finished_ids.extend(ids.iter().copied()); + + let new_cur = if let Some(min_unfinished) = state.min_unfinished() { + // The new position is the largest finished ID below the smallest + // unfinished one. The rest stay finished. + let mut new_cur = None; + state.finished_ids.retain(|&s| { + if s < min_unfinished { + new_cur = new_cur.max(Some(s)); + false + } else { + true + } + }); + new_cur + } else { + // Nothing is unfinished, so the largest finished ID is the new + // position. + let new_cur = state.finished_ids.iter().copied().max(); + state.finished_ids.clear(); + new_cur + }; + + // Python skipped a new position of 0 as falsy. IDs start at 2, so + // that only matters for the type. + if let Some(new_cur) = new_cur.filter(|&c| c != 0) { + let curr = state + .current_positions + .get(&self.instance_name) + .copied() + .unwrap_or(0); + state + .current_positions + .insert(self.instance_name.clone(), curr.max(new_cur)); + state.max_position_of_local_instance = + state.max_position_of_local_instance.max(curr.max(new_cur)); + } + + let mut log_lines = Vec::new(); + for &id in ids { + if let Some(line) = self.add_persisted_position(&mut state, id, log_issue_19795) { + log_lines.push(line); + } + } + drop(state); + // Logging can take the GIL, so it must not happen under the lock: a + // Python thread holding the GIL may be waiting for it. + for line in log_lines { + debug!(target: ISSUE_19795_LOG_TARGET, "{line}"); + } + } + + /// Another writer has persisted up to `new_id` (signed). + pub fn advance(&self, instance_name: &str, new_id: i64) { + let log_issue_19795 = self.wants_issue_19795_log(); + let new_id = new_id * self.return_factor; + let mut state = self.lock(); + let position = state + .current_positions + .entry(instance_name.to_owned()) + .or_insert(0); + *position = (*position).max(new_id); + state.max_seen_allocated_stream_id = state.max_seen_allocated_stream_id.max(new_id); + let line = self.add_persisted_position(&mut state, new_id, log_issue_19795); + drop(state); + if let Some(line) = line { + debug!(target: ISSUE_19795_LOG_TARGET, "{line}"); + } + } + + /// The highest position such that everything up to it has been + /// persisted by every writer. Signed. + pub fn get_persisted_upto_position(&self) -> i64 { + self.return_factor * self.lock().persisted_upto_position + } + + /// The position of `instance_name`, signed. For a writer we have not + /// heard from yet, the persisted position: assuming it is new saves a + /// table scan when it announces itself. + pub fn get_current_token_for_writer(&self, instance_name: &str) -> i64 { + let state = self.lock(); + if self.instance_name == instance_name { + return self.return_factor * state.max_position_of_local_instance; + } + let pos = state + .current_positions + .get(instance_name) + .copied() + .unwrap_or(state.persisted_upto_position) + .max(state.persisted_upto_position); + self.return_factor * pos + } + + /// Our last finished write, or the persisted position if we have none. + /// Signed. + pub fn get_minimal_local_current_token(&self) -> i64 { + let state = self.lock(); + self.return_factor + * state + .current_positions + .get(&self.instance_name) + .copied() + .unwrap_or(state.persisted_upto_position) + } + + /// Each writer's position, signed. + pub fn get_positions(&self) -> HashMap { + self.lock() + .current_positions + .iter() + .map(|(name, &pos)| (name.clone(), self.return_factor * pos)) + .collect() + } + + /// The persisted position as a positive number, also for streams that + /// count down. For the `synapse_storage_stream_current_position` gauge. + pub fn persisted_upto_position_abs(&self) -> i64 { + self.lock().persisted_upto_position + } + + fn wants_issue_19795_log(&self) -> bool { + self.stream_name == "to_device" + && log::log_enabled!(target: ISSUE_19795_LOG_TARGET, log::Level::Debug) + } + + /// Record that `new_id` (positive) has been persisted, and move the + /// persisted position up as far as that allows. Returns the issue 19795 + /// debug line if `log_issue_19795` and it applies. + fn add_persisted_position( + &self, + state: &mut State, + new_id: i64, + log_issue_19795: bool, + ) -> Option { + state.known_persisted_positions.push(Reverse(new_id)); + + // Every position below the lowest writer's is persisted. While we + // have nothing in flight our own position does not count: our next + // write will be above every other writer's position. + let our_current_position = state + .current_positions + .get(&self.instance_name) + .copied() + .unwrap_or(0); + let mut min_curr = state + .current_positions + .iter() + .filter(|(name, _)| **name != self.instance_name) + .map(|(_, &token)| token) + .min() + .unwrap_or(our_current_position); + let in_flight = !state.unfinished_ids.is_empty() || !state.in_flight_fetches.is_empty(); + if our_current_position != 0 && in_flight { + min_curr = min_curr.min(our_current_position); + } + + state.persisted_upto_position = state.persisted_upto_position.max(min_curr); + state.max_position_of_local_instance = state + .max_position_of_local_instance + .max(state.persisted_upto_position); + + if !in_flight { + // Nothing in flight, so our next write will be above everything + // allocated so far. + state.max_position_of_local_instance = state + .max_position_of_local_instance + .max(state.max_seen_allocated_stream_id); + } + + // Drop known positions at or below the persisted position, and move + // it up through any run without gaps. + while let Some(&Reverse(lowest)) = state.known_persisted_positions.peek() { + if lowest <= state.persisted_upto_position { + state.known_persisted_positions.pop(); + } else if lowest == state.persisted_upto_position + 1 { + state.known_persisted_positions.pop(); + state.persisted_upto_position += 1; + } else { + break; + } + } + + // Debugging for https://github.com/element-hq/synapse/issues/19795, + // on the to-device stream of a writer. + if log_issue_19795 && our_current_position > 0 { + let mut known: Vec = state + .known_persisted_positions + .iter() + .map(|Reverse(p)| *p) + .collect(); + known.sort_unstable(); + Some(format!( + "stream_id={new_id} now persisted for stream={}; _current_positions={:?} \ + _unfinished_ids={:?}, _known_persisted_positions={known:?} \ + _persisted_upto_position={} min_curr={min_curr}", + self.stream_name, + state.current_positions, + state.unfinished_ids, + state.persisted_upto_position, + )) + } else { + None + } + } +} + +impl State { + /// The smallest ID the current position must stay below, if any. + fn min_unfinished(&self) -> Option { + // An in-flight fetch records the highest safe ID, so the smallest + // unsafe one is one above it. + let in_flight = self.in_flight_fetches.keys().next().map(|&max| max + 1); + let unfinished = self.unfinished_ids.first().copied(); + match (unfinished, in_flight) { + (Some(a), Some(b)) => Some(a.min(b)), + (a, b) => a.or(b), + } + } +} + +/// The Python face of a [`StreamIdTracker`]. Shared with Rust code as an +/// `Arc` through [`PyStreamIdTracker::tracker`]. +#[pyclass(frozen, name = "StreamIdTracker")] +pub struct PyStreamIdTracker { + tracker: std::sync::Arc, +} + +impl PyStreamIdTracker { + pub fn tracker(&self) -> std::sync::Arc { + std::sync::Arc::clone(&self.tracker) + } +} + +#[pymethods] +impl PyStreamIdTracker { + #[new] + #[pyo3(signature = (stream_name, instance_name, positive))] + fn py_new(stream_name: String, instance_name: String, positive: bool) -> Self { + PyStreamIdTracker { + tracker: std::sync::Arc::new(StreamIdTracker::new( + stream_name, + instance_name, + positive, + )), + } + } + + fn load( + &self, + current_positions: HashMap, + persisted_upto_position: i64, + persisted_rows: Vec<(Option, i64)>, + writers: Vec, + ) { + self.tracker.load( + current_positions, + persisted_upto_position, + &persisted_rows, + &writers, + ) + } + + fn begin_fetch(&self) -> i64 { + self.tracker.begin_fetch() + } + + #[pyo3(signature = (token, allocated))] + fn end_fetch(&self, token: i64, allocated: Option>) { + self.tracker.end_fetch(token, allocated.as_deref()) + } + + fn mark_finished(&self, ids: Vec) { + self.tracker.mark_finished(&ids) + } + + fn advance(&self, instance_name: &str, new_id: i64) { + self.tracker.advance(instance_name, new_id) + } + + fn get_persisted_upto_position(&self) -> i64 { + self.tracker.get_persisted_upto_position() + } + + fn get_current_token_for_writer(&self, instance_name: &str) -> i64 { + self.tracker.get_current_token_for_writer(instance_name) + } + + fn get_minimal_local_current_token(&self) -> i64 { + self.tracker.get_minimal_local_current_token() + } + + fn get_positions(&self) -> HashMap { + self.tracker.get_positions() + } + + fn persisted_upto_position_abs(&self) -> i64 { + self.tracker.persisted_upto_position_abs() + } +} + +/// Called when registering modules with python. +pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { + let child_module = PyModule::new(py, "stream_ids")?; + child_module.add_class::()?; + + m.add_submodule(&child_module)?; + + // We need to manually add the module to sys.modules to make `from + // synapse.synapse_rust import stream_ids` work. + py.import("sys")? + .getattr("modules")? + .set_item("synapse.synapse_rust.stream_ids", child_module)?; + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn tracker(writers: &[&str]) -> StreamIdTracker { + let tracker = StreamIdTracker::new("events".into(), "first".into(), true); + let writers: Vec = writers.iter().map(|w| w.to_string()).collect(); + let positions = writers.iter().map(|w| (w.clone(), 7)).collect(); + tracker.load(positions, 7, &[], &writers); + tracker + } + + fn allocate(tracker: &StreamIdTracker, ids: &[i64]) { + let token = tracker.begin_fetch(); + tracker.end_fetch(token, Some(ids)); + } + + #[test] + fn finishing_out_of_order_waits_for_the_gap() { + let tracker = tracker(&["first"]); + allocate(&tracker, &[8]); + allocate(&tracker, &[9]); + + tracker.mark_finished(&[9]); + assert_eq!(tracker.get_persisted_upto_position(), 7); + assert_eq!(tracker.get_current_token_for_writer("first"), 7); + + tracker.mark_finished(&[8]); + assert_eq!(tracker.get_persisted_upto_position(), 9); + assert_eq!(tracker.get_current_token_for_writer("first"), 9); + } + + #[test] + fn in_flight_fetches_count_each_start() { + // Two fetches that start at the same maximum: the position must stay + // pinned until both have ended. + let tracker = tracker(&["first"]); + let a = tracker.begin_fetch(); + let b = tracker.begin_fetch(); + assert_eq!(a, b); + + tracker.end_fetch(a, Some(&[8])); + tracker.mark_finished(&[8]); + // Our own position waits for `b`. (The persisted position moves to + // 8, as 8 is the next ID after 7 and IDs are never reused.) + assert_eq!(tracker.get_minimal_local_current_token(), 7); + assert_eq!(tracker.get_persisted_upto_position(), 8); + + tracker.end_fetch(b, Some(&[9])); + tracker.mark_finished(&[9]); + assert_eq!(tracker.get_persisted_upto_position(), 9); + } + + #[test] + fn failed_fetch_releases_the_position() { + let tracker = tracker(&["first"]); + let token = tracker.begin_fetch(); + tracker.end_fetch(token, None); + allocate(&tracker, &[8]); + tracker.mark_finished(&[8]); + assert_eq!(tracker.get_persisted_upto_position(), 8); + } + + #[test] + fn unfinished_ids_hold_back_the_persisted_position() { + let tracker = StreamIdTracker::new("events".into(), "first".into(), true); + let writers = vec!["first".to_string(), "second".to_string()]; + tracker.load( + [("first".to_string(), 3), ("second".to_string(), 5)].into(), + 3, + &[], + &writers, + ); + assert_eq!(tracker.get_persisted_upto_position(), 3); + + // With nothing in flight, our own position does not hold the + // persisted position back. + tracker.advance("second", 6); + assert_eq!(tracker.get_persisted_upto_position(), 6); + + tracker.advance("second", 8); + allocate(&tracker, &[9]); + tracker.advance("second", 11); + // Our unfinished 9 counts our own position (still 3), so the + // persisted position stays where it was. + assert_eq!(tracker.get_persisted_upto_position(), 8); + tracker.mark_finished(&[9]); + assert_eq!(tracker.get_persisted_upto_position(), 11); + } + + #[test] + fn negative_streams_return_negative_ids() { + let tracker = StreamIdTracker::new("backfill".into(), "first".into(), false); + let writers = vec!["first".to_string()]; + tracker.load([("first".to_string(), 3)].into(), 3, &[], &writers); + allocate(&tracker, &[4]); + tracker.mark_finished(&[4]); + assert_eq!(tracker.get_persisted_upto_position(), -4); + assert_eq!(tracker.get_positions()["first"], -4); + assert_eq!(tracker.persisted_upto_position_abs(), 4); + + tracker.advance("other", -6); + assert_eq!(tracker.get_current_token_for_writer("other"), -6); + } + + #[test] + fn load_moves_our_position_to_our_later_rows() { + let tracker = StreamIdTracker::new("events".into(), "first".into(), true); + let writers = vec!["first".to_string(), "second".to_string()]; + tracker.load( + [("first".to_string(), 3), ("second".to_string(), 5)].into(), + 3, + &[ + (Some("first".to_string()), 4), + (None, 5), + (Some("first".to_string()), 6), + ], + &writers, + ); + assert_eq!(tracker.get_positions()["first"], 6); + assert_eq!(tracker.get_persisted_upto_position(), 6); + assert_eq!(tracker.get_current_token_for_writer("first"), 6); + } +} diff --git a/synapse/storage/util/id_generators.py b/synapse/storage/util/id_generators.py index d370f2f815..e215f7d4b7 100644 --- a/synapse/storage/util/id_generators.py +++ b/synapse/storage/util/id_generators.py @@ -20,7 +20,6 @@ # # import abc -import heapq import logging import threading from types import TracebackType @@ -38,9 +37,7 @@ from typing import ( import attr from immutabledict import immutabledict from prometheus_client import Gauge -from sortedcontainers import SortedList, SortedSet -from synapse.logging import issue9533_logger from synapse.metrics import SERVER_NAME_LABEL from synapse.metrics.background_process_metrics import run_as_background_process from synapse.storage.database import ( @@ -52,6 +49,7 @@ from synapse.storage.database import ( from synapse.storage.engines import BaseDatabaseEngine, PostgresEngine, Sqlite3Engine from synapse.storage.types import Cursor from synapse.storage.util.sequence import build_sequence_generator +from synapse.synapse_rust.stream_ids import StreamIdTracker from synapse.types import MultiWriterStreamToken if TYPE_CHECKING: @@ -246,76 +244,15 @@ class MultiWriterIdGenerator(AbstractStreamIdGenerator): self._writers = writers self._return_factor = 1 if positive else -1 - # The `instance_name` label is what separates one worker's metrics from - # another's when there are multiple Synapse workers for the same homeserver - # running in the same process. In production the scrape labels would do that - # too, but tests run several homeservers in one process. - self._current_position_gauge = stream_current_position_gauge.labels( - stream_name=stream_name, - **{SERVER_NAME_LABEL: server_name}, - ) - - # We lock as some functions may be called from DB threads. - self._lock = threading.Lock() - - # Note: If we are a negative stream then we still store all the IDs as - # positive to make life easier for us, and simply negate the IDs when we - # return them. - self._current_positions: dict[str, int] = {} - - # Set of local IDs that we're still processing. The current position - # should be less than the minimum of this set (if not empty). - self._unfinished_ids: SortedSet[int] = SortedSet() - - # We also need to track when we've requested some new stream IDs but - # they haven't yet been added to the `_unfinished_ids` set. Every time - # we request a new stream ID we add the current max stream ID to the - # list, and remove it once we've added the newly allocated IDs to the - # `_unfinished_ids` set. This means that we *may* be allocated stream - # IDs above those in the list, and so we can't advance the local current - # position beyond the minimum stream ID in this list. - self._in_flight_fetches: SortedList[int] = SortedList() - - # Set of local IDs that we've processed that are larger than the current - # position, due to there being smaller unpersisted IDs. - self._finished_ids: set[int] = set() - - # We track the max position where we know everything before has been - # persisted. This is done by a) looking at the min across all instances - # and b) noting that if we have seen a run of persisted positions - # without gaps (e.g. 5, 6, 7) then we can skip forward (e.g. to 7). - # - # Note: There is no guarantee that the IDs generated by the sequence - # will be gapless; gaps can form when e.g. a transaction was rolled - # back. This means that sometimes we won't be able to skip forward the - # position even though everything has been persisted. However, since - # gaps should be relatively rare it's still worth doing the book keeping - # that allows us to skip forwards when there are gapless runs of - # positions. - # - # We start at 1 here as a) the first generated stream ID will be 2, and - # b) other parts of the code assume that stream IDs are strictly greater - # than 0. - self._persisted_upto_position = ( - min(self._current_positions.values()) if self._current_positions else 1 - ) - self._known_persisted_positions: list[int] = [] - - # The maximum stream ID that we have seen been allocated across any writer. - # Since this defaults to 1, this means that ID 1 is assumed to have already - # been 'seen'. In other words, multi-writer streams start at 2. - # Unclear if this is desirable behaviour. - self._max_seen_allocated_stream_id = 1 - - # The maximum position of the local instance. This can be higher than - # the corresponding position in `current_positions` table when there are - # no active writes in progress. - self._max_position_of_local_instance = self._max_seen_allocated_stream_id + # The bookkeeping of which IDs are allocated, finished and persisted, + # by us and by the other writers. Rust code allocates IDs through it + # too. See `synapse_rust.stream_ids`. + self._tracker = StreamIdTracker(stream_name, instance_name, positive) self._sequence_gen = build_sequence_generator( db_conn=db_conn, database_engine=db.engine, - get_first_callback=lambda _: self._persisted_upto_position, + get_first_callback=lambda _: self._tracker.persisted_upto_position_abs(), sequence_name=sequence_name, # We only need to set the below if we want it to call # `check_consistency`, but we do that ourselves below so we can @@ -336,46 +273,25 @@ class MultiWriterIdGenerator(AbstractStreamIdGenerator): positive=positive, ) - # This goes and fills out the above state from the database. + # This goes and fills out the tracker from the database. # This may read on the PostgreSQL sequence, and # SequenceGenerator.check_consistency might have fixed up the sequence, which # means the SequenceGenerator needs to be setup before we read the value from # the sequence. self._load_current_ids(db_conn, tables, sequence_name) - self._max_seen_allocated_stream_id = max( - self._current_positions.values(), default=1 - ) - - # For the case where `stream_positions` is not up to date, - # `_persisted_upto_position` may be higher. - self._max_seen_allocated_stream_id = max( - self._max_seen_allocated_stream_id, self._persisted_upto_position - ) - - # Bump our local maximum position now that we've loaded things from the - # DB. - self._max_position_of_local_instance = self._max_seen_allocated_stream_id - - if not writers: - # If there have been no explicit writers given then any instance can - # write to the stream. In which case, let's pre-seed our own - # position with the current minimum. - self._current_positions[self._instance_name] = self._persisted_upto_position - - with self._lock: - self._report_current_position() - - def _report_current_position(self) -> None: - """Report the current position as a metric. - - Must be called with `_lock` held. - """ - assert self._lock.locked() - - # Note that we always report this as a positive value, so we don't - # multiply by the `_return_factor`. - self._current_position_gauge.set(self._persisted_upto_position) + # The `instance_name` label is what separates one worker's metrics from + # another's when there are multiple Synapse workers for the same homeserver + # running in the same process. In production the scrape labels would do that + # too, but tests run several homeservers in one process. + # + # The gauge reads the tracker when scraped, rather than being set on + # every change, since Rust code changes it too. Note that it is always + # positive, even for negative streams. + stream_current_position_gauge.labels( + stream_name=stream_name, + **{SERVER_NAME_LABEL: server_name}, + ).set_function(self._tracker.persisted_upto_position_abs) def _load_current_ids( self, @@ -385,7 +301,9 @@ class MultiWriterIdGenerator(AbstractStreamIdGenerator): ) -> None: cur = db_conn.cursor(txn_name="_load_current_ids") - # Load the current positions of all writers for the stream. + # Load the current positions of all writers for the stream. (Kept + # positive, as the tracker stores them, also for negative streams.) + current_positions: dict[str, int] = {} if self._writers: # We delete any stale entries in the positions table. This is # important if we add back a writer after a long time; we want to @@ -409,7 +327,7 @@ class MultiWriterIdGenerator(AbstractStreamIdGenerator): """ cur.execute(sql, (self._stream_name,)) - self._current_positions = { + current_positions = { instance: stream_id * self._return_factor for instance, stream_id in cur if instance in self._writers @@ -425,12 +343,13 @@ class MultiWriterIdGenerator(AbstractStreamIdGenerator): cur.execute(f"SELECT last_value FROM {sequence_name}") row = cur.fetchone() assert row is not None - self._current_positions[self._instance_name] = row[0] + current_positions[self._instance_name] = row[0] - # We set the `_persisted_upto_position` to be the minimum of all current + # We set the persisted position to be the minimum of all current # positions. If empty we use the max stream ID from the DB table. - min_stream_id = min(self._current_positions.values(), default=None) + min_stream_id = min(current_positions.values(), default=None) + rows: list[tuple[str | None, int]] = [] if min_stream_id is None: # We add a GREATEST here to ensure that the result is always # positive. (This can be a problem for e.g. backfill streams where @@ -456,22 +375,19 @@ class MultiWriterIdGenerator(AbstractStreamIdGenerator): max_stream_id = max(max_stream_id, stream_id) - self._persisted_upto_position = max_stream_id + persisted_upto_position = max_stream_id else: # If we have a min_stream_id then we pull out everything greater - # than it from the DB so that we can prefill - # `_known_persisted_positions` and get a more accurate - # `_persisted_upto_position`. + # than it from the DB so that the tracker can note them as + # persisted, and get a more accurate persisted position. # - # We also check if any of the later rows are from this instance, in - # which case we use that for this instance's current position. This - # is to handle the case where we didn't finish persisting to the - # stream positions table before restart (or the stream position - # table otherwise got out of date). + # The tracker also checks if any of the later rows are from this + # instance, in which case it uses that for this instance's current + # position. This is to handle the case where we didn't finish + # persisting to the stream positions table before restart (or the + # stream position table otherwise got out of date). + persisted_upto_position = min_stream_id - self._persisted_upto_position = min_stream_id - - rows: list[tuple[str, int]] = [] for table, instance_column, id_column in tables: sql = """ SELECT %(instance)s, %(id)s FROM %(table)s @@ -485,13 +401,16 @@ class MultiWriterIdGenerator(AbstractStreamIdGenerator): cur.execute(sql, (min_stream_id * self._return_factor,)) # Cast safety: this corresponds to the types returned by the query above. - rows.extend(cast(Iterable[tuple[str, int]], cur)) + rows.extend( + (instance, self._return_factor * stream_id) + for instance, stream_id in cast(Iterable[tuple[str | None, int]], cur) + ) # Sort by stream_id (ascending, lowest -> highest) so that we handle # rows in order for each instance because we don't want to overwrite # the current_position of an instance to a lower stream ID than # we're actually at. - def sort_by_stream_id_key_func(row: tuple[str, int]) -> int: + def sort_by_stream_id_key_func(row: tuple[str | None, int]) -> int: (instance, stream_id) = row # If `stream_id` is ever `None`, we will see a `TypeError: '<' # not supported between instances of 'NoneType' and 'X'` error. @@ -499,27 +418,12 @@ class MultiWriterIdGenerator(AbstractStreamIdGenerator): rows.sort(key=sort_by_stream_id_key_func) - with self._lock: - for ( - instance, - stream_id, - ) in rows: - stream_id = self._return_factor * stream_id - self._add_persisted_position(stream_id) - - if instance == self._instance_name: - self._current_positions[instance] = stream_id - - if self._writers: - # If we have explicit writers then make sure that each instance has - # a position. - for writer in self._writers: - self._current_positions.setdefault( - writer, self._persisted_upto_position - ) - cur.close() + self._tracker.load( + current_positions, persisted_upto_position, rows, self._writers + ) + def _load_next_id_txn(self, txn: Cursor) -> int: stream_ids = self._load_next_mult_id_txn(txn, 1) return stream_ids[0] @@ -527,24 +431,15 @@ class MultiWriterIdGenerator(AbstractStreamIdGenerator): def _load_next_mult_id_txn(self, txn: Cursor, n: int) -> list[int]: # We need to track that we've requested some more stream IDs, and what # the current max allocated stream ID is. This is to prevent a race - # where we've been allocated stream IDs but they have not yet been added - # to the `_unfinished_ids` set, allowing the current position to advance + # where we've been allocated stream IDs but they have not yet been + # recorded as unfinished, allowing the current position to advance # past them. - with self._lock: - current_max = self._max_seen_allocated_stream_id - self._in_flight_fetches.add(current_max) - + token = self._tracker.begin_fetch() + stream_ids = None try: stream_ids = self._sequence_gen.get_next_mult_txn(txn, n) - - with self._lock: - self._unfinished_ids.update(stream_ids) - self._max_seen_allocated_stream_id = max( - self._max_seen_allocated_stream_id, self._unfinished_ids[-1] - ) finally: - with self._lock: - self._in_flight_fetches.remove(current_max) + self._tracker.end_fetch(token, stream_ids) return stream_ids @@ -671,62 +566,7 @@ class MultiWriterIdGenerator(AbstractStreamIdGenerator): """These IDs have finished being processed so we should advance the current position if possible. """ - - with self._lock: - self._unfinished_ids.difference_update(next_ids) - self._finished_ids.update(next_ids) - - new_cur: int | None = None - - if self._unfinished_ids or self._in_flight_fetches: - # If there are unfinished IDs then the new position will be the - # largest finished ID strictly less than the minimum unfinished - # ID. - - # The minimum unfinished ID needs to take account of both - # `_unfinished_ids` and `_in_flight_fetches`. - if self._unfinished_ids and self._in_flight_fetches: - # `_in_flight_fetches` stores the maximum safe stream ID, so - # we add one to make it equivalent to the minimum unsafe ID. - min_unfinished = min( - self._unfinished_ids[0], self._in_flight_fetches[0] + 1 - ) - elif self._in_flight_fetches: - min_unfinished = self._in_flight_fetches[0] + 1 - else: - min_unfinished = self._unfinished_ids[0] - - finished = set() - for s in self._finished_ids: - if s < min_unfinished: - if new_cur is None or new_cur < s: - new_cur = s - else: - finished.add(s) - - # We clear these out since they're now all less than the new - # position. - self._finished_ids = finished - else: - # There are no unfinished IDs so the new position is simply the - # largest finished one. - new_cur = max(self._finished_ids) - - # We clear these out since they're now all less than the new - # position. - self._finished_ids.clear() - - if new_cur: - curr = self._current_positions.get(self._instance_name, 0) - self._current_positions[self._instance_name] = max(curr, new_cur) - self._max_position_of_local_instance = max( - curr, new_cur, self._max_position_of_local_instance - ) - - # TODO Can we call this for just the last position or somehow batch - # _add_persisted_position. - for next_id in next_ids: - self._add_persisted_position(next_id) + self._tracker.mark_finished(next_ids) def get_current_token(self) -> int: return self.get_persisted_upto_position() @@ -738,26 +578,10 @@ class MultiWriterIdGenerator(AbstractStreamIdGenerator): # For new writers we assume their initial position to be the current # persisted up to position. This stops Synapse from doing a full table # scan when a new writer announces itself over replication. - with self._lock: - if self._instance_name == instance_name: - return self._return_factor * self._max_position_of_local_instance - - pos = self._current_positions.get( - instance_name, self._persisted_upto_position - ) - - # We want to return the maximum "current token" that we can for a - # writer, this helps ensure that streams progress as fast as - # possible. - pos = max(pos, self._persisted_upto_position) - - return self._return_factor * pos + return self._tracker.get_current_token_for_writer(instance_name) def get_minimal_local_current_token(self) -> int: - with self._lock: - return self._return_factor * self._current_positions.get( - self._instance_name, self._persisted_upto_position - ) + return self._tracker.get_minimal_local_current_token() def get_positions(self) -> dict[str, int]: """Get a copy of the current positon map. @@ -765,26 +589,10 @@ class MultiWriterIdGenerator(AbstractStreamIdGenerator): Note that this won't necessarily include all configured writers if some writers haven't written anything yet. """ - - with self._lock: - return { - name: self._return_factor * i - for name, i in self._current_positions.items() - } + return self._tracker.get_positions() def advance(self, instance_name: str, new_id: int) -> None: - new_id *= self._return_factor - - with self._lock: - self._current_positions[instance_name] = max( - new_id, self._current_positions.get(instance_name, 0) - ) - - self._max_seen_allocated_stream_id = max( - self._max_seen_allocated_stream_id, new_id - ) - - self._add_persisted_position(new_id) + self._tracker.advance(instance_name, new_id) def get_persisted_upto_position(self) -> int: """Get the max position where all previous positions have been @@ -794,96 +602,7 @@ class MultiWriterIdGenerator(AbstractStreamIdGenerator): position across writers. This means that the returned position here can lag if one writer doesn't write very often. """ - - with self._lock: - return self._return_factor * self._persisted_upto_position - - def _add_persisted_position(self, new_id: int) -> None: - """Record that we have persisted a position. - - This is used to keep the `_current_positions` up to date. - """ - - # We require that the lock is locked by caller - assert self._lock.locked() - - heapq.heappush(self._known_persisted_positions, new_id) - - # We move the current min position up if the minimum current positions - # of all instances is higher (since by definition all positions less - # that that have been persisted). - # - # If we are one of several writers, then we don't need to factor our own - # `_current_position` into `_persisted_upto_position` unless we have unfinished - # writes (since we know that any future write that happens locally will have - # a higher stream ID than any of the other writers' current positions). In other - # words, when we have no outstanding writes, then the new `_persisted_upto_position` - # can be the minimum of all *other* writers' current positions, - # - our_current_position = self._current_positions.get(self._instance_name, 0) - min_curr = min( - ( - token - for name, token in self._current_positions.items() - if name != self._instance_name - ), - default=our_current_position, - ) - if our_current_position and (self._unfinished_ids or self._in_flight_fetches): - min_curr = min(min_curr, our_current_position) - - self._persisted_upto_position = max(min_curr, self._persisted_upto_position) - - # Advance our local max position. - self._max_position_of_local_instance = max( - self._max_position_of_local_instance, self._persisted_upto_position - ) - - if not self._unfinished_ids and not self._in_flight_fetches: - # If we don't have anything in flight, it's safe to advance to the - # max seen stream ID. - self._max_position_of_local_instance = max( - self._max_seen_allocated_stream_id, self._max_position_of_local_instance - ) - - # We now iterate through the seen positions, discarding those that are - # less than the current min positions, and incrementing the min position - # if its exactly one greater. - # - # This is also where we discard items from `_known_persisted_positions` - # (to ensure the list doesn't infinitely grow). - while self._known_persisted_positions: - if self._known_persisted_positions[0] <= self._persisted_upto_position: - heapq.heappop(self._known_persisted_positions) - elif ( - self._known_persisted_positions[0] == self._persisted_upto_position + 1 - ): - heapq.heappop(self._known_persisted_positions) - self._persisted_upto_position += 1 - else: - # There was a gap in seen positions, so there is nothing more to - # do. - break - - self._report_current_position() - - # Hacky debug logging to attempt to trace https://github.com/element-hq/synapse/issues/19795. - # If this is the to-device stream, and we are a writer for that stream, log some stats - if ( - issue9533_logger.isEnabledFor(logging.DEBUG) - and our_current_position > 0 - and self._stream_name == "to_device" - ): - issue9533_logger.debug( - "stream_id=%i now persisted for stream=%s; _current_positions=%s _unfinished_ids=%s, _known_persisted_positions=%s _persisted_upto_position=%i min_curr=%i", - new_id, - self._stream_name, - self._current_positions, - self._unfinished_ids, - self._known_persisted_positions, - self._persisted_upto_position, - min_curr, - ) + return self._tracker.get_persisted_upto_position() def _update_stream_positions_table_txn(self, txn: Cursor) -> None: """Update the `stream_positions` table with newly persisted position.""" @@ -973,7 +692,7 @@ class _MultiWriterCtxManager: # after the transaction allocated the IDs but before we returned, # then `__aexit__` will never run, because Python only invokes it # once `__aenter__` has returned. The allocated IDs would then be - # leaked into `_unfinished_ids` forever, permanently pinning the + # leaked into the unfinished IDs forever, permanently pinning the # persisted stream position and, e.g., wedging presence. # # So mark them as finished here to unblock the position. This mirrors diff --git a/synapse/synapse_rust/stream_ids.pyi b/synapse/synapse_rust/stream_ids.pyi new file mode 100644 index 0000000000..d391a692fa --- /dev/null +++ b/synapse/synapse_rust/stream_ids.pyi @@ -0,0 +1,54 @@ +# +# 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: +# . + +class StreamIdTracker: + """The in-memory bookkeeping of a multi-writer stream, for + `MultiWriterIdGenerator`: which IDs this instance has allocated and not + finished, each writer's position, and from those the current token. + + IDs are stored positive, also for streams that count down. `advance` and + the `get_*` methods take and return signed IDs. The bookkeeping methods + (`load`, `begin_fetch`, `end_fetch`, `mark_finished`) take positive IDs, + as the sequence returns them. + """ + + def __init__( + self, stream_name: str, instance_name: str, positive: bool + ) -> None: ... + def load( + self, + current_positions: dict[str, int], + persisted_upto_position: int, + persisted_rows: list[tuple[str | None, int]], + writers: list[str], + ) -> None: + """Set the start state from the database: each writer's position, the + persisted position, the rows above it sorted by stream ID, and the + configured writers.""" + + def begin_fetch(self) -> int: + """Start allocating IDs. Returns a token for `end_fetch`.""" + + def end_fetch(self, token: int, allocated: list[int] | None) -> None: + """Finish allocating IDs: the IDs allocated, or None on failure.""" + + def mark_finished(self, ids: list[int]) -> None: + """The transactions that wrote these IDs have finished.""" + + def advance(self, instance_name: str, new_id: int) -> None: ... + def get_persisted_upto_position(self) -> int: ... + def get_current_token_for_writer(self, instance_name: str) -> int: ... + def get_minimal_local_current_token(self) -> int: ... + def get_positions(self) -> dict[str, int]: ... + def persisted_upto_position_abs(self) -> int: + """The persisted position, positive also for streams that count down.""" diff --git a/tests/storage/test_id_generators.py b/tests/storage/test_id_generators.py index 8eb2befe9c..18384ec4a3 100644 --- a/tests/storage/test_id_generators.py +++ b/tests/storage/test_id_generators.py @@ -270,9 +270,9 @@ class MultiWriterIdGeneratorTestCase(MultiWriterIdGeneratorBase): If the `get_next()` async context manager is cancelled while `__aenter__` is allocating a stream ID, the DB interaction that runs the - sequence has already added the ID to `_unfinished_ids`, but `__aexit__` + sequence has already recorded the ID as unfinished, but `__aexit__` is never called (Python only invokes `__aexit__` if `__aenter__` - returned). The abandoned ID is therefore leaked into `_unfinished_ids` + returned). The abandoned ID is therefore leaked into the unfinished IDs forever, which permanently pins the persisted stream position: new rows keep getting higher IDs, but `get_current_token()` can never advance past `leaked_id - 1` until the process restarts. @@ -291,7 +291,7 @@ class MultiWriterIdGeneratorTestCase(MultiWriterIdGeneratorBase): # We model the cancellation at the seam it actually happens in # production: `__aenter__` awaits `runInteraction("_load_next_mult_id")`, # whose transaction runs in a thread pool and so *always* completes - - # allocating stream ID 8 and adding it to `_unfinished_ids` - but the + # allocating stream ID 8 and recording it as unfinished - but the # awaiting coroutine is handed a `CancelledError` because the enclosing # `/sync` request was cancelled. We reproduce that by letting the real # interaction run (applying its side effects) and then failing the @@ -302,7 +302,7 @@ class MultiWriterIdGeneratorTestCase(MultiWriterIdGeneratorBase): async def blocking_run_interaction(desc, func, *args, **kwargs): # type: ignore[no-untyped-def] result = await original_run_interaction(desc, func, *args, **kwargs) if desc == "_load_next_mult_id": - # Stream ID 8 is now allocated and recorded in `_unfinished_ids`. + # Stream ID 8 is now allocated and recorded as unfinished. # Deliver the cancellation here, exactly as a cancelled `/sync` # would land it on this `await`. await make_deferred_yieldable(cancel_enter) @@ -321,7 +321,7 @@ class MultiWriterIdGeneratorTestCase(MultiWriterIdGeneratorBase): write = ensureDeferred(presence_like_write()) # The write is now blocked inside `__aenter__`, i.e. after stream ID - # 8 has been allocated and added to `_unfinished_ids`. + # 8 has been allocated and recorded as unfinished. self.assertNoResult(write) # The client goes away and the `/sync` request is cancelled. @@ -340,7 +340,7 @@ class MultiWriterIdGeneratorTestCase(MultiWriterIdGeneratorBase): self.get_success(_successful_write()) # On the buggy code the token is still stuck at 7 (ID 8 is leaked in - # `_unfinished_ids`, blocking everything behind it). Once the leak is + # the unfinished IDs, blocking everything behind it). Once the leak is # fixed, the token advances to 9: ID 8 was allocated (and abandoned) by # the cancelled write, so the successful write above takes ID 9. self.assertEqual( @@ -893,6 +893,41 @@ class BackwardsMultiWriterIdGeneratorTestCase(MultiWriterIdGeneratorBase): self.assertEqual(id_gen_1.get_persisted_upto_position(), -2) self.assertEqual(id_gen_2.get_persisted_upto_position(), -2) + def test_restart_with_own_rows_past_stream_positions(self) -> None: + """On restart, rows of our own past `stream_positions` move our + position to the furthest of them, not the nearest. + + The loader once sorted the rows by their stored, negative stream ID + before negating them, so it visited ours furthest first and left our + position at the nearest. + """ + + def _insert(txn: LoggingTransaction) -> None: + # The table and `stream_positions` hold the negative IDs, the + # sequence the positive ones. + for stream_id, instance_name in ( + (-1, "second"), + (-2, "first"), + (-3, "first"), + ): + txn.execute( + "INSERT INTO foobar (stream_id, instance_name) VALUES (?, ?)", + (stream_id, instance_name), + ) + txn.execute("SELECT setval('foobar_seq', 3)") + # `second` is at 1, so the loader reads every row from 1 on. Our + # rows at 2 and 3 never reached `stream_positions`. + txn.execute( + "INSERT INTO stream_positions VALUES ('test_stream', 'second', -1)" + ) + + self.get_success(self.db_pool.runInteraction("_insert", _insert)) + + id_gen = self._create_id_generator("first", writers=["first", "second"]) + + self.assertEqual(id_gen.get_positions()["first"], -3) + self.assertEqual(id_gen.get_persisted_upto_position(), -3) + class MultiTableMultiWriterIdGeneratorTestCase(MultiWriterIdGeneratorBase): if not USE_POSTGRES_FOR_TESTS: