feat: snapshot restores through rpc
This commit is contained in:
Nicolas Meienberger
parent
8fedeef4d1
commit
44c4e80526
16 changed files with 1206 additions and 171 deletions
|
|
@ -2,7 +2,7 @@ import { z } from "zod";
|
|||
import { resticBackupProgressMetricsSchema, resticBackupRunSummarySchema } from "@zerobyte/core/restic";
|
||||
|
||||
const backupEventStatusSchema = z.enum(["success", "error", "stopped", "warning"]);
|
||||
const restoreEventStatusSchema = z.enum(["success", "error"]);
|
||||
const restoreEventStatusSchema = z.enum(["success", "error", "cancelled"]);
|
||||
|
||||
const backupEventBaseSchema = z.object({
|
||||
scheduleId: z.string(),
|
||||
|
|
@ -15,6 +15,7 @@ const organizationScopedSchema = z.object({
|
|||
});
|
||||
|
||||
const restoreEventBaseSchema = z.object({
|
||||
restoreId: z.string(),
|
||||
repositoryId: z.string(),
|
||||
snapshotId: z.string(),
|
||||
});
|
||||
|
|
@ -51,6 +52,8 @@ const restoreProgressEventSchema = restoreEventBaseSchema.extend(restoreProgress
|
|||
const restoreCompletedEventSchema = restoreEventBaseSchema.extend({
|
||||
status: restoreEventStatusSchema,
|
||||
error: z.string().optional(),
|
||||
filesRestored: z.number().optional(),
|
||||
filesSkipped: z.number().optional(),
|
||||
});
|
||||
|
||||
const serverBackupStartedEventSchema = organizationScopedSchema.extend(backupStartedEventSchema.shape);
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import { afterEach, expect, test, vi } from "vitest";
|
||||
import { Effect } from "effect";
|
||||
import { fromAny, fromPartial } from "@total-typescript/shoehorn";
|
||||
import type { BackupRunPayload } from "@zerobyte/contracts/agent-protocol";
|
||||
import type { BackupRunPayload, RestoreRunPayload } from "@zerobyte/contracts/agent-protocol";
|
||||
import type { AgentManagerEvent } from "../controller/server";
|
||||
import type { ProcessWithAgentRuntime } from "../agents-manager";
|
||||
|
||||
|
|
@ -9,6 +9,8 @@ const controllerMock = vi.hoisted(() => ({
|
|||
onEvent: null as null | ((event: AgentManagerEvent) => void),
|
||||
sendBackup: vi.fn(),
|
||||
cancelBackup: vi.fn(),
|
||||
sendRestore: vi.fn(),
|
||||
cancelRestore: vi.fn(),
|
||||
stop: vi.fn(),
|
||||
}));
|
||||
|
||||
|
|
@ -22,6 +24,8 @@ vi.mock("../controller/server", async () => {
|
|||
stop: Effect.sync(controllerMock.stop),
|
||||
sendBackup: controllerMock.sendBackup,
|
||||
cancelBackup: controllerMock.cancelBackup,
|
||||
sendRestore: controllerMock.sendRestore,
|
||||
cancelRestore: controllerMock.cancelRestore,
|
||||
};
|
||||
}),
|
||||
};
|
||||
|
|
@ -37,6 +41,7 @@ const resetAgentRuntime = () => {
|
|||
localAgentRestartTimeout: null,
|
||||
activeBackupsByScheduleId: new Map(),
|
||||
activeBackupScheduleIdsByJobId: new Map(),
|
||||
activeRestoresByRestoreId: new Map(),
|
||||
};
|
||||
};
|
||||
|
||||
|
|
@ -44,12 +49,17 @@ const backupPayload = fromPartial<BackupRunPayload>({
|
|||
jobId: "job-1",
|
||||
scheduleId: "schedule-1",
|
||||
});
|
||||
const restorePayload = fromPartial<RestoreRunPayload>({
|
||||
restoreId: "restore-1",
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
delete processWithAgentRuntime.__zerobyteAgentRuntime;
|
||||
controllerMock.onEvent = null;
|
||||
controllerMock.sendBackup.mockReset();
|
||||
controllerMock.cancelBackup.mockReset();
|
||||
controllerMock.sendRestore.mockReset();
|
||||
controllerMock.cancelRestore.mockReset();
|
||||
controllerMock.stop.mockReset();
|
||||
vi.resetModules();
|
||||
vi.restoreAllMocks();
|
||||
|
|
@ -73,7 +83,11 @@ test("backup progress is delivered to the running backup callback", async () =>
|
|||
type: "backup.progress",
|
||||
agentId: "local",
|
||||
agentName: "Local Agent",
|
||||
payload: fromAny({ jobId: "job-1", scheduleId: "schedule-1", progress: { percentDone: 0.5 } }),
|
||||
payload: fromAny({
|
||||
jobId: "job-1",
|
||||
scheduleId: "schedule-1",
|
||||
progress: { percentDone: 0.5 },
|
||||
}),
|
||||
});
|
||||
controllerMock.onEvent?.({
|
||||
type: "backup.completed",
|
||||
|
|
@ -108,7 +122,12 @@ test("backup failed and cancelled events resolve the matching running backup", a
|
|||
type: "backup.failed",
|
||||
agentId: "local",
|
||||
agentName: "Local Agent",
|
||||
payload: { jobId: "job-1", scheduleId: "schedule-1", error: "failed", errorDetails: "restic failed" },
|
||||
payload: {
|
||||
jobId: "job-1",
|
||||
scheduleId: "schedule-1",
|
||||
error: "failed",
|
||||
errorDetails: "restic failed",
|
||||
},
|
||||
});
|
||||
await expect(failedPromise).resolves.toEqual({ status: "failed", error: "restic failed" });
|
||||
|
||||
|
|
@ -124,7 +143,10 @@ test("backup failed and cancelled events resolve the matching running backup", a
|
|||
agentName: "Local Agent",
|
||||
payload: { jobId: "job-2", scheduleId: "schedule-2", message: "cancelled remotely" },
|
||||
});
|
||||
await expect(cancelledPromise).resolves.toEqual({ status: "cancelled", message: "cancelled remotely" });
|
||||
await expect(cancelledPromise).resolves.toEqual({
|
||||
status: "cancelled",
|
||||
message: "cancelled remotely",
|
||||
});
|
||||
await stopAgentController();
|
||||
});
|
||||
|
||||
|
|
@ -147,7 +169,11 @@ test("agent disconnect cancels only backups owned by that agent", async () => {
|
|||
onProgress: vi.fn(),
|
||||
});
|
||||
|
||||
controllerMock.onEvent?.({ type: "agent.disconnected", agentId: "local", agentName: "Local Agent" });
|
||||
controllerMock.onEvent?.({
|
||||
type: "agent.disconnected",
|
||||
agentId: "local",
|
||||
agentName: "Local Agent",
|
||||
});
|
||||
controllerMock.onEvent?.({
|
||||
type: "backup.completed",
|
||||
agentId: "remote",
|
||||
|
|
@ -230,6 +256,117 @@ test("runBackup requests cancellation when the abort signal fires while sending"
|
|||
});
|
||||
|
||||
expect(result).toEqual({ status: "cancelled" });
|
||||
expect(controllerMock.cancelBackup).toHaveBeenCalledWith("local", { jobId: "job-1", scheduleId: "schedule-1" });
|
||||
expect(controllerMock.cancelBackup).toHaveBeenCalledWith("local", {
|
||||
jobId: "job-1",
|
||||
scheduleId: "schedule-1",
|
||||
});
|
||||
await stopAgentController();
|
||||
});
|
||||
|
||||
test("restore events are delivered to the running restore callbacks", async () => {
|
||||
resetAgentRuntime();
|
||||
controllerMock.sendRestore.mockImplementation(() => Effect.succeed(true));
|
||||
const { agentManager, startAgentController, stopAgentController } = await import("../agents-manager");
|
||||
const onStarted = vi.fn();
|
||||
const onProgress = vi.fn();
|
||||
|
||||
await startAgentController();
|
||||
const started = await agentManager.startRestore("local", {
|
||||
payload: restorePayload,
|
||||
signal: new AbortController().signal,
|
||||
onStarted,
|
||||
onProgress,
|
||||
});
|
||||
if (started.status !== "started") {
|
||||
throw new Error("Expected restore to start");
|
||||
}
|
||||
|
||||
controllerMock.onEvent?.({
|
||||
type: "restore.started",
|
||||
agentId: "local",
|
||||
agentName: "Local Agent",
|
||||
payload: {
|
||||
restoreId: "restore-1",
|
||||
organizationId: "org-1",
|
||||
repositoryId: "repo-1",
|
||||
snapshotId: "snapshot-1",
|
||||
},
|
||||
});
|
||||
controllerMock.onEvent?.({
|
||||
type: "restore.progress",
|
||||
agentId: "local",
|
||||
agentName: "Local Agent",
|
||||
payload: fromAny({ restoreId: "restore-1", progress: { percent_done: 0.5 } }),
|
||||
});
|
||||
controllerMock.onEvent?.({
|
||||
type: "restore.completed",
|
||||
agentId: "local",
|
||||
agentName: "Local Agent",
|
||||
payload: {
|
||||
restoreId: "restore-1",
|
||||
organizationId: "org-1",
|
||||
repositoryId: "repo-1",
|
||||
snapshotId: "snapshot-1",
|
||||
result: { message_type: "summary", files_restored: 2, files_skipped: 1 },
|
||||
},
|
||||
});
|
||||
|
||||
await expect(started.result).resolves.toEqual({
|
||||
status: "completed",
|
||||
result: { message_type: "summary", files_restored: 2, files_skipped: 1 },
|
||||
});
|
||||
expect(onStarted).toHaveBeenCalledOnce();
|
||||
expect(onProgress).toHaveBeenCalledWith({ percent_done: 0.5 });
|
||||
await stopAgentController();
|
||||
});
|
||||
|
||||
test("agent disconnect cancels only restores owned by that agent", async () => {
|
||||
resetAgentRuntime();
|
||||
controllerMock.sendRestore.mockImplementation(() => Effect.succeed(true));
|
||||
const { agentManager, startAgentController, stopAgentController } = await import("../agents-manager");
|
||||
|
||||
await startAgentController();
|
||||
const localStarted = await agentManager.startRestore("local", {
|
||||
payload: restorePayload,
|
||||
signal: new AbortController().signal,
|
||||
onStarted: vi.fn(),
|
||||
onProgress: vi.fn(),
|
||||
});
|
||||
const remoteStarted = await agentManager.startRestore("remote", {
|
||||
payload: fromPartial<RestoreRunPayload>({ restoreId: "restore-2" }),
|
||||
signal: new AbortController().signal,
|
||||
onStarted: vi.fn(),
|
||||
onProgress: vi.fn(),
|
||||
});
|
||||
if (localStarted.status !== "started" || remoteStarted.status !== "started") {
|
||||
throw new Error("Expected restores to start");
|
||||
}
|
||||
|
||||
controllerMock.onEvent?.({
|
||||
type: "agent.disconnected",
|
||||
agentId: "local",
|
||||
agentName: "Local Agent",
|
||||
});
|
||||
controllerMock.onEvent?.({
|
||||
type: "restore.completed",
|
||||
agentId: "remote",
|
||||
agentName: "Remote Agent",
|
||||
payload: {
|
||||
restoreId: "restore-2",
|
||||
organizationId: "org-1",
|
||||
repositoryId: "repo-1",
|
||||
snapshotId: "snapshot-1",
|
||||
result: { message_type: "summary", files_restored: 1, files_skipped: 0 },
|
||||
},
|
||||
});
|
||||
|
||||
await expect(localStarted.result).resolves.toEqual({
|
||||
status: "cancelled",
|
||||
message: "The connection to the restore agent was lost. Restart the restore to ensure it completes.",
|
||||
});
|
||||
await expect(remoteStarted.result).resolves.toEqual({
|
||||
status: "completed",
|
||||
result: { message_type: "summary", files_restored: 1, files_skipped: 0 },
|
||||
});
|
||||
await stopAgentController();
|
||||
});
|
||||
|
|
|
|||
|
|
@ -215,6 +215,37 @@ test("websocket lifecycle updates agent connection status", async () => {
|
|||
expect(stop).toHaveBeenCalledWith(true);
|
||||
});
|
||||
|
||||
test("websocket restore events are forwarded with agent metadata", async () => {
|
||||
const serve = vi
|
||||
.spyOn(Bun, "serve")
|
||||
.mockReturnValue(fromPartial({ port: 3001, stop: vi.fn(() => Promise.resolve()) }));
|
||||
const { runtime, onEvent } = await startRuntime(vi.fn());
|
||||
const websocket = serve.mock.calls[0]?.[0].websocket;
|
||||
const socket = createSocket("connection-1");
|
||||
|
||||
await websocket?.open?.(fromPartial(socket));
|
||||
await websocket?.message?.(
|
||||
fromPartial(socket),
|
||||
createAgentMessage("restore.completed", {
|
||||
restoreId: "restore-1",
|
||||
organizationId: "org-1",
|
||||
repositoryId: "repo-1",
|
||||
snapshotId: "snapshot-1",
|
||||
result: { message_type: "summary", files_restored: 2, files_skipped: 0 },
|
||||
}),
|
||||
);
|
||||
await Effect.runPromise(runtime.stop);
|
||||
|
||||
expect(onEvent).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
type: "restore.completed",
|
||||
agentId: LOCAL_AGENT_ID,
|
||||
agentName: LOCAL_AGENT_NAME,
|
||||
payload: expect.objectContaining({ restoreId: "restore-1" }),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
test("websocket open failure closes the upgraded socket", async () => {
|
||||
agentsServiceMocks.markAgentConnecting.mockRejectedValueOnce(new Error("db unavailable"));
|
||||
const serve = vi
|
||||
|
|
|
|||
|
|
@ -1,5 +1,10 @@
|
|||
import { logger } from "@zerobyte/core/node";
|
||||
import type { BackupRunPayload, VolumeCommand, VolumeCommandResult } from "@zerobyte/contracts/agent-protocol";
|
||||
import type {
|
||||
BackupRunPayload,
|
||||
RestoreRunPayload,
|
||||
VolumeCommand,
|
||||
VolumeCommandResult,
|
||||
} from "@zerobyte/contracts/agent-protocol";
|
||||
import { Effect } from "effect";
|
||||
import { config } from "../../core/config";
|
||||
import { createAgentManagerRuntime, type AgentManagerEvent } from "./controller/server";
|
||||
|
|
@ -10,9 +15,16 @@ import {
|
|||
type AgentRuntimeState,
|
||||
type BackupExecutionProgress,
|
||||
type BackupExecutionResult,
|
||||
type RestoreExecutionProgress,
|
||||
type RestoreExecutionResult,
|
||||
} from "./helpers/runtime-state";
|
||||
import { getDevAgentRuntimeState } from "./helpers/runtime-state.dev";
|
||||
export type { BackupExecutionProgress, BackupExecutionResult } from "./helpers/runtime-state";
|
||||
export type {
|
||||
BackupExecutionProgress,
|
||||
BackupExecutionResult,
|
||||
RestoreExecutionProgress,
|
||||
RestoreExecutionResult,
|
||||
} from "./helpers/runtime-state";
|
||||
export type { ProcessWithAgentRuntime } from "./helpers/runtime-state.dev";
|
||||
|
||||
type ProcessWithProductionAgentRuntime = NodeJS.Process & {
|
||||
|
|
@ -26,6 +38,17 @@ type AgentRunBackupRequest = {
|
|||
onProgress: (progress: BackupExecutionProgress) => void;
|
||||
};
|
||||
|
||||
type AgentStartRestoreRequest = {
|
||||
payload: RestoreRunPayload;
|
||||
signal: AbortSignal;
|
||||
onStarted: () => void;
|
||||
onProgress: (progress: RestoreExecutionProgress) => void;
|
||||
};
|
||||
|
||||
type AgentRestoreStartResult =
|
||||
| { status: "started"; result: Promise<RestoreExecutionResult> }
|
||||
| { status: "unavailable"; error: Error };
|
||||
|
||||
const getProductionAgentRuntimeState = () => {
|
||||
// Nitro production builds can bundle startup plugins and API handlers into separate chunks.
|
||||
// Keep the live controller on process so both chunks see the same agent sessions.
|
||||
|
|
@ -41,6 +64,7 @@ const getAgentRuntimeState = () => (config.__prod__ ? getProductionAgentRuntimeS
|
|||
const getAgentManagerRuntime = () => getAgentRuntimeState().agentManager;
|
||||
const getActiveBackupsByScheduleId = () => getAgentRuntimeState().activeBackupsByScheduleId;
|
||||
const getActiveBackupScheduleIdsByJobId = () => getAgentRuntimeState().activeBackupScheduleIdsByJobId;
|
||||
const getActiveRestoresByRestoreId = () => getAgentRuntimeState().activeRestoresByRestoreId;
|
||||
|
||||
const clearActiveBackupRun = (scheduleId: number) => {
|
||||
const activeBackupsByScheduleId = getActiveBackupsByScheduleId();
|
||||
|
|
@ -65,6 +89,27 @@ const resolveActiveBackupRun = (scheduleId: number, result: BackupExecutionResul
|
|||
return true;
|
||||
};
|
||||
|
||||
const clearActiveRestoreRun = (restoreId: string) => {
|
||||
const activeRestoresByRestoreId = getActiveRestoresByRestoreId();
|
||||
const activeRestoreRun = activeRestoresByRestoreId.get(restoreId);
|
||||
if (!activeRestoreRun) {
|
||||
return null;
|
||||
}
|
||||
|
||||
activeRestoresByRestoreId.delete(restoreId);
|
||||
return activeRestoreRun;
|
||||
};
|
||||
|
||||
const resolveActiveRestoreRun = (restoreId: string, result: RestoreExecutionResult) => {
|
||||
const activeRestoreRun = clearActiveRestoreRun(restoreId);
|
||||
if (!activeRestoreRun) {
|
||||
return false;
|
||||
}
|
||||
|
||||
activeRestoreRun.resolve(result);
|
||||
return true;
|
||||
};
|
||||
|
||||
const cancelActiveBackupRunsForAgent = (agentId: string, message: string) => {
|
||||
const activeBackupsByScheduleId = getActiveBackupsByScheduleId();
|
||||
const matchingScheduleIds = [...activeBackupsByScheduleId.values()]
|
||||
|
|
@ -76,6 +121,17 @@ const cancelActiveBackupRunsForAgent = (agentId: string, message: string) => {
|
|||
}
|
||||
};
|
||||
|
||||
const cancelActiveRestoreRunsForAgent = (agentId: string, message: string) => {
|
||||
const activeRestoresByRestoreId = getActiveRestoresByRestoreId();
|
||||
const matchingRestoreIds = [...activeRestoresByRestoreId.values()]
|
||||
.filter((activeRestoreRun) => activeRestoreRun.agentId === agentId)
|
||||
.map((activeRestoreRun) => activeRestoreRun.restoreId);
|
||||
|
||||
for (const restoreId of matchingRestoreIds) {
|
||||
resolveActiveRestoreRun(restoreId, { status: "cancelled", message });
|
||||
}
|
||||
};
|
||||
|
||||
const getActiveBackupRun = (jobId: string, scheduleId: string, eventName: string, agentId: string) => {
|
||||
const trackedScheduleId = getActiveBackupScheduleIdsByJobId().get(jobId);
|
||||
if (trackedScheduleId === undefined) {
|
||||
|
|
@ -99,6 +155,21 @@ const getActiveBackupRun = (jobId: string, scheduleId: string, eventName: string
|
|||
return activeBackupRun;
|
||||
};
|
||||
|
||||
const getActiveRestoreRun = (restoreId: string, eventName: string, agentId: string) => {
|
||||
const activeRestoreRun = getActiveRestoresByRestoreId().get(restoreId);
|
||||
if (!activeRestoreRun) {
|
||||
logger.warn(`Received ${eventName} for unknown restore ${restoreId} from agent ${agentId}`);
|
||||
return null;
|
||||
}
|
||||
|
||||
if (activeRestoreRun.agentId !== agentId) {
|
||||
logger.warn(`Ignoring ${eventName} for restore ${restoreId} from unexpected agent ${agentId}`);
|
||||
return null;
|
||||
}
|
||||
|
||||
return activeRestoreRun;
|
||||
};
|
||||
|
||||
const requestBackupCancellation = async (agentId: string, scheduleId: number) => {
|
||||
const activeBackupRun = getActiveBackupsByScheduleId().get(scheduleId);
|
||||
if (!activeBackupRun) {
|
||||
|
|
@ -132,6 +203,32 @@ const requestBackupCancellation = async (agentId: string, scheduleId: number) =>
|
|||
return true;
|
||||
};
|
||||
|
||||
const requestRestoreCancellation = async (agentId: string, restoreId: string) => {
|
||||
const activeRestoreRun = getActiveRestoresByRestoreId().get(restoreId);
|
||||
if (!activeRestoreRun) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (activeRestoreRun.cancellationRequested) {
|
||||
return true;
|
||||
}
|
||||
|
||||
activeRestoreRun.cancellationRequested = true;
|
||||
|
||||
const runtime = getAgentManagerRuntime();
|
||||
if (!runtime) {
|
||||
resolveActiveRestoreRun(restoreId, { status: "cancelled" });
|
||||
return true;
|
||||
}
|
||||
|
||||
if (await Effect.runPromise(runtime.cancelRestore(agentId, { restoreId }))) {
|
||||
return true;
|
||||
}
|
||||
|
||||
resolveActiveRestoreRun(restoreId, { status: "cancelled" });
|
||||
return true;
|
||||
};
|
||||
|
||||
const handleAgentManagerEvent = (event: AgentManagerEvent) => {
|
||||
switch (event.type) {
|
||||
case "agent.disconnected": {
|
||||
|
|
@ -139,6 +236,10 @@ const handleAgentManagerEvent = (event: AgentManagerEvent) => {
|
|||
event.agentId,
|
||||
"The connection to the backup agent was lost. Restart the backup to ensure it completes.",
|
||||
);
|
||||
cancelActiveRestoreRunsForAgent(
|
||||
event.agentId,
|
||||
"The connection to the restore agent was lost. Restart the restore to ensure it completes.",
|
||||
);
|
||||
break;
|
||||
}
|
||||
case "backup.started": {
|
||||
|
|
@ -212,6 +313,60 @@ const handleAgentManagerEvent = (event: AgentManagerEvent) => {
|
|||
});
|
||||
break;
|
||||
}
|
||||
case "restore.started": {
|
||||
const activeRestoreRun = getActiveRestoreRun(event.payload.restoreId, event.type, event.agentId);
|
||||
if (!activeRestoreRun) {
|
||||
break;
|
||||
}
|
||||
|
||||
activeRestoreRun.onStarted();
|
||||
break;
|
||||
}
|
||||
case "restore.progress": {
|
||||
const activeRestoreRun = getActiveRestoreRun(event.payload.restoreId, event.type, event.agentId);
|
||||
if (!activeRestoreRun) {
|
||||
break;
|
||||
}
|
||||
|
||||
activeRestoreRun.onProgress(event.payload.progress);
|
||||
break;
|
||||
}
|
||||
case "restore.completed": {
|
||||
const activeRestoreRun = getActiveRestoreRun(event.payload.restoreId, event.type, event.agentId);
|
||||
if (!activeRestoreRun) {
|
||||
break;
|
||||
}
|
||||
|
||||
resolveActiveRestoreRun(activeRestoreRun.restoreId, {
|
||||
status: "completed",
|
||||
result: event.payload.result,
|
||||
});
|
||||
break;
|
||||
}
|
||||
case "restore.failed": {
|
||||
const activeRestoreRun = getActiveRestoreRun(event.payload.restoreId, event.type, event.agentId);
|
||||
if (!activeRestoreRun) {
|
||||
break;
|
||||
}
|
||||
|
||||
resolveActiveRestoreRun(activeRestoreRun.restoreId, {
|
||||
status: "failed",
|
||||
error: event.payload.errorDetails ?? event.payload.error,
|
||||
});
|
||||
break;
|
||||
}
|
||||
case "restore.cancelled": {
|
||||
const activeRestoreRun = getActiveRestoreRun(event.payload.restoreId, event.type, event.agentId);
|
||||
if (!activeRestoreRun) {
|
||||
break;
|
||||
}
|
||||
|
||||
resolveActiveRestoreRun(activeRestoreRun.restoreId, {
|
||||
status: "cancelled",
|
||||
message: activeRestoreRun.cancellationRequested ? undefined : event.payload.message,
|
||||
});
|
||||
break;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
|
|
@ -307,6 +462,52 @@ export const agentManager = {
|
|||
|
||||
return response.command;
|
||||
},
|
||||
startRestore: async (agentId: string, request: AgentStartRestoreRequest): Promise<AgentRestoreStartResult> => {
|
||||
const runtime = getAgentManagerRuntime();
|
||||
if (!runtime) {
|
||||
return {
|
||||
status: "unavailable",
|
||||
error: new Error(`Restore agent ${agentId} is not connected`),
|
||||
};
|
||||
}
|
||||
|
||||
if (request.signal.aborted) {
|
||||
throw request.signal.reason || new Error("Operation aborted");
|
||||
}
|
||||
|
||||
const completion = new Promise<RestoreExecutionResult>((resolve) => {
|
||||
getActiveRestoresByRestoreId().set(request.payload.restoreId, {
|
||||
agentId,
|
||||
restoreId: request.payload.restoreId,
|
||||
onStarted: request.onStarted,
|
||||
onProgress: request.onProgress,
|
||||
resolve,
|
||||
cancellationRequested: false,
|
||||
});
|
||||
});
|
||||
|
||||
try {
|
||||
if (!(await Effect.runPromise(runtime.sendRestore(agentId, request.payload)))) {
|
||||
clearActiveRestoreRun(request.payload.restoreId);
|
||||
return {
|
||||
status: "unavailable",
|
||||
error: new Error(`Failed to send restore command to agent ${agentId}`),
|
||||
};
|
||||
}
|
||||
|
||||
if (request.signal.aborted) {
|
||||
await requestRestoreCancellation(agentId, request.payload.restoreId);
|
||||
}
|
||||
|
||||
return { status: "started", result: completion };
|
||||
} catch (error) {
|
||||
clearActiveRestoreRun(request.payload.restoreId);
|
||||
throw error;
|
||||
}
|
||||
},
|
||||
cancelRestore: async (agentId: string, restoreId: string) => {
|
||||
return requestRestoreCancellation(agentId, restoreId);
|
||||
},
|
||||
};
|
||||
|
||||
export const startLocalAgent = async () => {
|
||||
|
|
|
|||
|
|
@ -5,6 +5,8 @@ import type {
|
|||
AgentMessage,
|
||||
BackupCancelPayload,
|
||||
BackupRunPayload,
|
||||
RestoreCancelPayload,
|
||||
RestoreRunPayload,
|
||||
VolumeCommand,
|
||||
VolumeCommandResponsePayload,
|
||||
} from "@zerobyte/contracts/agent-protocol";
|
||||
|
|
@ -22,16 +24,9 @@ type AgentEventContext = {
|
|||
agentName: string;
|
||||
};
|
||||
|
||||
type AgentBackupMessage = Extract<
|
||||
AgentMessage,
|
||||
{
|
||||
type: "backup.started" | "backup.progress" | "backup.completed" | "backup.failed" | "backup.cancelled";
|
||||
}
|
||||
>;
|
||||
|
||||
export type AgentManagerEvent =
|
||||
| (AgentEventContext & { type: "agent.disconnected" })
|
||||
| (AgentEventContext & AgentBackupMessage);
|
||||
| (AgentEventContext & AgentMessage);
|
||||
|
||||
type ControllerAgentSessionHandle = {
|
||||
agentId: string;
|
||||
|
|
@ -86,7 +81,7 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) =>
|
|||
return !!session && Effect.runSync(session.isReady());
|
||||
};
|
||||
|
||||
const handleSessionEvent = (params: { agentId: string; agentName: string; sessionId: string }) => {
|
||||
const handleSessionEvent = (params: { agentId: string; agentName: string }) => {
|
||||
const { agentId, agentName } = params;
|
||||
|
||||
return (event: ControllerAgentSessionEvent) => {
|
||||
|
|
@ -109,15 +104,6 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) =>
|
|||
case "agent.disconnected": {
|
||||
return Effect.sync(() => onEvent({ type: "agent.disconnected", agentId, agentName }));
|
||||
}
|
||||
case "restore.cancelled":
|
||||
case "restore.completed":
|
||||
case "restore.failed":
|
||||
case "restore.progress":
|
||||
case "restore.started": {
|
||||
// TODO: once we implement the app side
|
||||
return Effect.void;
|
||||
}
|
||||
|
||||
default: {
|
||||
return Effect.sync(() => onEvent({ ...event, agentId, agentName }));
|
||||
}
|
||||
|
|
@ -135,7 +121,6 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) =>
|
|||
handleSessionEvent({
|
||||
agentId: ws.data.agentId,
|
||||
agentName: ws.data.agentName,
|
||||
sessionId: ws.data.id,
|
||||
}),
|
||||
),
|
||||
scope,
|
||||
|
|
@ -368,6 +353,47 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) =>
|
|||
logger.info(`Sent backup cancel for command ${payload.jobId} to agent ${agentId}`);
|
||||
return true;
|
||||
}),
|
||||
sendRestore: (agentId: string, payload: RestoreRunPayload) =>
|
||||
Effect.gen(function* () {
|
||||
const session = getSession(agentId);
|
||||
|
||||
if (!session) {
|
||||
logger.warn(`Cannot send restore command. Agent ${agentId} is not connected.`);
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!(yield* session.isReady())) {
|
||||
logger.warn(`Cannot send restore command. Agent ${agentId} is not ready.`);
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!(yield* session.sendRestore(payload))) {
|
||||
logger.warn(`Cannot send restore command. Agent ${agentId} is no longer accepting commands.`);
|
||||
return false;
|
||||
}
|
||||
|
||||
logger.info(
|
||||
`Sent restore command ${payload.restoreId} to agent ${agentId} for snapshot ${payload.snapshotId}`,
|
||||
);
|
||||
return true;
|
||||
}),
|
||||
cancelRestore: (agentId: string, payload: RestoreCancelPayload) =>
|
||||
Effect.gen(function* () {
|
||||
const session = getSession(agentId);
|
||||
|
||||
if (!session) {
|
||||
logger.warn(`Cannot cancel restore command. Agent ${agentId} is not connected.`);
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!(yield* session.sendRestoreCancel(payload))) {
|
||||
logger.warn(`Cannot cancel restore command. Agent ${agentId} is no longer accepting commands.`);
|
||||
return false;
|
||||
}
|
||||
|
||||
logger.info(`Sent restore cancel for command ${payload.restoreId} to agent ${agentId}`);
|
||||
return true;
|
||||
}),
|
||||
runVolumeCommand: (
|
||||
agentId: string,
|
||||
command: VolumeCommand,
|
||||
|
|
|
|||
|
|
@ -7,6 +7,8 @@ import {
|
|||
type BackupCancelPayload,
|
||||
type BackupRunPayload,
|
||||
type ControllerWireMessage,
|
||||
type RestoreCancelPayload,
|
||||
type RestoreRunPayload,
|
||||
type VolumeCommand,
|
||||
type VolumeCommandResponsePayload,
|
||||
} from "@zerobyte/contracts/agent-protocol";
|
||||
|
|
@ -29,22 +31,19 @@ type SessionState = {
|
|||
lastPongAt: number | null;
|
||||
};
|
||||
|
||||
type PendingCommand = {
|
||||
deferred: Deferred.Deferred<VolumeCommandResponsePayload, Error>;
|
||||
description: string;
|
||||
};
|
||||
type PendingCommand = { deferred: Deferred.Deferred<VolumeCommandResponsePayload, Error>; description: string };
|
||||
|
||||
export type ControllerAgentSessionEvent =
|
||||
| Exclude<AgentMessage, { type: "volume.commandResult" }>
|
||||
| {
|
||||
type: "agent.disconnected";
|
||||
};
|
||||
| { type: "agent.disconnected" };
|
||||
|
||||
export type ControllerAgentSession = {
|
||||
readonly connectionId: string;
|
||||
handleMessage: (data: string) => Effect.Effect<void>;
|
||||
sendBackup: (payload: BackupRunPayload) => Effect.Effect<boolean>;
|
||||
sendBackupCancel: (payload: BackupCancelPayload) => Effect.Effect<boolean>;
|
||||
sendRestore: (payload: RestoreRunPayload) => Effect.Effect<boolean>;
|
||||
sendRestoreCancel: (payload: RestoreCancelPayload) => Effect.Effect<boolean>;
|
||||
runVolumeCommand: (command: VolumeCommand) => Effect.Effect<VolumeCommandResponsePayload, Error>;
|
||||
isReady: () => Effect.Effect<boolean>;
|
||||
run: Effect.Effect<void, never, Scope.Scope>;
|
||||
|
|
@ -68,7 +67,9 @@ export const createControllerAgentSession = (
|
|||
Queue.offer(outboundQueue, message).pipe(
|
||||
Effect.catchAllCause((cause) =>
|
||||
Effect.sync(() => {
|
||||
logger.error(`Failed to queue outbound message for agent ${socket.data.agentId}: ${toMessage(cause)}`);
|
||||
logger.error(
|
||||
`Failed to queue outbound message for agent ${socket.data.agentId}: ${toMessage(cause)}`,
|
||||
);
|
||||
return false;
|
||||
}),
|
||||
),
|
||||
|
|
@ -189,7 +190,11 @@ export const createControllerAgentSession = (
|
|||
}
|
||||
case "heartbeat.pong": {
|
||||
const seenAt = Date.now();
|
||||
yield* updateState((current) => ({ ...current, lastSeenAt: seenAt, lastPongAt: message.payload.sentAt }));
|
||||
yield* updateState((current) => ({
|
||||
...current,
|
||||
lastSeenAt: seenAt,
|
||||
lastPongAt: message.payload.sentAt,
|
||||
}));
|
||||
yield* onEvent(message);
|
||||
break;
|
||||
}
|
||||
|
|
@ -216,7 +221,9 @@ export const createControllerAgentSession = (
|
|||
}
|
||||
|
||||
if (!parsed.success) {
|
||||
yield* logger.effect.warn(`Invalid agent message from ${socket.data.agentId}: ${parsed.error.message}`);
|
||||
yield* logger.effect.warn(
|
||||
`Invalid agent message from ${socket.data.agentId}: ${parsed.error.message}`,
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
|
|
@ -225,6 +232,8 @@ export const createControllerAgentSession = (
|
|||
},
|
||||
sendBackup: (payload) => offerOutbound(createControllerMessage("backup.run", payload)),
|
||||
sendBackupCancel: (payload) => offerOutbound(createControllerMessage("backup.cancel", payload)),
|
||||
sendRestore: (payload) => offerOutbound(createControllerMessage("restore.run", payload)),
|
||||
sendRestoreCancel: (payload) => offerOutbound(createControllerMessage("restore.cancel", payload)),
|
||||
runVolumeCommand: (command) =>
|
||||
Effect.gen(function* () {
|
||||
const commandId = Bun.randomUUIDv7();
|
||||
|
|
@ -232,7 +241,9 @@ export const createControllerAgentSession = (
|
|||
const deferred = yield* Deferred.make<VolumeCommandResponsePayload, Error>();
|
||||
yield* setPendingCommand(commandId, { deferred, description });
|
||||
|
||||
const queued = yield* offerOutbound(createControllerMessage("volume.command", { commandId, command }));
|
||||
const queued = yield* offerOutbound(
|
||||
createControllerMessage("volume.command", { commandId, command }),
|
||||
);
|
||||
if (!queued) {
|
||||
yield* removePendingCommand(commandId);
|
||||
return yield* Effect.fail(new Error(`Failed to queue volume command ${command.name}`));
|
||||
|
|
|
|||
|
|
@ -1,20 +1,26 @@
|
|||
import { createAgentRuntimeState, type AgentRuntimeState } from "./runtime-state";
|
||||
|
||||
type LegacyAgentRuntimeState = Omit<AgentRuntimeState, "activeBackupsByScheduleId" | "activeBackupScheduleIdsByJobId"> &
|
||||
Partial<Pick<AgentRuntimeState, "activeBackupsByScheduleId" | "activeBackupScheduleIdsByJobId">>;
|
||||
type RuntimeMapKey = "activeBackupsByScheduleId" | "activeBackupScheduleIdsByJobId" | "activeRestoresByRestoreId";
|
||||
|
||||
type LegacyAgentRuntimeState = Omit<AgentRuntimeState, RuntimeMapKey> & Partial<Pick<AgentRuntimeState, RuntimeMapKey>>;
|
||||
|
||||
export type ProcessWithAgentRuntime = NodeJS.Process & {
|
||||
__zerobyteAgentRuntime?: LegacyAgentRuntimeState;
|
||||
};
|
||||
|
||||
const hasActiveBackupMaps = (runtime: LegacyAgentRuntimeState): runtime is AgentRuntimeState => {
|
||||
return runtime.activeBackupsByScheduleId instanceof Map && runtime.activeBackupScheduleIdsByJobId instanceof Map;
|
||||
const hasActiveRuntimeMaps = (runtime: LegacyAgentRuntimeState): runtime is AgentRuntimeState => {
|
||||
return (
|
||||
runtime.activeBackupsByScheduleId instanceof Map &&
|
||||
runtime.activeBackupScheduleIdsByJobId instanceof Map &&
|
||||
runtime.activeRestoresByRestoreId instanceof Map
|
||||
);
|
||||
};
|
||||
|
||||
const hydrateAgentRuntimeState = (runtime: LegacyAgentRuntimeState): AgentRuntimeState => ({
|
||||
...runtime,
|
||||
activeBackupsByScheduleId: runtime.activeBackupsByScheduleId ?? new Map(),
|
||||
activeBackupScheduleIdsByJobId: runtime.activeBackupScheduleIdsByJobId ?? new Map(),
|
||||
activeRestoresByRestoreId: runtime.activeRestoresByRestoreId ?? new Map(),
|
||||
});
|
||||
|
||||
export const getDevAgentRuntimeState = (): AgentRuntimeState => {
|
||||
|
|
@ -27,7 +33,7 @@ export const getDevAgentRuntimeState = (): AgentRuntimeState => {
|
|||
return runtime;
|
||||
}
|
||||
|
||||
if (hasActiveBackupMaps(existingRuntime)) {
|
||||
if (hasActiveRuntimeMaps(existingRuntime)) {
|
||||
return existingRuntime;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,12 +1,27 @@
|
|||
import type { ChildProcess } from "node:child_process";
|
||||
import type { ResticBackupOutputDto } from "@zerobyte/core/restic";
|
||||
import type { BackupProgressPayload } from "@zerobyte/contracts/agent-protocol";
|
||||
import type {
|
||||
BackupProgressPayload,
|
||||
RestoreCompletedPayload,
|
||||
RestoreProgressPayload,
|
||||
} from "@zerobyte/contracts/agent-protocol";
|
||||
import type { AgentManagerRuntime } from "../controller/server";
|
||||
|
||||
export type BackupExecutionProgress = BackupProgressPayload["progress"];
|
||||
export type BackupExecutionResult =
|
||||
| { status: "unavailable"; error: Error }
|
||||
| { status: "completed"; exitCode: number; result: ResticBackupOutputDto | null; warningDetails: string | null }
|
||||
| {
|
||||
status: "completed";
|
||||
exitCode: number;
|
||||
result: ResticBackupOutputDto | null;
|
||||
warningDetails: string | null;
|
||||
}
|
||||
| { status: "failed"; error: string }
|
||||
| { status: "cancelled"; message?: string };
|
||||
export type RestoreExecutionProgress = RestoreProgressPayload["progress"];
|
||||
export type RestoreExecutionResult =
|
||||
| { status: "unavailable"; error: Error }
|
||||
| { status: "completed"; result: RestoreCompletedPayload["result"] }
|
||||
| { status: "failed"; error: string }
|
||||
| { status: "cancelled"; message?: string };
|
||||
|
||||
|
|
@ -20,6 +35,15 @@ type ActiveBackupRun = {
|
|||
cancellationRequested: boolean;
|
||||
};
|
||||
|
||||
type ActiveRestoreRun = {
|
||||
agentId: string;
|
||||
restoreId: string;
|
||||
onStarted: () => void;
|
||||
onProgress: (progress: RestoreExecutionProgress) => void;
|
||||
resolve: (result: RestoreExecutionResult) => void;
|
||||
cancellationRequested: boolean;
|
||||
};
|
||||
|
||||
export type AgentRuntimeState = {
|
||||
agentManager: AgentManagerRuntime | null;
|
||||
localAgent: ChildProcess | null;
|
||||
|
|
@ -27,6 +51,7 @@ export type AgentRuntimeState = {
|
|||
localAgentRestartTimeout: ReturnType<typeof setTimeout> | null;
|
||||
activeBackupsByScheduleId: Map<number, ActiveBackupRun>;
|
||||
activeBackupScheduleIdsByJobId: Map<string, number>;
|
||||
activeRestoresByRestoreId: Map<string, ActiveRestoreRun>;
|
||||
};
|
||||
|
||||
export const createAgentRuntimeState = (): AgentRuntimeState => ({
|
||||
|
|
@ -36,4 +61,5 @@ export const createAgentRuntimeState = (): AgentRuntimeState => ({
|
|||
localAgentRestartTimeout: null,
|
||||
activeBackupsByScheduleId: new Map(),
|
||||
activeBackupScheduleIdsByJobId: new Map(),
|
||||
activeRestoresByRestoreId: new Map(),
|
||||
});
|
||||
|
|
|
|||
|
|
@ -4,22 +4,24 @@ import * as os from "node:os";
|
|||
import nodePath from "node:path";
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { Readable } from "node:stream";
|
||||
import { Effect } from "effect";
|
||||
import { afterEach, beforeAll, beforeEach, describe, expect, test, vi } from "vitest";
|
||||
import type { RepositoryConfig } from "@zerobyte/core/restic";
|
||||
import { REPOSITORY_BASE } from "~/server/core/constants";
|
||||
import { config } from "~/server/core/config";
|
||||
import { withContext } from "~/server/core/request-context";
|
||||
import { db } from "~/server/db/db";
|
||||
import { repositoriesTable } from "~/server/db/schema";
|
||||
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 { 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 { repositoriesService } from "../repositories.service";
|
||||
import { Effect } from "effect";
|
||||
|
||||
const createTestRepository = async (organizationId: string) => {
|
||||
const createTestRepository = async (organizationId: string, overrides: Partial<RepositoryInsert> = {}) => {
|
||||
const id = randomUUID();
|
||||
const shortId = generateShortId();
|
||||
const [repository] = await db
|
||||
|
|
@ -33,6 +35,7 @@ const createTestRepository = async (organizationId: string) => {
|
|||
compressionMode: "auto",
|
||||
status: "healthy",
|
||||
organizationId,
|
||||
...overrides,
|
||||
})
|
||||
.returning();
|
||||
return repository;
|
||||
|
|
@ -115,7 +118,11 @@ describe("repositoriesService.createRepository", () => {
|
|||
test("keeps an explicit local repository path unchanged when importing existing repository", async () => {
|
||||
// arrange
|
||||
const explicitPath = `${REPOSITORY_BASE}/custom-${randomUUID()}`;
|
||||
const config: RepositoryConfig = { backend: "local", path: explicitPath, isExistingRepository: true };
|
||||
const config: RepositoryConfig = {
|
||||
backend: "local",
|
||||
path: explicitPath,
|
||||
isExistingRepository: true,
|
||||
};
|
||||
|
||||
vi.spyOn(restic, "snapshots").mockImplementation(() => Effect.succeed([]));
|
||||
|
||||
|
|
@ -184,7 +191,9 @@ describe("repositoriesService repository stats", () => {
|
|||
expect(refreshed).toEqual(expectedStats);
|
||||
expect(statsSpy).toHaveBeenCalledTimes(1);
|
||||
|
||||
const persistedRepository = await db.query.repositoriesTable.findFirst({ where: { id: repository.id } });
|
||||
const persistedRepository = await db.query.repositoriesTable.findFirst({
|
||||
where: { id: repository.id },
|
||||
});
|
||||
expect(persistedRepository?.stats).toEqual(expectedStats);
|
||||
expect(typeof persistedRepository?.statsUpdatedAt).toBe("number");
|
||||
|
||||
|
|
@ -404,7 +413,19 @@ describe("repositoriesService.dumpSnapshot", () => {
|
|||
});
|
||||
|
||||
describe("repositoriesService.restoreSnapshot", () => {
|
||||
let originalEnableLocalAgent: boolean;
|
||||
const createPendingRestoreStart = () => ({
|
||||
status: "started" as const,
|
||||
result: new Promise<never>(() => {}),
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
originalEnableLocalAgent = config.flags.enableLocalAgent;
|
||||
config.flags.enableLocalAgent = true;
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
config.flags.enableLocalAgent = originalEnableLocalAgent;
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
|
|
@ -424,19 +445,8 @@ describe("repositoriesService.restoreSnapshot", () => {
|
|||
]),
|
||||
);
|
||||
|
||||
const restoreMock = vi.fn(() =>
|
||||
Effect.succeed({
|
||||
message_type: "summary" as const,
|
||||
seconds_elapsed: 1,
|
||||
percent_done: 100,
|
||||
files_skipped: 0,
|
||||
total_files: 1,
|
||||
files_restored: 1,
|
||||
total_bytes: 1,
|
||||
bytes_restored: 1,
|
||||
}),
|
||||
);
|
||||
vi.spyOn(restic, "restore").mockImplementation(restoreMock);
|
||||
const restoreMock = vi.fn(() => Promise.resolve(createPendingRestoreStart()));
|
||||
vi.spyOn(agentManager, "startRestore").mockImplementation(restoreMock);
|
||||
|
||||
return {
|
||||
organizationId,
|
||||
|
|
@ -446,13 +456,15 @@ describe("repositoriesService.restoreSnapshot", () => {
|
|||
};
|
||||
};
|
||||
|
||||
test("rejects restore targets inside protected roots", async () => {
|
||||
test("rejects protected targets even when the local agent is enabled", async () => {
|
||||
const { organizationId, userId, repositoryShortId, restoreMock } = await setupRestoreSnapshotScenario();
|
||||
const targetPath = nodePath.join(os.tmpdir(), "zerobyte-restore-target");
|
||||
|
||||
await expect(
|
||||
withContext({ organizationId, userId }, () =>
|
||||
repositoriesService.restoreSnapshot(repositoryShortId, "snapshot-restore", { targetPath }),
|
||||
repositoriesService.restoreSnapshot(repositoryShortId, "snapshot-restore", {
|
||||
targetPath,
|
||||
}),
|
||||
),
|
||||
).rejects.toThrow("Restore target path is not allowed");
|
||||
|
||||
|
|
@ -465,25 +477,197 @@ describe("repositoriesService.restoreSnapshot", () => {
|
|||
|
||||
try {
|
||||
await withContext({ organizationId, userId }, () =>
|
||||
repositoriesService.restoreSnapshot(repositoryShortId, "snapshot-restore", { targetPath }),
|
||||
repositoriesService.restoreSnapshot(repositoryShortId, "snapshot-restore", {
|
||||
targetPath,
|
||||
}),
|
||||
);
|
||||
} finally {
|
||||
await fs.rm(targetPath, { recursive: true, force: true });
|
||||
}
|
||||
|
||||
expect(restoreMock).toHaveBeenCalledWith(
|
||||
"local",
|
||||
expect.objectContaining({
|
||||
backend: "local",
|
||||
}),
|
||||
"snapshot-restore",
|
||||
targetPath,
|
||||
expect.objectContaining({
|
||||
organizationId,
|
||||
basePath: "/var/lib/zerobyte/volumes/vol123/_data",
|
||||
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 () => {
|
||||
const { organizationId, userId, repositoryShortId } = await setupRestoreSnapshotScenario();
|
||||
const targetPath = await fs.mkdtemp(nodePath.join(process.cwd(), "restore-target-"));
|
||||
|
||||
try {
|
||||
await withContext({ organizationId, userId }, () =>
|
||||
repositoriesService.restoreSnapshot(repositoryShortId, "snapshot-restore", {
|
||||
targetPath,
|
||||
}),
|
||||
);
|
||||
|
||||
await expect(
|
||||
withContext({ organizationId, userId }, () =>
|
||||
repositoriesService.restoreSnapshot(repositoryShortId, "snapshot-restore", {
|
||||
targetPath,
|
||||
}),
|
||||
),
|
||||
).rejects.toThrow("A restore is already running for this snapshot");
|
||||
} finally {
|
||||
await fs.rm(targetPath, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
test("routes restore to the requested target agent", async () => {
|
||||
const organizationId = session.organizationId;
|
||||
const agentId = `agent-${randomUUID()}`;
|
||||
const repository = await createTestRepository(organizationId, {
|
||||
type: "s3",
|
||||
config: {
|
||||
backend: "s3",
|
||||
endpoint: "https://s3.example.com",
|
||||
bucket: "bucket",
|
||||
accessKeyId: "access-key",
|
||||
secretAccessKey: "secret-key",
|
||||
},
|
||||
});
|
||||
await db.insert(agentsTable).values({
|
||||
id: agentId,
|
||||
organizationId,
|
||||
name: "Remote Agent",
|
||||
kind: "remote",
|
||||
status: "online",
|
||||
capabilities: {},
|
||||
updatedAt: Date.now(),
|
||||
});
|
||||
vi.spyOn(restic, "snapshots").mockReturnValue(
|
||||
Effect.succeed([
|
||||
{
|
||||
id: "snapshot-restore",
|
||||
short_id: "snapshot-restore",
|
||||
time: new Date().toISOString(),
|
||||
paths: ["/var/lib/zerobyte/volumes/vol123/_data"],
|
||||
hostname: "host",
|
||||
},
|
||||
]),
|
||||
);
|
||||
const restoreMock = vi.fn(() => Promise.resolve(createPendingRestoreStart()));
|
||||
vi.spyOn(agentManager, "startRestore").mockImplementation(restoreMock);
|
||||
const targetPath = await fs.mkdtemp(nodePath.join(process.cwd(), "restore-target-"));
|
||||
|
||||
try {
|
||||
await withContext({ organizationId, userId: session.user.id }, () =>
|
||||
repositoriesService.restoreSnapshot(repository.shortId, "snapshot-restore", {
|
||||
targetPath,
|
||||
targetAgentId: agentId,
|
||||
}),
|
||||
);
|
||||
} finally {
|
||||
await fs.rm(targetPath, { recursive: true, force: true });
|
||||
}
|
||||
|
||||
expect(restoreMock).toHaveBeenCalledWith(
|
||||
agentId,
|
||||
expect.objectContaining({
|
||||
payload: expect.objectContaining({
|
||||
target: targetPath,
|
||||
}),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
test("rejects a target agent outside the current organization", async () => {
|
||||
const organizationId = session.organizationId;
|
||||
const otherSession = await createTestSession();
|
||||
const otherAgentId = `agent-${randomUUID()}`;
|
||||
const repository = await createTestRepository(organizationId, {
|
||||
type: "s3",
|
||||
config: {
|
||||
backend: "s3",
|
||||
endpoint: "https://s3.example.com",
|
||||
bucket: "bucket",
|
||||
accessKeyId: "access-key",
|
||||
secretAccessKey: "secret-key",
|
||||
},
|
||||
});
|
||||
|
||||
await db.insert(agentsTable).values({
|
||||
id: otherAgentId,
|
||||
organizationId: otherSession.organizationId,
|
||||
name: "Other Org Agent",
|
||||
kind: "remote",
|
||||
status: "online",
|
||||
capabilities: {},
|
||||
updatedAt: Date.now(),
|
||||
});
|
||||
vi.spyOn(restic, "snapshots").mockReturnValue(
|
||||
Effect.succeed([
|
||||
{
|
||||
id: "snapshot-restore",
|
||||
short_id: "snapshot-restore",
|
||||
time: new Date().toISOString(),
|
||||
paths: ["/var/lib/zerobyte/volumes/vol123/_data"],
|
||||
hostname: "host",
|
||||
},
|
||||
]),
|
||||
);
|
||||
const restoreMock = vi.fn(() => Promise.resolve(createPendingRestoreStart()));
|
||||
vi.spyOn(agentManager, "startRestore").mockImplementation(restoreMock);
|
||||
|
||||
await expect(
|
||||
withContext({ organizationId, userId: session.user.id }, () =>
|
||||
repositoriesService.restoreSnapshot(repository.shortId, "snapshot-restore", {
|
||||
targetAgentId: otherAgentId,
|
||||
}),
|
||||
),
|
||||
).rejects.toThrow("Restore target agent not found");
|
||||
|
||||
expect(restoreMock).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
test("uses controller-local restore fallback when local agent supervision is disabled", async () => {
|
||||
config.flags.enableLocalAgent = false;
|
||||
const { organizationId, userId, repositoryShortId, restoreMock } = await setupRestoreSnapshotScenario();
|
||||
const resticRestoreMock = vi.spyOn(restic, "restore").mockReturnValue(
|
||||
Effect.succeed({
|
||||
message_type: "summary" as const,
|
||||
files_skipped: 0,
|
||||
files_restored: 1,
|
||||
}),
|
||||
);
|
||||
const targetPath = await fs.mkdtemp(nodePath.join(process.cwd(), "restore-target-"));
|
||||
|
||||
try {
|
||||
await withContext({ organizationId, userId }, () =>
|
||||
repositoriesService.restoreSnapshot(repositoryShortId, "snapshot-restore", {
|
||||
targetPath,
|
||||
}),
|
||||
);
|
||||
} finally {
|
||||
await fs.rm(targetPath, { recursive: true, force: true });
|
||||
}
|
||||
|
||||
expect(restoreMock).not.toHaveBeenCalled();
|
||||
await waitForExpect(() => {
|
||||
expect(resticRestoreMock).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ backend: "local" }),
|
||||
"snapshot-restore",
|
||||
targetPath,
|
||||
expect.objectContaining({
|
||||
organizationId,
|
||||
basePath: "/var/lib/zerobyte/volumes/vol123/_data",
|
||||
signal: expect.any(AbortSignal),
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
test("rejects original-location restore for snapshots with non-posix source paths", async () => {
|
||||
const { organizationId, userId, repositoryShortId, restoreMock } = await setupRestoreSnapshotScenario([
|
||||
"d:\\some\\path",
|
||||
|
|
@ -509,21 +693,26 @@ describe("repositoriesService.restoreSnapshot", () => {
|
|||
|
||||
try {
|
||||
await withContext({ organizationId, userId }, () =>
|
||||
repositoriesService.restoreSnapshot(repositoryShortId, "snapshot-restore", { targetPath }),
|
||||
repositoriesService.restoreSnapshot(repositoryShortId, "snapshot-restore", {
|
||||
targetPath,
|
||||
}),
|
||||
);
|
||||
} finally {
|
||||
await fs.rm(targetPath, { recursive: true, force: true });
|
||||
}
|
||||
|
||||
expect(restoreMock).toHaveBeenCalledWith(
|
||||
"local",
|
||||
expect.objectContaining({
|
||||
backend: "local",
|
||||
}),
|
||||
"snapshot-restore",
|
||||
targetPath,
|
||||
expect.objectContaining({
|
||||
organizationId,
|
||||
basePath: "/",
|
||||
payload: expect.objectContaining({
|
||||
snapshotId: "snapshot-restore",
|
||||
target: targetPath,
|
||||
repositoryConfig: expect.objectContaining({ backend: "local" }),
|
||||
options: expect.objectContaining({
|
||||
organizationId,
|
||||
basePath: "/",
|
||||
}),
|
||||
}),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
|
@ -536,9 +725,14 @@ describe("repositoriesService.getRetentionCategories", () => {
|
|||
|
||||
test("recomputes retention categories after repository cache invalidation", async () => {
|
||||
const organizationId = session.organizationId;
|
||||
const schedule = await createTestBackupSchedule({ organizationId, retentionPolicy: { keepLast: 1 } });
|
||||
const schedule = await createTestBackupSchedule({
|
||||
organizationId,
|
||||
retentionPolicy: { keepLast: 1 },
|
||||
});
|
||||
|
||||
const repository = await db.query.repositoriesTable.findFirst({ where: { id: schedule.repositoryId } });
|
||||
const repository = await db.query.repositoriesTable.findFirst({
|
||||
where: { id: schedule.repositoryId },
|
||||
});
|
||||
|
||||
expect(repository).toBeTruthy();
|
||||
if (!repository) {
|
||||
|
|
@ -622,7 +816,9 @@ describe("repositoriesService.deleteSnapshot", () => {
|
|||
expect(statsSpy).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
const updatedRepository = await db.query.repositoriesTable.findFirst({ where: { id: repository.id } });
|
||||
const updatedRepository = await db.query.repositoriesTable.findFirst({
|
||||
where: { id: repository.id },
|
||||
});
|
||||
expect(updatedRepository?.stats).toEqual(expectedStats);
|
||||
expect(typeof updatedRepository?.statsUpdatedAt).toBe("number");
|
||||
});
|
||||
|
|
|
|||
|
|
@ -182,7 +182,10 @@ export const repositoriesController = new Hono()
|
|||
const offset = Math.max(0, query.offset ?? 0);
|
||||
const limit = Math.min(1000, Math.max(1, query.limit ?? 500));
|
||||
|
||||
const result = await repositoriesService.listSnapshotFiles(shortId, snapshotId, decodedPath, { offset, limit });
|
||||
const result = await repositoriesService.listSnapshotFiles(shortId, snapshotId, decodedPath, {
|
||||
offset,
|
||||
limit,
|
||||
});
|
||||
|
||||
c.header("Cache-Control", "max-age=300, stale-while-revalidate=600");
|
||||
|
||||
|
|
@ -235,7 +238,7 @@ export const repositoriesController = new Hono()
|
|||
const { snapshotId, ...options } = c.req.valid("json");
|
||||
const result = await repositoriesService.restoreSnapshot(shortId, snapshotId, options);
|
||||
|
||||
return c.json<RestoreSnapshotDto>(result, 200);
|
||||
return c.json<RestoreSnapshotDto>(result, 202);
|
||||
})
|
||||
.post("/:shortId/doctor", startDoctorDto, async (c) => {
|
||||
const shortId = asShortId(c.req.param("shortId"));
|
||||
|
|
|
|||
|
|
@ -343,14 +343,13 @@ export const restoreSnapshotBody = z.object({
|
|||
excludeXattr: z.array(z.string()).optional(),
|
||||
delete: z.boolean().optional(),
|
||||
targetPath: z.string().optional(),
|
||||
targetAgentId: z.string().optional(),
|
||||
overwrite: overwriteModeSchema.optional(),
|
||||
});
|
||||
|
||||
const restoreSnapshotResponse = z.object({
|
||||
success: z.boolean(),
|
||||
message: z.string(),
|
||||
filesRestored: z.number(),
|
||||
filesSkipped: z.number(),
|
||||
restoreId: z.string(),
|
||||
status: z.literal("started"),
|
||||
});
|
||||
|
||||
export type RestoreSnapshotDto = z.infer<typeof restoreSnapshotResponse>;
|
||||
|
|
@ -360,8 +359,8 @@ export const restoreSnapshotDto = describeRoute({
|
|||
tags: ["Repositories"],
|
||||
operationId: "restoreSnapshot",
|
||||
responses: {
|
||||
200: {
|
||||
description: "Snapshot restored successfully",
|
||||
202: {
|
||||
description: "Snapshot restore started",
|
||||
content: {
|
||||
"application/json": {
|
||||
schema: resolver(restoreSnapshotResponse),
|
||||
|
|
|
|||
|
|
@ -19,7 +19,7 @@ import { logger } from "@zerobyte/core/node";
|
|||
import { parseRetentionCategories, type RetentionCategory } from "~/server/utils/retention-categories";
|
||||
import { repoMutex } from "../../core/repository-mutex";
|
||||
import { db } from "../../db/db";
|
||||
import { repositoriesTable, type Repository } from "../../db/schema";
|
||||
import { repositoriesTable, type Repository, type RepositoryInsert } from "../../db/schema";
|
||||
import { cache, cacheKeys } from "../../utils/cache";
|
||||
import { runEffectPromise, toMessage } from "../../utils/errors";
|
||||
import { generateShortId } from "../../utils/id";
|
||||
|
|
@ -30,12 +30,22 @@ import type { DumpPathKind, UpdateRepositoryBody } from "./repositories.dto";
|
|||
import { findCommonAncestor } from "@zerobyte/core/utils";
|
||||
import { prepareSnapshotDump } from "./helpers/dump";
|
||||
import { executeDoctor } from "./helpers/doctor";
|
||||
import { restoreExecutor } from "./restore-executor";
|
||||
import type { ShortId } from "~/server/utils/branded";
|
||||
import { decryptRepositoryConfig, encryptRepositoryConfig } from "./repository-config-secrets";
|
||||
import { getScheduleByIdOrShortId } from "../backups/helpers/backup-schedule-lookups";
|
||||
import type { RestoreExecutionProgress, RestoreExecutionResult } from "../agents/agents-manager";
|
||||
import { agentsService } from "../agents/agents.service";
|
||||
import { LOCAL_AGENT_ID } from "../agents/constants";
|
||||
import { taskStore } from "../tasks/tasks.store";
|
||||
import type { ParsedTask, TaskInput } from "../tasks/tasks.schemas";
|
||||
import { Effect } from "effect";
|
||||
|
||||
const runningDoctors = new Map<string, AbortController>();
|
||||
const RESTORE_TASK_RESOURCE_TYPE = "repository";
|
||||
|
||||
type RestoreTaskInput = Extract<TaskInput, { kind: "restore" }>;
|
||||
type RestoreTask = ParsedTask & { kind: "restore"; input: RestoreTaskInput };
|
||||
|
||||
const emptyRepositoryStats: ResticStatsDto = {
|
||||
total_size: 0,
|
||||
|
|
@ -56,6 +66,175 @@ const getBlockedRestoreTargets = () => {
|
|||
].filter((e) => e !== undefined);
|
||||
};
|
||||
|
||||
const assertAllowedControllerLocalRestoreTarget = (target: string) => {
|
||||
const resolvedTarget = nodePath.resolve(target);
|
||||
|
||||
for (const blockedTarget of getBlockedRestoreTargets()) {
|
||||
if (isPathWithin(blockedTarget, resolvedTarget)) {
|
||||
throw new BadRequestError(
|
||||
"Restore target path is not allowed. Restoring to this path could overwrite critical system files or application data.",
|
||||
);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
const isRestoreTask = (task: ParsedTask): task is RestoreTask =>
|
||||
task.kind === "restore" && task.input.kind === "restore";
|
||||
|
||||
const asRestoreTask = (task: ParsedTask, restoreId: string, eventName: string) => {
|
||||
if (!isRestoreTask(task)) {
|
||||
logger.warn(`Received ${eventName} for non-restore task ${restoreId}`);
|
||||
return null;
|
||||
}
|
||||
|
||||
return task;
|
||||
};
|
||||
|
||||
const updateActiveRestoreTask = (restoreId: string, eventName: string, update: () => ParsedTask) => {
|
||||
try {
|
||||
return asRestoreTask(update(), restoreId, eventName);
|
||||
} catch (error) {
|
||||
logger.warn(`Received ${eventName} for inactive restore ${restoreId}: ${toMessage(error)}`);
|
||||
return null;
|
||||
}
|
||||
};
|
||||
|
||||
const findActiveRestoreTask = (
|
||||
organizationId: string,
|
||||
repositoryShortId: string,
|
||||
snapshotId: string,
|
||||
): RestoreTask | null => {
|
||||
return (
|
||||
taskStore
|
||||
.listActiveByResource({
|
||||
organizationId,
|
||||
kind: "restore",
|
||||
resourceType: RESTORE_TASK_RESOURCE_TYPE,
|
||||
resourceId: repositoryShortId,
|
||||
})
|
||||
.find((task): task is RestoreTask => isRestoreTask(task) && task.input.snapshotId === snapshotId) ?? null
|
||||
);
|
||||
};
|
||||
|
||||
const emitRestoreStarted = (task: RestoreTask) => {
|
||||
serverEvents.emit("restore:started", {
|
||||
restoreId: task.id,
|
||||
organizationId: task.organizationId,
|
||||
repositoryId: task.input.repositoryId,
|
||||
snapshotId: task.input.snapshotId,
|
||||
});
|
||||
};
|
||||
|
||||
const emitRestoreProgress = (task: RestoreTask, progress: RestoreExecutionProgress) => {
|
||||
serverEvents.emit("restore:progress", {
|
||||
restoreId: task.id,
|
||||
organizationId: task.organizationId,
|
||||
repositoryId: task.input.repositoryId,
|
||||
snapshotId: task.input.snapshotId,
|
||||
...progress,
|
||||
});
|
||||
};
|
||||
|
||||
const emitRestoreCompleted = (
|
||||
task: RestoreTask,
|
||||
payload: {
|
||||
status: "success" | "error" | "cancelled";
|
||||
error?: string;
|
||||
filesRestored?: number;
|
||||
filesSkipped?: number;
|
||||
},
|
||||
) => {
|
||||
serverEvents.emit("restore:completed", {
|
||||
restoreId: task.id,
|
||||
organizationId: task.organizationId,
|
||||
repositoryId: task.input.repositoryId,
|
||||
snapshotId: task.input.snapshotId,
|
||||
...payload,
|
||||
});
|
||||
};
|
||||
|
||||
const markRestoreStarted = (restoreId: string) => {
|
||||
const task = updateActiveRestoreTask(restoreId, "restore.started", () => taskStore.markRunning(restoreId));
|
||||
if (!task) return;
|
||||
|
||||
emitRestoreStarted(task);
|
||||
};
|
||||
|
||||
const updateRestoreProgress = (restoreId: string, progress: RestoreExecutionProgress) => {
|
||||
const task = updateActiveRestoreTask(restoreId, "restore.progress", () =>
|
||||
taskStore.updateProgress(restoreId, { kind: "restore", progress }),
|
||||
);
|
||||
if (!task) return;
|
||||
|
||||
emitRestoreProgress(task, progress);
|
||||
};
|
||||
|
||||
const completeRestoreTask = (
|
||||
restoreId: string,
|
||||
result: Extract<RestoreExecutionResult, { status: "completed" }>["result"],
|
||||
) => {
|
||||
const task = updateActiveRestoreTask(restoreId, "restore.completed", () =>
|
||||
taskStore.complete(restoreId, { kind: "restore", result }),
|
||||
);
|
||||
if (!task) return;
|
||||
|
||||
emitRestoreCompleted(task, {
|
||||
status: "success",
|
||||
filesRestored: result.files_restored,
|
||||
filesSkipped: result.files_skipped,
|
||||
});
|
||||
};
|
||||
|
||||
const failRestoreTask = (restoreId: string, error: string) => {
|
||||
const task = updateActiveRestoreTask(restoreId, "restore.failed", () => taskStore.fail(restoreId, error));
|
||||
if (!task) return;
|
||||
|
||||
emitRestoreCompleted(task, { status: "error", error });
|
||||
};
|
||||
|
||||
const cancelRestoreTask = (restoreId: string, message?: string) => {
|
||||
const task = updateActiveRestoreTask(restoreId, "restore.cancelled", () =>
|
||||
taskStore.cancel(restoreId, message ?? null),
|
||||
);
|
||||
if (!task) return;
|
||||
|
||||
emitRestoreCompleted(task, { status: "cancelled", error: task.cancellationRequested ? undefined : message });
|
||||
};
|
||||
|
||||
const finishRestoreExecution = async (restoreId: string, resultPromise: Promise<RestoreExecutionResult>) => {
|
||||
try {
|
||||
const result = await resultPromise;
|
||||
|
||||
switch (result.status) {
|
||||
case "completed":
|
||||
completeRestoreTask(restoreId, result.result);
|
||||
return;
|
||||
case "failed":
|
||||
failRestoreTask(restoreId, result.error);
|
||||
return;
|
||||
case "cancelled":
|
||||
cancelRestoreTask(restoreId, result.message);
|
||||
return;
|
||||
case "unavailable":
|
||||
failRestoreTask(restoreId, result.error.message);
|
||||
return;
|
||||
}
|
||||
} catch (error) {
|
||||
failRestoreTask(restoreId, toMessage(error));
|
||||
}
|
||||
};
|
||||
|
||||
const assertAllowedRestoreAgent = async (agentId: string, organizationId: string) => {
|
||||
if (agentId === LOCAL_AGENT_ID) {
|
||||
return;
|
||||
}
|
||||
|
||||
const agent = await agentsService.getAgent(agentId);
|
||||
if (!agent || agent.organizationId !== organizationId) {
|
||||
throw new NotFoundError("Restore target agent not found");
|
||||
}
|
||||
};
|
||||
|
||||
const findRepository = async (shortId: ShortId) => {
|
||||
const organizationId = getOrganizationId();
|
||||
return await db.query.repositoriesTable.findFirst({
|
||||
|
|
@ -268,7 +447,13 @@ const listSnapshotFiles = async (
|
|||
|
||||
const cacheKey = cacheKeys.repository.ls(repository.id, snapshotId, path, offset, limit);
|
||||
type LsResult = {
|
||||
snapshot: { id: string; short_id: string; time: string; hostname: string; paths: string[] } | null;
|
||||
snapshot: {
|
||||
id: string;
|
||||
short_id: string;
|
||||
time: string;
|
||||
hostname: string;
|
||||
paths: string[];
|
||||
} | null;
|
||||
nodes: { name: string; type: string; path: string; size?: number; mode?: number }[];
|
||||
pagination: { offset: number; limit: number; total: number; hasMore: boolean };
|
||||
};
|
||||
|
|
@ -327,6 +512,7 @@ const restoreSnapshot = async (
|
|||
excludeXattr?: string[];
|
||||
delete?: boolean;
|
||||
targetPath?: string;
|
||||
targetAgentId?: string;
|
||||
overwrite?: OverwriteMode;
|
||||
},
|
||||
) => {
|
||||
|
|
@ -337,77 +523,71 @@ const restoreSnapshot = async (
|
|||
throw new NotFoundError("Repository not found");
|
||||
}
|
||||
|
||||
const target = options?.targetPath || "/";
|
||||
const resolvedTarget = nodePath.resolve(target);
|
||||
const blockedTargets = getBlockedRestoreTargets();
|
||||
const { targetAgentId, targetPath, ...restoreExecutionOptions } = options ?? {};
|
||||
const target = targetPath || "/";
|
||||
|
||||
for (const blockedTarget of blockedTargets) {
|
||||
if (isPathWithin(blockedTarget, resolvedTarget)) {
|
||||
throw new BadRequestError(
|
||||
"Restore target path is not allowed. Restoring to this path could overwrite critical system files or application data.",
|
||||
);
|
||||
}
|
||||
}
|
||||
const snapshot = await getSnapshotDetails(repository.shortId, snapshotId);
|
||||
const hasNonPosixSnapshotPaths = snapshot.paths.some((path) => !path.startsWith("/"));
|
||||
|
||||
const { paths } = await getSnapshotDetails(repository.shortId, snapshotId);
|
||||
const hasNonPosixSnapshotPaths = paths.some((path) => !path.startsWith("/"));
|
||||
|
||||
if (hasNonPosixSnapshotPaths && !options?.targetPath) {
|
||||
if (hasNonPosixSnapshotPaths && !targetPath) {
|
||||
throw new BadRequestError(
|
||||
"Original location restore is unavailable for this snapshot. Restore it to a custom location instead.",
|
||||
);
|
||||
}
|
||||
|
||||
const basePath = hasNonPosixSnapshotPaths ? "/" : findCommonAncestor(paths);
|
||||
const basePath = hasNonPosixSnapshotPaths ? "/" : findCommonAncestor(snapshot.paths);
|
||||
const executionAgentId = targetAgentId ?? LOCAL_AGENT_ID;
|
||||
const useControllerLocalRestoreFallback = executionAgentId === LOCAL_AGENT_ID && !appConfig.flags.enableLocalAgent;
|
||||
await assertAllowedRestoreAgent(executionAgentId, organizationId);
|
||||
|
||||
const releaseLock = await repoMutex.acquireShared(repository.id, `restore:${snapshotId}`);
|
||||
try {
|
||||
serverEvents.emit("restore:started", {
|
||||
organizationId,
|
||||
repositoryId: repository.shortId,
|
||||
snapshotId,
|
||||
});
|
||||
|
||||
const result = await runEffectPromise(
|
||||
restic.restore(repository.config, snapshotId, target, {
|
||||
basePath,
|
||||
...options,
|
||||
organizationId,
|
||||
onProgress: (progress) => {
|
||||
serverEvents.emit("restore:progress", {
|
||||
organizationId,
|
||||
repositoryId: repository.shortId,
|
||||
snapshotId,
|
||||
...progress,
|
||||
});
|
||||
},
|
||||
}),
|
||||
if (!useControllerLocalRestoreFallback && repository.type === "local" && executionAgentId !== LOCAL_AGENT_ID) {
|
||||
throw new BadRequestError(
|
||||
"Local repository restores must run on the agent that can access the repository path.",
|
||||
);
|
||||
}
|
||||
|
||||
serverEvents.emit("restore:completed", {
|
||||
if (executionAgentId === LOCAL_AGENT_ID) {
|
||||
assertAllowedControllerLocalRestoreTarget(target);
|
||||
}
|
||||
|
||||
const activeRestore = findActiveRestoreTask(organizationId, repository.shortId, snapshotId);
|
||||
if (activeRestore) {
|
||||
throw new ConflictError("A restore is already running for this snapshot");
|
||||
}
|
||||
|
||||
const task = taskStore.create({
|
||||
organizationId,
|
||||
resourceType: RESTORE_TASK_RESOURCE_TYPE,
|
||||
resourceId: repository.shortId,
|
||||
targetAgentId: useControllerLocalRestoreFallback ? null : executionAgentId,
|
||||
input: { kind: "restore", repositoryId: repository.shortId, snapshotId, target },
|
||||
});
|
||||
const restoreId = task.id;
|
||||
try {
|
||||
const repositoryConfig = await decryptRepositoryConfig(repository.config);
|
||||
const execution = await restoreExecutor.start({
|
||||
restoreId,
|
||||
organizationId,
|
||||
repositoryId: repository.shortId,
|
||||
repositoryId: repository.id,
|
||||
repositoryShortId: repository.shortId,
|
||||
repositoryConfig,
|
||||
snapshotId,
|
||||
status: "success",
|
||||
target,
|
||||
executionAgentId,
|
||||
options: {
|
||||
basePath,
|
||||
...restoreExecutionOptions,
|
||||
},
|
||||
onStarted: () => markRestoreStarted(restoreId),
|
||||
onProgress: (progress) => updateRestoreProgress(restoreId, progress),
|
||||
});
|
||||
|
||||
return {
|
||||
success: true,
|
||||
message: "Snapshot restored successfully",
|
||||
filesRestored: result.files_restored,
|
||||
filesSkipped: result.files_skipped,
|
||||
};
|
||||
void finishRestoreExecution(restoreId, execution.result);
|
||||
|
||||
return { restoreId, status: "started" as const };
|
||||
} catch (error) {
|
||||
serverEvents.emit("restore:completed", {
|
||||
organizationId,
|
||||
repositoryId: repository.shortId,
|
||||
snapshotId,
|
||||
status: "error",
|
||||
error: toMessage(error),
|
||||
});
|
||||
failRestoreTask(restoreId, toMessage(error));
|
||||
throw error;
|
||||
} finally {
|
||||
releaseLock();
|
||||
}
|
||||
};
|
||||
|
||||
|
|
@ -424,7 +604,11 @@ const dumpSnapshot = async (shortId: ShortId, snapshotId: string, path?: string,
|
|||
|
||||
try {
|
||||
const snapshot = await getSnapshotDetails(repository.shortId, snapshotId);
|
||||
const preparedDump = prepareSnapshotDump({ snapshotId, snapshotPaths: snapshot.paths, requestedPath: path });
|
||||
const preparedDump = prepareSnapshotDump({
|
||||
snapshotId,
|
||||
snapshotPaths: snapshot.paths,
|
||||
requestedPath: path,
|
||||
});
|
||||
const dumpOptions: Parameters<typeof restic.dump>[2] = {
|
||||
organizationId,
|
||||
path: preparedDump.path,
|
||||
|
|
@ -738,23 +922,22 @@ const updateRepository = async (shortId: ShortId, updates: UpdateRepositoryBody)
|
|||
const configChanged = updates.config && JSON.stringify(decryptedExisting) !== JSON.stringify(parsedConfig);
|
||||
const encryptedConfig = updates.config ? await encryptRepositoryConfig(parsedConfig) : existingConfig;
|
||||
const updatedAt = Date.now();
|
||||
const updatePayload = {
|
||||
const updatePayload: Partial<RepositoryInsert> = {
|
||||
name: newName,
|
||||
compressionMode: updates.compressionMode ?? existing.compressionMode,
|
||||
updatedAt,
|
||||
config: encryptedConfig,
|
||||
...(configChanged
|
||||
? {
|
||||
status: "unknown" as const,
|
||||
lastChecked: null,
|
||||
lastError: null,
|
||||
doctorResult: null,
|
||||
stats: null,
|
||||
statsUpdatedAt: null,
|
||||
}
|
||||
: {}),
|
||||
};
|
||||
|
||||
if (configChanged) {
|
||||
updatePayload.status = "unknown";
|
||||
updatePayload.lastChecked = null;
|
||||
updatePayload.lastError = null;
|
||||
updatePayload.doctorResult = null;
|
||||
updatePayload.stats = null;
|
||||
updatePayload.statsUpdatedAt = null;
|
||||
}
|
||||
|
||||
const [updated] = await db
|
||||
.update(repositoriesTable)
|
||||
.set(updatePayload)
|
||||
|
|
@ -816,7 +999,14 @@ const execResticCommand = async (
|
|||
addCommonArgs(resticArgs, env, repository.config);
|
||||
|
||||
try {
|
||||
const result = await safeSpawn({ command: "restic", args: resticArgs, env, signal, onStdout, onStderr });
|
||||
const result = await safeSpawn({
|
||||
command: "restic",
|
||||
args: resticArgs,
|
||||
env,
|
||||
signal,
|
||||
onStdout,
|
||||
onStderr,
|
||||
});
|
||||
return { exitCode: result.exitCode };
|
||||
} finally {
|
||||
await cleanupTemporaryKeys(env, resticDeps);
|
||||
|
|
|
|||
122
app/server/modules/repositories/restore-executor.ts
Normal file
122
app/server/modules/repositories/restore-executor.ts
Normal file
|
|
@ -0,0 +1,122 @@
|
|||
import type { RepositoryConfig } from "@zerobyte/core/restic";
|
||||
import type { RestoreRunPayload } from "@zerobyte/contracts/agent-protocol";
|
||||
import { config as appConfig } from "~/server/core/config";
|
||||
import { repoMutex } from "../../core/repository-mutex";
|
||||
import { restic, resticDeps } from "../../core/restic";
|
||||
import { runEffectPromise, toMessage } from "../../utils/errors";
|
||||
import { agentManager, type RestoreExecutionProgress, type RestoreExecutionResult } from "../agents/agents-manager";
|
||||
import { LOCAL_AGENT_ID } from "../agents/constants";
|
||||
|
||||
type RestoreExecutionOptions = Omit<Parameters<typeof restic.restore>[3], "organizationId" | "signal" | "onProgress">;
|
||||
|
||||
type RestoreExecutionRequest = {
|
||||
restoreId: string;
|
||||
organizationId: string;
|
||||
repositoryId: string;
|
||||
repositoryShortId: string;
|
||||
repositoryConfig: RepositoryConfig;
|
||||
snapshotId: string;
|
||||
target: string;
|
||||
executionAgentId: string;
|
||||
options: RestoreExecutionOptions;
|
||||
onStarted: () => void;
|
||||
onProgress: (progress: RestoreExecutionProgress) => void;
|
||||
};
|
||||
|
||||
type RestoreExecutionHandle = {
|
||||
result: Promise<RestoreExecutionResult>;
|
||||
};
|
||||
|
||||
const shouldRunInController = (agentId: string) => agentId === LOCAL_AGENT_ID && !appConfig.flags.enableLocalAgent;
|
||||
|
||||
const createRestoreRunPayload = async (request: RestoreExecutionRequest): Promise<RestoreRunPayload> => {
|
||||
const encryptedResticPassword = await resticDeps.getOrganizationResticPassword(request.organizationId);
|
||||
const resticPassword = await resticDeps.resolveSecret(encryptedResticPassword);
|
||||
|
||||
return {
|
||||
restoreId: request.restoreId,
|
||||
organizationId: request.organizationId,
|
||||
repositoryId: request.repositoryShortId,
|
||||
snapshotId: request.snapshotId,
|
||||
target: request.target,
|
||||
repositoryConfig: request.repositoryConfig,
|
||||
runtime: { password: resticPassword },
|
||||
options: {
|
||||
...request.options,
|
||||
organizationId: request.organizationId,
|
||||
},
|
||||
};
|
||||
};
|
||||
|
||||
const executeControllerRestore = async (
|
||||
request: RestoreExecutionRequest,
|
||||
signal: AbortSignal,
|
||||
): Promise<RestoreExecutionResult> => {
|
||||
if (signal.aborted) {
|
||||
return { status: "cancelled", message: "Restore was cancelled" };
|
||||
}
|
||||
|
||||
request.onStarted();
|
||||
|
||||
try {
|
||||
const result = await runEffectPromise(
|
||||
restic.restore(request.repositoryConfig, request.snapshotId, request.target, {
|
||||
...request.options,
|
||||
organizationId: request.organizationId,
|
||||
signal,
|
||||
onProgress: request.onProgress,
|
||||
}),
|
||||
);
|
||||
|
||||
return { status: "completed", result };
|
||||
} catch (error) {
|
||||
if (signal.aborted) {
|
||||
return { status: "cancelled", message: "Restore was cancelled" };
|
||||
}
|
||||
|
||||
return { status: "failed", error: toMessage(error) };
|
||||
}
|
||||
};
|
||||
|
||||
export const restoreExecutor = {
|
||||
start: async (request: RestoreExecutionRequest): Promise<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;
|
||||
}
|
||||
},
|
||||
};
|
||||
|
|
@ -6,14 +6,18 @@ import { generateBackupOutput } from "~/test/helpers/restic";
|
|||
import { taskStore } from "../tasks.store";
|
||||
import type { TaskInput, TaskProgress, TaskResult } from "../tasks.schemas";
|
||||
|
||||
const backupInput = (scheduleId = 1): TaskInput => ({
|
||||
type BackupTaskInput = Extract<TaskInput, { kind: "backup" }>;
|
||||
type BackupTaskProgress = Extract<TaskProgress, { kind: "backup" }>;
|
||||
type BackupTaskResult = Extract<TaskResult, { kind: "backup" }>;
|
||||
|
||||
const backupInput = (scheduleId = 1): BackupTaskInput => ({
|
||||
kind: "backup",
|
||||
scheduleId,
|
||||
scheduleShortId: `schedule-${scheduleId}`,
|
||||
manual: false,
|
||||
});
|
||||
|
||||
const backupProgress = (percentDone = 0.5): TaskProgress => ({
|
||||
const backupProgress = (percentDone = 0.5): BackupTaskProgress => ({
|
||||
kind: "backup",
|
||||
progress: {
|
||||
message_type: "status",
|
||||
|
|
@ -28,7 +32,7 @@ const backupProgress = (percentDone = 0.5): TaskProgress => ({
|
|||
},
|
||||
});
|
||||
|
||||
const backupResult = (): TaskResult => ({
|
||||
const backupResult = (): BackupTaskResult => ({
|
||||
kind: "backup",
|
||||
exitCode: 0,
|
||||
result: JSON.parse(generateBackupOutput()),
|
||||
|
|
@ -45,6 +49,16 @@ const createBackupTask = (overrides: Partial<Parameters<typeof taskStore.create>
|
|||
...overrides,
|
||||
});
|
||||
|
||||
const createRestoreTask = (overrides: Partial<Parameters<typeof taskStore.create>[0]> = {}) =>
|
||||
taskStore.create({
|
||||
organizationId: TEST_ORG_ID,
|
||||
resourceType: "repository",
|
||||
resourceId: "repo-short",
|
||||
targetAgentId: "local",
|
||||
input: { kind: "restore", repositoryId: "repo-short", snapshotId: "snapshot-1", target: "/tmp/restore" },
|
||||
...overrides,
|
||||
});
|
||||
|
||||
beforeEach(async () => {
|
||||
await ensureTestOrganization();
|
||||
await db.delete(tasksTable);
|
||||
|
|
@ -97,6 +111,10 @@ test("moves an active task through running, progress, cancellation request, and
|
|||
|
||||
const completed = taskStore.complete(task.id, backupResult());
|
||||
expect(completed.status).toBe("succeeded");
|
||||
expect(completed.result?.kind).toBe("backup");
|
||||
if (completed.result?.kind !== "backup") {
|
||||
throw new Error("Expected backup result");
|
||||
}
|
||||
expect(completed.result?.result?.snapshot_id).toBe("abcd1234");
|
||||
expect(completed.finishedAt).toEqual(expect.any(Number));
|
||||
expect(completed.cancellationRequested).toBe(true);
|
||||
|
|
@ -117,6 +135,44 @@ test("records failed and cancelled terminal task states", () => {
|
|||
expect(cancelled.cancellationRequested).toBe(true);
|
||||
});
|
||||
|
||||
test("moves restore tasks through progress and success", () => {
|
||||
const task = createRestoreTask({ id: "restore-task" });
|
||||
|
||||
const running = taskStore.markRunning(task.id);
|
||||
expect(running.kind).toBe("restore");
|
||||
expect(running.input).toMatchObject({ kind: "restore", snapshotId: "snapshot-1" });
|
||||
|
||||
const progressed = taskStore.updateProgress(task.id, {
|
||||
kind: "restore",
|
||||
progress: {
|
||||
message_type: "status",
|
||||
seconds_elapsed: 2,
|
||||
percent_done: 0.25,
|
||||
total_files: 4,
|
||||
files_restored: 1,
|
||||
total_bytes: 400,
|
||||
bytes_restored: 100,
|
||||
},
|
||||
});
|
||||
expect(progressed.progress?.progress.percent_done).toBe(0.25);
|
||||
|
||||
const completed = taskStore.complete(task.id, {
|
||||
kind: "restore",
|
||||
result: {
|
||||
message_type: "summary",
|
||||
total_files: 4,
|
||||
files_restored: 4,
|
||||
files_skipped: 0,
|
||||
},
|
||||
});
|
||||
expect(completed.status).toBe("succeeded");
|
||||
expect(completed.result?.kind).toBe("restore");
|
||||
if (completed.result?.kind !== "restore") {
|
||||
throw new Error("Expected restore result");
|
||||
}
|
||||
expect(completed.result?.result.files_restored).toBe(4);
|
||||
});
|
||||
|
||||
test("finds the newest active task for a resource and marks only matching active tasks stale", async () => {
|
||||
createBackupTask({ id: "task-a-old", resourceId: "shared" });
|
||||
const newest = createBackupTask({ id: "task-z-new", resourceId: "shared" });
|
||||
|
|
|
|||
|
|
@ -1,4 +1,9 @@
|
|||
import { resticBackupOutputSchema, resticBackupProgressSchema } from "@zerobyte/core/restic";
|
||||
import {
|
||||
resticBackupOutputSchema,
|
||||
resticBackupProgressSchema,
|
||||
resticRestoreOutputSchema,
|
||||
restoreProgressSchema,
|
||||
} from "@zerobyte/core/restic";
|
||||
import { z } from "zod";
|
||||
|
||||
export const taskStatuses = ["queued", "running", "cancelling", "cancelled", "succeeded", "failed", "stale"] as const;
|
||||
|
|
@ -6,7 +11,7 @@ export const activeTaskStatuses = ["queued", "running", "cancelling"] as const;
|
|||
|
||||
export const taskStatusSchema = z.enum(taskStatuses);
|
||||
export const activeTaskStatusSchema = z.enum(activeTaskStatuses);
|
||||
export const taskKindSchema = z.enum(["backup"]);
|
||||
export const taskKindSchema = z.enum(["backup", "restore"]);
|
||||
|
||||
export const taskInputSchema = z.discriminatedUnion("kind", [
|
||||
z.object({
|
||||
|
|
@ -15,6 +20,12 @@ export const taskInputSchema = z.discriminatedUnion("kind", [
|
|||
scheduleShortId: z.string(),
|
||||
manual: z.boolean(),
|
||||
}),
|
||||
z.object({
|
||||
kind: z.literal("restore"),
|
||||
repositoryId: z.string(),
|
||||
snapshotId: z.string(),
|
||||
target: z.string(),
|
||||
}),
|
||||
]);
|
||||
|
||||
export const taskProgressSchema = z.discriminatedUnion("kind", [
|
||||
|
|
@ -22,6 +33,10 @@ export const taskProgressSchema = z.discriminatedUnion("kind", [
|
|||
kind: z.literal("backup"),
|
||||
progress: resticBackupProgressSchema,
|
||||
}),
|
||||
z.object({
|
||||
kind: z.literal("restore"),
|
||||
progress: restoreProgressSchema,
|
||||
}),
|
||||
]);
|
||||
|
||||
export const taskResultSchema = z.discriminatedUnion("kind", [
|
||||
|
|
@ -31,6 +46,10 @@ export const taskResultSchema = z.discriminatedUnion("kind", [
|
|||
result: resticBackupOutputSchema.nullable(),
|
||||
warningDetails: z.string().nullable(),
|
||||
}),
|
||||
z.object({
|
||||
kind: z.literal("restore"),
|
||||
result: resticRestoreOutputSchema,
|
||||
}),
|
||||
]);
|
||||
|
||||
export const taskSchema = z
|
||||
|
|
|
|||
|
|
@ -30,9 +30,7 @@ type CreateTaskParams = {
|
|||
input: TaskInput;
|
||||
};
|
||||
|
||||
type MarkActiveStaleParams = Partial<TaskResource> & {
|
||||
error?: string;
|
||||
};
|
||||
type MarkActiveStaleParams = Partial<TaskResource> & { error?: string };
|
||||
|
||||
const parseTask = (row: unknown): ParsedTask => taskSchema.parse(row);
|
||||
|
||||
|
|
@ -188,6 +186,17 @@ export const taskStore = {
|
|||
return row ? parseTask(row) : null;
|
||||
},
|
||||
|
||||
listActiveByResource: (params: TaskResource): ParsedTask[] => {
|
||||
const rows = db
|
||||
.select()
|
||||
.from(tasksTable)
|
||||
.where(and(...buildActiveConditions(params)))
|
||||
.orderBy(desc(tasksTable.createdAt), desc(tasksTable.id))
|
||||
.all();
|
||||
|
||||
return rows.map(parseTask);
|
||||
},
|
||||
|
||||
markActiveStale: (params: MarkActiveStaleParams = {}): ParsedTask[] => {
|
||||
const now = Date.now();
|
||||
const rows = db
|
||||
|
|
|
|||
Loading…
Reference in a new issue