From d528836d6137d9eb263e72c2f41d1a2012b023ec Mon Sep 17 00:00:00 2001 From: Nicolas Meienberger Date: Sun, 31 May 2026 21:29:39 +0200 Subject: [PATCH] fix(restore): do not wait for lock before returning response --- .../__tests__/repositories.service.test.ts | 160 ++++++++++++++---- .../modules/repositories/restore-executor.ts | 97 +++++++---- 2 files changed, 187 insertions(+), 70 deletions(-) diff --git a/app/server/modules/repositories/__tests__/repositories.service.test.ts b/app/server/modules/repositories/__tests__/repositories.service.test.ts index 1381543e..b5c6ccc4 100644 --- a/app/server/modules/repositories/__tests__/repositories.service.test.ts +++ b/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 = {}) => { @@ -416,8 +418,26 @@ describe("repositoriesService.restoreSnapshot", () => { let originalEnableLocalAgent: boolean; const createPendingRestoreStart = () => ({ status: "started" as const, - result: new Promise(() => {}), + result: new Promise(() => {}), }); + const resolveWithin = async (promise: Promise, timeoutMs: number) => { + return await new Promise((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( - "local", - expect.objectContaining({ - payload: expect.objectContaining({ - snapshotId: "snapshot-restore", - target: targetPath, - repositoryConfig: expect.objectContaining({ backend: "local" }), - options: expect.objectContaining({ - organizationId, - basePath: "/var/lib/zerobyte/volumes/vol123/_data", + await waitForExpect(() => { + expect(restoreMock).toHaveBeenCalledWith( + "local", + expect.objectContaining({ + payload: expect.objectContaining({ + snapshotId: "snapshot-restore", + target: targetPath, + repositoryConfig: expect.objectContaining({ backend: "local" }), + options: expect.objectContaining({ + 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((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( - agentId, - expect.objectContaining({ - payload: expect.objectContaining({ - target: targetPath, + await waitForExpect(() => { + expect(restoreMock).toHaveBeenCalledWith( + agentId, + expect.objectContaining({ + 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( - "local", - expect.objectContaining({ - payload: expect.objectContaining({ - snapshotId: "snapshot-restore", - target: targetPath, - repositoryConfig: expect.objectContaining({ backend: "local" }), - options: expect.objectContaining({ - organizationId, - basePath: "/", + await waitForExpect(() => { + expect(restoreMock).toHaveBeenCalledWith( + "local", + expect.objectContaining({ + payload: expect.objectContaining({ + snapshotId: "snapshot-restore", + target: targetPath, + repositoryConfig: expect.objectContaining({ backend: "local" }), + options: expect.objectContaining({ + organizationId, + basePath: "/", + }), }), }), - }), - ); + ); + }); }); }); diff --git a/app/server/modules/repositories/restore-executor.ts b/app/server/modules/repositories/restore-executor.ts index ae8b7114..d631938f 100644 --- a/app/server/modules/repositories/restore-executor.ts +++ b/app/server/modules/repositories/restore-executor.ts @@ -78,45 +78,68 @@ const executeControllerRestore = async ( } }; +const executeAgentRestore = async ( + request: RestoreExecutionRequest, + signal: AbortSignal, +): Promise => { + 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 => { + 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 => { + start: (request: RestoreExecutionRequest): RestoreExecutionHandle => { const abortController = new AbortController(); - let releaseLock: (() => void) | null = null; - try { - releaseLock = await repoMutex.acquireShared( - request.repositoryId, - `restore:${request.restoreId}`, - abortController.signal, - ); - - let result: Promise; - 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; - } + return { + result: executeRestoreWithRepositoryLock(request, abortController.signal), + }; }, };