Calculate the forward extremities in Rust

Port `_calculate_new_extremities`, `_calculate_new_state_dag_extremities`
and `_prune_extremities` from the persist events controller to
`rust/src/persist/extremities.rs`. They read through an `ExtremityReads`
trait, which the `PersistReader` implements and the unit tests implement
over an in-memory DAG. Metrics go to a `PersistCounts` the caller owns,
so shadow mode can drop them.

Missing prev state events fail with `MissingPrevStateEvents`, which
`persist_error_to_py` turns into the same 500 `SynapseError` as Python.

`PersistReaderForTest` gains the three calculations, and
`test_persist_reader.py` compares each with the Python method on the
same database and the same new events.

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

rust/src/events/mod.rs

@@ -727,6 +727,15 @@ impl Event {
         self.parsed_event.common_fields.state_key.as_deref_opt()
     }
 
+    /// The event's MSC4242 `prev_state_events`, or `None` if its format
+    /// has none.
+    pub fn prev_state_event_ids(&self) -> Option<&[String]> {
+        match &*self.parsed_event.specific_fields {
+            EventFormatEnum::VMSC4242(format) => Some(&format.prev_state_events),
+            _ => None,
+        }
+    }
+
     /// The event's `content`.
     pub fn content(&self) -> &JsonObject {
         &self.parsed_event.common_fields.content

rust/src/persist/context.rs

@@ -86,6 +86,16 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PersistContext {
     }
 }
 
+impl PersistContext {
+    /// Whether the event was rejected, as Python's `if context.rejected`
+    /// reads it: an empty reason does not count.
+    pub fn is_rejected(&self) -> bool {
+        self.rejected
+            .as_deref()
+            .is_some_and(|reason| !reason.is_empty())
+    }
+}
+
 /// Convert a dict of `(type, state_key)` to event ID.
 fn state_map_from_py(dict: &Bound<'_, PyDict>) -> PyResult<StateMap> {
     let mut map = StateMap::with_capacity(dict.len());

rust/src/persist/extremities.rs

@@ -0,0 +1,995 @@
+/*
+ * 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 room's forward extremities after a chunk of events: ports of
+//! `_calculate_new_extremities`, `_calculate_new_state_dag_extremities` and
+//! `_prune_extremities` from `synapse/storage/controllers/persist_events.py`.
+//!
+//! The reads go through [`ExtremityReads`], which [`PersistReader`]
+//! implements, so that the unit tests can run over an in-memory DAG. The
+//! metrics go to a [`PersistCounts`] the caller owns. Shadow mode drops it,
+//! as the Python path has already counted the same chunk.
+
+use std::{
+    collections::{HashMap, HashSet},
+    fmt,
+    sync::Arc,
+};
+
+use anyhow::{anyhow, bail};
+use async_trait::async_trait;
+use log::{error, info};
+
+use super::{context::PersistContext, metrics::PersistCounts, reader::PersistReader};
+use crate::{events::Event, storage::events::RedactBehaviour};
+
+/// `EventTypes.Dummy`.
+const DUMMY: &str = "org.matrix.dummy_event";
+
+/// What the extremity calculations read from an event being persisted.
+#[derive(Clone, Debug, Default)]
+pub struct NewEvent {
+    pub event_id: Arc<str>,
+    pub sender: Arc<str>,
+    pub depth: i64,
+    pub prev_event_ids: Vec<Arc<str>>,
+    /// The MSC4242 `prev_state_events`, if the event's format has them.
+    pub prev_state_events: Option<Vec<Arc<str>>>,
+    /// Whether the event has a state key, which in an MSC4242 room is
+    /// what puts it in the state DAG.
+    pub is_state: bool,
+    pub outlier: bool,
+    pub soft_failed: bool,
+    pub rejected: bool,
+}
+
+impl NewEvent {
+    pub fn new(event: &Event, context: &PersistContext) -> anyhow::Result<Self> {
+        let metadata = event.internal_metadata();
+        Ok(NewEvent {
+            event_id: Arc::from(event.event_id()),
+            sender: Arc::from(event.sender()),
+            depth: event.depth(),
+            prev_event_ids: event.prev_event_ids().into_iter().map(Arc::from).collect(),
+            prev_state_events: event
+                .prev_state_event_ids()
+                .map(|ids| ids.iter().map(|id| Arc::from(id.as_str())).collect()),
+            is_state: event.state_key().is_some(),
+            outlier: metadata.is_outlier()?,
+            soft_failed: metadata.soft_failed()?,
+            rejected: context.is_rejected(),
+        })
+    }
+
+    /// The `prev_state_events`, which an event in the state DAG must have.
+    fn prev_state(&self) -> anyhow::Result<&[Arc<str>]> {
+        self.prev_state_events
+            .as_deref()
+            .ok_or_else(|| anyhow!("{} has no prev_state_events", self.event_id))
+    }
+}
+
+/// What `_prune_extremities` reads from an extremity it might drop, and
+/// from the prev events of a local dummy event.
+#[derive(Clone, Debug, Default)]
+pub struct DagEvent {
+    pub sender: Arc<str>,
+    pub event_type: Arc<str>,
+    pub prev_event_ids: Vec<Arc<str>>,
+    pub origin_server_ts: i64,
+    pub depth: i64,
+}
+
+impl From<&Event> for DagEvent {
+    fn from(event: &Event) -> Self {
+        DagEvent {
+            sender: Arc::from(event.sender()),
+            event_type: Arc::from(event.type_()),
+            prev_event_ids: event.prev_event_ids().into_iter().map(Arc::from).collect(),
+            origin_server_ts: event.origin_server_ts(),
+            depth: event.depth(),
+        }
+    }
+}
+
+/// `_calculate_new_state_dag_extremities` found prev state events that are
+/// neither in the chunk nor in the database. Python raises a 500
+/// `SynapseError` for this, and so does the facade.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct MissingPrevStateEvents {
+    pub room_id: String,
+    pub count: usize,
+}
+
+impl fmt::Display for MissingPrevStateEvents {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "missing {} prev_state_events in room {}",
+            self.count, self.room_id
+        )
+    }
+}
+
+impl std::error::Error for MissingPrevStateEvents {}
+
+/// The reads the extremity calculations make.
+#[async_trait]
+pub trait ExtremityReads: Sync {
+    /// See [`PersistReader::events_which_are_prevs`].
+    async fn events_which_are_prevs(&self, event_ids: &[Arc<str>])
+        -> anyhow::Result<Vec<Arc<str>>>;
+
+    /// See [`PersistReader::prevs_before_rejected`].
+    async fn prevs_before_rejected(
+        &self,
+        event_ids: &[Arc<str>],
+        include_soft_failed: bool,
+    ) -> anyhow::Result<HashSet<Arc<str>>>;
+
+    /// See [`PersistReader::have_seen_events`].
+    async fn have_seen_events(&self, event_ids: &[Arc<str>]) -> anyhow::Result<HashSet<Arc<str>>>;
+
+    /// The events found among `event_ids`, rejected ones included, as
+    /// `get_events(allow_rejected=True, redact_behaviour=as_is)` returns
+    /// them.
+    async fn dag_events(&self, event_ids: &[Arc<str>]) -> anyhow::Result<Vec<DagEvent>>;
+}
+
+#[async_trait]
+impl ExtremityReads for PersistReader<'_> {
+    async fn events_which_are_prevs(
+        &self,
+        event_ids: &[Arc<str>],
+    ) -> anyhow::Result<Vec<Arc<str>>> {
+        PersistReader::events_which_are_prevs(self, event_ids).await
+    }
+
+    async fn prevs_before_rejected(
+        &self,
+        event_ids: &[Arc<str>],
+        include_soft_failed: bool,
+    ) -> anyhow::Result<HashSet<Arc<str>>> {
+        PersistReader::prevs_before_rejected(self, event_ids, include_soft_failed).await
+    }
+
+    async fn have_seen_events(&self, event_ids: &[Arc<str>]) -> anyhow::Result<HashSet<Arc<str>>> {
+        PersistReader::have_seen_events(self, event_ids).await
+    }
+
+    async fn dag_events(&self, event_ids: &[Arc<str>]) -> anyhow::Result<Vec<DagEvent>> {
+        let events = self
+            .events(event_ids, RedactBehaviour::AsIs, true, false)
+            .await?;
+        Ok(events.iter().map(DagEvent::from).collect())
+    }
+}
+
+/// `_calculate_new_extremities`: the forward extremities after persisting
+/// `events` into a room whose extremities were `latest_event_ids`.
+pub async fn calculate_new_extremities(
+    reads: &dyn ExtremityReads,
+    events: &[NewEvent],
+    latest_event_ids: &HashSet<Arc<str>>,
+    counts: &mut PersistCounts,
+) -> anyhow::Result<HashSet<Arc<str>>> {
+    // Only the new events that are not outliers, rejected or soft failed can
+    // become extremities.
+    let new_events: Vec<&NewEvent> = events
+        .iter()
+        .filter(|event| !event.outlier && !event.rejected && !event.soft_failed)
+        .collect();
+    let new_prevs: Vec<Arc<str>> = new_events
+        .iter()
+        .flat_map(|event| event.prev_event_ids.iter().cloned())
+        .collect();
+
+    let mut result = latest_event_ids.clone();
+    result.extend(new_events.iter().map(|event| Arc::clone(&event.event_id)));
+    for prev in &new_prevs {
+        result.remove(prev);
+    }
+
+    // Drop the events that are already prev events of persisted events.
+    let candidates: Vec<Arc<str>> = result.iter().cloned().collect();
+    for prev in reads.events_which_are_prevs(&candidates).await? {
+        result.remove(&prev);
+    }
+
+    // Drop the extremities behind soft failed or rejected prev events of
+    // the new events. Otherwise they would dangle.
+    for prev in reads.prevs_before_rejected(&new_prevs, true).await? {
+        result.remove(&prev);
+    }
+
+    // Only chunks that change the extremities count, which leaves out
+    // backfill and outliers.
+    if result != *latest_event_ids {
+        *counts
+            .forward_extremities
+            .entry(result.len() as u64)
+            .or_default() += 1;
+        let stale = latest_event_ids.intersection(&result).count();
+        *counts
+            .stale_forward_extremities
+            .entry(stale as u64)
+            .or_default() += 1;
+    }
+
+    Ok(result)
+}
+
+/// `_calculate_new_state_dag_extremities`: the MSC4242 state DAG forward
+/// extremities after persisting `events` into a room whose state DAG
+/// extremities were `existing`.
+///
+/// Fails with [`MissingPrevStateEvents`] if a prev state event is neither
+/// in the chunk nor persisted. Also fails if no extremity is left, where
+/// Python's `assert` raises an `AssertionError` instead.
+pub async fn calculate_new_state_dag_extremities(
+    reads: &dyn ExtremityReads,
+    room_id: &str,
+    existing: &HashSet<Arc<str>>,
+    events: &[NewEvent],
+    counts: &mut PersistCounts,
+) -> anyhow::Result<HashSet<Arc<str>>> {
+    let state_events: Vec<&NewEvent> = events.iter().filter(|event| event.is_state).collect();
+    if state_events.is_empty() {
+        return Ok(existing.clone());
+    }
+
+    // Unlike the room DAG, this ignores outliers, keeps soft failed events
+    // and does not look for persisted children: the state DAG is persisted
+    // in causal order, so no persisted event can point at a new one.
+    let accepted: Vec<&NewEvent> = state_events
+        .iter()
+        .copied()
+        .filter(|event| !event.rejected)
+        .collect();
+
+    // Every prev state event must be known. Rejected events count as known,
+    // as other events may point at them.
+    let mut missing: HashSet<Arc<str>> = HashSet::new();
+    for event in &state_events {
+        missing.extend(event.prev_state()?.iter().cloned());
+    }
+    for event in &state_events {
+        missing.remove(&event.event_id);
+    }
+    let to_check: Vec<Arc<str>> = missing.iter().cloned().collect();
+    for seen in reads.have_seen_events(&to_check).await? {
+        missing.remove(&seen);
+    }
+    if !missing.is_empty() {
+        let mut missing: Vec<&str> = missing.iter().map(|id| &**id).collect();
+        missing.sort_unstable();
+        error!(
+            "_calculate_new_state_dag_extremities: missing the following prev_state_events in room {room_id} : {missing:?}"
+        );
+        let handling: Vec<&str> = state_events.iter().map(|e| &*e.event_id).collect();
+        error!("_calculate_new_state_dag_extremities: was handling {handling:?}");
+        return Err(MissingPrevStateEvents {
+            room_id: room_id.to_owned(),
+            count: missing.len(),
+        }
+        .into());
+    }
+
+    let mut new_prevs: Vec<Arc<str>> = Vec::new();
+    for event in &accepted {
+        new_prevs.extend(event.prev_state()?.iter().cloned());
+    }
+
+    let mut result = existing.clone();
+    result.extend(accepted.iter().map(|event| Arc::clone(&event.event_id)));
+    for prev in &new_prevs {
+        result.remove(prev);
+    }
+
+    // Drop the extremities behind rejected prev state events of the new
+    // events, so that they do not dangle.
+    for prev in reads.prevs_before_rejected(&new_prevs, false).await? {
+        result.remove(&prev);
+    }
+
+    if result != *existing {
+        *counts
+            .state_dag_forward_extremities
+            .entry(result.len() as u64)
+            .or_default() += 1;
+    }
+
+    if result.is_empty() {
+        bail!("No state dag forward extremities left in room {room_id}!");
+    }
+    Ok(result)
+}
+
+/// `_prune_extremities`: after state resolution, drop the extremities
+/// whose state group is not `resolved_state_group`, if they all look stuck.
+/// Either all of them go or none do, as state resolution over a subset
+/// might come out differently.
+///
+/// An extremity looks stuck when it is not being persisted, is not a local
+/// event (a local dummy event counts as local if it leads to a local event
+/// through other local dummy events) and is either over a day old and more
+/// than 100 events deep behind the chunk, or from a server that sent an
+/// event in the chunk and more than 20 events deep behind it.
+///
+/// Fails if an extremity has no state group in `event_id_to_state_group`
+/// (a `KeyError` in Python) or a sender is not a valid ID (a 400
+/// `SynapseError` in Python, though event validation rules that out). As in
+/// Python, the attempt is counted even then, so the caller should keep
+/// `counts` on failure too.
+#[allow(clippy::too_many_arguments)]
+pub async fn prune_extremities(
+    reads: &dyn ExtremityReads,
+    room_id: &str,
+    server_name: &str,
+    new_latest_event_ids: HashSet<Arc<str>>,
+    resolved_state_group: i64,
+    event_id_to_state_group: &HashMap<Arc<str>, i64>,
+    events: &[NewEvent],
+    now_ms: i64,
+    counts: &mut PersistCounts,
+) -> anyhow::Result<HashSet<Arc<str>>> {
+    counts.potential_times_prune_extremities += 1;
+
+    let mut kept = HashSet::new();
+    let mut dropped = Vec::new();
+    for event_id in &new_latest_event_ids {
+        let state_group = event_id_to_state_group
+            .get(event_id)
+            .ok_or_else(|| anyhow!("no state group for extremity {event_id}"))?;
+        if *state_group == resolved_state_group {
+            kept.insert(Arc::clone(event_id));
+        } else {
+            dropped.push(Arc::clone(event_id));
+        }
+    }
+
+    let dropped_set: HashSet<&Arc<str>> = dropped.iter().collect();
+    if events
+        .iter()
+        .any(|event| dropped_set.contains(&event.event_id))
+    {
+        return Ok(new_latest_event_ids);
+    }
+
+    let dropped_events = reads.dag_events(&dropped).await?;
+
+    let mut new_senders = HashSet::new();
+    for event in events {
+        new_senders.insert(domain_from_id(&event.sender)?);
+    }
+
+    let one_day_ago = now_ms - 24 * 60 * 60 * 1000;
+    let current_depth = events
+        .iter()
+        .map(|event| event.depth)
+        .max()
+        .ok_or_else(|| anyhow!("pruning extremities for an empty chunk"))?;
+
+    for event in &dropped_events {
+        // A local dummy event can go only if it does not lead to a local
+        // event through other local dummy events: a new local event should
+        // reference those.
+        let mut to_check = vec![event.clone()];
+        while !to_check.is_empty() {
+            let mut next: HashSet<Arc<str>> = HashSet::new();
+            for checking in &to_check {
+                if is_mine_id(&checking.sender, server_name) {
+                    if &*checking.event_type != DUMMY {
+                        return Ok(new_latest_event_ids);
+                    }
+                    next.extend(checking.prev_event_ids.iter().cloned());
+                }
+            }
+            if next.is_empty() {
+                break;
+            }
+            let next: Vec<Arc<str>> = next.into_iter().collect();
+            to_check = reads.dag_events(&next).await?;
+        }
+
+        if event.origin_server_ts < one_day_ago && event.depth < current_depth - 100 {
+            continue;
+        }
+
+        // Extremities from a server in the chunk can go sooner, but not at
+        // once, or every extremity of that server would go.
+        if new_senders.contains(domain_from_id(&event.sender)?) && event.depth < current_depth - 20
+        {
+            continue;
+        }
+
+        return Ok(new_latest_event_ids);
+    }
+
+    counts.times_pruned_extremities += 1;
+    info!(
+        "Pruning forward extremities in room {room_id}: from {:?} -> {:?}",
+        new_latest_event_ids, kept
+    );
+    Ok(kept)
+}
+
+/// `get_domain_from_id`.
+fn domain_from_id(id: &str) -> anyhow::Result<&str> {
+    id.split_once(':')
+        .map(|(_, domain)| domain)
+        .ok_or_else(|| anyhow!("Invalid ID: {id:?}"))
+}
+
+/// `HomeServer.is_mine_id`: false for a malformed ID.
+fn is_mine_id(id: &str, server_name: &str) -> bool {
+    id.split_once(':')
+        .is_some_and(|(_, domain)| domain == server_name)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// A persisted event in the in-memory DAG.
+    #[derive(Clone, Default)]
+    struct Stored {
+        prevs: Vec<&'static str>,
+        outlier: bool,
+        rejected: bool,
+        soft_failed: bool,
+        sender: &'static str,
+        event_type: &'static str,
+        ts: i64,
+        depth: i64,
+    }
+
+    /// The reads over an in-memory DAG, with the same semantics as the SQL.
+    #[derive(Default)]
+    struct FakeReads {
+        events: HashMap<&'static str, Stored>,
+    }
+
+    impl FakeReads {
+        fn add(&mut self, id: &'static str, stored: Stored) {
+            self.events.insert(id, stored);
+        }
+    }
+
+    #[async_trait]
+    impl ExtremityReads for FakeReads {
+        async fn events_which_are_prevs(
+            &self,
+            event_ids: &[Arc<str>],
+        ) -> anyhow::Result<Vec<Arc<str>>> {
+            let mut out = Vec::new();
+            for child in self.events.values() {
+                if child.outlier || child.rejected || child.soft_failed {
+                    continue;
+                }
+                for prev in &child.prevs {
+                    if event_ids.iter().any(|id| &**id == *prev) {
+                        out.push(Arc::from(*prev));
+                    }
+                }
+            }
+            Ok(out)
+        }
+
+        async fn prevs_before_rejected(
+            &self,
+            event_ids: &[Arc<str>],
+            include_soft_failed: bool,
+        ) -> anyhow::Result<HashSet<Arc<str>>> {
+            let mut found: HashSet<Arc<str>> = HashSet::new();
+            let mut to_check: Vec<Arc<str>> = event_ids.to_vec();
+            while let Some(id) = to_check.pop() {
+                let Some(event) = self.events.get(&*id) else {
+                    continue;
+                };
+                // As in the SQL, the flags are those of the event being
+                // walked, and what is found is its prev events.
+                let bad = event.rejected || (include_soft_failed && event.soft_failed);
+                if event.outlier || !bad {
+                    continue;
+                }
+                for prev in &event.prevs {
+                    if found.insert(Arc::from(*prev)) {
+                        to_check.push(Arc::from(*prev));
+                    }
+                }
+            }
+            Ok(found)
+        }
+
+        async fn have_seen_events(
+            &self,
+            event_ids: &[Arc<str>],
+        ) -> anyhow::Result<HashSet<Arc<str>>> {
+            Ok(event_ids
+                .iter()
+                .filter(|id| self.events.contains_key(&***id))
+                .cloned()
+                .collect())
+        }
+
+        async fn dag_events(&self, event_ids: &[Arc<str>]) -> anyhow::Result<Vec<DagEvent>> {
+            Ok(event_ids
+                .iter()
+                .filter_map(|id| self.events.get(&**id))
+                .map(|stored| DagEvent {
+                    sender: Arc::from(stored.sender),
+                    event_type: Arc::from(stored.event_type),
+                    prev_event_ids: stored.prevs.iter().map(|p| Arc::from(*p)).collect(),
+                    origin_server_ts: stored.ts,
+                    depth: stored.depth,
+                })
+                .collect())
+        }
+    }
+
+    fn ids(ids: &[&str]) -> HashSet<Arc<str>> {
+        ids.iter().map(|id| Arc::from(*id)).collect()
+    }
+
+    fn new_event(id: &str, prevs: &[&str]) -> NewEvent {
+        NewEvent {
+            event_id: Arc::from(id),
+            sender: Arc::from("@a:remote"),
+            depth: 1,
+            prev_event_ids: prevs.iter().map(|p| Arc::from(*p)).collect(),
+            ..Default::default()
+        }
+    }
+
+    fn state_event(id: &str, prev_state: &[&str]) -> NewEvent {
+        NewEvent {
+            is_state: true,
+            prev_state_events: Some(prev_state.iter().map(|p| Arc::from(*p)).collect()),
+            ..new_event(id, &[])
+        }
+    }
+
+    fn block_on<F: std::future::Future>(f: F) -> F::Output {
+        futures::executor::block_on(f)
+    }
+
+    #[test]
+    fn new_event_replaces_its_prev() {
+        let mut reads = FakeReads::default();
+        reads.add("$a", Stored::default());
+        let mut counts = PersistCounts::default();
+        let result = block_on(calculate_new_extremities(
+            &reads,
+            &[new_event("$b", &["$a"])],
+            &ids(&["$a"]),
+            &mut counts,
+        ))
+        .unwrap();
+        assert_eq!(result, ids(&["$b"]));
+        assert_eq!(counts.forward_extremities[&1], 1);
+        assert_eq!(counts.stale_forward_extremities[&0], 1);
+    }
+
+    #[test]
+    fn outliers_rejected_and_soft_failed_do_not_count() {
+        let reads = FakeReads::default();
+        let events = [
+            NewEvent {
+                outlier: true,
+                ..new_event("$o", &["$a"])
+            },
+            NewEvent {
+                rejected: true,
+                ..new_event("$r", &["$a"])
+            },
+            NewEvent {
+                soft_failed: true,
+                ..new_event("$s", &["$a"])
+            },
+        ];
+        let mut counts = PersistCounts::default();
+        let result = block_on(calculate_new_extremities(
+            &reads,
+            &events,
+            &ids(&["$a"]),
+            &mut counts,
+        ))
+        .unwrap();
+        assert_eq!(result, ids(&["$a"]));
+        // Nothing changed, so nothing is counted.
+        assert!(counts.forward_extremities.is_empty());
+        assert!(counts.stale_forward_extremities.is_empty());
+    }
+
+    #[test]
+    fn events_with_persisted_children_are_dropped() {
+        // $c was persisted (as by backfill) before its prev $b arrived.
+        let mut reads = FakeReads::default();
+        reads.add(
+            "$c",
+            Stored {
+                prevs: vec!["$b"],
+                ..Default::default()
+            },
+        );
+        // A soft failed child does not count.
+        reads.add(
+            "$sf",
+            Stored {
+                prevs: vec!["$d"],
+                soft_failed: true,
+                ..Default::default()
+            },
+        );
+        let mut counts = PersistCounts::default();
+        let result = block_on(calculate_new_extremities(
+            &reads,
+            &[new_event("$b", &["$a"]), new_event("$d", &["$a"])],
+            &ids(&["$a", "$c"]),
+            &mut counts,
+        ))
+        .unwrap();
+        assert_eq!(result, ids(&["$c", "$d"]));
+        assert_eq!(counts.stale_forward_extremities[&1], 1);
+    }
+
+    #[test]
+    fn soft_failed_prevs_take_their_ancestors_with_them() {
+        // F <- SF1 <- SF2 <- N: F stops being an extremity.
+        let mut reads = FakeReads::default();
+        reads.add("$f", Stored::default());
+        reads.add(
+            "$sf1",
+            Stored {
+                prevs: vec!["$f"],
+                soft_failed: true,
+                ..Default::default()
+            },
+        );
+        reads.add(
+            "$sf2",
+            Stored {
+                prevs: vec!["$sf1"],
+                soft_failed: true,
+                ..Default::default()
+            },
+        );
+        let mut counts = PersistCounts::default();
+        let result = block_on(calculate_new_extremities(
+            &reads,
+            &[new_event("$n", &["$sf2"])],
+            &ids(&["$f"]),
+            &mut counts,
+        ))
+        .unwrap();
+        assert_eq!(result, ids(&["$n"]));
+    }
+
+    #[test]
+    fn state_dag_without_state_events_is_unchanged() {
+        let reads = FakeReads::default();
+        let mut counts = PersistCounts::default();
+        let result = block_on(calculate_new_state_dag_extremities(
+            &reads,
+            "!r:x",
+            &ids(&["$a"]),
+            &[new_event("$m", &["$a"])],
+            &mut counts,
+        ))
+        .unwrap();
+        assert_eq!(result, ids(&["$a"]));
+        assert!(counts.state_dag_forward_extremities.is_empty());
+    }
+
+    #[test]
+    fn state_dag_missing_prev_state_events_fail() {
+        let mut reads = FakeReads::default();
+        reads.add("$a", Stored::default());
+        let mut counts = PersistCounts::default();
+        let err = block_on(calculate_new_state_dag_extremities(
+            &reads,
+            "!r:x",
+            &ids(&["$a"]),
+            &[
+                state_event("$b", &["$a", "$gone1"]),
+                state_event("$c", &["$b", "$gone2"]),
+            ],
+            &mut counts,
+        ))
+        .unwrap_err();
+        assert_eq!(
+            err.downcast_ref::<MissingPrevStateEvents>(),
+            Some(&MissingPrevStateEvents {
+                room_id: "!r:x".to_owned(),
+                count: 2,
+            })
+        );
+    }
+
+    #[test]
+    fn state_dag_checks_the_prevs_of_rejected_events() {
+        let mut reads = FakeReads::default();
+        reads.add("$a", Stored::default());
+        let events = [NewEvent {
+            rejected: true,
+            ..state_event("$r", &["$gone"])
+        }];
+        let mut counts = PersistCounts::default();
+        let err = block_on(calculate_new_state_dag_extremities(
+            &reads,
+            "!r:x",
+            &ids(&["$a"]),
+            &events,
+            &mut counts,
+        ))
+        .unwrap_err();
+        assert_eq!(
+            err.downcast_ref::<MissingPrevStateEvents>()
+                .map(|e| e.count),
+            Some(1)
+        );
+    }
+
+    #[test]
+    fn state_dag_unchanged_is_not_counted() {
+        let mut reads = FakeReads::default();
+        reads.add("$a", Stored::default());
+        let events = [NewEvent {
+            rejected: true,
+            ..state_event("$r", &["$a"])
+        }];
+        let mut counts = PersistCounts::default();
+        let result = block_on(calculate_new_state_dag_extremities(
+            &reads,
+            "!r:x",
+            &ids(&["$a"]),
+            &events,
+            &mut counts,
+        ))
+        .unwrap();
+        assert_eq!(result, ids(&["$a"]));
+        assert!(counts.state_dag_forward_extremities.is_empty());
+    }
+
+    #[test]
+    fn state_dag_keeps_soft_failed_and_knows_rejected() {
+        let mut reads = FakeReads::default();
+        reads.add("$a", Stored::default());
+        let events = [
+            // Rejected, but $c may point at it.
+            NewEvent {
+                rejected: true,
+                ..state_event("$r", &["$a"])
+            },
+            NewEvent {
+                soft_failed: true,
+                ..state_event("$c", &["$r"])
+            },
+            // Not in the state DAG at all.
+            new_event("$m", &["$c"]),
+        ];
+        let mut counts = PersistCounts::default();
+        let result = block_on(calculate_new_state_dag_extremities(
+            &reads,
+            "!r:x",
+            &ids(&["$a"]),
+            &events,
+            &mut counts,
+        ))
+        .unwrap();
+        assert_eq!(result, ids(&["$a", "$c"]));
+        assert_eq!(counts.state_dag_forward_extremities[&2], 1);
+    }
+
+    #[test]
+    fn state_dag_drops_extremities_behind_rejected_prevs() {
+        // F <- R (rejected, persisted) <- N.
+        let mut reads = FakeReads::default();
+        reads.add("$f", Stored::default());
+        reads.add(
+            "$r",
+            Stored {
+                prevs: vec!["$f"],
+                rejected: true,
+                ..Default::default()
+            },
+        );
+        let mut counts = PersistCounts::default();
+        let result = block_on(calculate_new_state_dag_extremities(
+            &reads,
+            "!r:x",
+            &ids(&["$f"]),
+            &[state_event("$n", &["$r"])],
+            &mut counts,
+        ))
+        .unwrap();
+        assert_eq!(result, ids(&["$n"]));
+    }
+
+    #[test]
+    fn state_dag_with_no_extremities_left_fails() {
+        let mut reads = FakeReads::default();
+        reads.add("$a", Stored::default());
+        let mut counts = PersistCounts::default();
+        let events = [NewEvent {
+            rejected: true,
+            ..state_event("$r", &["$a"])
+        }];
+        let err = block_on(calculate_new_state_dag_extremities(
+            &reads,
+            "!r:x",
+            &HashSet::new(),
+            &events,
+            &mut counts,
+        ))
+        .unwrap_err();
+        assert!(err.to_string().contains("No state dag forward extremities"));
+    }
+
+    const NOW: i64 = 10 * 24 * 60 * 60 * 1000;
+    const TWO_DAYS_AGO: i64 = NOW - 2 * 24 * 60 * 60 * 1000;
+
+    /// Prune with extremities $keep (state group 1, the resolved one) and
+    /// $drop (state group 2), persisting `events`.
+    fn prune(reads: &FakeReads, events: &[NewEvent]) -> (HashSet<Arc<str>>, PersistCounts) {
+        let groups: HashMap<Arc<str>, i64> =
+            [(Arc::from("$keep"), 1), (Arc::from("$drop"), 2)].into();
+        let mut counts = PersistCounts::default();
+        let result = block_on(prune_extremities(
+            reads,
+            "!r:x",
+            "local",
+            ids(&["$keep", "$drop"]),
+            1,
+            &groups,
+            events,
+            NOW,
+            &mut counts,
+        ))
+        .unwrap();
+        (result, counts)
+    }
+
+    fn deep_event(depth: i64) -> NewEvent {
+        NewEvent {
+            depth,
+            sender: Arc::from("@n:other"),
+            ..new_event("$n", &["$keep"])
+        }
+    }
+
+    fn stored(sender: &'static str, ts: i64, depth: i64) -> Stored {
+        Stored {
+            sender,
+            event_type: "m.room.message",
+            ts,
+            depth,
+            ..Default::default()
+        }
+    }
+
+    #[test]
+    fn prune_drops_old_deep_extremities() {
+        let mut reads = FakeReads::default();
+        reads.add("$drop", stored("@x:remote", TWO_DAYS_AGO, 10));
+        let (result, counts) = prune(&reads, &[deep_event(111)]);
+        assert_eq!(result, ids(&["$keep"]));
+        assert_eq!(counts.potential_times_prune_extremities, 1);
+        assert_eq!(counts.times_pruned_extremities, 1);
+    }
+
+    #[test]
+    fn prune_keeps_extremities_that_are_too_new_or_shallow() {
+        let mut reads = FakeReads::default();
+        // Old, but exactly 100 deep: kept.
+        reads.add("$drop", stored("@x:remote", TWO_DAYS_AGO, 10));
+        let (result, counts) = prune(&reads, &[deep_event(110)]);
+        assert_eq!(result, ids(&["$keep", "$drop"]));
+        assert_eq!(counts.potential_times_prune_extremities, 1);
+        assert_eq!(counts.times_pruned_extremities, 0);
+
+        // Deep, but recent: kept.
+        reads.add("$drop", stored("@x:remote", NOW, 10));
+        let (result, _) = prune(&reads, &[deep_event(500)]);
+        assert_eq!(result, ids(&["$keep", "$drop"]));
+    }
+
+    #[test]
+    fn prune_drops_sooner_for_a_server_in_the_chunk() {
+        let mut reads = FakeReads::default();
+        reads.add("$drop", stored("@x:other", NOW, 10));
+        let (result, _) = prune(&reads, &[deep_event(31)]);
+        assert_eq!(result, ids(&["$keep"]));
+        let (result, _) = prune(&reads, &[deep_event(30)]);
+        assert_eq!(result, ids(&["$keep", "$drop"]));
+    }
+
+    #[test]
+    fn prune_keeps_extremities_being_persisted() {
+        let mut reads = FakeReads::default();
+        reads.add("$drop", stored("@x:remote", TWO_DAYS_AGO, 10));
+        let (result, _) = prune(
+            &reads,
+            &[
+                deep_event(500),
+                NewEvent {
+                    depth: 10,
+                    ..new_event("$drop", &[])
+                },
+            ],
+        );
+        assert_eq!(result, ids(&["$keep", "$drop"]));
+    }
+
+    #[test]
+    fn prune_keeps_local_events() {
+        let mut reads = FakeReads::default();
+        reads.add("$drop", stored("@me:local", TWO_DAYS_AGO, 10));
+        let (result, _) = prune(&reads, &[deep_event(500)]);
+        assert_eq!(result, ids(&["$keep", "$drop"]));
+    }
+
+    #[test]
+    fn prune_follows_local_dummy_events() {
+        let dummy = |prevs: Vec<&'static str>| Stored {
+            event_type: DUMMY,
+            prevs,
+            ..stored("@me:local", TWO_DAYS_AGO, 10)
+        };
+
+        // dummy -> dummy -> remote: can go.
+        let mut reads = FakeReads::default();
+        reads.add("$drop", dummy(vec!["$d2"]));
+        reads.add("$d2", dummy(vec!["$remote"]));
+        reads.add("$remote", stored("@x:remote", TWO_DAYS_AGO, 8));
+        let (result, _) = prune(&reads, &[deep_event(500)]);
+        assert_eq!(result, ids(&["$keep"]));
+
+        // dummy -> dummy -> local message: stays.
+        reads.add("$local", stored("@me:local", TWO_DAYS_AGO, 8));
+        reads.add("$d2", dummy(vec!["$remote", "$local"]));
+        let (result, _) = prune(&reads, &[deep_event(500)]);
+        assert_eq!(result, ids(&["$keep", "$drop"]));
+    }
+
+    #[test]
+    fn prune_needs_every_state_group() {
+        let reads = FakeReads::default();
+        let mut counts = PersistCounts::default();
+        let err = block_on(prune_extremities(
+            &reads,
+            "!r:x",
+            "local",
+            ids(&["$a"]),
+            1,
+            &HashMap::new(),
+            &[deep_event(1)],
+            NOW,
+            &mut counts,
+        ))
+        .unwrap_err();
+        assert!(err.to_string().contains("no state group"));
+    }
+
+    #[test]
+    fn ids_and_domains() {
+        assert_eq!(domain_from_id("@a:b.c:8448").unwrap(), "b.c:8448");
+        assert!(domain_from_id("nope").is_err());
+        assert!(is_mine_id("@a:local", "local"));
+        assert!(!is_mine_id("@a:other", "local"));
+        assert!(!is_mine_id("nope", "local"));
+    }
+}

rust/src/persist/mod.rs

@@ -21,6 +21,7 @@
 //! - [`context`]: what persisting needs from an `EventContext`.
 //! - [`effects`]: the Python cache updates a Rust transaction asks for once
 //!   it commits.
+//! - [`extremities`]: the room's forward extremities after a chunk.
 //! - [`metrics`]: counts for the Python metrics, flushed when scraped.
 //! - [`python`]: the [`python::Persister`] and test helpers.
 //! - [`reader`]: what persisting reads before its write transaction.
@@ -32,6 +33,7 @@ use pyo3::prelude::*;
 
 pub mod context;
 pub mod effects;
+pub mod extremities;
 pub mod metrics;
 pub mod python;
 pub mod reader;

rust/src/persist/python.rs

@@ -18,10 +18,14 @@
 
 use std::sync::Arc;
 
+use http::StatusCode;
 use pyo3::{exceptions::PyRuntimeError, intern, prelude::*, types::PyDict};
 
-use super::{context::PersistContext, metrics::PersistMetrics};
+use super::{
+    context::PersistContext, extremities::MissingPrevStateEvents, metrics::PersistMetrics,
+};
 use crate::{
+    errors::SynapseError,
     events::Event,
     homeserver::HomeServer,
     runtime::RustRuntime,
@@ -122,6 +126,28 @@ impl Persister {
     }
 }
 
+/// Convert an error from persisting to the Python exception the Python
+/// path would raise: a Python exception as it was raised, a 500
+/// `SynapseError` for missing prev state events, as
+/// `_calculate_new_state_dag_extremities` raises, and a `RuntimeError`
+/// otherwise.
+pub fn persist_error_to_py(e: anyhow::Error) -> PyErr {
+    let e = match e.downcast::<PyErr>() {
+        Ok(err) => return err,
+        Err(e) => e,
+    };
+    if let Some(missing) = e.downcast_ref::<MissingPrevStateEvents>() {
+        return SynapseError::new(
+            StatusCode::INTERNAL_SERVER_ERROR,
+            missing.to_string(),
+            "M_UNKNOWN",
+            None,
+            None,
+        );
+    }
+    PyRuntimeError::new_err(format!("{e:#}"))
+}
+
 /// `(internal_metadata, json)`: the two `event_json` columns as Rust would
 /// write them for `event`, for tests to compare with Python's.
 #[pyfunction]

rust/src/persist/reader.rs

@@ -39,13 +39,20 @@ use std::{
 
 use anyhow::Context;
 use futures::{future::BoxFuture, FutureExt};
-use pyo3::prelude::*;
+use pyo3::{prelude::*, types::PyDict};
 use serde::Deserialize;
 use serde_json::Value;
 use tokio::runtime::Handle;
 use tokio_postgres::Row;
 
+use super::{
+    context::PersistContext,
+    extremities::{self, NewEvent},
+    metrics::PersistCounts,
+    python::persist_error_to_py,
+};
 use crate::database::postgres::pool::{ConnectionPool, PyConnectionPool};
+use crate::events::Event;
 use crate::storage::events::Querier;
 use crate::storage::sync_store::{
     PersistedEventPosition, SyncReader, SyncStore, SELECT_MANY_BATCH_SIZE,
@@ -669,10 +676,6 @@ fn is_truthy(value: &Value) -> bool {
 // Test-only entry points, so trial can compare the reads against the Python
 // store on a real schema.
 
-fn anyhow_to_py(e: anyhow::Error) -> PyErr {
-    pyo3::exceptions::PyRuntimeError::new_err(format!("{e:#}"))
-}
-
 fn arcs(ids: Vec<String>) -> Vec<Arc<str>> {
     ids.into_iter().map(Arc::from).collect()
 }
@@ -702,7 +705,7 @@ impl PyPersistReaderForTest {
             self.handle
                 .block_on(f(PersistReader::new(self.store.reader(self.pool.clone()))))
         })
-        .map_err(anyhow_to_py)
+        .map_err(persist_error_to_py)
     }
 }
 
@@ -982,6 +985,109 @@ impl PyPersistReaderForTest {
             async move { max_sequence_numbers(&reader.borrow().await?, &chain_ids).await }.boxed()
         })
     }
+
+    // The extremity calculations over the reader, each returning its result
+    // and what it counted.
+
+    fn calculate_new_extremities<'py>(
+        &self,
+        py: Python<'py>,
+        events_and_contexts: Vec<(Py<Event>, PersistContext)>,
+        latest_event_ids: Vec<String>,
+    ) -> PyResult<(HashSet<String>, Bound<'py, PyDict>)> {
+        let events = new_events(py, &events_and_contexts)?;
+        let latest: HashSet<Arc<str>> = arcs(latest_event_ids).into_iter().collect();
+        let (result, counts) = self.run(py, |reader| {
+            async move {
+                let mut counts = PersistCounts::default();
+                let result =
+                    extremities::calculate_new_extremities(&reader, &events, &latest, &mut counts)
+                        .await?;
+                Ok((result, counts))
+            }
+            .boxed()
+        })?;
+        Ok((strings(result), counts.to_py(py)?))
+    }
+
+    fn calculate_new_state_dag_extremities<'py>(
+        &self,
+        py: Python<'py>,
+        room_id: String,
+        existing: Vec<String>,
+        events_and_contexts: Vec<(Py<Event>, PersistContext)>,
+    ) -> PyResult<(HashSet<String>, Bound<'py, PyDict>)> {
+        let events = new_events(py, &events_and_contexts)?;
+        let existing: HashSet<Arc<str>> = arcs(existing).into_iter().collect();
+        let (result, counts) = self.run(py, |reader| {
+            async move {
+                let mut counts = PersistCounts::default();
+                let result = extremities::calculate_new_state_dag_extremities(
+                    &reader,
+                    &room_id,
+                    &existing,
+                    &events,
+                    &mut counts,
+                )
+                .await?;
+                Ok((result, counts))
+            }
+            .boxed()
+        })?;
+        Ok((strings(result), counts.to_py(py)?))
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    fn prune_extremities<'py>(
+        &self,
+        py: Python<'py>,
+        room_id: String,
+        server_name: String,
+        new_latest_event_ids: Vec<String>,
+        resolved_state_group: i64,
+        event_id_to_state_group: HashMap<String, i64>,
+        events_and_contexts: Vec<(Py<Event>, PersistContext)>,
+        now_ms: i64,
+    ) -> PyResult<(HashSet<String>, Bound<'py, PyDict>)> {
+        let events = new_events(py, &events_and_contexts)?;
+        let latest: HashSet<Arc<str>> = arcs(new_latest_event_ids).into_iter().collect();
+        let groups: HashMap<Arc<str>, i64> = event_id_to_state_group
+            .into_iter()
+            .map(|(id, group)| (Arc::from(id), group))
+            .collect();
+        let (result, counts) = self.run(py, |reader| {
+            async move {
+                let mut counts = PersistCounts::default();
+                let result = extremities::prune_extremities(
+                    &reader,
+                    &room_id,
+                    &server_name,
+                    latest,
+                    resolved_state_group,
+                    &groups,
+                    &events,
+                    now_ms,
+                    &mut counts,
+                )
+                .await?;
+                Ok((result, counts))
+            }
+            .boxed()
+        })?;
+        Ok((strings(result), counts.to_py(py)?))
+    }
+}
+
+fn new_events(
+    py: Python<'_>,
+    events_and_contexts: &[(Py<Event>, PersistContext)],
+) -> PyResult<Vec<NewEvent>> {
+    events_and_contexts
+        .iter()
+        .map(|(event, context)| {
+            NewEvent::new(event.bind(py).get(), context).map_err(persist_error_to_py)
+        })
+        .collect()
 }
 
 #[cfg(test)]

synapse/synapse_rust/persist.pyi

@@ -143,6 +143,32 @@ class PersistReaderForTest:
         """`(origin_chain_id, origin_sequence_number, target_chain_id,
         target_sequence_number)` rows."""
     def max_sequence_numbers(self, chain_ids: list[int]) -> dict[int, int]: ...
+    # The extremity calculations (`rust/src/persist/extremities.rs`), each
+    # returning its result and what it counted, in the shape
+    # `Persister.take_metrics` returns.
+    def calculate_new_extremities(
+        self,
+        events_and_contexts: list[tuple[EventBase, EventContext]],
+        latest_event_ids: list[str],
+    ) -> tuple[set[str], dict[str, Any]]: ...
+    def calculate_new_state_dag_extremities(
+        self,
+        room_id: str,
+        existing: list[str],
+        events_and_contexts: list[tuple[EventBase, EventContext]],
+    ) -> tuple[set[str], dict[str, Any]]:
+        """Raises a 500 `SynapseError` for missing prev state events, as
+        Python does."""
+    def prune_extremities(
+        self,
+        room_id: str,
+        server_name: str,
+        new_latest_event_ids: list[str],
+        resolved_state_group: int,
+        event_id_to_state_group: dict[str, int],
+        events_and_contexts: list[tuple[EventBase, EventContext]],
+        now_ms: int,
+    ) -> tuple[set[str], dict[str, Any]]: ...
 
 def persist_gil_stats() -> dict[str, tuple[int, float]]:
     """`{site: (calls, seconds held)}` for every time event persistence took

tests/storage/test_persist_reader.py

@@ -26,7 +26,10 @@ from unittest import SkipTest
 from twisted.internet.testing import MemoryReactor
 
 from synapse.api.constants import EventTypes
-from synapse.api.errors import NotFoundError
+from synapse.api.errors import NotFoundError, SynapseError
+from synapse.api.room_versions import RoomVersions
+from synapse.events import EventBase
+from synapse.events.snapshot import EventContext
 from synapse.rest import admin
 from synapse.rest.client import login, room
 from synapse.server import HomeServer
@@ -38,6 +41,7 @@ from synapse.types.handlers import SLIDING_SYNC_DEFAULT_BUMP_EVENT_TYPES
 from synapse.util.clock import Clock
 
 from tests import unittest
+from tests.test_utils.event_builders import make_test_event
 from tests.test_utils.event_injection import inject_event
 from tests.utils import RUST_PERSIST_FOR_TESTS
 
@@ -110,6 +114,51 @@ class PersistReaderTestBase(unittest.HomeserverTestCase):
             )
         ]
 
+    def _dag_with_bad_events(self) -> list[str]:
+        """A chain of events in the room, of which the third is rejected,
+        the fourth soft failed, the fifth an outlier and the seventh both
+        rejected and an outlier."""
+        chain = [self._send(f"m{i}") for i in range(3)]
+        chain.append(
+            self.get_success(
+                inject_event(
+                    self.hs,
+                    prev_event_ids=[chain[-1]],
+                    internal_metadata={"soft_failed": True},
+                    room_id=self.room_id,
+                    sender=self.alice,
+                    type=EventTypes.Message,
+                    content={"msgtype": "m.text", "body": "soft failed"},
+                )
+            ).event_id
+        )
+        chain.extend(
+            self.get_success(
+                inject_event(
+                    self.hs,
+                    prev_event_ids=[chain[-1]],
+                    room_id=self.room_id,
+                    sender=self.alice,
+                    type=EventTypes.Message,
+                    content={"msgtype": "m.text", "body": body},
+                )
+            ).event_id
+            for body in ("outlier", "last", "rejected outlier")
+        )
+        for event_id in (chain[2], chain[6]):
+            self._reject(event_id)
+        for event_id in (chain[4], chain[6]):
+            self._write("UPDATE events SET outlier = TRUE WHERE event_id = ?", event_id)
+        return chain
+
+    def _reject(self, event_id: str) -> None:
+        self.get_success(
+            self.store.db_pool.simple_insert(
+                "rejections",
+                {"event_id": event_id, "reason": "test", "last_check": "now"},
+            )
+        )
+
 
 class PersistReaderTestCase(PersistReaderTestBase):
     """Each test runs a Rust read and the Python store method against the
@@ -157,51 +206,6 @@ class PersistReaderTestCase(PersistReaderTestBase):
             )
         self.assertEqual(len(self.rust.state_dag_extremities(self.room_id)), 2)
 
-    def _dag_with_bad_events(self) -> list[str]:
-        """A chain of events in the room, of which the third is rejected,
-        the fourth soft failed, the fifth an outlier and the seventh both
-        rejected and an outlier."""
-        chain = [self._send(f"m{i}") for i in range(3)]
-        chain.append(
-            self.get_success(
-                inject_event(
-                    self.hs,
-                    prev_event_ids=[chain[-1]],
-                    internal_metadata={"soft_failed": True},
-                    room_id=self.room_id,
-                    sender=self.alice,
-                    type=EventTypes.Message,
-                    content={"msgtype": "m.text", "body": "soft failed"},
-                )
-            ).event_id
-        )
-        chain.extend(
-            self.get_success(
-                inject_event(
-                    self.hs,
-                    prev_event_ids=[chain[-1]],
-                    room_id=self.room_id,
-                    sender=self.alice,
-                    type=EventTypes.Message,
-                    content={"msgtype": "m.text", "body": body},
-                )
-            ).event_id
-            for body in ("outlier", "last", "rejected outlier")
-        )
-        for event_id in (chain[2], chain[6]):
-            self._reject(event_id)
-        for event_id in (chain[4], chain[6]):
-            self._write("UPDATE events SET outlier = TRUE WHERE event_id = ?", event_id)
-        return chain
-
-    def _reject(self, event_id: str) -> None:
-        self.get_success(
-            self.store.db_pool.simple_insert(
-                "rejections",
-                {"event_id": event_id, "reason": "test", "last_check": "now"},
-            )
-        )
-
     def test_events_which_are_prevs(self) -> None:
         chain = self._dag_with_bad_events()
         # A second child of "last", so it is a prev event twice.
@@ -668,3 +672,286 @@ class ChainCoverReadsTestCase(PersistReaderTestBase):
             self.assertEqual(
                 self.rust.max_sequence_numbers(chain_ids), dict(rows), chain_ids
             )
+
+
+class ExtremitiesTestCase(PersistReaderTestBase):
+    """The extremity calculations in `rust/src/persist/extremities.rs`
+    against the controller's Python methods, over the same database and
+    the same new events."""
+
+    def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
+        super().prepare(reactor, clock, hs)
+        persistence = hs.get_storage_controllers().persistence
+        assert persistence is not None
+        self.persistence = persistence
+        self.room_version = self.get_success(self.store.get_room_version(self.room_id))
+        self._counter = 0
+
+    def _new(
+        self,
+        prev_event_ids: list[str],
+        *,
+        sender: str | None = None,
+        depth: int = 10,
+        prev_state_events: list[str] | None = None,
+        rejected: bool = False,
+        soft_failed: bool = False,
+        outlier: bool = False,
+    ) -> tuple[EventBase, EventContext]:
+        """A new event for the chunk, not persisted. With
+        `prev_state_events` it is an MSC4242 state event."""
+        self._counter += 1
+        fields: dict[str, Any] = {
+            "room_id": self.room_id,
+            "sender": sender or self.alice,
+            "type": EventTypes.Message,
+            "content": {"msgtype": "m.text", "body": f"new {self._counter}"},
+            "prev_events": prev_event_ids,
+            "depth": depth,
+        }
+        room_version = self.room_version
+        if prev_state_events is not None:
+            room_version = RoomVersions.MSC4242v12
+            fields["type"] = "m.test.state"
+            fields["state_key"] = str(self._counter)
+            fields["prev_state_events"] = prev_state_events
+        event = make_test_event(fields, room_version)
+        event.internal_metadata.outlier = outlier
+        if soft_failed:
+            event.internal_metadata.soft_failed = True
+        context = EventContext.for_outlier(self.hs.get_storage_controllers())
+        if rejected:
+            context.rejected = "test"
+        return event, context
+
+    def _new_extremities(
+        self,
+        events: list[tuple[EventBase, EventContext]],
+        latest: set[str],
+    ) -> tuple[set[str], dict[str, Any]]:
+        expected = self.get_success(
+            self.persistence._calculate_new_extremities(self.room_id, events, latest)
+        )
+        result, counts = self.rust.calculate_new_extremities(events, list(latest))
+        self.assertEqual(result, expected)
+        return result, counts
+
+    def test_calculate_new_extremities(self) -> None:
+        chain = self._dag_with_bad_events()
+        # The fixture flags its events after persisting them, so the stored
+        # extremities are out of date. Use the one they would be.
+        latest = {chain[5]}
+
+        # A message on top, with new events that cannot be extremities.
+        new = self._new([chain[5]])
+        result, counts = self._new_extremities(
+            [
+                new,
+                self._new([new[0].event_id], rejected=True),
+                self._new([new[0].event_id], soft_failed=True),
+                self._new([new[0].event_id], outlier=True),
+            ],
+            latest,
+        )
+        self.assertEqual(result, {new[0].event_id})
+        self.assertEqual(counts["forward_extremities"], {1: 1})
+        self.assertEqual(counts["stale_forward_extremities"], {0: 1})
+
+        # Nothing changes, so nothing is counted.
+        result, counts = self._new_extremities(
+            [self._new([chain[5]], outlier=True)], latest
+        )
+        self.assertEqual(result, latest)
+        self.assertEqual(counts["forward_extremities"], {})
+
+        # Stale extremities with persisted children go. The new event's prev
+        # is soft failed, and behind it a rejected event: everything up to
+        # chain[1] goes, while chain[0], whose child chain[1] is persisted,
+        # goes anyway.
+        on_soft_failed = self._new([chain[3]])
+        result, counts = self._new_extremities(
+            [on_soft_failed], {chain[0], chain[1], chain[5]}
+        )
+        self.assertEqual(result, {chain[5], on_soft_failed[0].event_id})
+        self.assertEqual(counts["forward_extremities"], {2: 1})
+        self.assertEqual(counts["stale_forward_extremities"], {1: 1})
+
+    def _new_extremities_state_dag(
+        self,
+        existing: set[str],
+        events: list[tuple[EventBase, EventContext]],
+    ) -> set[str]:
+        expected = self.get_success(
+            self.persistence._calculate_new_state_dag_extremities(
+                self.room_id,
+                frozenset(existing),
+                events,  # type: ignore[arg-type]
+            )
+        )
+        result, _ = self.rust.calculate_new_state_dag_extremities(
+            self.room_id, list(existing), events
+        )
+        self.assertEqual(result, expected)
+        return result
+
+    def test_calculate_new_state_dag_extremities(self) -> None:
+        # The function reads the room DAG tables, so the room's messages
+        # stand in for persisted state events.
+        chain = self._dag_with_bad_events()
+
+        # No state events: unchanged.
+        self.assertEqual(
+            self._new_extremities_state_dag({chain[5]}, [self._new([chain[5]])]),
+            {chain[5]},
+        )
+
+        # A rejected new event is known to the next one, and a soft failed
+        # one can be an extremity.
+        rejected = self._new([], prev_state_events=[chain[5]], rejected=True)
+        soft_failed = self._new(
+            [], prev_state_events=[rejected[0].event_id], soft_failed=True
+        )
+        fork = self._new([], prev_state_events=[chain[1]])
+        # A rejected event nothing points at is never an extremity.
+        rejected_leaf = self._new([], prev_state_events=[chain[5]], rejected=True)
+        self.assertEqual(
+            self._new_extremities_state_dag(
+                {chain[5], chain[1]}, [rejected, soft_failed, fork, rejected_leaf]
+            ),
+            {chain[5], soft_failed[0].event_id, fork[0].event_id},
+        )
+
+        # A prev state event that is rejected takes its ancestors with it:
+        # chain[2] is rejected, so chain[1] goes.
+        on_rejected = self._new([], prev_state_events=[chain[2]])
+        self.assertEqual(
+            self._new_extremities_state_dag({chain[1], chain[5]}, [on_rejected]),
+            {chain[5], on_rejected[0].event_id},
+        )
+
+        # Soft failed events do not count on the state DAG: chain[3] is soft
+        # failed, so the walk stops there and chain[1] stays.
+        on_soft_failed = self._new([], prev_state_events=[chain[3]])
+        self.assertEqual(
+            self._new_extremities_state_dag({chain[1], chain[5]}, [on_soft_failed]),
+            {chain[1], chain[5], on_soft_failed[0].event_id},
+        )
+
+    def test_missing_prev_state_events(self) -> None:
+        message = self._send("known")
+        events = [
+            self._new([], prev_state_events=[message, "$gone1"]),
+            self._new([], prev_state_events=["$gone2"]),
+        ]
+        failure = self.get_failure(
+            self.persistence._calculate_new_state_dag_extremities(
+                self.room_id,
+                frozenset({message}),
+                events,  # type: ignore[arg-type]
+            ),
+            SynapseError,
+        )
+        with self.assertRaises(SynapseError) as cm:
+            self.rust.calculate_new_state_dag_extremities(
+                self.room_id, [message], events
+            )
+        self.assertEqual(cm.exception.code, 500)
+        self.assertEqual(cm.exception.msg, failure.value.msg)
+        self.assertEqual(cm.exception.errcode, failure.value.errcode)
+
+    def _inject(
+        self, sender: str, prev: str, event_type: str = EventTypes.Message
+    ) -> str:
+        content: dict[str, Any] = (
+            {} if event_type == EventTypes.Dummy else {"msgtype": "m.text", "body": "x"}
+        )
+        self._counter += 1
+        content["n"] = self._counter
+        return self.get_success(
+            inject_event(
+                self.hs,
+                prev_event_ids=[prev],
+                room_id=self.room_id,
+                sender=sender,
+                type=event_type,
+                content=content,
+            )
+        ).event_id
+
+    def _prune(
+        self,
+        dropped: str,
+        events: list[tuple[EventBase, EventContext]],
+        keep: str,
+    ) -> set[str]:
+        """Prune with `keep` in the resolved state group and `dropped` in
+        another."""
+        latest = {keep, dropped}
+        groups = {keep: 1, dropped: 2}
+        expected = self.get_success(
+            self.persistence._prune_extremities(
+                self.room_id, set(latest), 1, groups, events
+            )
+        )
+        result, counts = self.rust.prune_extremities(
+            self.room_id,
+            self.hs.hostname,
+            list(latest),
+            1,
+            groups,
+            events,
+            self.clock.time_msec(),
+        )
+        self.assertEqual(result, expected)
+        self.assertEqual(counts["potential_times_prune_extremities"], 1)
+        self.assertEqual(
+            counts["times_pruned_extremities"], 1 if result == {keep} else 0
+        )
+        return result
+
+    def test_prune_extremities(self) -> None:
+        keep = self._send("keep")
+        old_remote = self._inject("@x:remote", keep)
+        old_local = self._inject(self.alice, keep)
+        old_dummy_to_remote = self._inject(self.alice, old_remote, EventTypes.Dummy)
+        old_dummy_to_local = self._inject(
+            self.alice,
+            self._inject(self.alice, old_local, EventTypes.Dummy),
+            EventTypes.Dummy,
+        )
+        # Make the events above more than a day old.
+        self.reactor.advance(2 * 24 * 60 * 60)
+        new_remote = self._inject("@x:other", keep)
+        depth = self.get_success(self.store.get_event(new_remote)).depth
+
+        def chunk(
+            depth_ahead: int, sender: str = "@y:elsewhere"
+        ) -> list[tuple[EventBase, EventContext]]:
+            return [self._new([keep], sender=sender, depth=depth + depth_ahead)]
+
+        # Old and more than 100 deep: goes.
+        self.assertEqual(self._prune(old_remote, chunk(101), keep), {keep})
+        self.assertEqual(self._prune(old_remote, chunk(100), keep), {keep, old_remote})
+        # Local events stay, unless dummy events that lead only to remote
+        # events.
+        self.assertEqual(self._prune(old_local, chunk(500), keep), {keep, old_local})
+        self.assertEqual(self._prune(old_dummy_to_remote, chunk(500), keep), {keep})
+        self.assertEqual(
+            self._prune(old_dummy_to_local, chunk(500), keep),
+            {keep, old_dummy_to_local},
+        )
+        # Recent: goes only if its server sent an event in the chunk, more
+        # than 20 deep.
+        self.assertEqual(self._prune(new_remote, chunk(500), keep), {keep, new_remote})
+        self.assertEqual(self._prune(new_remote, chunk(21, "@z:other"), keep), {keep})
+        self.assertEqual(
+            self._prune(new_remote, chunk(20, "@z:other"), keep), {keep, new_remote}
+        )
+        # Never an extremity being persisted.
+        being_persisted = self._new([keep], depth=depth + 500)
+        self.assertEqual(
+            self._prune(
+                being_persisted[0].event_id, chunk(500) + [being_persisted], keep
+            ),
+            {keep, being_persisted[0].event_id},
+        )