fix(rivetkit): remove the connection of an aborted HTTP action

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

@@ -25,6 +25,7 @@ use crate::actor::persist::{
 };
 use crate::actor::state::RequestSaveOpts;
 use crate::error::ActorRuntime;
+use crate::runtime::RuntimeSpawner;
 use crate::time::timeout;
 use crate::types::ConnId;
 
@@ -153,6 +154,46 @@ pub(crate) fn decode_persisted_connection(payload: &[u8]) -> Result<PersistedCon
 	Ok(connection)
 }
 
+/// Disconnects a connection when dropped, unless it was disconnected or
+/// disarmed first. A caller that goes away drops its future at any await, and
+/// a connection left registered keeps the actor awake forever.
+pub(crate) struct DisconnectOnDrop {
+	conn: Option<ConnHandle>,
+}
+
+impl DisconnectOnDrop {
+	pub(crate) fn new(conn: ConnHandle) -> Self {
+		Self { conn: Some(conn) }
+	}
+
+	pub(crate) async fn disconnect(mut self) -> Result<()> {
+		match self.conn.take() {
+			Some(conn) => conn.disconnect(None).await,
+			None => Ok(()),
+		}
+	}
+
+	pub(crate) fn disarm(mut self) {
+		self.conn = None;
+	}
+}
+
+impl Drop for DisconnectOnDrop {
+	fn drop(&mut self) {
+		if let Some(conn) = self.conn.take() {
+			RuntimeSpawner::spawn(async move {
+				if let Err(error) = conn.disconnect(None).await {
+					tracing::warn!(
+						conn_id = conn.id(),
+						?error,
+						"failed to disconnect a connection whose caller went away"
+					);
+				}
+			});
+		}
+	}
+}
+
 #[derive(Clone)]
 pub struct ConnHandle(Arc<ConnHandleInner>);
 
@@ -724,10 +765,14 @@ impl ActorContext {
 			.await?;
 		self.insert_existing(conn.clone());
 
+		// The caller can go away while onConnect runs.
+		let opening = DisconnectOnDrop::new(conn.clone());
 		if let Err(error) = self.emit_connection_open(&conn, request).await {
+			opening.disarm();
 			self.remove_existing(conn.id());
 			return Err(error);
 		}
+		opening.disarm();
 		self.0.metrics.inc_connections_total();
 		self.record_connections_updated();
 		self.reset_sleep_timer();

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

@@ -241,6 +241,7 @@ impl RegistryDispatcher {
 				);
 			}
 		};
+		let request_connection = DisconnectOnDrop::new(conn.clone());
 
 		let dispatch_result = with_action_dispatch_timeout(
 			config.action_timeout,
@@ -253,7 +254,7 @@ impl RegistryDispatcher {
 			),
 		)
 		.await;
-		let disconnect_result = conn.disconnect(None).await;
+		let disconnect_result = request_connection.disconnect().await;
 
 		match dispatch_result {
 			Ok(output) => {
@@ -362,6 +363,7 @@ impl RegistryDispatcher {
 				);
 			}
 		};
+		let request_connection = DisconnectOnDrop::new(conn.clone());
 
 		let incoming =
 			crate::telemetry::IncomingInvocationContext::from_http_headers(request.headers());
@@ -391,7 +393,7 @@ impl RegistryDispatcher {
 			}
 			Err(error) => Err(error),
 		};
-		let disconnect_result = conn.disconnect(None).await;
+		let disconnect_result = request_connection.disconnect().await;
 
 		match queue_result {
 			Ok(result) => {

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

@@ -35,7 +35,7 @@ use vbare::OwnedVersionedData;
 
 use crate::actor::action::ActionDispatchError;
 use crate::actor::config::CanHibernateWebSocket;
-use crate::actor::connection::{ConnHandle, HibernatableConnectionMetadata};
+use crate::actor::connection::{ConnHandle, DisconnectOnDrop, HibernatableConnectionMetadata};
 use crate::actor::context::{ActorContext, InspectorAttachGuard};
 use crate::actor::factory::ActorFactory;
 use crate::actor::kv::LegacyActorKv;

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/actors/sleepWithSlowConnect.ts

@@ -0,0 +1,3 @@
+import { sleepWithSlowConnect } from "../sleep";
+
+export default sleepWithSlowConnect;

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/registry-static.ts

@@ -126,6 +126,7 @@ import {
 	sleepWithNoSleepOption,
 	sleepWithRawHttp,
 	sleepWithRawWebSocket,
+	sleepWithSlowConnect,
 	sleepWithWaitUntilInOnWake,
 	sleepWithWaitUntilMessage,
 } from "./sleep";
@@ -225,6 +226,7 @@ export const registry = setup({
 		// From sleep.ts
 		sleep,
 		sleepWithLongRpc,
+		sleepWithSlowConnect,
 		sleepWithRawHttp,
 		sleepWithRawWebSocket,
 		sleepWithNoSleepOption,

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/sleep.ts

@@ -176,6 +176,42 @@ export const sleepWithLongRpc = actor({
 	},
 });
 
+export const sleepWithSlowConnect = actor({
+	state: { startCount: 0, sleepCount: 0 },
+	createVars: () => ({}) as { releaseConnect?: () => void },
+	onWake: (c) => {
+		c.state.startCount += 1;
+	},
+	onSleep: (c) => {
+		c.state.sleepCount += 1;
+	},
+	onConnect: async (c, conn) => {
+		if (
+			!(conn.params as { holdConnect?: boolean } | undefined)?.holdConnect
+		)
+			return;
+		const { promise, resolve } = promiseWithResolvers<void>((reason) =>
+			c.log.warn({ msg: "unhandled held connect rejection", reason }),
+		);
+		c.vars.releaseConnect = resolve;
+		c.broadcast("connecting");
+		await promise;
+	},
+	actions: {
+		getCounts: (c) => {
+			return {
+				startCount: c.state.startCount,
+				sleepCount: c.state.sleepCount,
+			};
+		},
+		ping: () => "pong",
+		releaseConnect: (c) => c.vars.releaseConnect?.(),
+	},
+	options: {
+		sleepTimeout: SLEEP_TIMEOUT,
+	},
+});
+
 export const sleepWithWaitUntilMessage = actor({
 	state: {
 		startCount: 0,

rivetkit-typescript/packages/rivetkit/tests/driver/actor-sleep.test.ts

@@ -1070,3 +1070,82 @@ describeDriverMatrix("Actor Sleep", (driverTestConfig) => {
 		);
 	});
 });
+
+// Gateway3 tells the actor when a caller disconnects before the response starts.
+describeDriverMatrix(
+	"Actor Sleep with Gateway3",
+	(driverTestConfig) => {
+		test("an aborted rpc lets the actor sleep once it finishes", async (c) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+			const key = [crypto.randomUUID()];
+
+			// A connection only observes that the rpc started, then leaves.
+			const observer = client.sleepWithLongRpc.getOrCreate(key).connect();
+			const started = new Promise((resolve) =>
+				observer.once("waiting", resolve),
+			);
+			// The subscription and this call share one ordered connection, so the subscription is active before the rpc starts.
+			await observer.getCounts();
+			const caller = new AbortController();
+			const aborted = client.sleepWithLongRpc.getOrCreate(key).action({
+				name: "longRunningRpc",
+				args: [],
+				signal: caller.signal,
+			});
+			await started;
+			await observer.dispose();
+			caller.abort();
+			await expect(aborted).rejects.toThrow();
+
+			// The rpc keeps running after the abort. Once it finishes, nothing is left to keep the actor awake.
+			await client.sleepWithLongRpc
+				.getOrCreate(key)
+				.finishLongRunningRpc();
+			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 250);
+
+			const { startCount, sleepCount } = await client.sleepWithLongRpc
+				.getOrCreate(key)
+				.getCounts();
+			expect(sleepCount).toBe(1);
+			expect(startCount).toBe(2);
+		});
+
+		test("a caller that leaves during onConnect lets the actor sleep", async (c) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+			const key = [crypto.randomUUID()];
+
+			// A connection only observes that onConnect started, then leaves.
+			const observer = client.sleepWithSlowConnect
+				.getOrCreate(key)
+				.connect();
+			const connecting = new Promise((resolve) =>
+				observer.once("connecting", resolve),
+			);
+			// The subscription and this call share one ordered connection, so the subscription is active before the call below connects.
+			await observer.ping();
+			const caller = new AbortController();
+			const aborted = client.sleepWithSlowConnect
+				.getOrCreate(key, { params: { holdConnect: true } })
+				.action({ name: "ping", args: [], signal: caller.signal });
+			await connecting;
+			await observer.dispose();
+			caller.abort();
+			await expect(aborted).rejects.toThrow();
+
+			// onConnect finishes after the caller left. Nothing is left to keep the actor awake.
+			await client.sleepWithSlowConnect.getOrCreate(key).releaseConnect();
+			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 250);
+
+			const { startCount, sleepCount } = await client.sleepWithSlowConnect
+				.getOrCreate(key)
+				.getCounts();
+			expect(sleepCount).toBe(1);
+			expect(startCount).toBe(2);
+		});
+	},
+	{
+		runtimes: ["native"],
+		encodings: ["bare"],
+		config: { engine: { gateway3: true } },
+	},
+);