fix(rivetkit): prevent connection snapshot stalls

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,57 +419,56 @@ pub(crate) struct PendingHibernationChanges {
 	pub removed: BTreeSet<ConnId>,
 }
 
-/// 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<ConnId, ConnHandle>>,
-	next_after: Option<ConnId>,
+	inner: std::vec::IntoIter<ConnHandle>,
+	_marker: PhantomData<&'a ()>,
 }
 
 impl<'a> ConnHandles<'a> {
-	fn new(guard: RwLockReadGuard<'a, BTreeMap<ConnId, ConnHandle>>) -> Self {
+	fn new(connections: Vec<ConnHandle>) -> 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
 	}
 }
 
 impl Iterator for ConnHandles<'_> {
 	type Item = ConnHandle;
 
 	fn next(&mut self) -> Option<Self::Item> {
-		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<usize>) {
+		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 {

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()
 	}

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::<Vec<_>>();
+		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(