feat: US-012 - Bound register_task shutdown drain against shutdown deadline

Author: NathanFlurry

rivetkit-rust/packages/rivetkit-core/CLAUDE.md

@@ -9,7 +9,7 @@
 
 - Any mutation that changes a `can_sleep` input must call `ActorContext::reset_sleep_timer()` so the `ActorTask` sleep deadline is re-evaluated. Inputs are: `ready`/`started`, `no_sleep`, `active_http_request_count`, `sleep_keep_awake_count`, `sleep_internal_keep_awake_count`, `pending_disconnect_count`, `conns()`, and `websocket_callback_count`. Missing this call leaves the sleep timer armed against stale state and triggers the `"sleep idle deadline elapsed but actor stayed awake"` warning on the next tick.
 - `ActorContext::set_prevent_sleep(...)` / `prevent_sleep()` are deprecated no-ops kept for NAPI bridge compatibility. Use `keep_awake(future)` (holds counter while awaited) or `wait_until(future)` (tracked shutdown task) instead. Do not reintroduce a `prevent_sleep` field, a `CanSleep::PreventSleep` variant, or branches that read it.
-- Runtime-owned promises that must drain during shutdown should use `ActorContext::register_task(...)`, not public `wait_until(...)`, so metrics and runtime intent stay distinct.
+- Runtime-owned promises that must drain during shutdown should use `ActorContext::register_task(...)`, not public `wait_until(...)`, so metrics and runtime intent stay distinct. Registered tasks must race user work against `shutdown_deadline_token()` so shutdown cannot hang forever.
 - `ctx.sleep()` and `ctx.destroy()` return `Result<()>`. They error with `ActorLifecycleError::Starting` when called before startup completes and `ActorLifecycleError::Stopping` if the requested flag has already been set this generation (atomic `swap(true, ...)`). Internal idle-timer paths log and suppress the already-requested error.
 - The grace deadline path (`on_sleep_grace_deadline`) aborts the user `run` handle and cancels `shutdown_deadline_token()`. Foreign-runtime adapters running `onSleep` / `onDestroy` must observe that token via `tokio::select!` so SQLite teardown does not race user cleanup work.
 - Counter `register_zero_notify(&idle_notify)` hooks only drive shutdown drain waits. They are not a substitute for the activity-dirty notification, so any new sleep-affecting counter must also notify on transitions that change `can_sleep`.

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

@@ -556,10 +556,7 @@ impl ActorContext {
 	pub fn register_task(&self, future: impl Future<Output = ()> + Send + 'static) {
 		let ctx = self.clone();
 		self.track_shutdown_task(async move {
-			ctx.record_user_task_started(UserTaskKind::RegisteredTask);
-			let started_at = Instant::now();
-			future.await;
-			ctx.record_user_task_finished(UserTaskKind::RegisteredTask, started_at.elapsed());
+			Self::run_registered_task(ctx, future).await;
 		});
 	}
 
@@ -583,13 +580,30 @@ impl ActorContext {
 	pub fn register_task(&self, future: impl Future<Output = ()> + 'static) {
 		let ctx = self.clone();
 		self.track_shutdown_task(async move {
-			ctx.record_user_task_started(UserTaskKind::RegisteredTask);
-			let started_at = Instant::now();
-			future.await;
-			ctx.record_user_task_finished(UserTaskKind::RegisteredTask, started_at.elapsed());
+			Self::run_registered_task(ctx, future).await;
 		});
 	}
 
+	async fn run_registered_task<F>(ctx: ActorContext, future: F)
+	where
+		F: Future<Output = ()>,
+	{
+		let shutdown_deadline = ctx.shutdown_deadline_token();
+		ctx.record_user_task_started(UserTaskKind::RegisteredTask);
+		let started_at = Instant::now();
+		tokio::select! {
+			_ = future => {}
+			_ = shutdown_deadline.cancelled() => {
+				tracing::warn!(
+					actor_id = %ctx.actor_id(),
+					reason = "shutdown_deadline_elapsed",
+					"registered task cancelled by shutdown deadline"
+				);
+			}
+		}
+		ctx.record_user_task_finished(UserTaskKind::RegisteredTask, started_at.elapsed());
+	}
+
 	pub async fn keep_awake<F>(&self, future: F) -> F::Output
 	where
 		F: Future,

rivetkit-rust/packages/rivetkit-core/tests/sleep.rs

@@ -19,18 +19,27 @@ mod moved_tests {
 	#[derive(Default)]
 	struct MessageVisitor {
 		message: Option<String>,
+		actor_id: Option<String>,
+		reason: Option<String>,
 	}
 
 	impl Visit for MessageVisitor {
 		fn record_str(&mut self, field: &Field, value: &str) {
-			if field.name() == "message" {
-				self.message = Some(value.to_owned());
+			match field.name() {
+				"message" => self.message = Some(value.to_owned()),
+				"actor_id" => self.actor_id = Some(value.to_owned()),
+				"reason" => self.reason = Some(value.to_owned()),
+				_ => {}
 			}
 		}
 
 		fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
-			if field.name() == "message" {
-				self.message = Some(format!("{value:?}").trim_matches('"').to_owned());
+			let value = format!("{value:?}").trim_matches('"').to_owned();
+			match field.name() {
+				"message" => self.message = Some(value),
+				"actor_id" => self.actor_id = Some(value),
+				"reason" => self.reason = Some(value),
+				_ => {}
 			}
 		}
 	}
@@ -40,6 +49,11 @@ mod moved_tests {
 		count: Arc<AtomicUsize>,
 	}
 
+	#[derive(Clone)]
+	struct RegisteredTaskDeadlineLayer {
+		count: Arc<AtomicUsize>,
+	}
+
 	impl<S> Layer<S> for ShutdownTaskRefusedLayer
 	where
 		S: Subscriber,
@@ -59,6 +73,27 @@ mod moved_tests {
 		}
 	}
 
+	impl<S> Layer<S> for RegisteredTaskDeadlineLayer
+	where
+		S: Subscriber,
+	{
+		fn on_event(&self, event: &Event<'_>, _ctx: LayerContext<'_, S>) {
+			if *event.metadata().level() != tracing::Level::WARN {
+				return;
+			}
+
+			let mut visitor = MessageVisitor::default();
+			event.record(&mut visitor);
+			if visitor.message.as_deref()
+				== Some("registered task cancelled by shutdown deadline")
+				&& visitor.actor_id.as_deref() == Some("actor-register-task-deadline")
+				&& visitor.reason.as_deref() == Some("shutdown_deadline_elapsed")
+			{
+				self.count.fetch_add(1, Ordering::SeqCst);
+			}
+		}
+	}
+
 	struct NotifyOnDrop(DropMutex<Option<oneshot::Sender<()>>>);
 
 	impl NotifyOnDrop {
@@ -162,6 +197,33 @@ mod moved_tests {
 		assert_eq!(warning_count.load(Ordering::SeqCst), 1);
 	}
 
+	#[tokio::test(start_paused = true)]
+	async fn register_task_exits_when_shutdown_deadline_cancels() {
+		let ctx = ActorContext::new_for_sleep_tests("actor-register-task-deadline");
+		let warning_count = Arc::new(AtomicUsize::new(0));
+		let subscriber = Registry::default().with(RegisteredTaskDeadlineLayer {
+			count: warning_count.clone(),
+		});
+		let _guard = tracing::subscriber::set_default(subscriber);
+
+		ctx.register_task(futures::future::pending::<()>());
+		assert_eq!(ctx.shutdown_task_count(), 1);
+
+		ctx.cancel_shutdown_deadline();
+
+		assert!(
+			ctx.0
+				.sleep
+				.work
+				.shutdown_counter
+				.wait_zero(Instant::now() + Duration::from_millis(1))
+				.await,
+			"registered task should stop waiting after the shutdown deadline"
+		);
+		assert_eq!(ctx.shutdown_task_count(), 0);
+		assert_eq!(warning_count.load(Ordering::SeqCst), 1);
+	}
+
 	#[tokio::test(start_paused = true)]
 	async fn sleep_then_destroy_signal_tasks_do_not_leak_after_teardown() {
 		let ctx = ActorContext::new_for_sleep_tests("actor-sleep-destroy");

scripts/ralph/prd.json

@@ -195,8 +195,8 @@
         "Tests pass"
       ],
       "priority": 12,
-      "passes": false,
-      "notes": ""
+      "passes": true,
+      "notes": "Core register_task now races registered futures against the shutdown deadline token for both native and wasm cfgs. A focused core test verifies a never-completing task drains after deadline cancellation and logs the stable reason."
     },
     {
       "id": "US-013",

scripts/ralph/progress.txt

@@ -19,6 +19,7 @@
 - NAPI actor-event synthetic `RivetErrorSchema` values should be `static` when fields are compile-time constants, or process-interned by `(group, code)` with `scc::HashMap` when the default message is dynamic.
 - Wasm websocket callback regions should be tracked in a map keyed by monotonic region IDs and removed on end so callback churn does not retain empty slots.
 - NAPI `Ref::unref` needs an `Env`; cleanup from worker or `Drop` paths should be routed through an Env-bearing TSF and must tolerate addon shutdown by falling back to a bounded leak.
+- Core `ActorContext::register_task(...)` must race registered runtime promises against `shutdown_deadline_token()` so shutdown drain cannot hang forever.
 
 Started: Sat May  2 02:13:32 AM PDT 2026
 ---
@@ -133,3 +134,13 @@ Started: Sat May  2 02:13:32 AM PDT 2026
   - Use unique `(group, code)` namespaces when asserting `BRIDGE_RIVET_ERROR_SCHEMAS.len()` deltas because the intern map is process-global.
   - Schema defaults are only the first-seen fallback for a `(group, code)`; callers should assert live bridged messages through `RivetTransportError.message()`.
 ---
+## 2026-05-02 03:30:40 PDT - US-012
+- Implemented bounded shutdown drain for core `ActorContext::register_task(...)` by racing registered futures against `shutdown_deadline_token()` in one shared native/wasm helper.
+- Added a core sleep regression test that registers a never-completing future, cancels the shutdown deadline, verifies the shutdown counter drains, and checks the warning includes `actor_id` plus `reason = "shutdown_deadline_elapsed"`.
+- Files changed: `rivetkit-rust/packages/rivetkit-core/src/actor/context.rs`, `rivetkit-rust/packages/rivetkit-core/tests/sleep.rs`, `rivetkit-rust/packages/rivetkit-core/AGENTS.md`, `scripts/ralph/prd.json`, `scripts/ralph/progress.txt`.
+- Checks: `cargo test -p rivetkit-core register_task_exits_when_shutdown_deadline_cancels`, `cargo check -p rivetkit-core`, `scripts/cargo/check-rivetkit-core-wasm.sh`, `cargo test -p rivetkit-core actor_task_logs_lifecycle_dispatch_and_actor_event_flow`.
+- Full `cargo test -p rivetkit-core` was stopped after it hung in the existing `save_tick_cancels_pending_inspector_deadline_and_broadcasts_overlay` test; that test also hung when run in isolation.
+- **Learnings for future iterations:**
+  - Keep native and wasm `register_task(...)` behavior in a shared helper where possible so the two cfg-specific public signatures cannot diverge.
+  - The package-wide core test suite currently has an unrelated hanging inspector debounce test, so use focused tests plus `check-rivetkit-core-wasm.sh` for this shutdown path.
+---