fix(rivetkit-core): close drained connections with ws.going_away

.claude/skills/driver-test-runner/SKILL.md

@@ -133,6 +133,7 @@ Each row: `[native] [wasm] <file> | <suite description>`
 - [ ] [ ] 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 `<FILE>` 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 |

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

@@ -32,6 +32,10 @@ use crate::types::ConnId;
 pub(crate) type EventSendCallback = Arc<dyn Fn(OutgoingEvent) -> Result<()> + Send + Sync>;
 pub(crate) type DisconnectCallback =
 	Arc<dyn Fn(Option<String>) -> 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<dyn Fn(&ConnHandle) + Send + Sync>;
 
 #[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<String>) -> Result<()> {
 		let Some(handler) = self.transport_disconnect_handler() else {
 			return Ok(());
 		};
-		handler(None).await
+		handler(reason).await
 	}
 
 	fn transport_disconnect_handler(&self) -> Option<DisconnectCallback> {
@@ -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(),

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(())
 		})

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",
+	},
+});

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

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<typeof registry>({
+				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<void>((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"],
+	},
+);

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);