refactor(mutex): improve multi locks operations to wait for all

Wait for all locks to be available instead of locking one side and
waiting

app/server/core/__tests__/repository-mutex.test.ts

@@ -232,9 +232,10 @@ describe("RepositoryMutex", () => {
 		releaseS3();
 	});
 
-	test("should acquire many locks in repository order", async () => {
+	test("should wait to acquire all many-lock requests before locking any repository", async () => {
 		const releaseSharedB = await repoMutex.acquireShared("repo-b", "holder-b");
-		const results: string[] = [];
+		let manyAcquired = false;
+		let exclusiveAAcquired = false;
 
 		const manyPromise = repoMutex
 			.acquireMany([
@@ -242,34 +243,34 @@ describe("RepositoryMutex", () => {
 				{ repositoryId: "repo-a", type: "shared", operation: "many-a" },
 			])
 			.then((release) => {
-				results.push("many");
+				manyAcquired = true;
 				return release;
 			});
 
 		await new Promise((resolve) => setTimeout(resolve, 10));
 
 		const exclusiveAPromise = repoMutex.acquireExclusive("repo-a", "exclusive-a").then((release) => {
-			results.push("exclusive-a");
+			exclusiveAAcquired = true;
 			return release;
 		});
 
 		await new Promise((resolve) => setTimeout(resolve, 10));
-		expect(results).toEqual([]);
+		expect(exclusiveAAcquired).toBe(true);
+		expect(manyAcquired).toBe(false);
 
 		releaseSharedB();
-
+		await new Promise((resolve) => setTimeout(resolve, 10));
-		const releaseMany = await manyPromise;
+		expect(manyAcquired).toBe(false);
-		expect(results).toEqual(["many"]);
-
-		releaseMany();
 
 		const releaseExclusiveA = await exclusiveAPromise;
-		expect(results).toEqual(["many", "exclusive-a"]);
-
 		releaseExclusiveA();
+
+		const releaseMany = await manyPromise;
+		expect(manyAcquired).toBe(true);
+		releaseMany();
 	});
 
-	test("should release earlier locks when acquireMany is aborted", async () => {
+	test("should abort acquireMany without leaving partial locks behind", async () => {
 		const releaseExclusiveB = await repoMutex.acquireExclusive("repo-b", "holder-b");
 		const controller = new AbortController();
 
@@ -290,7 +291,7 @@ describe("RepositoryMutex", () => {
 		});
 
 		await new Promise((resolve) => setTimeout(resolve, 10));
-		expect(exclusiveAAcquired).toBe(false);
+		expect(exclusiveAAcquired).toBe(true);
 
 		controller.abort(new Error("stop"));
 		await expect(manyPromise).rejects.toThrow("stop");
@@ -302,6 +303,30 @@ describe("RepositoryMutex", () => {
 		releaseExclusiveB();
 	});
 
+	test("should abort acquireMany if the signal aborts after waiting resolves", async () => {
+		const repoA = "many-abort-after-wake-a";
+		const repoB = "many-abort-after-wake-b";
+		const releaseExclusiveB = await repoMutex.acquireExclusive(repoB, "holder-b");
+		const controller = new AbortController();
+
+		const manyPromise = repoMutex.acquireMany(
+			[
+				{ repositoryId: repoB, type: "exclusive", operation: "many-b" },
+				{ repositoryId: repoA, type: "shared", operation: "many-a" },
+			],
+			controller.signal,
+		);
+
+		await new Promise((resolve) => setTimeout(resolve, 10));
+
+		releaseExclusiveB();
+		controller.abort(new Error("stop after wake"));
+
+		await expect(manyPromise).rejects.toThrow("stop after wake");
+		expect(repoMutex.isLocked(repoA)).toBe(false);
+		expect(repoMutex.isLocked(repoB)).toBe(false);
+	});
+
 	test("should safely handle multiple calls to the release function", async () => {
 		const repoId = "idempotent-release";
 

app/server/core/repository-mutex.ts

@@ -26,6 +26,7 @@ interface RepositoryLockState {
 
 class RepositoryMutex {
 	private locks = new Map<string, RepositoryLockState>();
+	private changeListeners = new Map<string, Set<() => void>>();
 	private lockIdCounter = 0;
 
 	private getOrCreateState(repositoryId: string): RepositoryLockState {
@@ -52,6 +53,129 @@ class RepositoryMutex {
 		}
 	}
 
+	private notifyChange(repositoryId: string): void {
+		const listeners = this.changeListeners.get(repositoryId);
+		if (!listeners) {
+			return;
+		}
+
+		for (const listener of listeners) {
+			listener();
+		}
+	}
+
+	private canAcquireImmediately(state: RepositoryLockState | undefined, type: LockType): boolean {
+		if (!state) {
+			return true;
+		}
+
+		if (type === "shared") {
+			const hasExclusiveInQueue = state.waitQueue.some((item) => item.type === "exclusive");
+			return !state.exclusiveHolder && !hasExclusiveInQueue;
+		}
+
+		return !state.exclusiveHolder && state.sharedHolders.size === 0 && state.waitQueue.length === 0;
+	}
+
+	private waitForChange(repositoryIds: string[], signal?: AbortSignal) {
+		if (signal?.aborted) {
+			throw signal.reason || new Error("Operation aborted");
+		}
+
+		return new Promise<void>((resolve, reject) => {
+			const uniqueRepositoryIds = [...new Set(repositoryIds)];
+			const cleanupCallbacks: Array<() => void> = [];
+			let settled = false;
+
+			const settle = (callback: () => void) => {
+				if (settled) {
+					return;
+				}
+
+				settled = true;
+				for (const cleanup of cleanupCallbacks) {
+					cleanup();
+				}
+				callback();
+			};
+
+			for (const repositoryId of uniqueRepositoryIds) {
+				let listeners = this.changeListeners.get(repositoryId);
+				if (!listeners) {
+					listeners = new Set();
+					this.changeListeners.set(repositoryId, listeners);
+				}
+
+				const listener = () => settle(resolve);
+				listeners.add(listener);
+
+				cleanupCallbacks.push(() => {
+					const currentListeners = this.changeListeners.get(repositoryId);
+					if (!currentListeners) {
+						return;
+					}
+
+					currentListeners.delete(listener);
+					if (currentListeners.size === 0) {
+						this.changeListeners.delete(repositoryId);
+					}
+				});
+			}
+
+			if (signal) {
+				const onAbort = () => {
+					settle(() => reject(signal.reason || new Error("Operation aborted")));
+				};
+
+				signal.addEventListener("abort", onAbort);
+				cleanupCallbacks.push(() => {
+					signal.removeEventListener("abort", onAbort);
+				});
+			}
+		});
+	}
+
+	private tryAcquireMany(requests: LockRequest[]) {
+		for (const request of requests) {
+			if (!this.canAcquireImmediately(this.locks.get(request.repositoryId), request.type)) {
+				return null;
+			}
+		}
+
+		const releases = requests.map((request) => {
+			const state = this.getOrCreateState(request.repositoryId);
+			const lockId = this.generateLockId();
+
+			if (request.type === "shared") {
+				state.sharedHolders.set(lockId, {
+					id: lockId,
+					operation: request.operation,
+					acquiredAt: Date.now(),
+				});
+			} else {
+				state.exclusiveHolder = {
+					id: lockId,
+					operation: request.operation,
+					acquiredAt: Date.now(),
+				};
+			}
+
+			return this.createRelease(request.type, request.repositoryId, lockId);
+		});
+
+		let released = false;
+		return () => {
+			if (released) {
+				return;
+			}
+
+			released = true;
+			for (const release of releases.toReversed()) {
+				release();
+			}
+		};
+	}
+
 	async acquireShared(repositoryId: string, operation: string, signal?: AbortSignal): Promise<() => void> {
 		if (signal?.aborted) {
 			throw signal.reason || new Error("Operation aborted");
@@ -88,6 +212,7 @@ class RepositoryMutex {
 						if (index !== -1) {
 							state.waitQueue.splice(index, 1);
 							this.cleanupStateIfEmpty(repositoryId);
+							this.notifyChange(repositoryId);
 							reject(signal.reason || new Error("Operation aborted"));
 						}
 					};
@@ -142,6 +267,7 @@ class RepositoryMutex {
 						if (index !== -1) {
 							state.waitQueue.splice(index, 1);
 							this.cleanupStateIfEmpty(repositoryId);
+							this.notifyChange(repositoryId);
 							reject(signal.reason || new Error("Operation aborted"));
 						}
 					};
@@ -182,34 +308,21 @@ class RepositoryMutex {
 		}
 
 		const sortedRequests = [...requests].sort((a, b) => a.repositoryId.localeCompare(b.repositoryId));
-		const releases: Array<() => void> = [];
+		while (true) {
+			const releaseLocks = this.tryAcquireMany(sortedRequests);
+			if (releaseLocks) {
+				return releaseLocks;
+			}
 
-		try {
+			await this.waitForChange(
-			for (const request of sortedRequests) {
+				sortedRequests.map((request) => request.repositoryId),
-				const release =
+				signal,
-					request.type === "shared"
+			);
-						? await this.acquireShared(request.repositoryId, request.operation, signal)
+
-						: await this.acquireExclusive(request.repositoryId, request.operation, signal);
+			if (signal?.aborted) {
-				releases.push(release);
+				throw signal.reason || new Error("Operation aborted");
 			}
-		} catch (error) {
-			for (const release of releases.toReversed()) {
-				release();
-			}
-			throw error;
 		}
-
-		let released = false;
-		return () => {
-			if (released) {
-				return;
-			}
-
-			released = true;
-			for (const release of releases.toReversed()) {
-				release();
-			}
-		};
 	}
 
 	private releaseShared(repositoryId: string, lockId: string): void {
@@ -229,6 +342,7 @@ class RepositoryMutex {
 
 		this.processWaitQueue(repositoryId);
 		this.cleanupStateIfEmpty(repositoryId);
+		this.notifyChange(repositoryId);
 	}
 
 	private releaseExclusive(repositoryId: string, lockId: string): void {
@@ -249,6 +363,7 @@ class RepositoryMutex {
 
 		this.processWaitQueue(repositoryId);
 		this.cleanupStateIfEmpty(repositoryId);
+		this.notifyChange(repositoryId);
 	}
 
 	private processWaitQueue(repositoryId: string): void {