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

This commit is contained in:
Nicolas Meienberger 2026-05-31 21:29:39 +02:00
parent 44c4e80526
commit d528836d61
No known key found for this signature in database
2 changed files with 187 additions and 70 deletions

View file

@ -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(
"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<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(
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: "/",
}),
}),
}),
}),
);
);
});
});
});

View file

@ -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;
try {
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;
}
return {
result: executeRestoreWithRepositoryLock(request, abortController.signal),
};
},
};