feat: add max request id handler

Author: itzmanish

lib/transport/connection.ts

@@ -105,13 +105,19 @@ export class Connection {
 	async #recv(msg: Control.MessageWithType) {
 		if (msg.type === Control.ControlMessageType.GoAway) {
 			await this.#handleGoAway(msg.message)
+		} else if (msg.type === Control.ControlMessageType.MaxRequestId) {
+			await this.#handleMaxRequestId(msg.message)
 		} else if (Control.isPublisher(msg.type)) {
 			await this.#subscriber.recv(msg)
 		} else {
 			await this.#publisher.recv(msg)
 		}
 	}
 
+	async #handleMaxRequestId(msg: Control.MaxRequestId) {
+		this.#controlStream.setRemoteMaxRequestId(msg.id)
+	}
+
 	async #handleGoAway(msg: Control.GoAway) {
 		console.log("preparing for migration, got go_away message:", msg)
 		if (this.#migrationState === "in_progress") {

lib/transport/control/index.ts

@@ -21,6 +21,7 @@ import { FetchCancel } from "./fetch_cancel"
 import { GoAway } from "./go_away"
 import { ClientSetup } from "./client_setup"
 import { ServerSetup } from "./server_setup"
+import { MaxRequestId } from "./max_request_id"
 
 
 enum Version {
@@ -60,6 +61,7 @@ type MessageWithType =
     | { type: ControlMessageType.SubscribeNamespaceError; message: SubscribeNamespaceError }
     | { type: ControlMessageType.Unsubscribe; message: Unsubscribe }
     | { type: ControlMessageType.GoAway; message: GoAway }
+    | { type: ControlMessageType.MaxRequestId; message: MaxRequestId }
 
 type Message = Subscriber | Publisher
 
@@ -185,6 +187,7 @@ export {
     FetchError,
     FetchCancel,
     GoAway,
+    MaxRequestId,
     ClientSetup,
     ServerSetup,
 

lib/transport/stream.ts

@@ -9,15 +9,18 @@ import {
 	Subscribe, SubscribeOk, SubscribeError,
 	SubscribeUpdate, SubscribeNamespace,
 	SubscribeNamespaceOk, SubscribeNamespaceError,
-	GoAway,
+	GoAway, MaxRequestId,
 } from "./control"
 import { debug } from "./utils"
 import { ImmutableBytesBuffer, ReadableWritableStreamBuffer, Reader, Writer } from "./buffer"
 
 export class ControlStream {
 	private decoder: Decoder
 	private encoder: Encoder
+	// Our next request ID to use when sending requests
 	#nextRequestId = 0n
+	// Remote's maximum request ID (first invalid ID). Requests we send must be < this value.
+	#remoteMaxRequestId?: bigint
 
 	#mutex = Promise.resolve()
 
@@ -61,11 +64,41 @@ export class ControlStream {
 		return lock
 	}
 
+	/**
+	 * Returns the next request ID or throws if max request ID is reached.
+	 * Per spec: "If a Request ID equal to or larger than this [MAX_REQUEST_ID] is received
+	 * by the endpoint that sent the MAX_REQUEST_ID in any request message, the endpoint
+	 * MUST close the session with an error of TOO_MANY_REQUESTS."
+	 * @param incr number at which the next request ID should be incremented (default 2 for client)
+	 * @returns next request ID
+	 */
 	nextRequestId(incr: bigint = 2n): bigint {
 		const id = this.#nextRequestId
+		
+		// Check if we're about to exceed the remote's max request ID
+		if (this.#remoteMaxRequestId !== undefined && id >= this.#remoteMaxRequestId) {
+			throw new Error(`TOO_MANY_REQUESTS: Request ID ${id} >= remote max ${this.#remoteMaxRequestId}`)
+		}
+		
 		this.#nextRequestId += incr
 		return id
 	}
+
+	/**
+	 * Sets the remote's maximum request ID. Per spec:
+	 * "The Maximum Request ID MUST only increase within a session, and receipt of a
+	 * MAX_REQUEST_ID message with an equal or smaller Request ID value is a PROTOCOL_VIOLATION."
+	 * 
+	 * Note: The id parameter is the "Maximum Request ID + 1" (first invalid ID).
+	 * @param id The new maximum request ID for the session plus 1
+	 */
+	setRemoteMaxRequestId(id: bigint) {
+		// Validate that MAX_REQUEST_ID only increases
+		if (this.#remoteMaxRequestId !== undefined && id <= this.#remoteMaxRequestId) {
+			throw new Error(`PROTOCOL_VIOLATION: MAX_REQUEST_ID must only increase (received ${id}, current ${this.#remoteMaxRequestId})`)
+		}
+		this.#remoteMaxRequestId = id
+	}
 }
 
 export class Decoder {
@@ -221,6 +254,12 @@ export class Decoder {
 					message: SubscribeNamespaceError.deserialize(payload),
 				}
 				break
+			case ControlMessageType.MaxRequestId:
+				res = {
+					type: t,
+					message: MaxRequestId.deserialize(payload),
+				}
+				break
 			default:
 				throw new Error(`unknown message kind: ${t}`)
 		}
@@ -280,6 +319,8 @@ export class Encoder {
 				return FetchOk.serialize(message as FetchOk)
 			case ControlMessageType.FetchError:
 				return FetchError.serialize(message as FetchError)
+			case ControlMessageType.MaxRequestId:
+				return MaxRequestId.serialize(message as MaxRequestId)
 			default:
 				throw new Error(`unknown message kind in encoder`)
 		}