fix(rivetkit): release HTTP cancellation resources

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

@@ -111,6 +111,13 @@ impl Request {
 		self.cancel_token.clone()
 	}
 
+	/// Replaces the request token with a child of `parent`, so cancelling the
+	/// parent (for example the actor abort signal) also cancels this request.
+	/// Call it before any clone of the token is handed out.
+	pub(crate) fn set_cancellation_parent(&mut self, parent: &CancellationToken) {
+		self.cancel_token = parent.child_token();
+	}
+
 	pub async fn into_buffered(mut self) -> Result<Self> {
 		if let Some(body_stream) = &self.body_stream {
 			let mut body_stream = body_stream.lock().take();

rivetkit-rust/packages/rivetkit-core/src/registry/http.rs

@@ -10,6 +10,12 @@ const HEADER_RIVET_ACTOR: &str = "x-rivet-actor";
 const HEADER_RIVET_ACTOR_GENERATION: &str = "x-rivet-actor-generation";
 const HEADER_RIVET_ACTOR_KEY: &str = "x-rivet-actor-key";
 
+/// Cancels a user request's token when the request finishes, unless a
+/// streaming response body takes ownership of it.
+///
+/// Every terminal path must cancel the token. Foreign runtimes park a task per
+/// request on it to deliver `abortSignal`, and a token that is dropped without
+/// being cancelled strands that task and its JS callback until the actor stops.
 struct RequestCancellationGuard {
 	token: Option<tokio_util::sync::CancellationToken>,
 }
@@ -18,10 +24,6 @@ impl RequestCancellationGuard {
 	fn new(token: tokio_util::sync::CancellationToken) -> Self {
 		Self { token: Some(token) }
 	}
-
-	fn disarm(&mut self) {
-		self.token = None;
-	}
 }
 
 impl Drop for RequestCancellationGuard {
@@ -100,10 +102,11 @@ impl RegistryDispatcher {
 	async fn handle_user_request_fetch(
 		&self,
 		instance: &ActorTaskHandle,
-		request: Request,
+		mut request: Request,
 	) -> Result<HttpResponse> {
 		let encoding = request_encoding(request.headers());
 		let request_method = request.method().clone();
+		request.set_cancellation_parent(&instance.ctx.actor_abort_signal());
 		let mut cancellation_guard = RequestCancellationGuard::new(request.cancellation_token());
 		let (reply_tx, reply_rx) = oneshot::channel();
 		try_send_dispatch_command(
@@ -123,14 +126,14 @@ impl RegistryDispatcher {
 			Ok(response) => {
 				rearm_sleep_after_request(instance.ctx.clone());
 				let mut response = build_envoy_response(response, &request_method)?;
+				// A streaming body outlives this call, so it cancels the token
+				// when it is dropped. Otherwise the guard cancels it on return.
 				if let Some(body_stream) = response.body_stream.as_mut() {
 					let token = cancellation_guard
 						.token
 						.take()
 						.expect("request cancellation token must be armed");
 					body_stream.set_on_drop(move || token.cancel());
-				} else {
-					cancellation_guard.disarm();
 				}
 				Ok(response)
 			}

rivetkit-rust/packages/rivetkit-core/src/serverless_http.rs

@@ -194,6 +194,8 @@ async fn forward_application_request(
 		Ok(response) => into_application_response(response, request_token),
 		Err(error) => {
 			tracing::error!(?error, "application request handler failed");
+			// The handler may have registered cancellation listeners on the token.
+			request_token.cancel();
 			(
 				StatusCode::INTERNAL_SERVER_ERROR,
 				[("content-type", "text/plain; charset=utf-8")],
@@ -258,6 +260,8 @@ async fn forward_request(
 		Ok(response) => into_application_response(response, request_token),
 		Err(error) => {
 			tracing::error!(?error, "application request handler failed");
+			// The handler may have registered cancellation listeners on the token.
+			request_token.cancel();
 			(
 				StatusCode::INTERNAL_SERVER_ERROR,
 				[("content-type", "text/plain; charset=utf-8")],

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

@@ -1,7 +1,10 @@
 use std::sync::Arc;
+use std::time::Duration;
 
 use super::*;
 use crate::actor::config::ActorConfig;
+use crate::actor::context::ActorContext;
+use crate::actor::messages::{ActorEvent, Response};
 
 #[test]
 fn actor_metadata_does_not_request_legacy_kv_preload_probes() {
@@ -276,3 +279,168 @@ async fn stop_for_current_generation_stops_it() {
 		"a stop for the current generation must stop it",
 	);
 }
+
+/// Builds a dispatcher whose actor answers every raw request with a buffered
+/// response. For each request it parks a task on the request cancellation
+/// token, the way foreign runtimes wait to deliver `abortSignal`, and hands
+/// that task's handle to the test.
+fn request_waiter_dispatcher(
+	waiters: tokio::sync::mpsc::UnboundedSender<tokio::task::JoinHandle<()>>,
+	hang_until_cancelled: bool,
+) -> Arc<RegistryDispatcher> {
+	let mut factories = HashMap::new();
+	factories.insert(
+		"counter".to_owned(),
+		Arc::new(ActorFactory::new(ActorConfig::default(), move |start| {
+			let waiters = waiters.clone();
+			Box::pin(async move {
+				let mut start = start;
+				if let Some(ready) = start.startup_ready.take() {
+					let _ = ready.send(Ok(()));
+				}
+				let mut events = start.events;
+				while let Some(event) = events.recv().await {
+					let ActorEvent::HttpRequest { request, reply, .. } = event else {
+						continue;
+					};
+					let token = request.cancellation_token();
+					let waiter = tokio::spawn({
+						let token = token.clone();
+						async move { token.cancelled().await }
+					});
+					let _ = waiters.send(waiter);
+					if hang_until_cancelled {
+						tokio::spawn(async move {
+							token.cancelled().await;
+							reply.send(Ok(Response::from(
+								::http::Response::builder()
+									.status(::http::StatusCode::SERVICE_UNAVAILABLE)
+									.body(Vec::new())
+									.expect("build response"),
+							)
+							.into()));
+						});
+					} else {
+						reply.send(Ok(Response::from(
+							::http::Response::builder()
+								.status(::http::StatusCode::OK)
+								.body(b"pong".to_vec())
+								.expect("build response"),
+						)
+						.into()));
+					}
+				}
+				Ok(())
+			})
+		})),
+	);
+	Arc::new(RegistryDispatcher::new(factories, false))
+}
+
+fn raw_ping_request() -> rivet_envoy_client::config::HttpRequest {
+	rivet_envoy_client::config::HttpRequest {
+		method: "GET".to_owned(),
+		path: "/request/ping".to_owned(),
+		headers: HashMap::new(),
+		body: None,
+		body_stream: None,
+	}
+}
+
+/// Regression test for a per-request task leak. A buffered response used to
+/// drop the request token without cancelling it, so every task waiting on it
+/// (one per `onRequest` in the TypeScript runtime) stayed alive until the
+/// actor stopped and the live task count grew by one per request.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn buffered_raw_requests_release_request_cancellation_waiters() {
+	const REQUESTS: usize = 500;
+
+	let (waiters_tx, mut waiters_rx) = tokio::sync::mpsc::unbounded_channel();
+	let dispatcher = request_waiter_dispatcher(waiters_tx, false);
+	let actor_id = "actor-request-leak";
+	dispatcher
+		.start_actor(StartActorRequest {
+			actor_id: actor_id.to_owned(),
+			generation: 1,
+			actor_name: "counter".to_owned(),
+			input: None,
+			ctx: ActorContext::new(actor_id, "counter", Vec::new(), "local"),
+		})
+		.await
+		.expect("actor should start");
+
+	let metrics = tokio::runtime::Handle::current().metrics();
+	let baseline_tasks = metrics.num_alive_tasks();
+
+	for _ in 0..REQUESTS {
+		let response = dispatcher
+			.handle_fetch(actor_id, raw_ping_request())
+			.await
+			.expect("raw request should succeed");
+		assert_eq!(response.status, ::http::StatusCode::OK.as_u16());
+		assert_eq!(response.body.as_deref(), Some(&b"pong"[..]));
+	}
+
+	for _ in 0..REQUESTS {
+		let waiter = waiters_rx.recv().await.expect("one waiter per request");
+		tokio::time::timeout(Duration::from_secs(5), waiter)
+			.await
+			.expect("request token must be cancelled once the response is returned")
+			.expect("waiter task should not panic");
+	}
+
+	let live_tasks = metrics.num_alive_tasks();
+	assert!(
+		live_tasks <= baseline_tasks + 2,
+		"live tasks should return to baseline after {REQUESTS} requests: baseline={baseline_tasks} live={live_tasks}",
+	);
+}
+
+/// Request tokens are children of the actor abort signal, so stopping the
+/// actor cancels requests that are still in flight.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn actor_abort_signal_cancels_in_flight_raw_requests() {
+	let (waiters_tx, mut waiters_rx) = tokio::sync::mpsc::unbounded_channel();
+	let dispatcher = request_waiter_dispatcher(waiters_tx, true);
+	let actor_id = "actor-request-abort";
+	let ctx = ActorContext::new(actor_id, "counter", Vec::new(), "local");
+	dispatcher
+		.start_actor(StartActorRequest {
+			actor_id: actor_id.to_owned(),
+			generation: 1,
+			actor_name: "counter".to_owned(),
+			input: None,
+			ctx: ctx.clone(),
+		})
+		.await
+		.expect("actor should start");
+
+	let fetch = tokio::spawn({
+		let dispatcher = dispatcher.clone();
+		async move { dispatcher.handle_fetch(actor_id, raw_ping_request()).await }
+	});
+	let waiter = waiters_rx
+		.recv()
+		.await
+		.expect("request should reach the actor");
+	assert!(
+		!waiter.is_finished(),
+		"request must stay in flight until aborted"
+	);
+
+	ctx.cancel_actor_abort_signal();
+
+	tokio::time::timeout(Duration::from_secs(5), waiter)
+		.await
+		.expect("actor abort must cancel the in-flight request token")
+		.expect("waiter task should not panic");
+	let response = tokio::time::timeout(Duration::from_secs(5), fetch)
+		.await
+		.expect("aborted request should complete")
+		.expect("fetch task should not panic")
+		.expect("fetch should return a response");
+	assert_eq!(
+		response.status,
+		::http::StatusCode::SERVICE_UNAVAILABLE.as_u16()
+	);
+}

rivetkit-typescript/packages/rivetkit-napi/index.d.ts

@@ -440,7 +440,28 @@ export declare class CancellationToken {
   constructor()
   aborted(): boolean
   cancel(): void
-  onCancelled(callback: (...args: any[]) => any): void
+  /**
+   * Calls `callback` once when the token is cancelled.
+   *
+   * The returned listener must be released when the caller stops caring
+   * about cancellation. Releasing ends the waiting task and frees the
+   * threadsafe function, which otherwise stays alive (with its libuv handle
+   * and JS closure) until the token is cancelled.
+   */
+  onCancelled(callback: (...args: any[]) => any): CancellationListener
+}
+/**
+ * Registration returned by `CancellationToken.onCancelled`.
+ *
+ * Dropping the JS object does not release the listener, because the callback
+ * must still fire if the caller keeps no reference to the registration.
+ */
+export declare class CancellationListener {
+  /**
+   * Stops waiting for cancellation without calling the callback.
+   * Idempotent, and a no-op after the callback has fired.
+   */
+  release(): void
 }
 export declare class ConnHandle {
   id(): string

rivetkit-typescript/packages/rivetkit-napi/index.js

@@ -310,7 +310,7 @@ if (!nativeBinding) {
   throw new Error(`Failed to load native binding`)
 }
 
-const { ActorContext, decodeInspectorRequest, encodeInspectorResponse, OutboundCall, WorkflowSpan, WorkflowStepSpan, NapiActorFactory, CancellationToken, ConnHandle, JsNativeDatabase, JsSqliteTransaction, JsActorStateTransaction, HttpResponseBodyStream, HttpRequestBodyStream, Kv, Queue, QueueMessage, CoreRegistry, Schedule, WebSocket, setTelemetryLogSink, shutdownTelemetry } = nativeBinding
+const { ActorContext, decodeInspectorRequest, encodeInspectorResponse, OutboundCall, WorkflowSpan, WorkflowStepSpan, NapiActorFactory, CancellationToken, CancellationListener, ConnHandle, JsNativeDatabase, JsSqliteTransaction, JsActorStateTransaction, HttpResponseBodyStream, HttpRequestBodyStream, Kv, Queue, QueueMessage, CoreRegistry, Schedule, WebSocket, setTelemetryLogSink, shutdownTelemetry } = nativeBinding
 
 module.exports.ActorContext = ActorContext
 module.exports.decodeInspectorRequest = decodeInspectorRequest
@@ -320,6 +320,7 @@ module.exports.WorkflowSpan = WorkflowSpan
 module.exports.WorkflowStepSpan = WorkflowStepSpan
 module.exports.NapiActorFactory = NapiActorFactory
 module.exports.CancellationToken = CancellationToken
+module.exports.CancellationListener = CancellationListener
 module.exports.ConnHandle = ConnHandle
 module.exports.JsNativeDatabase = JsNativeDatabase
 module.exports.JsSqliteTransaction = JsSqliteTransaction