Fix the prefetched cache stream IDs on retry and while in flight

A write transaction that failed after prefetching cache stream IDs
finished all of them, but left the unused ones spare. The retry handed
those out first, so it wrote cache invalidations at IDs below the
stream's position, which other workers never fetch. Finishing the
pending IDs now drops the spare ones too.

The prefetch also began its fetch only once the `nextval` rows were
back, so the position could pass an ID still on its way. It now begins
before the reads are sent, as `next_ids` does, and ends unallocated if
they fail.

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

rust/src/persist/txn/mod.rs

@@ -338,15 +338,14 @@ async fn persist_events_txn_steps<'a>(
             ),
         );
     }
+    let prefetch = input.cache_stream.begin_prefetch();
     let mut results = txn.query_many(&queries).await?.into_iter();
     if new_room.is_none() && results.next().is_some_and(|rows| rows.is_empty()) {
         return Err(anyhow!("Room does not exist {room_id}").into());
     }
     let outlier_rows = results.next().unwrap_or_default();
     let redaction_rows = results.next().unwrap_or_default();
-    input
-        .cache_stream
-        .add_prefetched(&results.next().unwrap_or_default())?;
+    prefetch.allocated(&results.next().unwrap_or_default())?;
 
     let events = events::update_outliers(txn, input, events, outlier_rows, effects).await?;
 

rust/src/persist/txn/streams.rs

@@ -100,17 +100,19 @@ impl TxnStream {
     }
 
     /// The Postgres sequence the IDs come from, for a [`NEXTVAL_SQL`] sent
-    /// with other reads, whose rows go to [`TxnStream::add_prefetched`].
+    /// with other reads.
     pub fn sequence(&self) -> &'static str {
         self.sequence
     }
 
-    /// Keep the IDs `rows` of a [`NEXTVAL_SQL`] allocated, for
-    /// [`TxnStream::next_ids`] to hand out without a round trip of its own.
-    pub fn add_prefetched(&self, rows: &[Row]) -> anyhow::Result<()> {
-        let ids = Fetch::begin(self).allocated_rows(rows)?;
-        self.lock_spare().extend(ids);
-        Ok(())
+    /// Start fetching IDs with a [`NEXTVAL_SQL`] sent with other reads.
+    /// Call this before sending it, as `next_ids` begins its fetch before
+    /// its query: otherwise the stream's position can pass the IDs while
+    /// they are on their way back.
+    pub fn begin_prefetch(&self) -> Prefetch<'_> {
+        Prefetch {
+            fetch: Fetch::begin(self),
+        }
     }
 
     /// `_load_next_id_txn`: the next ID on the stream, unfinished.
@@ -154,6 +156,9 @@ impl TxnStream {
     /// Mark finished the IDs of an attempt that failed, as Python's
     /// `call_on_exception` does.
     pub fn finish_pending(&self) {
+        // The spare IDs are pending too, so they are finished with the rest
+        // and must not be handed out again.
+        self.lock_spare().clear();
         let ids = std::mem::take(&mut *self.lock_pending());
         if !ids.is_empty() {
             self.tracker.mark_finished(&ids);
@@ -167,6 +172,31 @@ impl TxnStream {
     }
 }
 
+/// A fetch begun by [`TxnStream::begin_prefetch`]. It ends without
+/// allocating if dropped before [`Prefetch::allocated`].
+pub struct Prefetch<'a> {
+    fetch: Fetch<'a>,
+}
+
+impl Prefetch<'_> {
+    /// Keep the IDs `rows` of the [`NEXTVAL_SQL`] allocated, for
+    /// [`TxnStream::next_ids`] to hand out without a round trip of its own.
+    pub fn allocated(self, rows: &[Row]) -> anyhow::Result<()> {
+        let ids = rows
+            .iter()
+            .map(|row| row.try_get(0))
+            .collect::<Result<Vec<i64>, _>>()?;
+        self.allocate_ids(&ids);
+        Ok(())
+    }
+
+    fn allocate_ids(self, ids: &[i64]) {
+        let stream = self.fetch.stream;
+        self.fetch.allocated(ids);
+        stream.lock_spare().extend_from_slice(ids);
+    }
+}
+
 /// A fetch of IDs from the sequence, between the tracker's `begin_fetch`
 /// and `end_fetch`. Python's `try`/`finally` cannot be interrupted, but a
 /// future can be dropped at any await, so the fetch ends on drop if it did
@@ -270,4 +300,34 @@ mod tests {
         // A fetch still in flight would hold the position at 1.
         assert_eq!(tracker.get_persisted_upto_position(), 2);
     }
+
+    #[test]
+    fn a_failed_attempt_leaves_no_spare_ids() {
+        let tracker = tracker();
+        let stream = TxnStream::caches(Arc::clone(&tracker));
+        stream.begin_prefetch().allocate_ids(&[2, 3, 4]);
+        // The attempt failed: its IDs, spare ones included, are finished.
+        stream.finish_pending();
+        assert_eq!(tracker.get_persisted_upto_position(), 4);
+        // So the next attempt must not hand them out.
+        assert!(stream.lock_spare().is_empty());
+    }
+
+    #[test]
+    fn a_prefetch_holds_the_position_until_it_allocates() {
+        let tracker = tracker();
+        let stream = TxnStream::caches(Arc::clone(&tracker));
+        let other = TxnStream::caches(Arc::clone(&tracker));
+        // Sent with the reads, and the sequence has answered 2.
+        let prefetch = stream.begin_prefetch();
+        // Meanwhile another fetch on this instance takes 3 and finishes it.
+        allocate(&other, &[3]);
+        other.finish_pending();
+        // The prefetch may still be holding 2, so the position must wait.
+        assert_eq!(tracker.get_persisted_upto_position(), 1);
+        prefetch.allocate_ids(&[2]);
+        assert_eq!(tracker.get_persisted_upto_position(), 1);
+        stream.finish_pending();
+        assert_eq!(tracker.get_persisted_upto_position(), 3);
+    }
 }