feat: implement session migration with GoAway message handling

Author: itzmanish

lib/transport/buffer.ts

@@ -143,6 +143,7 @@ export class ImmutableBytesBuffer {
 
     getUtf8String(): string {
         const len = this.getNumberVarInt()
+        if (len === 0) return ""
         const val = this.getBytes(len)
         return new TextDecoder().decode(val)
     }

lib/transport/client.ts

@@ -26,17 +26,20 @@ export class Client {
 		})
 	}
 
-	async connect(): Promise<Connection> {
+	async createQuic(sessionUri?: string): Promise<WebTransport> {
 		const options: WebTransportOptions = {}
 
 		const fingerprint = await this.#fingerprint
 		if (fingerprint) options.serverCertificateHashes = [fingerprint]
 
-		const quic = new WebTransport(this.config.url, options)
+		const uri = sessionUri ?? this.config.url
+		const quic = new WebTransport(uri, options)
 		await quic.ready
+		return quic
+	}
 
+	async prepareConnection(quic: WebTransport): Promise<{ control: Stream.ControlStream, objects: Objects }> {
 		const stream = await quic.createBidirectionalStream({ sendOrder: Number.MAX_SAFE_INTEGER })
-
 		const buffer = new ReadableWritableStreamBuffer(stream.readable, stream.writable)
 
 		const msg: Control.ClientSetup = {
@@ -57,7 +60,21 @@ export class Client {
 		const control = new Stream.ControlStream(buffer)
 		const objects = new Objects(quic)
 
-		return new Connection(quic, control, objects)
+		return { control, objects }
+	}
+
+	async migrateSession(sessionUri?: string): Promise<{ quic: WebTransport, control: Stream.ControlStream, objects: Objects }> {
+		const quic = await this.createQuic(sessionUri)
+		const { control, objects } = await this.prepareConnection(quic)
+		return { quic, control, objects }
+	}
+
+	async connect(sessionUri?: string): Promise<Connection> {
+		const quic = await this.createQuic(sessionUri)
+		const { control, objects } = await this.prepareConnection(quic)
+		const conn = new Connection(quic, control, objects)
+		conn.onMigration = this.migrateSession.bind(this)
+		return conn
 	}
 
 	async #fetchFingerprint(url?: string): Promise<WebTransportHash | undefined> {

lib/transport/connection.ts

@@ -6,7 +6,10 @@ import { ControlStream } from "./stream"
 import { Publisher } from "./publisher"
 import { Subscriber } from "./subscriber"
 
+export type MigrationState = "none" | "in_progress" | "done"
+
 export class Connection {
+	#migrationState: MigrationState = "none"
 	// The established WebTransport session.
 	#quic: WebTransport
 
@@ -36,6 +39,11 @@ export class Connection {
 		this.#running = this.#run()
 	}
 
+	// Callback, should be set when creating connection in the client
+	onMigration: (sessionUri?: string) => Promise<{ quic: WebTransport, control: ControlStream, objects: Objects }> = async () => {
+		throw new Error("not implemented")
+	}
+
 	close(code = 0, reason = "") {
 		this.#quic.close({ closeCode: code, reason })
 	}
@@ -95,13 +103,43 @@ export class Connection {
 	}
 
 	async #recv(msg: Control.MessageWithType) {
-		if (Control.isPublisher(msg.type)) {
+		if (msg.type === Control.ControlMessageType.GoAway) {
+			await this.#handleGoAway(msg.message)
+		} else if (Control.isPublisher(msg.type)) {
 			await this.#subscriber.recv(msg)
 		} else {
 			await this.#publisher.recv(msg)
 		}
 	}
 
+	async #handleGoAway(msg: Control.GoAway) {
+		console.log("preparing for migration, got go_away message:", msg)
+		if (this.#migrationState === "in_progress") {
+			throw new Error("go away received twice")
+		}
+		this.#migrationState = "in_progress"
+		await this.#subscriber.startMigration()
+		await this.#publisher.startMigration()
+
+		// FIXME(itzmanish): is this how we close the quic connection for go_away?
+		this.#quic.close({ closeCode: 0, reason: "going_away" })
+
+		const { quic, control, objects } = await this.onMigration(msg.session_uri)
+		this.#quic = quic
+		this.#controlStream = control
+		this.#objects = objects
+		this.#migrationState = "done"
+
+		await this.#publisher.migrationDone(control, objects)
+		await this.#subscriber.migrationDone(control, objects)
+	}
+
+	async migrateSession(quic: WebTransport, control: ControlStream, objects: Objects) {
+		this.#quic = quic
+		this.#controlStream = control
+		this.#objects = objects
+	}
+
 	async closed(): Promise<Error> {
 		try {
 			await this.#running

lib/transport/control/go_away.ts

@@ -21,6 +21,10 @@ export namespace GoAway {
 
     export function deserialize(reader: ImmutableBytesBuffer): GoAway {
         const session_uri = reader.getUtf8String()
+        if (session_uri.length > 8192) {
+            // TODO(itzmanish): should be PROTOCOL_ERROR
+            throw new Error("session_uri too long")
+        }
         return {
             session_uri
         }

lib/transport/control/index.ts

@@ -59,6 +59,7 @@ type MessageWithType =
     | { type: ControlMessageType.SubscribeNamespaceOk; message: SubscribeNamespaceOk }
     | { type: ControlMessageType.SubscribeNamespaceError; message: SubscribeNamespaceError }
     | { type: ControlMessageType.Unsubscribe; message: Unsubscribe }
+    | { type: ControlMessageType.GoAway; message: GoAway }
 
 type Message = Subscriber | Publisher
 

lib/transport/publisher.ts

@@ -3,11 +3,14 @@ import { ControlStream } from "./stream"
 import { Queue, Watch } from "../common/async"
 import { Objects, TrackWriter, ObjectDatagramType } from "./objects"
 import { SubgroupType, SubgroupWriter } from "./subgroup"
+import { MigrationState } from "./connection"
 
 export class Publisher {
 	// Used to send control messages
 	#control: ControlStream
 
+	#migrationState: MigrationState = "none"
+
 	// Use to send objects.
 	#objects: Objects
 
@@ -27,7 +30,21 @@ export class Publisher {
 		this.#objects = objects
 	}
 
+	async startMigration() {
+		this.#migrationState = "in_progress"
+	}
+
+	async migrationDone(control: ControlStream, objects: Objects) {
+		this.#migrationState = "done"
+		this.#control = control
+		this.#objects = objects
+		// NOTE(itzmanish): should we republish all the tracks?
+	}
+
 	async publish_namespace(namespace: string[]): Promise<PublishNamespaceSend> {
+		if (this.#migrationState === "in_progress") {
+			throw new Error(`migration in progress`)
+		}
 		if (this.#publishedNamespaces.has(namespace.join("/"))) {
 			throw new Error(`already announced: ${namespace.join("/")}`)
 		}
@@ -103,6 +120,9 @@ export class Publisher {
 	}
 
 	async recvSubscribe(msg: Control.Subscribe) {
+		if (this.#migrationState === "in_progress") {
+			throw new Error(`migration in progress`)
+		}
 		try {
 			if (this.#subscribe.has(msg.id)) {
 				throw new Error(`duplicate subscribe for id: ${msg.id}`)

lib/transport/stream.ts

@@ -9,6 +9,7 @@ import {
 	Subscribe, SubscribeOk, SubscribeError,
 	SubscribeUpdate, SubscribeNamespace,
 	SubscribeNamespaceOk, SubscribeNamespaceError,
+	GoAway,
 } from "./control"
 import { debug } from "./utils"
 import { ImmutableBytesBuffer, ReadableWritableStreamBuffer, Reader, Writer } from "./buffer"
@@ -94,6 +95,12 @@ export class Decoder {
 
 		let res: MessageWithType
 		switch (t) {
+			case ControlMessageType.GoAway:
+				res = {
+					type: t,
+					message: GoAway.deserialize(payload),
+				}
+				break
 			case ControlMessageType.Subscribe:
 				res = {
 					type: t,

lib/transport/subscriber.ts

@@ -5,6 +5,7 @@ import type { TrackReader } from "./objects"
 import { debug } from "./utils"
 import { ControlStream } from "./stream"
 import { SubgroupReader } from "./subgroup"
+import { MigrationState } from "./connection"
 
 export interface TrackInfo {
 	track_alias: bigint
@@ -15,6 +16,9 @@ export class Subscriber {
 	// Use to send control messages.
 	#control: ControlStream
 
+	#migrationState: MigrationState = "none"
+	#tracksToMigrate = new Set<[string[], string]>()
+
 	// Use to send objects.
 	#objects: Objects
 
@@ -38,6 +42,24 @@ export class Subscriber {
 		return this.#publishedNamespacesQueue
 	}
 
+	async startMigration() {
+		this.#migrationState = "in_progress"
+		// close all the subscription
+		this.#trackToIDMap.forEach(async (id, track) => {
+			await this.unsubscribe(track, true);
+		})
+	}
+
+	async migrationDone(control: ControlStream, objects: Objects) {
+		this.#migrationState = "done"
+		this.#control = control
+		this.#objects = objects
+		this.#tracksToMigrate.forEach(async (track) => {
+			await this.subscribe(track[0], track[1]);
+		})
+		this.#tracksToMigrate.clear()
+	}
+
 	async recv(msg: Control.MessageWithType) {
 		const { type, message } = msg;
 		switch (type) {
@@ -82,8 +104,10 @@ export class Subscriber {
 	}
 
 	async subscribe_namespace(namespace: string[]) {
+		if (this.#migrationState === "in_progress") {
+			throw new Error(`migration in progress`)
+		}
 		const id = this.#control.nextRequestId()
-		// TODO(itzmanish): implement this
 		const msg: Control.MessageWithType = {
 			type: Control.ControlMessageType.SubscribeNamespace,
 			message: {
@@ -95,6 +119,9 @@ export class Subscriber {
 	}
 
 	async subscribe(namespace: string[], track: string) {
+		if (this.#migrationState === "in_progress") {
+			throw new Error(`migration in progress`)
+		}
 		const id = this.#control.nextRequestId()
 
 		const subscribe = new SubscribeSend(this.#control, id, namespace, track)
@@ -122,22 +149,29 @@ export class Subscriber {
 		return subscribe
 	}
 
-	async unsubscribe(track: string) {
-		if (this.#trackToIDMap.has(track)) {
-			const trackID = this.#trackToIDMap.get(track)
-			if (trackID === undefined) {
-				console.warn(`Exception track ${track} not found in trackToIDMap.`)
-				return
-			}
-			try {
-				await this.#control.send({ type: Control.ControlMessageType.Unsubscribe, message: { id: trackID } })
-				this.#trackToIDMap.delete(track)
-			} catch (error) {
-				console.error(`Failed to unsubscribe from track ${track}:`, error)
+	async unsubscribe(track: string, isMigrating = false) {
+		const trackID = this.#trackToIDMap.get(track)
+
+		if (trackID === undefined) {
+			console.warn(`Exception track ${track} not found in trackToIDMap.`)
+			return
+		}
+
+		try {
+			const subscription = this.#subscribe.get(trackID)
+			if (subscription) {
+				this.#subscribe.delete(trackID)
+				await subscription.close()
+				if (isMigrating) {
+					this.#tracksToMigrate.add([subscription.namespace, track])
+				}
 			}
-		} else {
-			console.warn(`During unsubscribe request initiation attempt track ${track} not found in trackToIDMap.`)
+			this.#trackToIDMap.delete(track)
+
+		} catch (error) {
+			console.error(`Failed to unsubscribe from track ${track}:`, error)
 		}
+
 	}
 
 	recvSubscribeOk(msg: Control.SubscribeOk) {
@@ -265,8 +299,10 @@ export class SubscribeSend {
 	}
 
 	async close(_code = 0n, _reason = "") {
-		// TODO implement unsubscribe
-		// await this.#inner.sendReset(code, reason)
+		this.#control.send({
+			type: Control.ControlMessageType.Unsubscribe,
+			message: { id: this.#id }
+		})
 	}
 
 	onOk(trackAlias: bigint) {