From b5798b7b5c8332ae2afcf08f47188fc3beeb90f5 Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Wed, 28 Oct 2026 22:37:12 -0700 Subject: [PATCH] fix(rivetkit): release HTTP cancellation resources --- .../rivetkit-core/src/actor/messages.rs | 7 + .../rivetkit-core/src/registry/http.rs | 17 +- .../rivetkit-core/src/serverless_http.rs | 4 + .../packages/rivetkit-core/tests/registry.rs | 168 ++++++++++++++++++ .../packages/rivetkit-napi/index.d.ts | 23 ++- .../packages/rivetkit-napi/index.js | 3 +- .../rivetkit-napi/src/cancellation_token.rs | 71 ++++++-- .../packages/rivetkit-wasm/src/lib.rs | 38 +++- .../rivetkit/src/registry/napi-runtime.ts | 5 +- .../packages/rivetkit/src/registry/native.ts | 57 ++++-- .../packages/rivetkit/src/registry/runtime.ts | 7 +- .../rivetkit/src/registry/wasm-runtime.ts | 9 +- .../native-cancellation-listener.test.ts | 91 ++++++++++ .../rivetkit/tests/runtime-parity.test.ts | 47 ++++- .../rivetkit/tests/wasm-host-smoke.test.ts | 9 +- .../rivetkit/tests/wasm-runtime.test.ts | 9 +- 16 files changed, 515 insertions(+), 50 deletions(-) create mode 100644 rivetkit-typescript/packages/rivetkit/tests/native-cancellation-listener.test.ts diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs index 0fd20d2afb..f39585a6ff 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/messages.rs +++ b/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 { if let Some(body_stream) = &self.body_stream { let mut body_stream = body_stream.lock().take(); diff --git a/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs b/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs index 42dd5ad2b9..ac00afe52d 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/registry/http.rs +++ b/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, } @@ -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 { 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) } diff --git a/rivetkit-rust/packages/rivetkit-core/src/serverless_http.rs b/rivetkit-rust/packages/rivetkit-core/src/serverless_http.rs index 0a8849e526..31cd34160c 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/serverless_http.rs +++ b/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")], diff --git a/rivetkit-rust/packages/rivetkit-core/tests/registry.rs b/rivetkit-rust/packages/rivetkit-core/tests/registry.rs index 05564197f3..a385002088 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/registry.rs +++ b/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>, + hang_until_cancelled: bool, +) -> Arc { + 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() + ); +} diff --git a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts index af10f8fe0d..bf6b34f4d0 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts +++ b/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 diff --git a/rivetkit-typescript/packages/rivetkit-napi/index.js b/rivetkit-typescript/packages/rivetkit-napi/index.js index 0d39d47667..d0040232d9 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.js +++ b/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 diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/cancellation_token.rs b/rivetkit-typescript/packages/rivetkit-napi/src/cancellation_token.rs index 3ac71fb7c4..4e58056fb5 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/cancellation_token.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/cancellation_token.rs @@ -42,33 +42,72 @@ impl CancellationToken { self.inner.cancel(); } + /// 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. #[napi] - pub fn on_cancelled(&self, callback: JsFunction) -> napi::Result<()> { + pub fn on_cancelled(&self, callback: JsFunction) -> napi::Result { let token = self.inner.clone(); + let released = CoreCancellationToken::new(); + let listener = CancellationListener { + released: released.clone(), + }; let tsfn: ThreadsafeFunction<(), ErrorStrategy::CalleeHandled> = callback .create_threadsafe_function(0, |_ctx: ThreadSafeCallContext<()>| { Ok(Vec::::new()) })?; napi::bindgen_prelude::spawn(async move { - token.cancelled().await; - tracing::debug!( - kind = "cancellationToken.onCancelled", - payload_summary = "cancelled=true", - "invoking napi TSF callback" - ); - let status = tsfn.call(Ok(()), ThreadsafeFunctionCallMode::NonBlocking); - tracing::debug!( - kind = "cancellationToken.onCancelled", - ?status, - "napi TSF callback returned" - ); - if status != napi::Status::Ok { - tracing::warn!(?status, "failed to deliver cancellation callback"); + tokio::select! { + biased; + _ = released.cancelled() => { + tracing::debug!( + kind = "cancellationToken.onCancelled", + "cancellation listener released before cancellation" + ); + } + _ = token.cancelled() => { + tracing::debug!( + kind = "cancellationToken.onCancelled", + payload_summary = "cancelled=true", + "invoking napi TSF callback" + ); + let status = tsfn.call(Ok(()), ThreadsafeFunctionCallMode::NonBlocking); + tracing::debug!( + kind = "cancellationToken.onCancelled", + ?status, + "napi TSF callback returned" + ); + if status != napi::Status::Ok { + tracing::warn!(?status, "failed to deliver cancellation callback"); + } + } } }); - Ok(()) + Ok(listener) + } +} + +/// 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. +#[napi] +pub struct CancellationListener { + released: CoreCancellationToken, +} + +#[napi] +impl CancellationListener { + /// Stops waiting for cancellation without calling the callback. + /// Idempotent, and a no-op after the callback has fired. + #[napi] + pub fn release(&self) { + self.released.cancel(); } } diff --git a/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs b/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs index 57b8eebe37..c3faf99216 100644 --- a/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs +++ b/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs @@ -1166,13 +1166,43 @@ impl WasmCancellationToken { self.inner.cancel(); } + /// Calls `callback` once when the token is cancelled. Release the returned + /// listener when cancellation no longer matters; otherwise the waiting task + /// and the callback stay alive until the token is cancelled. #[wasm_bindgen(js_name = onCancelled)] - pub fn on_cancelled(&self, callback: Function) { - let token = self.inner.clone(); + pub fn on_cancelled(&self, callback: Function) -> WasmCancellationListener { + // The waiter wakes when either the token or the listener's release cancels + // its child token, and the release flag tells the two apart. + let wake = self.inner.child_token(); + let released = Rc::new(Cell::new(false)); + let listener = WasmCancellationListener { + wake: wake.clone(), + released: released.clone(), + }; spawn_local(async move { - token.cancelled().await; - let _ = callback.call0(&JsValue::UNDEFINED); + wake.cancelled().await; + if !released.get() { + let _ = callback.call0(&JsValue::UNDEFINED); + } }); + listener + } +} + +/// Registration returned by `CancellationToken.onCancelled`. +#[wasm_bindgen(js_name = CancellationListener)] +pub struct WasmCancellationListener { + wake: CoreCancellationToken, + released: Rc>, +} + +#[wasm_bindgen(js_class = CancellationListener)] +impl WasmCancellationListener { + /// Stops waiting for cancellation without calling the callback. + #[wasm_bindgen] + pub fn release(&self) { + self.released.set(true); + self.wake.cancel(); } } diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts index f822e70a0e..f4cef6aaa4 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts @@ -481,8 +481,9 @@ export class NapiCoreRuntime implements CoreRuntime { onCancellationTokenCancelled( token: CancellationTokenHandle, callback: (...args: unknown[]) => unknown, - ): void { - asNativeCancellationToken(token).onCancelled(callback); + ): () => void { + const listener = asNativeCancellationToken(token).onCancelled(callback); + return () => listener.release(); } actorState(ctx: ActorContextHandle): Buffer { diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts index 0c5c55648e..bc791c6abc 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts @@ -347,6 +347,11 @@ type NativeDestroyGate = { type NativeDatabaseClientState = { client: unknown; }; +type NativeActorAbort = { + signal: AbortSignal; + listeners: Set<() => void>; +}; + type NativeActorRuntimeState = { sql?: ReturnType; databaseClient?: NativeDatabaseClientState; @@ -354,6 +359,12 @@ type NativeActorRuntimeState = { vars?: unknown; destroyGate?: NativeDestroyGate; persistState?: NativePersistActorState; + // One native actor abort signal shared by every context of the actor. Each + // native signal parks a task and a callback handle until the actor stops, + // so creating one per request or action leaks both. Per-context listeners + // live in `listeners` so the shared signal never collects one listener per + // in-flight request. + actorAbort?: NativeActorAbort; }; // Keep JS-only actor caches on the NAPI ActorContext runtime-state bag instead @@ -3022,22 +3033,28 @@ export class ActorContextHandleAdapter { get abortSignal(): AbortSignal { if (!this.#abortSignal) { - const actorSignal = this.#createActorAbortSignal(); + const actorAbort = this.#actorAbort(); + const actorSignal = actorAbort.signal; if (this.#dispatchCancelToken === undefined) { this.#abortSignal = actorSignal; } else { const controller = new AbortController(); let cleanedUp = false; + let releaseDispatchListener: (() => void) | undefined; const onActorAbort = () => { cleanup(); controller.abort(); }; + // Releasing the native listener on dispose frees it even when the + // dispatch token is never cancelled. const cleanup = () => { if (cleanedUp) { return; } cleanedUp = true; - actorSignal.removeEventListener("abort", onActorAbort); + actorAbort.listeners.delete(onActorAbort); + releaseDispatchListener?.(); + releaseDispatchListener = undefined; this.#abortSignalCleanup = undefined; }; @@ -3051,10 +3068,8 @@ export class ActorContextHandleAdapter { } else { const dispatchCancelToken = this.#dispatchCancelToken; this.#abortSignalCleanup = cleanup; - actorSignal.addEventListener("abort", onActorAbort, { - once: true, - }); - callNativeSync(() => + actorAbort.listeners.add(onActorAbort); + releaseDispatchListener = callNativeSync(() => this.#runtime.onCancellationTokenCancelled( dispatchCancelToken, () => { @@ -3511,19 +3526,37 @@ export class ActorContextHandleAdapter { this.#abortSignalCleanup?.(); } - #createActorAbortSignal(): AbortSignal { + #actorAbort(): NativeActorAbort { + // An aborted cached signal may belong to a previous generation, so + // only a live one is reused. + const runtimeState = getNativeRuntimeState(this.#runtime, this.#ctx); + const cached = runtimeState.actorAbort; + if (cached && !cached.signal.aborted) { + return cached; + } + const nativeSignal = callNativeSync(() => this.#runtime.actorAbortSignal(this.#ctx), ); const controller = new AbortController(); - if (nativeSignal.aborted) { + const actorAbort: NativeActorAbort = { + signal: controller.signal, + listeners: new Set(), + }; + const abort = () => { controller.abort(); + for (const listener of actorAbort.listeners) { + listener(); + } + actorAbort.listeners.clear(); + }; + if (nativeSignal.aborted) { + abort(); } else { - nativeSignal.addEventListener("abort", () => controller.abort(), { - once: true, - }); + nativeSignal.addEventListener("abort", abort, { once: true }); } - return controller.signal; + runtimeState.actorAbort = actorAbort; + return actorAbort; } #readState(): unknown { diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts index 286868ff45..a84544234b 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts @@ -532,10 +532,15 @@ export interface CoreRuntime { createCancellationToken(): CancellationTokenHandle; cancellationTokenAborted(token: CancellationTokenHandle): boolean; cancelCancellationToken(token: CancellationTokenHandle): void; + /** + * Calls `callback` once when `token` is cancelled. Returns a release + * function that drops the registration without calling `callback`; call it + * once cancellation no longer matters so the native listener is freed. + */ onCancellationTokenCancelled( token: CancellationTokenHandle, callback: (...args: unknown[]) => unknown, - ): void; + ): () => void; actorState(ctx: ActorContextHandle): RuntimeBytes; actorBeginOnStateChange(ctx: ActorContextHandle): void; diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts index 052c0fb591..bb17bff648 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts @@ -66,7 +66,9 @@ type WasmWebSocketHandle = object; interface WasmCancellationToken { aborted(): boolean; cancel(): void; - onCancelled(callback: (...args: unknown[]) => unknown): void; + onCancelled(callback: (...args: unknown[]) => unknown): { + release(): void; + }; } interface WasmCoreRegistry { @@ -390,10 +392,11 @@ export class WasmCoreRuntime implements CoreRuntime { onCancellationTokenCancelled( token: CancellationTokenHandle, callback: (...args: unknown[]) => unknown, - ): void { - callWasmSync(() => + ): () => void { + const listener = callWasmSync(() => asWasmCancellationToken(token).onCancelled(callback), ); + return () => callWasmSync(() => listener.release()); } actorState(ctx: ActorContextHandle): RuntimeBytes { diff --git a/rivetkit-typescript/packages/rivetkit/tests/native-cancellation-listener.test.ts b/rivetkit-typescript/packages/rivetkit/tests/native-cancellation-listener.test.ts new file mode 100644 index 0000000000..449b4a6a77 --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/tests/native-cancellation-listener.test.ts @@ -0,0 +1,91 @@ +import { spawn } from "node:child_process"; +import { createRequire } from "node:module"; +import { describe, expect, test } from "vitest"; + +const require = createRequire(import.meta.url); +const NAPI_BINDING = require.resolve("@rivetkit/rivetkit-napi"); + +interface NativeCancellationBindings { + CancellationToken: new () => { + cancel(): void; + onCancelled(callback: () => void): { release(): void }; + }; +} + +function loadBindings(): NativeCancellationBindings { + return require(NAPI_BINDING) as NativeCancellationBindings; +} + +/** Waits for one turn of native callbacks to be delivered to JS. */ +async function flushNativeCallbacks(): Promise { + const { CancellationToken } = loadBindings(); + const token = new CancellationToken(); + await new Promise((resolve) => { + token.onCancelled(resolve); + token.cancel(); + }); +} + +describe("native cancellation listeners", () => { + test("a cancelled token calls its listener once", async () => { + const { CancellationToken } = loadBindings(); + const token = new CancellationToken(); + let calls = 0; + const fired = new Promise((resolve) => { + token.onCancelled(() => { + calls += 1; + resolve(); + }); + }); + + token.cancel(); + token.cancel(); + await fired; + await flushNativeCallbacks(); + expect(calls).toBe(1); + }); + + test("a released listener is never called", async () => { + const { CancellationToken } = loadBindings(); + const token = new CancellationToken(); + let calls = 0; + const listener = token.onCancelled(() => { + calls += 1; + }); + + listener.release(); + listener.release(); + token.cancel(); + await flushNativeCallbacks(); + expect(calls).toBe(0); + }); + + // Regression test for a per-request leak: every request registered a + // listener whose native callback handle stayed alive until the token was + // cancelled. Each live handle keeps the event loop referenced, so a process + // that releases all of its listeners must still exit on its own. + test("released listeners free their native callback handles", async () => { + const script = ` + const { CancellationToken } = require(${JSON.stringify(NAPI_BINDING)}); + const tokens = []; + for (let i = 0; i < 5000; i++) { + const token = new CancellationToken(); + token.onCancelled(() => {}).release(); + tokens.push(token); + } + `; + const child = spawn(process.execPath, ["-e", script], { + stdio: "ignore", + }); + const timeout = setTimeout(() => child.kill("SIGKILL"), 20_000); + const exit = await new Promise<{ + code: number | null; + signal: NodeJS.Signals | null; + }>((resolve) => { + child.on("exit", (code, signal) => resolve({ code, signal })); + }); + clearTimeout(timeout); + + expect(exit).toEqual({ code: 0, signal: null }); + }); +}); diff --git a/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts b/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts index 867a9a6743..a4dcb9391d 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts @@ -105,6 +105,7 @@ class ParityScenario { class FakeActorContext { stateBytes = Buffer.alloc(0); + abortSignalCalls = 0; readonly runtimeBag = {}; readonly registeredTasks: Array> = []; readonly abortController = new AbortController(); @@ -187,6 +188,7 @@ class FakeActorContext { } abortSignal(): AbortSignal { + this.abortSignalCalls += 1; return this.abortController.signal; } @@ -209,6 +211,10 @@ class FakeCancellationToken { #cancelled = false; #callbacks: Array<() => void> = []; + get listenerCount(): number { + return this.#callbacks.length; + } + aborted(): boolean { return this.#cancelled; } @@ -220,8 +226,15 @@ class FakeCancellationToken { } } - onCancelled(callback: () => void): void { + onCancelled(callback: () => void): { release(): void } { this.#callbacks.push(callback); + return { + release: () => { + this.#callbacks = this.#callbacks.filter( + (registered) => registered !== callback, + ); + }, + }; } } @@ -461,6 +474,38 @@ async function invokePromotedStatus( } describe("CoreRuntime NAPI and wasm parity", () => { + test.each([ + "napi", + "wasm", + ] as const)("%s releases completed action listeners and shares one actor abort signal", async (kind) => { + const runtimeCase = createRuntimeCase(kind); + const definition = actor({ + actions: { + inspectSignal: (c) => c.abortSignal.aborted, + }, + }); + const factory = buildNativeFactory( + runtimeCase.runtime, + registryConfig(definition), + definition, + ) as unknown as FakeActorFactory; + const ctx = new FakeActorContext(runtimeCase.scenario); + + for (let i = 0; i < 100; i++) { + const cancelToken = new FakeCancellationToken(); + await factory.callbacks.actions.inspectSignal(null, { + ctx, + conn: null, + name: "inspectSignal", + args: encodeValue([]), + cancelToken, + }); + expect(cancelToken.listenerCount).toBe(0); + } + + expect(ctx.abortSignalCalls).toBe(1); + }); + test.each([ ["napi", 42n, 42], ["wasm", 42n, 42], diff --git a/rivetkit-typescript/packages/rivetkit/tests/wasm-host-smoke.test.ts b/rivetkit-typescript/packages/rivetkit/tests/wasm-host-smoke.test.ts index 1020028149..b4908f8aa1 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/wasm-host-smoke.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/wasm-host-smoke.test.ts @@ -305,8 +305,15 @@ class FakeCancellationToken { } } - onCancelled(callback: () => void): void { + onCancelled(callback: () => void): { release(): void } { this.#callbacks.push(callback); + return { + release: () => { + this.#callbacks = this.#callbacks.filter( + (registered) => registered !== callback, + ); + }, + }; } } diff --git a/rivetkit-typescript/packages/rivetkit/tests/wasm-runtime.test.ts b/rivetkit-typescript/packages/rivetkit/tests/wasm-runtime.test.ts index c8209a363d..8a712c027a 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/wasm-runtime.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/wasm-runtime.test.ts @@ -213,8 +213,15 @@ class FakeCancellationToken { } } - onCancelled(callback: () => void): void { + onCancelled(callback: () => void): { release(): void } { this.#callbacks.push(callback); + return { + release: () => { + this.#callbacks = this.#callbacks.filter( + (registered) => registered !== callback, + ); + }, + }; } }