fix(restore): do not wait for lock before returning response

app/server/modules/repositories/__tests__/repositories.service.test.ts

@@ -14,11 +14,13 @@ import { db } from "~/server/db/db";
 import { agentsTable, repositoriesTable, type RepositoryInsert } from "~/server/db/schema";
 import { generateShortId } from "~/server/utils/id";
 import { restic } from "~/server/core/restic";
-import { agentManager } from "~/server/modules/agents/agents-manager";
+import { agentManager, type RestoreExecutionResult } from "~/server/modules/agents/agents-manager";
 import { createTestSession } from "~/test/helpers/auth";
 import { createTestBackupSchedule } from "~/test/helpers/backup";
 import { cache, cacheKeys } from "~/server/utils/cache";
 import { ResticError } from "@zerobyte/core/restic/server";
+import { repoMutex } from "~/server/core/repository-mutex";
+import { taskStore } from "~/server/modules/tasks/tasks.store";
 import { repositoriesService } from "../repositories.service";
 
 const createTestRepository = async (organizationId: string, overrides: Partial<RepositoryInsert> = {}) => {
@@ -416,8 +418,26 @@ describe("repositoriesService.restoreSnapshot", () => {
 	let originalEnableLocalAgent: boolean;
 	const createPendingRestoreStart = () => ({
 		status: "started" as const,
-		result: new Promise<never>(() => {}),
+		result: new Promise<RestoreExecutionResult>(() => {}),
 	});
+	const resolveWithin = async <T>(promise: Promise<T>, timeoutMs: number) => {
+		return await new Promise<T>((resolve, reject) => {
+			const timeout = setTimeout(() => {
+				reject(new Error(`Expected promise to resolve within ${timeoutMs}ms`));
+			}, timeoutMs);
+
+			promise.then(
+				(value) => {
+					clearTimeout(timeout);
+					resolve(value);
+				},
+				(error) => {
+					clearTimeout(timeout);
+					reject(error);
+				},
+			);
+		});
+	};
 
 	beforeEach(() => {
 		originalEnableLocalAgent = config.flags.enableLocalAgent;
@@ -451,6 +471,7 @@ describe("repositoriesService.restoreSnapshot", () => {
 		return {
 			organizationId,
 			userId: session.user.id,
+			repositoryId: repository.id,
 			repositoryShortId: repository.shortId,
 			restoreMock,
 		};
@@ -485,20 +506,22 @@ describe("repositoriesService.restoreSnapshot", () => {
 			await fs.rm(targetPath, { recursive: true, force: true });
 		}
 
-		expect(restoreMock).toHaveBeenCalledWith(
+		await waitForExpect(() => {
-			"local",
+			expect(restoreMock).toHaveBeenCalledWith(
-			expect.objectContaining({
+				"local",
-				payload: expect.objectContaining({
+				expect.objectContaining({
-					snapshotId: "snapshot-restore",
+					payload: expect.objectContaining({
-					target: targetPath,
+						snapshotId: "snapshot-restore",
-					repositoryConfig: expect.objectContaining({ backend: "local" }),
+						target: targetPath,
-					options: expect.objectContaining({
+						repositoryConfig: expect.objectContaining({ backend: "local" }),
-						organizationId,
+						options: expect.objectContaining({
-						basePath: "/var/lib/zerobyte/volumes/vol123/_data",
+							organizationId,
+							basePath: "/var/lib/zerobyte/volumes/vol123/_data",
+						}),
 					}),
 				}),
-			}),
+			);
-		);
+		});
 	});
 
 	test("rejects starting a second active restore for the same snapshot", async () => {
@@ -524,6 +547,73 @@ describe("repositoriesService.restoreSnapshot", () => {
 		}
 	});
 
+	test("returns a restore id while waiting for the repository mutex", async () => {
+		const { organizationId, userId, repositoryId, repositoryShortId, restoreMock } =
+			await setupRestoreSnapshotScenario();
+		const targetPath = await fs.mkdtemp(nodePath.join(process.cwd(), "restore-target-"));
+		await withContext({ organizationId, userId }, () =>
+			repositoriesService.getSnapshotDetails(repositoryShortId, "snapshot-restore"),
+		);
+
+		let finishRestore: (result: RestoreExecutionResult) => void = () => {};
+		const restoreResult = new Promise<RestoreExecutionResult>((resolve) => {
+			finishRestore = resolve;
+		});
+		restoreMock.mockResolvedValueOnce({ status: "started", result: restoreResult });
+
+		const releaseExclusive = await repoMutex.acquireExclusive(repositoryId, "check");
+		let restoreId = "";
+		try {
+			const restoreStart = withContext({ organizationId, userId }, () =>
+				repositoriesService.restoreSnapshot(repositoryShortId, "snapshot-restore", {
+					targetPath,
+				}),
+			);
+
+			const result = await resolveWithin(restoreStart, 1000);
+			restoreId = result.restoreId;
+
+			expect(result.status).toBe("started");
+			expect(restoreMock).not.toHaveBeenCalled();
+
+			const task = taskStore.findActiveByResource({
+				organizationId,
+				kind: "restore",
+				resourceType: "repository",
+				resourceId: repositoryShortId,
+			});
+			expect(task?.id).toBe(restoreId);
+			expect(task?.status).toBe("queued");
+		} finally {
+			releaseExclusive();
+			await fs.rm(targetPath, { recursive: true, force: true });
+		}
+
+		await waitForExpect(() => {
+			expect(restoreMock).toHaveBeenCalledTimes(1);
+		});
+
+		finishRestore({
+			status: "completed",
+			result: {
+				message_type: "summary",
+				files_skipped: 0,
+				files_restored: 1,
+			},
+		});
+
+		await waitForExpect(() => {
+			expect(
+				taskStore.findActiveByResource({
+					organizationId,
+					kind: "restore",
+					resourceType: "repository",
+					resourceId: repositoryShortId,
+				}),
+			).toBeNull();
+		});
+	});
+
 	test("routes restore to the requested target agent", async () => {
 		const organizationId = session.organizationId;
 		const agentId = `agent-${randomUUID()}`;
@@ -572,14 +662,16 @@ describe("repositoriesService.restoreSnapshot", () => {
 			await fs.rm(targetPath, { recursive: true, force: true });
 		}
 
-		expect(restoreMock).toHaveBeenCalledWith(
+		await waitForExpect(() => {
-			agentId,
+			expect(restoreMock).toHaveBeenCalledWith(
-			expect.objectContaining({
+				agentId,
-				payload: expect.objectContaining({
+				expect.objectContaining({
-					target: targetPath,
+					payload: expect.objectContaining({
+						target: targetPath,
+					}),
 				}),
-			}),
+			);
-		);
+		});
 	});
 
 	test("rejects a target agent outside the current organization", async () => {
@@ -701,20 +793,22 @@ describe("repositoriesService.restoreSnapshot", () => {
 			await fs.rm(targetPath, { recursive: true, force: true });
 		}
 
-		expect(restoreMock).toHaveBeenCalledWith(
+		await waitForExpect(() => {
-			"local",
+			expect(restoreMock).toHaveBeenCalledWith(
-			expect.objectContaining({
+				"local",
-				payload: expect.objectContaining({
+				expect.objectContaining({
-					snapshotId: "snapshot-restore",
+					payload: expect.objectContaining({
-					target: targetPath,
+						snapshotId: "snapshot-restore",
-					repositoryConfig: expect.objectContaining({ backend: "local" }),
+						target: targetPath,
-					options: expect.objectContaining({
+						repositoryConfig: expect.objectContaining({ backend: "local" }),
-						organizationId,
+						options: expect.objectContaining({
-						basePath: "/",
+							organizationId,
+							basePath: "/",
+						}),
 					}),
 				}),
-			}),
+			);
-		);
+		});
 	});
 });
 

app/server/modules/repositories/restore-executor.ts

@@ -78,45 +78,68 @@ const executeControllerRestore = async (
 	}
 };
 
+const executeAgentRestore = async (
+	request: RestoreExecutionRequest,
+	signal: AbortSignal,
+): Promise<RestoreExecutionResult> => {
+	if (signal.aborted) {
+		return { status: "cancelled", message: "Restore was cancelled" };
+	}
+
+	try {
+		const payload = await createRestoreRunPayload(request);
+		const started = await agentManager.startRestore(request.executionAgentId, {
+			payload,
+			signal,
+			onStarted: request.onStarted,
+			onProgress: request.onProgress,
+		});
+
+		if (started.status === "unavailable") {
+			return started;
+		}
+
+		return await started.result;
+	} catch (error) {
+		if (signal.aborted) {
+			return { status: "cancelled", message: "Restore was cancelled" };
+		}
+
+		return { status: "failed", error: toMessage(error) };
+	}
+};
+
+const executeRestoreWithRepositoryLock = async (
+	request: RestoreExecutionRequest,
+	signal: AbortSignal,
+): Promise<RestoreExecutionResult> => {
+	let releaseLock: (() => void) | null = null;
+
+	try {
+		releaseLock = await repoMutex.acquireShared(request.repositoryId, `restore:${request.restoreId}`, signal);
+
+		if (shouldRunInController(request.executionAgentId)) {
+			return await executeControllerRestore(request, signal);
+		}
+
+		return await executeAgentRestore(request, signal);
+	} catch (error) {
+		if (signal.aborted) {
+			return { status: "cancelled", message: "Restore was cancelled" };
+		}
+
+		return { status: "failed", error: toMessage(error) };
+	} finally {
+		releaseLock?.();
+	}
+};
+
 export const restoreExecutor = {
-	start: async (request: RestoreExecutionRequest): Promise<RestoreExecutionHandle> => {
+	start: (request: RestoreExecutionRequest): RestoreExecutionHandle => {
 		const abortController = new AbortController();
 
-		let releaseLock: (() => void) | null = null;
+		return {
-		try {
+			result: executeRestoreWithRepositoryLock(request, abortController.signal),
-			releaseLock = await repoMutex.acquireShared(
+		};
-				request.repositoryId,
-				`restore:${request.restoreId}`,
-				abortController.signal,
-			);
-
-			let result: Promise<RestoreExecutionResult>;
-			if (shouldRunInController(request.executionAgentId)) {
-				result = executeControllerRestore(request, abortController.signal);
-			} else {
-				const payload = await createRestoreRunPayload(request);
-				const started = await agentManager.startRestore(request.executionAgentId, {
-					payload,
-					signal: abortController.signal,
-					onStarted: request.onStarted,
-					onProgress: request.onProgress,
-				});
-
-				if (started.status === "unavailable") {
-					throw started.error;
-				}
-
-				result = started.result;
-			}
-
-			return {
-				result: result.finally(() => {
-					releaseLock?.();
-				}),
-			};
-		} catch (error) {
-			releaseLock?.();
-			throw error;
-		}
 	},
 };