From 4e957623325bd9683213c549f5ab8c25be2fbb7f Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Wed, 19 Aug 2026 22:49:08 -0700 Subject: [PATCH] fix(rivetkit): prevent connection snapshot stalls --- .../rivetkit-core/src/actor/connection.rs | 48 +++++++++---------- .../rivetkit-core/src/actor/context.rs | 7 ++- .../rivetkit-core/tests/connection.rs | 42 ++++++++++++++++ 3 files changed, 69 insertions(+), 28 deletions(-) diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/connection.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/connection.rs index a1a7ba1428..49c08b89b8 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/connection.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/connection.rs @@ -1,13 +1,14 @@ use std::collections::{BTreeMap, BTreeSet}; use std::fmt; -use std::ops::Bound::{Excluded, Unbounded}; +use std::iter::FusedIterator; +use std::marker::PhantomData; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; use anyhow::{Context, Result}; use futures::future::BoxFuture; -use parking_lot::{RwLock, RwLockReadGuard}; +use parking_lot::RwLock; use rivet_error::RivetError; use rivetkit_actor_persist::{generated::v4 as persist_v4, versioned as persist_versioned}; use serde::Serialize; @@ -418,31 +419,30 @@ pub(crate) struct PendingHibernationChanges { pub removed: BTreeSet, } -/// Lock-backed iterator over live connection handles. +/// Point-in-time snapshot iterator over live connection handles. /// -/// Do not hold this iterator across `.await`. It keeps a read lock on the -/// connection map until dropped, which blocks writers such as add/remove or -/// connection reconfiguration. -#[must_use = "connection iterators hold a read lock until dropped"] +/// The snapshot captures membership only. Its handles remain live and may be +/// disconnected or removed while the iterator exists. +#[must_use] pub struct ConnHandles<'a> { - guard: RwLockReadGuard<'a, BTreeMap>, - next_after: Option, + inner: std::vec::IntoIter, + _marker: PhantomData<&'a ()>, } impl<'a> ConnHandles<'a> { - fn new(guard: RwLockReadGuard<'a, BTreeMap>) -> Self { + fn new(connections: Vec) -> Self { Self { - guard, - next_after: None, + inner: connections.into_iter(), + _marker: PhantomData, } } pub fn len(&self) -> usize { - self.guard.len() + self.inner.len() } pub fn is_empty(&self) -> bool { - self.guard.is_empty() + self.inner.len() == 0 } } @@ -450,25 +450,25 @@ impl Iterator for ConnHandles<'_> { type Item = ConnHandle; fn next(&mut self) -> Option { - let (conn_id, conn) = match self.next_after.as_ref() { - Some(conn_id) => self - .guard - .range((Excluded(conn_id.clone()), Unbounded)) - .next()?, - None => self.guard.iter().next()?, - }; - self.next_after = Some(conn_id.clone()); - Some(conn.clone()) + self.inner.next() + } + + fn size_hint(&self) -> (usize, Option) { + self.inner.size_hint() } } +impl ExactSizeIterator for ConnHandles<'_> {} + +impl FusedIterator for ConnHandles<'_> {} + impl ActorContext { pub(crate) fn configure_connection_storage(&self, config: ActorConfig) { *self.0.connection_config.write() = config; } pub(crate) fn iter_connections(&self) -> ConnHandles<'_> { - ConnHandles::new(self.0.connections.read()) + ConnHandles::new(self.0.connections.read().values().cloned().collect()) } pub(crate) fn active_connection_count(&self) -> u32 { diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs index fbe40d9740..6326bdc3f4 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs @@ -790,11 +790,10 @@ impl ActorContext { } } - /// Returns a lock-backed iterator over live connections. + /// Returns a point-in-time snapshot iterator over live connections. /// - /// Do not hold the returned iterator across `.await`. It keeps a read lock - /// on the connection map until dropped, which blocks connection writers. - #[must_use] + /// The snapshot captures membership only. Its handles remain live and may be + /// disconnected or removed while the iterator exists. pub fn conns(&self) -> ConnHandles<'_> { self.iter_connections() } diff --git a/rivetkit-rust/packages/rivetkit-core/tests/connection.rs b/rivetkit-rust/packages/rivetkit-core/tests/connection.rs index 5162f6c80f..e05fabc4ff 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/connection.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/connection.rs @@ -4,6 +4,7 @@ mod moved_tests { use std::collections::BTreeSet; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::mpsc as std_mpsc; use std::time::Duration; use parking_lot::Mutex; @@ -31,6 +32,47 @@ mod moved_tests { assert_eq!(make_connection_key("conn-1"), b"\x02conn-1".to_vec()); } + #[test] + fn retained_connection_snapshot_does_not_block_disconnect_writer() { + let ctx = ActorContext::new_with_kv( + "actor-connection-snapshot", + "actor", + Vec::new(), + "local", + Kv::new_in_memory(), + ); + ctx.insert_existing(super::ConnHandle::new( + "conn-snapshot", + Vec::new(), + Vec::new(), + false, + )); + let snapshot = ctx.conns(); + let (started_tx, started_rx) = std_mpsc::channel(); + let (removed_tx, removed_rx) = std_mpsc::channel(); + let writer = std::thread::spawn({ + let ctx = ctx.clone(); + move || { + started_tx.send(()).expect("test receiver should stay open"); + let removed = ctx.remove_existing("conn-snapshot").is_some(); + removed_tx + .send(removed) + .expect("test receiver should stay open"); + } + }); + + started_rx.recv().expect("disconnect writer should start"); + let removed = removed_rx.recv_timeout(Duration::from_secs(2)); + let snapshot_ids = snapshot + .map(|conn| conn.id().to_owned()) + .collect::>(); + writer.join().expect("disconnect writer should join"); + + assert!(removed.expect("retained snapshot must not block connection writer")); + assert_eq!(snapshot_ids, vec!["conn-snapshot".to_owned()]); + assert!(ctx.conns().is_empty()); + } + #[tokio::test] async fn pending_connection_is_invisible_until_preflight_succeeds() { let ctx = ActorContext::new_with_kv(