Move the multi-writer stream ID bookkeeping to Rust

`MultiWriterIdGenerator` tracks in memory which IDs this instance has
allocated and not finished, and each writer's position, and derives the
current token from them. That state now lives in a Rust `StreamIdTracker`,
so that Rust code can allocate and finish stream IDs without the GIL. The
generator keeps what touches the database or the reactor: loading the start
positions, the sequence, the `stream_positions` table and notifying
replication.

The logic is a line-by-line port. Two changes in form: the in-flight fetches
are a counted map (they were a sorted multiset), and the position gauge reads
the tracker when scraped instead of being set on each change. The issue 19795
debug line is logged after the lock is released, since logging can take the
GIL.

One behaviour change, a fix: on restart, the loader sorted the rows past
`stream_positions` by their stored ID before negating them. On the backfill
stream that visited our own rows furthest first, so our position ended at
the nearest of them. The rows are now negated first.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UHKYPSc37kWqGACn9aKxKL

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.

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)?;

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;

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:
+ * <https://www.gnu.org/licenses/agpl-3.0.html>.
+ *
+ */
+
+//! 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<State>,
+}
+
+#[derive(Debug)]
+struct State {
+    /// Each writer's position. Includes our own once we have written.
+    current_positions: HashMap<String, i64>,
+
+    /// IDs this instance has allocated and not yet finished. The current
+    /// position must stay below the smallest.
+    unfinished_ids: BTreeSet<i64>,
+
+    /// 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<i64, usize>,
+
+    /// IDs this instance has finished that are above the current position,
+    /// because a lower ID is still unfinished.
+    finished_ids: HashSet<i64>,
+
+    /// 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<Reverse<i64>>,
+
+    /// 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<String, i64>,
+        persisted_upto_position: i64,
+        persisted_rows: &[(Option<String>, 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<String, i64> {
+        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<String> {
+        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<i64> = 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<i64> {
+        // 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<StreamIdTracker>,
+}
+
+impl PyStreamIdTracker {
+    pub fn tracker(&self) -> std::sync::Arc<StreamIdTracker> {
+        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<String, i64>,
+        persisted_upto_position: i64,
+        persisted_rows: Vec<(Option<String>, i64)>,
+        writers: Vec<String>,
+    ) {
+        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<Vec<i64>>) {
+        self.tracker.end_fetch(token, allocated.as_deref())
+    }
+
+    fn mark_finished(&self, ids: Vec<i64>) {
+        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<String, i64> {
+        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::<PyStreamIdTracker>()?;
+
+    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<String> = 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);
+    }
+}

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

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:
+# <https://www.gnu.org/licenses/agpl-3.0.html>.
+
+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."""

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: