Add WebSocket support (#12)

frontend/sync-client/src/sync-operations/syncer.ts

@@ -1,30 +1,40 @@
-import type { Database, RelativePath } from "../persistence/database";
+import type {
+	Database,
+	DocumentId,
+	RelativePath
+} from "../persistence/database";
 import type { SyncService } from "../services/sync-service";
 import type { Logger } from "../tracing/logger";
 import PQueue from "p-queue";
 import { hash } from "../utils/hash";
 import { v4 as uuidv4 } from "uuid";
 import type { components } from "../services/types";
-import type { Settings } from "../persistence/settings";
+import type { Settings, SyncSettings } from "../persistence/settings";
 import type { FileOperations } from "../file-operations/file-operations";
 import { findMatchingFile } from "../utils/find-matching-file";
 import type { UnrestrictedSyncer } from "./unrestricted-syncer";
 import { createPromise } from "../utils/create-promise";
 import { SyncResetError } from "../services/sync-reset-error";
+import { Locks } from "../utils/locks";
 
 export class Syncer {
+	private readonly remoteDocumentsLock: Locks<DocumentId>;
 	private readonly remainingOperationsListeners: ((
 		remainingOperations: number
 	) => void)[] = [];
+	private readonly webSocketStatusChangeListeners: (() => void)[] = [];
 	private readonly syncQueue: PQueue;
 
 	private runningScheduleSyncForOfflineChanges: Promise<void> | undefined;
-	private runningApplyRemoteChangesLocally: Promise<void> | undefined;
+	private refreshApplyRemoteChangesWebSocketInterval:
+		| NodeJS.Timeout
+		| undefined;
+	private applyRemoteChangesWebSocket: WebSocket | undefined;
 
 	public constructor(
 		private readonly logger: Logger,
 		private readonly database: Database,
-		settings: Settings,
+		private readonly settings: Settings,
 		private readonly syncService: SyncService,
 		private readonly operations: FileOperations,
 		private readonly internalSyncer: UnrestrictedSyncer
@@ -33,11 +43,23 @@ export class Syncer {
 			concurrency: settings.getSettings().syncConcurrency
 		});
 
+		this.updateWebSocket(settings.getSettings());
+
+		this.remoteDocumentsLock = new Locks<DocumentId>(this.logger);
+
 		settings.addOnSettingsChangeListener((newSettings, oldSettings) => {
-			if (newSettings.syncConcurrency === oldSettings.syncConcurrency) {
-				return;
+			if (
+				newSettings.remoteUri !== oldSettings.remoteUri ||
+				newSettings.vaultName !== oldSettings.vaultName ||
+				newSettings.token !== oldSettings.token ||
+				newSettings.isSyncEnabled !== oldSettings.isSyncEnabled
+			) {
+				this.updateWebSocket(newSettings);
+			}
+
+			if (newSettings.syncConcurrency !== oldSettings.syncConcurrency) {
+				this.syncQueue.concurrency = newSettings.syncConcurrency;
 			}
-			this.syncQueue.concurrency = newSettings.syncConcurrency;
 		});
 
 		this.syncQueue.on("active", () => {
@@ -45,6 +67,12 @@ export class Syncer {
 				listener(this.syncQueue.size);
 			});
 		});
+
+		this.setWebSocketRefreshInterval();
+	}
+
+	public get isWebSocketConnected(): boolean {
+		return this.applyRemoteChangesWebSocket?.readyState === WebSocket.OPEN;
 	}
 
 	public addRemainingOperationsListener(
@@ -53,6 +81,10 @@ export class Syncer {
 		this.remainingOperationsListeners.push(listener);
 	}
 
+	public addWebSocketStatusChangeListener(listener: () => void): void {
+		this.webSocketStatusChangeListeners.push(listener);
+	}
+
 	public async syncLocallyCreatedFile(
 		relativePath: RelativePath
 	): Promise<void> {
@@ -206,109 +238,139 @@ export class Syncer {
 		}
 	}
 
-	public async applyRemoteChangesLocally(): Promise<void> {
-		if (this.runningApplyRemoteChangesLocally !== undefined) {
-			this.logger.debug(
-				"Applying remote changes locally is already in progress"
-			);
-			return this.runningApplyRemoteChangesLocally;
-		}
-
-		try {
-			this.runningApplyRemoteChangesLocally =
-				this.internalApplyRemoteChangesLocally();
-			await this.runningApplyRemoteChangesLocally;
-			this.logger.info("All remote changes have been applied locally");
-		} catch (e) {
-			if (e instanceof SyncResetError) {
-				this.logger.info(
-					"Failed to apply remote changes locally due to a reset"
-				);
-				return;
-			}
-			this.logger.error(`Failed to apply remote changes locally: ${e}`);
-			throw e;
-		} finally {
-			this.runningApplyRemoteChangesLocally = undefined;
-		}
+	public async waitUntilFinished(): Promise<void> {
+		await this.runningScheduleSyncForOfflineChanges;
+		return this.syncQueue.onEmpty();
 	}
 
 	public async reset(): Promise<void> {
 		await this.waitUntilFinished();
-		this.internalSyncer.reset();
+		this.setWebSocketRefreshInterval();
+		this.updateWebSocket(this.settings.getSettings());
 	}
 
-	public async waitUntilFinished(): Promise<void> {
-		await Promise.allSettled([
-			this.runningScheduleSyncForOfflineChanges,
-			this.runningApplyRemoteChangesLocally
-		]);
-		return this.syncQueue.onEmpty();
+	public stop(): void {
+		clearInterval(this.refreshApplyRemoteChangesWebSocketInterval);
+		this.applyRemoteChangesWebSocket?.close();
 	}
 
-	private async internalApplyRemoteChangesLocally(): Promise<void> {
-		const remote = await this.syncQueue.add(async () =>
-			this.syncService.getAll(this.database.getLastSeenUpdateId())
-		);
+	private updateWebSocket(settings: SyncSettings): void {
+		this.applyRemoteChangesWebSocket?.close();
 
-		if (!remote) {
-			throw new Error("Failed to fetch remote changes");
-		}
-
-		if (remote.latestDocuments.length === 0) {
-			this.logger.debug("No remote changes to apply");
+		if (!settings.isSyncEnabled) {
+			this.applyRemoteChangesWebSocket = undefined;
 			return;
 		}
 
-		this.logger.info("Applying remote changes locally");
+		const wsUri = new URL(settings.remoteUri);
+		wsUri.protocol = wsUri.protocol === "https" ? "wss" : "ws";
+		wsUri.pathname = `/vaults/${settings.vaultName}/ws`;
 
-		await Promise.all(
-			remote.latestDocuments.map(this.syncRemotelyUpdatedFile.bind(this))
-		);
-
-		const lastSeenUpdateId = this.database.getLastSeenUpdateId();
 		if (
-			lastSeenUpdateId === undefined ||
-			lastSeenUpdateId < remote.lastUpdateId
+			typeof globalThis !== "undefined" &&
+			typeof globalThis.WebSocket === "undefined"
 		) {
-			this.database.setLastSeenUpdateId(remote.lastUpdateId);
+			// polyfill for WebSocket in Node.js
+			// eslint-disable-next-line
+			globalThis.WebSocket = require("ws");
 		}
+
+		this.applyRemoteChangesWebSocket = new WebSocket(wsUri);
+
+		this.applyRemoteChangesWebSocket.onmessage = (event): void =>
+			void this.syncRemotelyUpdatedFile(event.data).catch(
+				(e: unknown) => {
+					this.logger.error(
+						`Failed to sync remotely updated file: ${e}`
+					);
+				}
+			);
+
+		// The JS WebSocket API doesn't support setting headers, so we have to send the token as a message
+		this.applyRemoteChangesWebSocket.onopen = (): void => {
+			this.applyRemoteChangesWebSocket?.send(settings.token);
+			this.webSocketStatusChangeListeners.forEach((listener) => {
+				listener();
+			});
+		};
+
+		this.applyRemoteChangesWebSocket.onclose = (event): void => {
+			this.logger.warn(
+				`WebSocket closed with code ${event.code}: ${event.reason}`
+			);
+			this.webSocketStatusChangeListeners.forEach((listener) => {
+				listener();
+			});
+		};
 	}
 
-	private async syncRemotelyUpdatedFile(
-		remoteVersion: components["schemas"]["DocumentVersionWithoutContent"]
-	): Promise<void> {
+	private setWebSocketRefreshInterval(): void {
+		this.refreshApplyRemoteChangesWebSocketInterval = setInterval(() => {
+			if (
+				this.applyRemoteChangesWebSocket?.readyState === WebSocket.OPEN
+			) {
+				return;
+			}
+			this.updateWebSocket(this.settings.getSettings());
+		}, 5000);
+	}
+
+	private async syncRemotelyUpdatedFile(message: string): Promise<void> {
+		// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
+		const remoteVersion = JSON.parse(
+			message
+		) as components["schemas"]["DocumentVersionWithoutContent"];
+
 		let document = this.database.getDocumentByDocumentId(
 			remoteVersion.documentId
 		);
 
-		const [promise, resolve, reject] = createPromise();
-
+		let hasLockToRelease = false;
 		if (document === undefined) {
-			await this.syncQueue.add(async () =>
-				this.internalSyncer.unrestrictedSyncRemotelyUpdatedFile(
-					remoteVersion
-				)
+			// Let's avoid the same documents getting created in parallel multiple times
+			await this.remoteDocumentsLock.waitForLock(
+				remoteVersion.documentId
 			);
-		} else {
-			document = await this.database.getResolvedDocumentByRelativePath(
-				document.relativePath,
-				promise
+			hasLockToRelease = true;
+			document = this.database.getDocumentByDocumentId(
+				remoteVersion.documentId
 			);
+		}
 
-			try {
+		try {
+			if (document === undefined) {
 				await this.syncQueue.add(async () =>
 					this.internalSyncer.unrestrictedSyncRemotelyUpdatedFile(
-						remoteVersion,
-						document
+						remoteVersion
 					)
 				);
+			} else {
+				const [promise, resolve, reject] = createPromise();
 
-				resolve();
-			} catch (e) {
-				reject(e);
-			} finally {
-				this.database.removeDocumentPromise(promise);
+				document =
+					await this.database.getResolvedDocumentByRelativePath(
+						document.relativePath,
+						promise
+					);
+
+				try {
+					await this.syncQueue.add(async () =>
+						this.internalSyncer.unrestrictedSyncRemotelyUpdatedFile(
+							remoteVersion,
+							document
+						)
+					);
+
+					resolve();
+				} catch (e) {
+					reject(e);
+				} finally {
+					this.database.removeDocumentPromise(promise);
+				}
+			}
+		} finally {
+			if (hasLockToRelease) {
+				this.remoteDocumentsLock.unlock(remoteVersion.documentId);
 			}
 		}
 	}