fix(rivetkit): drain shutdown work before sleep

Author: NathanFlurry

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

@@ -32,12 +32,10 @@ use crate::actor::queue::{QueueInspectorUpdateCallback, QueueMetadata, QueueWait
 use crate::actor::schedule::{InternalKeepAwakeCallback, LocalAlarmCallback};
 use crate::actor::sleep::{CanSleep, SleepState};
 use crate::actor::state::{PendingSave, PersistedActor, RequestSaveOpts};
+use crate::actor::task::LifecycleEvent;
 #[cfg(not(target_arch = "wasm32"))]
 use crate::actor::task::{LIFECYCLE_EVENT_INBOX_CHANNEL, actor_channel_overloaded_error};
-use crate::actor::task::LifecycleEvent;
 use crate::actor::task_types::UserTaskKind;
-#[cfg(feature = "wasm-runtime")]
-use crate::actor::work_registry::CountGuard;
 use crate::actor::work_registry::RegionGuard;
 use crate::error::{ActorLifecycle as ActorLifecycleError, ActorRuntime};
 use crate::inspector::{Inspector, InspectorSnapshot};
@@ -107,7 +105,7 @@ pub(crate) struct ActorContextInner {
 	// Forced-sync: queue config is read from sync public methods before blocking
 	// on async queue work.
 	pub(super) queue_config: Mutex<ActorConfig>,
-	pub(super) queue_abort_signal: Option<CancellationToken>,
+	pub(super) queue_abort_signal: Mutex<CancellationToken>,
 	pub(super) queue_initialize: OnceCell<()>,
 	// Forced-sync: startup installs preload before any queue method awaits init.
 	pub(super) queue_preloaded_kv: Mutex<Option<PreloadedKv>>,
@@ -135,7 +133,7 @@ pub(crate) struct ActorContextInner {
 	destroy_requested: AtomicBool,
 	destroy_completed: AtomicBool,
 	destroy_completion_notify: Notify,
-	abort_signal: CancellationToken,
+	abort_signal: Mutex<CancellationToken>,
 	shutdown_deadline: CancellationToken,
 	// Forced-sync: runtime wiring slots are configured through synchronous
 	// lifecycle setup and cloned before sending events.
@@ -280,7 +278,7 @@ impl ActorContext {
 			#[cfg(test)]
 			schedule_driver_alarm_cancel_count: AtomicUsize::new(0),
 			queue_config: Mutex::new(config.clone()),
-			queue_abort_signal: Some(abort_signal.clone()),
+			queue_abort_signal: Mutex::new(abort_signal.clone()),
 			queue_initialize: OnceCell::new(),
 			queue_preloaded_kv: Mutex::new(None),
 			queue_preloaded_message_entries: Mutex::new(None),
@@ -303,7 +301,7 @@ impl ActorContext {
 			destroy_requested: AtomicBool::new(false),
 			destroy_completed: AtomicBool::new(false),
 			destroy_completion_notify: Notify::new(),
-			abort_signal,
+			abort_signal: Mutex::new(abort_signal),
 			shutdown_deadline,
 			inspector: RwLock::new(None),
 			inspector_attach_count: RwLock::new(None),
@@ -481,30 +479,41 @@ impl ActorContext {
 		self.flush_on_shutdown();
 		self.0.destroy_requested.store(true, Ordering::SeqCst);
 		self.0.destroy_completed.store(false, Ordering::SeqCst);
-		self.0.abort_signal.cancel();
+		self.0.abort_signal.lock().cancel();
 	}
 
 	#[cfg(feature = "wasm-runtime")]
 	fn mark_destroy_requested_without_spawn(&self) {
 		self.cancel_sleep_timer();
 		self.0.destroy_requested.store(true, Ordering::SeqCst);
 		self.0.destroy_completed.store(false, Ordering::SeqCst);
-		self.0.abort_signal.cancel();
+		self.0.abort_signal.lock().cancel();
 	}
 
 	#[doc(hidden)]
 	pub fn cancel_abort_signal_for_sleep(&self) {
-		self.0.abort_signal.cancel();
+		self.0.abort_signal.lock().cancel();
+	}
+
+	pub(crate) fn reset_abort_signal_for_start(&self) {
+		let mut abort_signal = self.0.abort_signal.lock();
+		if !abort_signal.is_cancelled() {
+			return;
+		}
+
+		let next_signal = CancellationToken::new();
+		*abort_signal = next_signal.clone();
+		*self.0.queue_abort_signal.lock() = next_signal;
 	}
 
 	#[doc(hidden)]
 	pub fn actor_abort_signal(&self) -> CancellationToken {
-		self.0.abort_signal.clone()
+		self.0.abort_signal.lock().clone()
 	}
 
 	#[doc(hidden)]
 	pub fn actor_aborted(&self) -> bool {
-		self.0.abort_signal.is_cancelled()
+		self.0.abort_signal.lock().is_cancelled()
 	}
 
 	/// Fires when the shutdown grace deadline has elapsed and core is forcing
@@ -562,12 +571,8 @@ impl ActorContext {
 
 	#[cfg(feature = "wasm-runtime")]
 	pub fn wait_until(&self, future: impl Future<Output = ()> + 'static) {
-		let counter = self.0.sleep.work.shutdown_counter.clone();
-		counter.increment();
-		let guard = CountGuard::from_incremented(counter);
 		let ctx = self.clone();
-		wasm_bindgen_futures::spawn_local(async move {
-			let _guard = guard;
+		self.track_shutdown_task(async move {
 			ctx.record_user_task_started(UserTaskKind::WaitUntil);
 			let started_at = Instant::now();
 			future.await;
@@ -1314,6 +1319,10 @@ impl ActorContext {
 		self.0.sleep_requested.load(Ordering::SeqCst)
 	}
 
+	pub(crate) fn clear_sleep_requested(&self) {
+		self.0.sleep_requested.store(false, Ordering::SeqCst);
+	}
+
 	fn keep_awake_guard(&self) -> KeepAwakeGuard {
 		let region = self
 			.keep_awake_region()
@@ -1484,15 +1493,15 @@ impl ActorContext {
 		};
 
 		match sender.try_reserve() {
-				Ok(permit) => {
-					permit.send(event);
-				}
-				#[cfg(target_arch = "wasm32")]
-				Err(_) => {}
-				#[cfg(not(target_arch = "wasm32"))]
-				Err(_) => {
-					let _ = actor_channel_overloaded_error(
-						LIFECYCLE_EVENT_INBOX_CHANNEL,
+			Ok(permit) => {
+				permit.send(event);
+			}
+			#[cfg(target_arch = "wasm32")]
+			Err(_) => {}
+			#[cfg(not(target_arch = "wasm32"))]
+			Err(_) => {
+				let _ = actor_channel_overloaded_error(
+					LIFECYCLE_EVENT_INBOX_CHANNEL,
 					self.0.lifecycle_event_inbox_capacity,
 					operation,
 					Some(&self.0.metrics),

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

@@ -818,25 +818,17 @@ impl ActorContext {
 		timeout: Option<Duration>,
 		signal: Option<&CancellationToken>,
 	) -> WaitOutcome {
+		let actor_abort_signal = self.0.queue_abort_signal.lock().clone();
 		if signal.is_some_and(CancellationToken::is_cancelled) {
 			return WaitOutcome::Aborted;
 		}
-		if self
-			.0
-			.queue_abort_signal
-			.as_ref()
-			.is_some_and(CancellationToken::is_cancelled)
-		{
+		if actor_abort_signal.is_cancelled() {
 			return WaitOutcome::Aborted;
 		}
 
 		let notified = self.0.queue_notify.notified();
 		let actor_aborted = async {
-			if let Some(signal) = &self.0.queue_abort_signal {
-				signal.cancelled().await;
-			} else {
-				pending::<()>().await;
-			}
+			actor_abort_signal.cancelled().await;
 		};
 		let external_aborted = async {
 			if let Some(signal) = signal {

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

@@ -14,14 +14,21 @@ use tracing::Instrument;
 
 use crate::actor::config::ActorConfig;
 use crate::actor::context::ActorContext;
+use crate::actor::task_types::ShutdownKind;
 use crate::actor::work_registry::{CountGuard, RegionGuard, WorkRegistry};
 #[cfg(feature = "wasm-runtime")]
+use crate::actor::work_registry::LocalShutdownTask;
+#[cfg(feature = "wasm-runtime")]
 use crate::runtime::RuntimeSpawner;
-use crate::time::{Instant, sleep};
 #[cfg(test)]
 use crate::time::sleep_until;
+use crate::time::{Instant, sleep};
 #[cfg(test)]
 use crate::types::ActorKey;
+#[cfg(feature = "wasm-runtime")]
+use futures::channel::oneshot as futures_oneshot;
+#[cfg(feature = "wasm-runtime")]
+use futures::future::{AbortHandle, Abortable};
 
 /// Per-actor sleep state.
 ///
@@ -250,8 +257,9 @@ impl ActorContext {
 		CanSleep::Yes
 	}
 
-	pub(crate) fn can_finalize_sleep(&self) -> bool {
+	pub(crate) fn can_finalize_shutdown(&self, reason: ShutdownKind) -> bool {
 		self.0.sleep.work.core_dispatched_hooks.load() == 0
+			&& (matches!(reason, ShutdownKind::Destroy) || !self.run_handler_active())
 			&& self.shutdown_task_count() == 0
 			&& self.sleep_keep_awake_count() == 0
 			&& self.sleep_internal_keep_awake_count() == 0
@@ -359,7 +367,7 @@ impl ActorContext {
 			tokio::pin!(idle);
 			idle.as_mut().enable();
 
-			if self.can_finalize_sleep() {
+			if self.can_finalize_shutdown(ShutdownKind::Sleep) {
 				return true;
 			}
 
@@ -472,12 +480,16 @@ impl ActorContext {
 		let ctx = self.clone();
 		shutdown_tasks.spawn(
 			async move {
-				let _guard = guard;
-				fut.await;
+				{
+					let _guard = guard;
+					fut.await;
+				}
 				ctx.reset_sleep_timer();
 			}
 			.in_current_span(),
 		);
+		drop(shutdown_tasks);
+		self.reset_sleep_timer();
 		true
 	}
 
@@ -486,6 +498,7 @@ impl ActorContext {
 	where
 		F: Future<Output = ()> + 'static,
 	{
+		let mut local_shutdown_tasks = self.0.sleep.work.local_shutdown_tasks.lock();
 		if self.0.sleep.work.teardown_started.load(Ordering::Acquire) {
 			tracing::warn!("shutdown task spawned after teardown; aborting immediately");
 			return false;
@@ -494,21 +507,46 @@ impl ActorContext {
 		counter.increment();
 		let guard = CountGuard::from_incremented(counter);
 		let ctx = self.clone();
+		let (complete_tx, complete_rx) = futures_oneshot::channel();
+		let (abort_handle, abort_registration) = AbortHandle::new_pair();
+		local_shutdown_tasks.push(LocalShutdownTask {
+			abort_handle,
+			complete_rx,
+		});
+		drop(local_shutdown_tasks);
+		let ctx_for_task = ctx.clone();
 		wasm_bindgen_futures::spawn_local(
 			async move {
-				let _guard = guard;
-				fut.await;
-				ctx.reset_sleep_timer();
+				let task = async move {
+					{
+						let _guard = guard;
+						fut.await;
+					}
+					let _ = complete_tx.send(());
+					ctx_for_task.reset_sleep_timer();
+				};
+				if Abortable::new(task, abort_registration).await.is_err() {
+					ctx.reset_sleep_timer();
+				}
 			}
 			.in_current_span(),
 		);
+		self.reset_sleep_timer();
 		true
 	}
 
 	pub(crate) fn shutdown_task_count(&self) -> usize {
 		self.0.sleep.work.shutdown_counter.load()
 	}
 
+	pub(crate) fn mark_shutdown_deadline_reached(&self) {
+		self.0
+			.sleep
+			.work
+			.shutdown_deadline_reached
+			.store(true, Ordering::Release);
+	}
+
 	pub(crate) fn begin_core_dispatched_hook(&self) {
 		self.0.sleep.work.core_dispatched_hooks.increment();
 		self.reset_sleep_timer();
@@ -524,17 +562,88 @@ impl ActorContext {
 	}
 
 	pub(crate) async fn teardown_sleep_state(&self) {
-		self.0
+		let abort_remaining = self
+			.0
 			.sleep
 			.work
-			.teardown_started
-			.store(true, Ordering::Release);
-		let mut shutdown_tasks = {
-			let mut guard = self.0.sleep.work.shutdown_tasks.lock();
-			std::mem::take(&mut *guard)
-		};
-		shutdown_tasks.shutdown().await;
-		*self.0.sleep.work.shutdown_tasks.lock() = shutdown_tasks;
+			.shutdown_deadline_reached
+			.swap(false, Ordering::AcqRel);
+		if abort_remaining {
+			self.0
+				.sleep
+				.work
+				.teardown_started
+				.store(true, Ordering::Release);
+		}
+
+		#[cfg(feature = "wasm-runtime")]
+		{
+			loop {
+				let local_shutdown_tasks = {
+					let mut guard = self.0.sleep.work.local_shutdown_tasks.lock();
+					let taken = std::mem::take(&mut *guard);
+					if taken.is_empty() {
+						self.0
+							.sleep
+							.work
+							.teardown_started
+							.store(true, Ordering::Release);
+						return;
+					}
+					taken
+				};
+
+				if abort_remaining {
+					for task in local_shutdown_tasks {
+						task.abort_handle.abort();
+						if task.complete_rx.await.is_err() {
+							tracing::debug!("aborted shutdown task during teardown");
+						}
+					}
+					self.0
+						.sleep
+						.work
+						.teardown_started
+						.store(true, Ordering::Release);
+					return;
+				}
+
+				for task in local_shutdown_tasks {
+					if task.complete_rx.await.is_err() {
+						tracing::debug!("shutdown task completion dropped during teardown");
+					}
+				}
+			}
+		}
+
+		#[cfg(not(feature = "wasm-runtime"))]
+		loop {
+			let mut shutdown_tasks = {
+				let mut guard = self.0.sleep.work.shutdown_tasks.lock();
+				let taken = std::mem::take(&mut *guard);
+				if taken.is_empty() {
+					self.0
+						.sleep
+						.work
+						.teardown_started
+						.store(true, Ordering::Release);
+					return;
+				}
+				taken
+			};
+
+			if abort_remaining {
+				shutdown_tasks.shutdown().await;
+			} else {
+				while let Some(result) = shutdown_tasks.join_next().await {
+					if let Err(error) = result
+						&& !error.is_cancelled()
+					{
+						tracing::error!(?error, "shutdown task join failed during teardown");
+					}
+				}
+			}
+		}
 	}
 
 	pub(crate) fn sleep_state_config(&self) -> ActorConfig {