feat: process update message in the order that they are received (#1058)

* feat: process update message in the order that they are received

* feat: beforeSync can be made async again, as the server now processes messages in order

Author: janthurau

packages/extension-redis/src/Redis.ts

@@ -369,16 +369,13 @@ export class Redis implements Extension {
 			return;
 		}
 
-		new MessageReceiver(message, this.redisTransactionOrigin).apply(
-			document,
-			undefined,
-			(reply) => {
-				return this.pub.publish(
-					this.pubKey(document.name),
-					this.encodeMessage(reply),
-				);
-			},
-		);
+		const receiver = new MessageReceiver(message, this.redisTransactionOrigin);
+		await receiver.apply(document, undefined, (reply) => {
+			return this.pub.publish(
+				this.pubKey(document.name),
+				this.encodeMessage(reply),
+			);
+		});
 	};
 
 	/**

packages/server/src/ClientConnection.ts

@@ -164,6 +164,12 @@ export class ClientConnection<Context = any> {
 		);
 
 		instance.onClose(async (document, event) => {
+			// Wait for any pending message processing to complete before running
+			// disconnect hooks. This ensures that document updates from queued messages
+			// are applied (and their debounced onStoreDocument scheduled) before the
+			// disconnect handler checks whether to call executeNow.
+			await instance.waitForPendingMessages();
+
 			const disconnectHookPayload: onDisconnectPayload = {
 				instance: this.documentProvider as Hocuspocus, // TODO, this will be removed when we use events instead of hooks for this class
 				clientsCount: document.getConnectionsCount(),

packages/server/src/Connection.ts

@@ -9,10 +9,7 @@ import type Document from "./Document.ts";
 import { IncomingMessage } from "./IncomingMessage.ts";
 import { MessageReceiver } from "./MessageReceiver.ts";
 import { OutgoingMessage } from "./OutgoingMessage.ts";
-import type {
-	beforeSyncPayload,
-	onStatelessPayload,
-} from "./types.ts";
+import type { beforeSyncPayload, onStatelessPayload } from "./types.ts";
 
 export class Connection<Context = any> {
 	webSocket: WebSocket;
@@ -39,6 +36,10 @@ export class Connection<Context = any> {
 
 	readOnly: boolean;
 
+	private messageQueue: Uint8Array[] = [];
+
+	private processingPromise: Promise<void> = Promise.resolve();
+
 	/**
 	 * Constructor.
 	 */
@@ -121,6 +122,13 @@ export class Connection<Context = any> {
 		return this;
 	}
 
+	/**
+	 * Returns a promise that resolves when all queued messages have been processed.
+	 */
+	waitForPendingMessages(): Promise<void> {
+		return this.processingPromise;
+	}
+
 	/**
 	 * Send the given message
 	 */
@@ -204,30 +212,34 @@ export class Connection<Context = any> {
 	 * @public
 	 */
 	public handleMessage(data: Uint8Array): void {
-		const message = new IncomingMessage(data);
-		const documentName = message.readVarString();
-
-		if (documentName !== this.document.name) return;
-
-		message.writeVarString(documentName);
-
-		this.callbacks
-			.beforeHandleMessage(this, data)
-			.then(() => {
-				try {
-					new MessageReceiver(message).apply(this.document, this);
-				} catch (e: any) {
-					console.error(
-						`closing connection ${this.socketId} (while handling ${documentName}) because of exception`,
-						e,
-					);
-					this.close({
-						code: "code" in e ? e.code : ResetConnection.code,
-						reason: "reason" in e ? e.reason : ResetConnection.reason,
-					});
-				}
-			})
-			.catch((e: any) => {
+		this.messageQueue.push(data);
+
+		if (this.messageQueue.length === 1) {
+			this.processingPromise = this.processMessages();
+		}
+	}
+
+	private async processMessages() {
+		while (this.messageQueue.length > 0) {
+			const rawUpdate = this.messageQueue.at(0) as Uint8Array;
+
+			const message = new IncomingMessage(rawUpdate);
+			const documentName = message.readVarString();
+
+			if (documentName !== this.document.name) {
+				this.messageQueue.shift();
+				continue;
+			}
+
+			message.writeVarString(documentName);
+
+			try {
+				await this.callbacks.beforeHandleMessage(this, rawUpdate);
+				const receiver = new MessageReceiver(message);
+
+				await receiver.apply(this.document, this);
+				// biome-ignore lint/suspicious/noExplicitAny: <explanation>
+			} catch (e: any) {
 				console.error(
 					`closing connection ${this.socketId} (while handling ${documentName}) because of exception`,
 					e,
@@ -236,7 +248,10 @@ export class Connection<Context = any> {
 					code: "code" in e ? e.code : ResetConnection.code,
 					reason: "reason" in e ? e.reason : ResetConnection.reason,
 				});
-			});
+			}
+
+			this.messageQueue.shift();
+		}
 	}
 }
 

packages/server/src/IncomingMessage.ts

@@ -1,17 +1,17 @@
 import type { Decoder } from "lib0/decoding";
 import {
 	createDecoder,
+	readVarString,
 	readVarUint,
 	readVarUint8Array,
-	readVarString,
 } from "lib0/decoding";
 import type { Encoder } from "lib0/encoding";
 import {
 	createEncoder,
+	length,
 	toUint8Array,
-	writeVarUint,
 	writeVarString,
-	length,
+	writeVarUint,
 } from "lib0/encoding";
 import type { MessageType } from "./types.ts";
 

packages/server/src/MessageReceiver.ts

@@ -27,7 +27,7 @@ export class MessageReceiver {
 		this.defaultTransactionOrigin = defaultTransactionOrigin;
 	}
 
-	public apply(
+	public async apply(
 		document: Document,
 		connection?: Connection,
 		reply?: (message: Uint8Array) => void,
@@ -40,7 +40,7 @@ export class MessageReceiver {
 			case MessageType.Sync:
 			case MessageType.SyncReply: {
 				message.writeVarUint(MessageType.Sync);
-				this.readSyncMessage(
+				await this.readSyncMessage(
 					message,
 					document,
 					connection,
@@ -126,7 +126,7 @@ export class MessageReceiver {
 		}
 	}
 
-	readSyncMessage(
+	async readSyncMessage(
 		message: IncomingMessage,
 		document: Document,
 		connection?: Connection,
@@ -136,7 +136,7 @@ export class MessageReceiver {
 		const type = message.readVarUint();
 
 		if (connection) {
-			connection.callbacks.beforeSync(connection, {
+			await connection.callbacks.beforeSync(connection, {
 				type,
 				payload: message.peekVarUint8Array(),
 			});