feat: [US-008] - [Explicit wake-time liveness query for hibernated conns (follow-up from US-007 audit)]

Author: NathanFlurry

rivetkit-rust/packages/rivetkit-core/src/actor/context.rs

@@ -1,3 +1,4 @@
+use std::collections::BTreeSet;
 use std::future::Future;
 use std::sync::Arc;
 use std::sync::Weak;
@@ -66,6 +67,8 @@ pub(crate) struct ActorContextInner {
 	inspector: std::sync::RwLock<Option<Inspector>>,
 	actor_events: std::sync::RwLock<Option<mpsc::Sender<ActorEvent>>>,
 	lifecycle_events: std::sync::RwLock<Option<mpsc::Sender<LifecycleEvent>>>,
+	hibernated_connection_liveness_override:
+		std::sync::RwLock<Option<BTreeSet<(Vec<u8>, Vec<u8>)>>>,
 	lifecycle_event_inbox_capacity: usize,
 	metrics: ActorMetrics,
 	diagnostics: ActorDiagnostics,
@@ -206,6 +209,7 @@ impl ActorContext {
 			inspector: std::sync::RwLock::new(None),
 			actor_events: std::sync::RwLock::new(None),
 			lifecycle_events: std::sync::RwLock::new(None),
+			hibernated_connection_liveness_override: std::sync::RwLock::new(None),
 			lifecycle_event_inbox_capacity,
 			metrics,
 			diagnostics,
@@ -801,6 +805,47 @@ impl ActorContext {
 		self.request_save(false);
 	}
 
+	pub(crate) fn hibernated_connection_is_live(
+		&self,
+		gateway_id: &[u8],
+		request_id: &[u8],
+	) -> Result<bool> {
+		if let Some(override_pairs) = self
+			.0
+			.hibernated_connection_liveness_override
+			.read()
+			.expect("hibernated connection liveness override lock poisoned")
+			.as_ref()
+		{
+			return Ok(
+				override_pairs.contains(&(gateway_id.to_vec(), request_id.to_vec()))
+			);
+		}
+
+		// TODO(hibernation-liveness): Replace this with an envoy-client query for
+		// active gateway_id/request_id membership once that API exists. `EnvoyHandle`
+		// currently exposes actor metadata and hibernation restore/ack hooks, but not
+		// a way to ask whether a persisted hibernating request is still live.
+		todo!(
+			"explicit wake-time hibernated connection liveness query requires envoy-client gateway_id/request_id membership support"
+		);
+	}
+
+	#[cfg(test)]
+	pub(crate) fn set_hibernated_connection_liveness_override<I>(
+		&self,
+		pairs: I,
+	) where
+		I: IntoIterator<Item = (Vec<u8>, Vec<u8>)>,
+	{
+		*self
+			.0
+			.hibernated_connection_liveness_override
+			.write()
+			.expect("hibernated connection liveness override lock poisoned") =
+			Some(pairs.into_iter().collect());
+	}
+
 	fn prepare_state_deltas(
 		&self,
 		deltas: Vec<StateDelta>,

rivetkit-rust/packages/rivetkit-core/src/actor/task.rs

@@ -5,7 +5,7 @@ use std::sync::Arc;
 use anyhow::{Context, Result, anyhow};
 use futures::FutureExt;
 use tokio::sync::{mpsc, oneshot};
-use tokio::task::{JoinError, JoinHandle, yield_now};
+use tokio::task::{JoinError, JoinHandle};
 use tokio::time::{Duration, Instant, sleep, sleep_until, timeout};
 
 use crate::actor::action::ActionDispatchError;
@@ -506,7 +506,10 @@ impl ActorTask {
 			.restore_hibernatable_connections()
 			.await
 			.context("restore hibernatable connections")?;
-		self.settle_hibernated_connections().await;
+		self
+			.settle_hibernated_connections()
+			.await
+			.context("settle hibernated connections")?;
 		self.ctx.schedule().sync_future_alarm_logged();
 
 		self.transition_to(LifecycleState::Started);
@@ -565,30 +568,39 @@ impl ActorTask {
 		}));
 	}
 
-	async fn settle_hibernated_connections(&self) {
-		let mut last_count = self
+	async fn settle_hibernated_connections(&self) -> Result<()> {
+		let mut dead_conn_ids = Vec::new();
+
+		for conn in self
 			.ctx
 			.conns()
 			.into_iter()
 			.filter(|conn| conn.is_hibernatable())
-			.count();
-		if last_count == 0 {
-			return;
+		{
+			let Some(hibernation) = conn.hibernation() else {
+				dead_conn_ids.push(conn.id().to_owned());
+				continue;
+			};
+			if self
+				.ctx
+				.hibernated_connection_is_live(
+					&hibernation.gateway_id,
+					&hibernation.request_id,
+				)?
+			{
+				continue;
+			}
+			dead_conn_ids.push(conn.id().to_owned());
 		}
 
-		for _ in 0..8 {
-			yield_now().await;
-			let count = self
+		for conn_id in dead_conn_ids {
+			self
 				.ctx
-				.conns()
-				.into_iter()
-				.filter(|conn| conn.is_hibernatable())
-				.count();
-			if count == last_count {
-				break;
-			}
-			last_count = count;
+				.request_hibernation_transport_removal(conn_id.clone());
+			self.ctx.remove_conn(&conn_id);
 		}
+
+		Ok(())
 	}
 
 	async fn fire_due_alarms(&mut self) -> Result<()> {

rivetkit-rust/packages/rivetkit-core/tests/modules/task.rs

@@ -99,6 +99,17 @@ mod moved_tests {
 		conn
 	}
 
+	fn configure_live_hibernated_pairs(
+		ctx: &ActorContext,
+		pairs: impl IntoIterator<Item = (&'static [u8], &'static [u8])>,
+	) {
+		ctx.set_hibernated_connection_liveness_override(
+			pairs
+				.into_iter()
+				.map(|(gateway_id, request_id)| (gateway_id.to_vec(), request_id.to_vec())),
+		);
+	}
+
 	#[tokio::test]
 	async fn save_tick_respects_debounce_and_immediate_requests() {
 		let ctx = new_with_kv("actor-1", "task-save", Vec::new(), "local", new_in_memory());
@@ -272,6 +283,7 @@ mod moved_tests {
 			request_headers: BTreeMap::from([("x-test".to_owned(), "true".to_owned())]),
 		}));
 		ctx.add_conn(hibernating_conn.clone());
+		configure_live_hibernated_pairs(&ctx, [(b"gateway".as_slice(), b"request".as_slice())]);
 
 		let hibernating_conn_id = hibernating_conn.id().to_owned();
 		let factory = Arc::new(ActorFactory::new(
@@ -381,6 +393,7 @@ mod moved_tests {
 			request_headers: BTreeMap::new(),
 		}));
 		ctx.add_conn(hibernating_conn.clone());
+		configure_live_hibernated_pairs(&ctx, [(b"gateway".as_slice(), b"request".as_slice())]);
 
 		let hibernating_conn_id = hibernating_conn.id().to_owned();
 		let factory = Arc::new(ActorFactory::new(
@@ -493,6 +506,7 @@ mod moved_tests {
 		let (_dispatch_tx, dispatch_rx) = mpsc::channel(4);
 		let (events_tx, events_rx) = mpsc::channel(4);
 		ctx.configure_lifecycle_events(Some(events_tx));
+		configure_live_hibernated_pairs(&ctx, [(b"gateway".as_slice(), b"request".as_slice())]);
 		let (started_tx, started_rx) = oneshot::channel();
 		let started_tx = Arc::new(Mutex::new(Some(started_tx)));
 		let factory = Arc::new(ActorFactory::new(Default::default(), move |start| {
@@ -871,23 +885,10 @@ mod moved_tests {
 		let (_dispatch_tx, dispatch_rx) = mpsc::channel(4);
 		let (events_tx, events_rx) = mpsc::channel(4);
 		ctx.configure_lifecycle_events(Some(events_tx));
-
-		let stale_ctx = ctx.clone();
-		tokio::spawn(async move {
-			for _ in 0..50 {
-				if stale_ctx.conns().len() == 2 {
-					break;
-				}
-				sleep(Duration::from_millis(5)).await;
-			}
-			if let Some(conn) = stale_ctx
-				.conns()
-				.into_iter()
-				.find(|conn| conn.id() == "conn-stale")
-			{
-				let _ = conn.disconnect(Some("stale")).await;
-			}
-		});
+		configure_live_hibernated_pairs(
+			&ctx,
+			[(b"gateway-live".as_slice(), b"request-live".as_slice())],
+		);
 
 		let (started_tx, started_rx) = oneshot::channel();
 		let started_tx = Arc::new(Mutex::new(Some(started_tx)));
@@ -971,4 +972,113 @@ mod moved_tests {
 			.await
 			.expect("sleep stop should succeed");
 	}
+
+	#[tokio::test]
+	async fn wake_start_reaps_dead_hibernated_connections_without_engine_registration() {
+		let kv = new_in_memory();
+		let seed_ctx =
+			new_with_kv("actor-wake-dead", "task-wake", Vec::new(), "local", kv.clone());
+		let dead_conn = ConnHandle::new("conn-dead", Vec::new(), Vec::new(), true);
+		dead_conn.configure_hibernation(Some(HibernatableConnectionMetadata {
+			gateway_id: b"gateway-dead".to_vec(),
+			request_id: b"request-dead".to_vec(),
+			server_message_index: 7,
+			client_message_index: 11,
+			request_path: "/ws".to_owned(),
+			request_headers: BTreeMap::new(),
+		}));
+		seed_ctx.add_conn(dead_conn.clone());
+		seed_ctx
+			.save_state(vec![StateDelta::ConnHibernation {
+				conn: dead_conn.id().into(),
+				bytes: vec![9, 8, 7],
+			}])
+			.await
+			.expect("seed hibernation should persist");
+
+		let ctx =
+			new_with_kv("actor-wake-dead", "task-wake", Vec::new(), "local", kv.clone());
+		let (_lifecycle_tx, lifecycle_rx) = mpsc::channel(4);
+		let (_dispatch_tx, dispatch_rx) = mpsc::channel(4);
+		let (events_tx, events_rx) = mpsc::channel(4);
+		ctx.configure_lifecycle_events(Some(events_tx));
+		ctx.set_hibernated_connection_liveness_override(std::iter::empty());
+
+		let (started_tx, started_rx) = oneshot::channel();
+		let started_tx = Arc::new(Mutex::new(Some(started_tx)));
+		let factory = Arc::new(ActorFactory::new(Default::default(), move |start| {
+			let started_tx = started_tx.clone();
+			Box::pin(async move {
+				let mut events = start.events;
+				started_tx
+					.lock()
+					.expect("started sender lock poisoned")
+					.take()
+					.expect("started sender should exist")
+					.send(start.hibernated.into_iter().map(|(conn, _)| conn.id().to_owned()).collect::<Vec<_>>())
+					.expect("started info should send");
+				while let Some(event) = events.recv().await {
+					match event {
+						ActorEvent::SaveTick { reply } => {
+							reply.send(Ok(Vec::new()));
+						}
+						ActorEvent::Sleep { reply } | ActorEvent::Destroy { reply } => {
+							reply.send(Ok(Vec::new()));
+							break;
+						}
+						ActorEvent::ConnectionOpen { .. } => {
+							panic!("dead hibernated connection should not refire ConnectionOpen");
+						}
+						_ => {}
+					}
+				}
+				Ok(())
+			})
+		}));
+
+		let mut task = ActorTask::new(
+			"actor-wake-dead".into(),
+			0,
+			lifecycle_rx,
+			dispatch_rx,
+			events_rx,
+			factory,
+			ctx.clone(),
+			None,
+			None,
+		);
+		let (start_tx, start_rx) = oneshot::channel();
+		task
+			.handle_lifecycle(LifecycleCommand::Start { reply: start_tx })
+			.await;
+		start_rx
+			.await
+			.expect("start reply should send")
+			.expect("start should succeed");
+
+		assert_eq!(started_rx.await.expect("start info should send"), Vec::<String>::new());
+		assert!(ctx.conns().is_empty());
+
+		task
+			.handle_event(crate::actor::task::LifecycleEvent::SaveRequested {
+				immediate: false,
+			})
+			.await;
+		task.on_state_save_tick().await;
+
+		let last_batch = kv
+			.test_last_apply_batch()
+			.expect("last apply batch should be recorded");
+		assert_eq!(last_batch.deletes, vec![make_connection_key("conn-dead")]);
+		assert!(
+			kv.get(&make_connection_key("conn-dead"))
+				.await
+				.expect("persisted connection lookup should succeed")
+				.is_none()
+		);
+
+		task.handle_stop(StopReason::Sleep)
+			.await
+			.expect("sleep stop should succeed");
+	}
 }