chore(rivetkit): debounce ws message ack

Author: NathanFlurry

rivetkit-typescript/packages/rivetkit/src/drivers/engine/actor-driver.ts

@@ -76,6 +76,10 @@ export class EngineActorDriver implements ActorDriver {
 	#runnerStarted: PromiseWithResolvers<undefined> = promiseWithResolvers();
 	#runnerStopped: PromiseWithResolvers<undefined> = promiseWithResolvers();
 
+	// WebSocket message acknowledgment debouncing
+	#wsAckQueue: Map<string, bigint> = new Map();
+	#wsAckFlushInterval?: NodeJS.Timeout;
+
 	constructor(
 		registryConfig: RegistryConfig,
 		runConfig: RunnerConfig,
@@ -284,6 +288,15 @@ export class EngineActorDriver implements ActorDriver {
 			namespace: runConfig.namespace,
 			runnerName: runConfig.runnerName,
 		});
+
+		// Start WebSocket ack flush interval
+		//
+		// Decreasing this reduces the amount of buffered messages on the
+		// gateway
+		//
+		// Gateway timeout configured to 30s
+		// https://github.com/rivet-dev/rivet/blob/222dae87e3efccaffa2b503de40ecf8afd4e31eb/engine/packages/pegboard-gateway/src/shared_state.rs#L17
+		this.#wsAckFlushInterval = setInterval(() => this.#flushWsAcks(), 1000);
 	}
 
 	async #loadActorHandler(actorId: string): Promise<ActorHandler> {
@@ -302,6 +315,20 @@ export class EngineActorDriver implements ActorDriver {
 		return handler.actor;
 	}
 
+	#flushWsAcks(): void {
+		if (this.#wsAckQueue.size === 0) return;
+
+		for (const [requestIdStr, messageIndex] of this.#wsAckQueue.entries()) {
+			// Convert string back to ArrayBuffer
+			const requestId = new Uint8Array(
+				requestIdStr.split(",").map((x) => Number.parseInt(x)),
+			).buffer;
+			this.#runner.sendWebsocketMessageAck(requestId, messageIndex);
+		}
+
+		this.#wsAckQueue.clear();
+	}
+
 	getContext(actorId: string): DriverContext {
 		return {};
 	}
@@ -554,13 +581,28 @@ export class EngineActorDriver implements ActorDriver {
 
 			invariant(event.rivetRequestId, "missing rivetRequestId");
 			invariant(event.rivetMessageIndex, "missing rivetMessageIndex");
-			this.#runner.sendWebsocketMessageAck(
-				event.rivetRequestId,
-				event.rivetMessageIndex,
-			);
+
+			// Track only the highest seen message index per request
+			// Convert ArrayBuffer to string for Map key
+			const currentMax = this.#wsAckQueue.get(requestId) ?? -1n;
+			if (event.rivetMessageIndex > currentMax) {
+				this.#wsAckQueue.set(
+					requestId,
+					BigInt(event.rivetMessageIndex),
+				);
+			} else {
+				logger().warn({
+					msg: "received lower index than ack queue for message",
+					requestId: requestId,
+					queuedMessageIndex: currentMax,
+					eventMessageIndex: event.rivetMessageIndex,
+				});
+			}
 		});
 
 		websocket.addEventListener("close", (event) => {
+			// Flush any pending acks before closing
+			this.#flushWsAcks();
 			wsHandlerPromise.then((x) => x.onClose?.(event, wsContext));
 		});
 
@@ -575,6 +617,16 @@ export class EngineActorDriver implements ActorDriver {
 
 	async shutdownRunner(immediate: boolean): Promise<void> {
 		logger().info({ msg: "stopping engine actor driver" });
+
+		// Clear the ack flush interval
+		if (this.#wsAckFlushInterval) {
+			clearInterval(this.#wsAckFlushInterval);
+			this.#wsAckFlushInterval = undefined;
+		}
+
+		// Flush any remaining acks
+		this.#flushWsAcks();
+
 		await this.#runner.shutdown(immediate);
 	}