From 924f9e059b45df41f0f8a522f0b845d49adf87a8 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Tue, 13 Oct 2026 00:58:57 +0400 Subject: [PATCH] fix(rivetkit-core): close drained connections with ws.going_away --- .claude/skills/driver-test-runner/SKILL.md | 2 + .../rivetkit-core/src/actor/connection.rs | 15 ++++-- .../rivetkit-core/src/registry/websocket.rs | 7 ++- .../fixtures/driver-test-suite/conn-drain.ts | 15 ++++++ .../driver-test-suite/registry-static.ts | 3 ++ .../tests/driver/actor-conn-drain.test.ts | 53 +++++++++++++++++++ .../fixtures/driver-test-suite-runtime.ts | 6 +++ 7 files changed, 97 insertions(+), 4 deletions(-) create mode 100644 rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/conn-drain.ts create mode 100644 rivetkit-typescript/packages/rivetkit/tests/driver/actor-conn-drain.test.ts diff --git a/.claude/skills/driver-test-runner/SKILL.md b/.claude/skills/driver-test-runner/SKILL.md index 8338072bcd..ea0d9bb4d8 100644 --- a/.claude/skills/driver-test-runner/SKILL.md +++ b/.claude/skills/driver-test-runner/SKILL.md @@ -133,6 +133,7 @@ Each row: `[native] [wasm] | ` - [ ] [ ] actor-db-pragma-migration | Actor Database Pragma Migration - [ ] [ ] actor-state-zod-coercion | Actor State Zod Coercion - [ ] [ ] actor-conn-status | Connection Status Changes +- [ ] [ ] actor-conn-drain | Connection Drain - [ ] [ ] gateway-routing | Gateway Routing - [ ] [ ] lifecycle-hooks | Lifecycle Hooks - [ ] [ ] serverless-handler | Serverless Handler Tests @@ -233,6 +234,7 @@ Replace `` with the file name stem (part before the `|` in the progress fi | actor-db-pragma-migration | Actor Database Pragma Migration | | actor-state-zod-coercion | Actor State Zod Coercion | | actor-conn-status | Connection Status Changes | +| actor-conn-drain | Connection Drain | | gateway-routing | Gateway Routing | | lifecycle-hooks | Lifecycle Hooks | | serverless-handler | Serverless Handler Tests | diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/connection.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/connection.rs index 028f6c9700..76f5f84b64 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/connection.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/connection.rs @@ -32,6 +32,10 @@ use crate::types::ConnId; pub(crate) type EventSendCallback = Arc Result<()> + Send + Sync>; pub(crate) type DisconnectCallback = Arc) -> BoxFuture<'static, Result<()>> + Send + Sync>; + +/// Close reason for connections closed while the actor shuts down. Clients +/// classify it as reconnect-only. +pub(crate) const WS_GOING_AWAY_CLOSE_REASON: &str = "ws.going_away"; type StateChangeCallback = Arc; #[derive(Clone, Debug, PartialEq, Eq)] @@ -426,11 +430,11 @@ impl ConnHandle { self.disconnect_handler() } - pub(crate) async fn disconnect_transport_only(&self) -> Result<()> { + pub(crate) async fn disconnect_transport_only(&self, reason: Option) -> Result<()> { let Some(handler) = self.transport_disconnect_handler() else { return Ok(()); }; - handler(None).await + handler(reason).await } fn transport_disconnect_handler(&self) -> Option { @@ -952,9 +956,14 @@ impl ActorContext { .collect(); let mut disconnected_ids = Vec::new(); let mut failures = Vec::new(); + let reason = if self.actor_aborted() { + Some(WS_GOING_AWAY_CLOSE_REASON.to_owned()) + } else { + None + }; for conn in &connections { - match conn.disconnect_transport_only().await { + match conn.disconnect_transport_only(reason.clone()).await { Ok(()) => { tracing::debug!( actor_id = %self.actor_id(), diff --git a/rivetkit-rust/packages/rivetkit-core/src/registry/websocket.rs b/rivetkit-rust/packages/rivetkit-core/src/registry/websocket.rs index efe223b905..ab2b779266 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/registry/websocket.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/registry/websocket.rs @@ -2,6 +2,7 @@ use super::actor_connect::*; use super::dispatch::*; use super::inspector::encode_json_as_cbor; use super::*; +use crate::actor::connection::WS_GOING_AWAY_CLOSE_REASON; use crate::error::ProtocolError; use crate::telemetry::IncomingInvocationContext; use crate::time::timeout; @@ -712,7 +713,11 @@ fn configure_actor_connect_connection( let transport_disconnect_sender = transport_disconnect_sender.clone(); Box::pin(async move { if !transport_closed.swap(true, Ordering::SeqCst) { - transport_disconnect_sender.close(Some(1000), reason); + let code = match reason.as_deref() { + Some(WS_GOING_AWAY_CLOSE_REASON) => 1001, + _ => 1000, + }; + transport_disconnect_sender.close(Some(code), reason); } Ok(()) }) diff --git a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/conn-drain.ts b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/conn-drain.ts new file mode 100644 index 0000000000..3eb6a0dc16 --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/conn-drain.ts @@ -0,0 +1,15 @@ +import { actor } from "rivetkit"; + +export const CONN_DRAIN_DISCONNECT_MARKER = "conn-drain onDisconnect"; + +export const connDrainActor = actor({ + onDisconnect: async (c) => { + console.log(`${CONN_DRAIN_DISCONNECT_MARKER} ${c.key.join("/")}`); + // A slow hook keeps the connection registered while the runtime + // closes its socket, which is when a duplicate hook can fire. + await new Promise((resolve) => setTimeout(resolve, 100)); + }, + actions: { + ping: () => "pong", + }, +}); diff --git a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/registry-static.ts b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/registry-static.ts index ac165dd402..bc1c3bad13 100644 --- a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/registry-static.ts +++ b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/registry-static.ts @@ -33,6 +33,7 @@ import { } from "./actor-db-raw"; import { onStateChangeActor } from "./actor-onstatechange"; import { connErrorSerializationActor } from "./conn-error-serialization"; +import { connDrainActor } from "./conn-drain"; import { counterWithParams } from "./conn-params"; import { connPreflightVisibilityActor } from "./conn-preflight-visibility"; import { connStateActor } from "./conn-state"; @@ -293,6 +294,8 @@ export const registry = setup({ connStateActor, // From conn-preflight-visibility.ts connPreflightVisibilityActor, + // From conn-drain.ts + connDrainActor, // From metadata.ts metadataActor, // From vars.ts diff --git a/rivetkit-typescript/packages/rivetkit/tests/driver/actor-conn-drain.test.ts b/rivetkit-typescript/packages/rivetkit/tests/driver/actor-conn-drain.test.ts new file mode 100644 index 0000000000..316e1e7692 --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/tests/driver/actor-conn-drain.test.ts @@ -0,0 +1,53 @@ +import { expect, test } from "vitest"; +import { CONN_DRAIN_DISCONNECT_MARKER } from "../../fixtures/driver-test-suite/conn-drain"; +import type { registry } from "../../fixtures/driver-test-suite/registry-static"; +import { createClient } from "@/client/mod"; +import type { ActorError } from "@/client/errors"; +import { describeDriverMatrix } from "./shared-matrix"; + +describeDriverMatrix( + "Connection Drain", + (driverTestConfig) => { + test("closes connections with ws.going_away and runs onDisconnect once", async (c) => { + const runtime = await driverTestConfig.start({ + env: { RIVETKIT_TEST_DRAIN_ON_SIGTERM: "1" }, + }); + const client = createClient({ + endpoint: runtime.endpoint, + namespace: runtime.namespace, + poolName: runtime.runnerName, + encoding: driverTestConfig.encoding, + disableMetadataLookup: true, + }); + c.onTestFinished(async () => { + await client.dispose(); + await runtime.cleanup(); + }); + + const key = crypto.randomUUID(); + const conn = client.connDrainActor.getOrCreate([key]).connect(); + const errors: ActorError[] = []; + conn.onError((error) => errors.push(error)); + const closed = new Promise((resolve) => + conn.onClose(resolve), + ); + expect(await conn.ping()).toBe("pong"); + + await runtime.cleanup(); + await closed; + + expect( + errors.map((error) => `${error.group}.${error.code}`), + ).toEqual(["ws.going_away"]); + const output = runtime.getRuntimeOutput?.() ?? ""; + const hookRuns = + output.split(`${CONN_DRAIN_DISCONNECT_MARKER} ${key}`).length - 1; + expect(hookRuns).toBe(1); + }); + }, + { + encodings: ["bare"], + runtimes: ["native"], + sqliteBackends: ["local"], + }, +); diff --git a/rivetkit-typescript/packages/rivetkit/tests/fixtures/driver-test-suite-runtime.ts b/rivetkit-typescript/packages/rivetkit/tests/fixtures/driver-test-suite-runtime.ts index 9a70734ef9..b9524ce6f6 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/fixtures/driver-test-suite-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/fixtures/driver-test-suite-runtime.ts @@ -49,4 +49,10 @@ const { registry: nativeRegistry, serveConfig } = await buildConfiguredRegistry( registry.parseConfig(), ); +// Production hosts drain on SIGTERM through `registry.start()`. This is opt-in +// so other driver tests keep a fast hard stop on cleanup. +if (process.env.RIVETKIT_TEST_DRAIN_ON_SIGTERM === "1") { + process.once("SIGTERM", () => void nativeRegistry.shutdown()); +} + await nativeRegistry.serve(serveConfig);