diff --git a/app/server.ts b/app/server.ts index 34364053..abfe59eb 100644 --- a/app/server.ts +++ b/app/server.ts @@ -1,5 +1,3 @@ -import { logger } from "@zerobyte/core/node"; -import { shutdown } from "./server/modules/lifecycle/shutdown"; import { runCLI } from "./server/cli"; import { createStartHandler, defaultStreamHandler, defineHandlerCallback } from "@tanstack/react-start/server"; import { createServerEntry } from "@tanstack/react-start/server-entry"; @@ -18,25 +16,3 @@ const fetch = createStartHandler(customHandler); export default createServerEntry({ fetch, }); - -process.on("SIGTERM", async () => { - logger.info("SIGTERM received, starting graceful shutdown..."); - try { - await shutdown(); - } catch (err) { - logger.error("Error during shutdown", err); - } finally { - process.exit(0); - } -}); - -process.on("SIGINT", async () => { - logger.info("SIGINT received, starting graceful shutdown..."); - try { - await shutdown(); - } catch (err) { - logger.error("Error during shutdown", err); - } finally { - process.exit(0); - } -}); diff --git a/app/server/core/__tests__/repository-mutex.test.ts b/app/server/core/__tests__/repository-mutex.test.ts index 6eba0172..4ff80e5f 100644 --- a/app/server/core/__tests__/repository-mutex.test.ts +++ b/app/server/core/__tests__/repository-mutex.test.ts @@ -2,8 +2,20 @@ import { describe, expect, test, vi } from "vitest"; import { eq } from "drizzle-orm"; import { db } from "~/server/db/db"; import { repositoryLocksTable, repositoryLockWaitersTable } from "~/server/db/schema"; +import { + createDeferred, + holdExclusiveLock as holdExclusive, + holdManyLocks as holdMany, + holdSharedLock as holdShared, +} from "~/test/helpers/repository-mutex"; import { repoMutex } from "../repository-mutex"; +const loadRepositoryMutexModule = async () => { + const moduleUrl = new URL("../repository-mutex.ts", import.meta.url); + moduleUrl.searchParams.set("test", crypto.randomUUID()); + return import(moduleUrl.href) as Promise; +}; + const acquireWithin = (promise: Promise, ms = 500) => Promise.race([ promise, @@ -17,15 +29,15 @@ describe("RepositoryMutex", () => { const repoId = "test-repo"; const results: string[] = []; - const releaseShared1 = await repoMutex.acquireShared(repoId, "backup-1"); + const releaseShared1 = await holdShared(repoId, "backup-1"); results.push("acquired-shared-1"); - const exclusivePromise = repoMutex.acquireExclusive(repoId, "unlock").then((release) => { + const exclusivePromise = holdExclusive(repoId, "unlock").then((release) => { results.push("acquired-exclusive"); return release; }); - const shared2Promise = repoMutex.acquireShared(repoId, "backup-2").then((release) => { + const shared2Promise = holdShared(repoId, "backup-2").then((release) => { results.push("acquired-shared-2"); return release; }); @@ -34,33 +46,33 @@ describe("RepositoryMutex", () => { expect(results).toEqual(["acquired-shared-1"]); - releaseShared1(); + await releaseShared1(); const releaseExclusive = await exclusivePromise; expect(results).toEqual(["acquired-shared-1", "acquired-exclusive"]); - releaseExclusive(); + await releaseExclusive(); const releaseShared2 = await shared2Promise; expect(results).toEqual(["acquired-shared-1", "acquired-exclusive", "acquired-shared-2"]); - releaseShared2(); + await releaseShared2(); }); test("should remove aborted acquisitions from the wait queue", async () => { const repoId = "abort-test"; const results: string[] = []; - const releaseShared1 = await repoMutex.acquireShared(repoId, "backup-1"); + const releaseShared1 = await holdShared(repoId, "backup-1"); results.push("acquired-shared-1"); const controller = new AbortController(); - const exclusivePromise = repoMutex.acquireExclusive(repoId, "unlock", controller.signal).catch((err) => { + const exclusivePromise = holdExclusive(repoId, "unlock", controller.signal).catch((err) => { results.push("aborted-exclusive"); throw err; }); - const shared2Promise = repoMutex.acquireShared(repoId, "backup-2").then((release) => { + const shared2Promise = holdShared(repoId, "backup-2").then((release) => { results.push("acquired-shared-2"); return release; }); @@ -73,25 +85,25 @@ describe("RepositoryMutex", () => { await expect(exclusivePromise).rejects.toThrow(); expect(results).toEqual(["acquired-shared-1", "aborted-exclusive"]); - releaseShared1(); + await releaseShared1(); // After exclusive is aborted, shared-2 should be next const releaseShared2 = await shared2Promise; expect(results).toEqual(["acquired-shared-1", "aborted-exclusive", "acquired-shared-2"]); - releaseShared2(); + await releaseShared2(); }); test("should handle multiple aborts correctly", async () => { const repoId = "multi-abort"; - const releaseShared1 = await repoMutex.acquireShared(repoId, "backup-1"); + const releaseShared1 = await holdShared(repoId, "backup-1"); const controller1 = new AbortController(); const controller2 = new AbortController(); - const p1 = repoMutex.acquireExclusive(repoId, "ex-1", controller1.signal); - const p2 = repoMutex.acquireExclusive(repoId, "ex-2", controller2.signal); - const p3 = repoMutex.acquireExclusive(repoId, "ex-3"); + const p1 = holdExclusive(repoId, "ex-1", controller1.signal); + const p2 = holdExclusive(repoId, "ex-2", controller2.signal); + const p3 = holdExclusive(repoId, "ex-3"); controller2.abort(); await expect(p2).rejects.toThrow(); @@ -99,11 +111,11 @@ describe("RepositoryMutex", () => { controller1.abort(); await expect(p1).rejects.toThrow(); - releaseShared1(); + await releaseShared1(); const releaseEx3 = await p3; expect(releaseEx3).toBeDefined(); - releaseEx3(); + await releaseEx3(); }); test("should remove signal listener and not release unacquired lock on abort", async () => { @@ -119,9 +131,9 @@ describe("RepositoryMutex", () => { }; // Hold the lock so the next one has to wait - const release = await repoMutex.acquireExclusive(repoId, "holder"); + const release = await holdExclusive(repoId, "holder"); - const abortedAcquisition = repoMutex.acquireShared(repoId, "waiter", signal); + const abortedAcquisition = holdShared(repoId, "waiter", signal); // Trigger abort while it's waiting controller.abort(); @@ -130,21 +142,21 @@ describe("RepositoryMutex", () => { expect(removed).toBe(true); expect(repoMutex.isLocked(repoId)).toBe(true); - release(); + await release(); expect(repoMutex.isLocked(repoId)).toBe(false); }); test("should allow concurrent shared locks", async () => { const repoId = "concurrent-shared"; - const release1 = await repoMutex.acquireShared(repoId, "op1"); - const release2 = await repoMutex.acquireShared(repoId, "op2"); - const release3 = await repoMutex.acquireShared(repoId, "op3"); + const release1 = await holdShared(repoId, "op1"); + const release2 = await holdShared(repoId, "op2"); + const release3 = await holdShared(repoId, "op3"); expect(repoMutex.isLocked(repoId)).toBe(true); - release1(); - release2(); - release3(); + await release1(); + await release2(); + await release3(); expect(repoMutex.isLocked(repoId)).toBe(false); }); @@ -153,10 +165,10 @@ describe("RepositoryMutex", () => { const repoId = "shared-blocks-exclusive"; let exclusiveAcquired = false; - const releaseShared1 = await repoMutex.acquireShared(repoId, "s1"); - const releaseShared2 = await repoMutex.acquireShared(repoId, "s2"); + const releaseShared1 = await holdShared(repoId, "s1"); + const releaseShared2 = await holdShared(repoId, "s2"); - const exclusivePromise = repoMutex.acquireExclusive(repoId, "e1").then((release) => { + const exclusivePromise = holdExclusive(repoId, "e1").then((release) => { exclusiveAcquired = true; return release; }); @@ -164,29 +176,29 @@ describe("RepositoryMutex", () => { await new Promise((resolve) => setTimeout(resolve, 10)); expect(exclusiveAcquired).toBe(false); - releaseShared1(); + await releaseShared1(); await new Promise((resolve) => setTimeout(resolve, 10)); expect(exclusiveAcquired).toBe(false); // still waiting for s2 - releaseShared2(); + await releaseShared2(); const releaseExclusive = await exclusivePromise; expect(exclusiveAcquired).toBe(true); - releaseExclusive(); + await releaseExclusive(); }); test("should block all locks while exclusive lock is held", async () => { const repoId = "exclusive-blocks-all"; const results: string[] = []; - const releaseExclusive = await repoMutex.acquireExclusive(repoId, "e1"); + const releaseExclusive = await holdExclusive(repoId, "e1"); results.push("e1-acquired"); - const s1Promise = repoMutex.acquireShared(repoId, "s1").then((release) => { + const s1Promise = holdShared(repoId, "s1").then((release) => { results.push("s1-acquired"); return release; }); - const e2Promise = repoMutex.acquireExclusive(repoId, "e2").then((release) => { + const e2Promise = holdExclusive(repoId, "e2").then((release) => { results.push("e2-acquired"); return release; }); @@ -194,34 +206,34 @@ describe("RepositoryMutex", () => { await new Promise((resolve) => setTimeout(resolve, 10)); expect(results).toEqual(["e1-acquired"]); - releaseExclusive(); + await releaseExclusive(); const releaseS1 = await s1Promise; expect(results).toEqual(["e1-acquired", "s1-acquired"]); - releaseS1(); + await releaseS1(); const releaseE2 = await e2Promise; expect(results).toEqual(["e1-acquired", "s1-acquired", "e2-acquired"]); - releaseE2(); + await releaseE2(); }); test("should grant all waiting shared locks at once when exclusive lock is released", async () => { const repoId = "batch-shared"; const results: string[] = []; - const releaseExclusive = await repoMutex.acquireExclusive(repoId, "e1"); + const releaseExclusive = await holdExclusive(repoId, "e1"); - const s1Promise = repoMutex.acquireShared(repoId, "s1").then((release) => { + const s1Promise = holdShared(repoId, "s1").then((release) => { results.push("s1"); return release; }); - const s2Promise = repoMutex.acquireShared(repoId, "s2").then((release) => { + const s2Promise = holdShared(repoId, "s2").then((release) => { results.push("s2"); return release; }); - const s3Promise = repoMutex.acquireShared(repoId, "s3").then((release) => { + const s3Promise = holdShared(repoId, "s3").then((release) => { results.push("s3"); return release; }); @@ -229,7 +241,7 @@ describe("RepositoryMutex", () => { await new Promise((resolve) => setTimeout(resolve, 10)); expect(results).toEqual([]); - releaseExclusive(); + await releaseExclusive(); const [releaseS1, releaseS2, releaseS3] = await Promise.all([s1Promise, s2Promise, s3Promise]); @@ -238,29 +250,27 @@ describe("RepositoryMutex", () => { expect(results).toContain("s2"); expect(results).toContain("s3"); - releaseS1(); - releaseS2(); - releaseS3(); + await releaseS1(); + await releaseS2(); + await releaseS3(); }); test("should wait to acquire all many-lock requests before locking any repository", async () => { - const releaseSharedB = await repoMutex.acquireShared("repo-b", "holder-b"); + const releaseSharedB = await holdShared("repo-b", "holder-b"); let manyAcquired = false; let exclusiveAAcquired = false; - const manyPromise = repoMutex - .acquireMany([ - { repositoryId: "repo-b", type: "exclusive", operation: "many-b" }, - { repositoryId: "repo-a", type: "shared", operation: "many-a" }, - ]) - .then((release) => { - manyAcquired = true; - return release; - }); + const manyPromise = holdMany([ + { repositoryId: "repo-b", type: "exclusive", operation: "many-b" }, + { repositoryId: "repo-a", type: "shared", operation: "many-a" }, + ]).then((release) => { + manyAcquired = true; + return release; + }); await new Promise((resolve) => setTimeout(resolve, 10)); - const exclusiveAPromise = repoMutex.acquireExclusive("repo-a", "exclusive-a").then((release) => { + const exclusiveAPromise = holdExclusive("repo-a", "exclusive-a").then((release) => { exclusiveAAcquired = true; return release; }); @@ -269,23 +279,23 @@ describe("RepositoryMutex", () => { expect(exclusiveAAcquired).toBe(true); expect(manyAcquired).toBe(false); - releaseSharedB(); + await releaseSharedB(); await new Promise((resolve) => setTimeout(resolve, 10)); expect(manyAcquired).toBe(false); const releaseExclusiveA = await exclusiveAPromise; - releaseExclusiveA(); + await releaseExclusiveA(); const releaseMany = await manyPromise; expect(manyAcquired).toBe(true); - releaseMany(); + await releaseMany(); }); - test("should abort acquireMany without leaving partial locks behind", async () => { - const releaseExclusiveB = await repoMutex.acquireExclusive("repo-b", "holder-b"); + test("should abort runMany without leaving partial locks behind", async () => { + const releaseExclusiveB = await holdExclusive("repo-b", "holder-b"); const controller = new AbortController(); - const manyPromise = repoMutex.acquireMany( + const manyPromise = holdMany( [ { repositoryId: "repo-b", type: "exclusive", operation: "many-b" }, { repositoryId: "repo-a", type: "shared", operation: "many-a" }, @@ -296,7 +306,7 @@ describe("RepositoryMutex", () => { await new Promise((resolve) => setTimeout(resolve, 10)); let exclusiveAAcquired = false; - const exclusiveAPromise = repoMutex.acquireExclusive("repo-a", "exclusive-a").then((release) => { + const exclusiveAPromise = holdExclusive("repo-a", "exclusive-a").then((release) => { exclusiveAAcquired = true; return release; }); @@ -310,17 +320,17 @@ describe("RepositoryMutex", () => { const releaseExclusiveA = await exclusiveAPromise; expect(exclusiveAAcquired).toBe(true); - releaseExclusiveA(); - releaseExclusiveB(); + await releaseExclusiveA(); + await releaseExclusiveB(); }); - test("should abort acquireMany if the signal aborts after waiting resolves", async () => { + test("should abort runMany if the signal aborts after waiting resolves", async () => { const repoA = "many-abort-after-wake-a"; const repoB = "many-abort-after-wake-b"; - const releaseExclusiveB = await repoMutex.acquireExclusive(repoB, "holder-b"); + const releaseExclusiveB = await holdExclusive(repoB, "holder-b"); const controller = new AbortController(); - const manyPromise = repoMutex.acquireMany( + const manyPromise = holdMany( [ { repositoryId: repoB, type: "exclusive", operation: "many-b" }, { repositoryId: repoA, type: "shared", operation: "many-a" }, @@ -330,7 +340,7 @@ describe("RepositoryMutex", () => { await new Promise((resolve) => setTimeout(resolve, 10)); - releaseExclusiveB(); + await releaseExclusiveB(); controller.abort(new Error("stop after wake")); await expect(manyPromise).rejects.toThrow("stop after wake"); @@ -341,17 +351,17 @@ describe("RepositoryMutex", () => { test("should not leave a promoted shared lock behind if that waiter aborts before observing acquisition", async () => { vi.useFakeTimers(); const repoId = "shared-abort-after-promotion"; - const releaseExclusive = await repoMutex.acquireExclusive(repoId, "holder"); - const firstSharedPromise = repoMutex.acquireShared(repoId, "shared-1"); + const releaseExclusive = await holdExclusive(repoId, "holder"); + const firstSharedPromise = holdShared(repoId, "shared-1"); try { await vi.advanceTimersByTimeAsync(100); const controller = new AbortController(); - const secondSharedPromise = repoMutex.acquireShared(repoId, "shared-2", controller.signal); + const secondSharedPromise = holdShared(repoId, "shared-2", controller.signal); await vi.advanceTimersByTimeAsync(140); - releaseExclusive(); + await releaseExclusive(); // The first shared waiter wakes first and promotes both shared waiters. await vi.advanceTimersByTimeAsync(10); @@ -360,7 +370,7 @@ describe("RepositoryMutex", () => { controller.abort(new Error("abort after promotion")); await expect(secondSharedPromise).rejects.toThrow("abort after promotion"); - releaseShared1(); + await releaseShared1(); const remainingLocks = await db.query.repositoryLocksTable.findMany({ where: { repositoryId: { eq: repoId } }, @@ -378,14 +388,14 @@ describe("RepositoryMutex", () => { test("should safely handle multiple calls to the release function", async () => { const repoId = "idempotent-release"; - const releaseShared = await repoMutex.acquireShared(repoId, "s1"); - releaseShared(); - releaseShared(); // Should not throw or cause issues - releaseShared(); + const releaseShared = await holdShared(repoId, "s1"); + await releaseShared(); + await releaseShared(); // Should not throw or cause issues + await releaseShared(); - const releaseExclusive = await repoMutex.acquireExclusive(repoId, "e1"); - releaseExclusive(); - releaseExclusive(); // Should not throw + const releaseExclusive = await holdExclusive(repoId, "e1"); + await releaseExclusive(); + await releaseExclusive(); // Should not throw expect(repoMutex.isLocked(repoId)).toBe(false); }); @@ -395,8 +405,8 @@ describe("RepositoryMutex", () => { const controller = new AbortController(); controller.abort(new Error("pre-aborted")); - await expect(repoMutex.acquireShared(repoId, "s1", controller.signal)).rejects.toThrow("pre-aborted"); - await expect(repoMutex.acquireExclusive(repoId, "e1", controller.signal)).rejects.toThrow("pre-aborted"); + await expect(holdShared(repoId, "s1", controller.signal)).rejects.toThrow("pre-aborted"); + await expect(holdExclusive(repoId, "e1", controller.signal)).rejects.toThrow("pre-aborted"); expect(repoMutex.isLocked(repoId)).toBe(false); }); @@ -406,22 +416,22 @@ describe("RepositoryMutex", () => { expect(repoMutex.isLocked(repoId)).toBe(false); - const releaseShared1 = await repoMutex.acquireShared(repoId, "s1"); + const releaseShared1 = await holdShared(repoId, "s1"); expect(repoMutex.isLocked(repoId)).toBe(true); - const releaseShared2 = await repoMutex.acquireShared(repoId, "s2"); + const releaseShared2 = await holdShared(repoId, "s2"); expect(repoMutex.isLocked(repoId)).toBe(true); - releaseShared1(); + await releaseShared1(); expect(repoMutex.isLocked(repoId)).toBe(true); // still locked by s2 - releaseShared2(); + await releaseShared2(); expect(repoMutex.isLocked(repoId)).toBe(false); // all shared released - const releaseExclusive = await repoMutex.acquireExclusive(repoId, "e1"); + const releaseExclusive = await holdExclusive(repoId, "e1"); expect(repoMutex.isLocked(repoId)).toBe(true); - releaseExclusive(); + await releaseExclusive(); expect(repoMutex.isLocked(repoId)).toBe(false); }); @@ -441,7 +451,7 @@ describe("RepositoryMutex", () => { heartbeatAt: now - 10_000, }); - const releaseShared = await acquireWithin(repoMutex.acquireShared(repoId, "backup")); + const releaseShared = await acquireWithin(holdShared(repoId, "backup")); try { const expiredLock = await db.query.repositoryLocksTable.findFirst({ @@ -451,7 +461,7 @@ describe("RepositoryMutex", () => { expect(expiredLock).toBeUndefined(); expect(repoMutex.isLocked(repoId)).toBe(true); } finally { - releaseShared(); + await releaseShared(); await db.delete(repositoryLocksTable).where(eq(repositoryLocksTable.repositoryId, repoId)); } }); @@ -472,7 +482,7 @@ describe("RepositoryMutex", () => { heartbeatAt: now - 10_000, }); - const releaseShared = await acquireWithin(repoMutex.acquireShared(repoId, "backup")); + const releaseShared = await acquireWithin(holdShared(repoId, "backup")); try { const expiredWaiter = await db.query.repositoryLockWaitersTable.findFirst({ @@ -482,7 +492,7 @@ describe("RepositoryMutex", () => { expect(expiredWaiter).toBeUndefined(); expect(repoMutex.isLocked(repoId)).toBe(true); } finally { - releaseShared(); + await releaseShared(); await db.delete(repositoryLockWaitersTable).where(eq(repositoryLockWaitersTable.repositoryId, repoId)); await db.delete(repositoryLocksTable).where(eq(repositoryLocksTable.repositoryId, repoId)); } @@ -491,7 +501,7 @@ describe("RepositoryMutex", () => { test("should release only the caller lock row", async () => { const repoId = "release-own-row"; const foreignLockId = "foreign-release-own-row"; - const releaseShared = await repoMutex.acquireShared(repoId, "owned-shared"); + const releaseShared = await holdShared(repoId, "owned-shared"); const now = Date.now(); try { @@ -506,7 +516,7 @@ describe("RepositoryMutex", () => { heartbeatAt: now, }); - releaseShared(); + await releaseShared(); const remainingLocks = await db.query.repositoryLocksTable.findMany({ where: { repositoryId: { eq: repoId } }, @@ -516,8 +526,108 @@ describe("RepositoryMutex", () => { expect(remainingLocks.map((lock) => lock.operation)).toEqual(["foreign-shared"]); expect(repoMutex.isLocked(repoId)).toBe(true); } finally { - releaseShared(); + await releaseShared(); await db.delete(repositoryLocksTable).where(eq(repositoryLocksTable.repositoryId, repoId)); } }); + + test("shutdown should abort an active managed operation and release its lock", async () => { + const repoId = "shutdown-active"; + const started = createDeferred(); + const operation = repoMutex.runShared(repoId, "active", async ({ signal }) => { + started.resolve(signal); + await new Promise((resolve) => signal.addEventListener("abort", () => resolve(), { once: true })); + }); + + const signal = await started.promise; + expect(repoMutex.isLocked(repoId)).toBe(true); + + await repoMutex.shutdown({ timeoutMs: 100 }); + await operation; + + expect(signal.aborted).toBe(true); + expect(repoMutex.isLocked(repoId)).toBe(false); + }); + + test("shutdown should abort queued managed operations before they start", async () => { + const repoId = "shutdown-queued"; + const holderStarted = createDeferred(); + const holder = repoMutex.runExclusive(repoId, "holder", async ({ signal }) => { + holderStarted.resolve(signal); + await new Promise((resolve) => signal.addEventListener("abort", () => resolve(), { once: true })); + }); + await holderStarted.promise; + + let queuedStarted = false; + const queued = repoMutex.runShared(repoId, "queued", async () => { + queuedStarted = true; + }); + const queuedResult = queued.then( + () => null, + (error: unknown) => error, + ); + + await new Promise((resolve) => setTimeout(resolve, 20)); + await repoMutex.shutdown({ timeoutMs: 100 }); + await holder; + + const queuedError = await queuedResult; + expect(queuedError).toBeInstanceOf(Error); + expect(String((queuedError as Error).message)).toContain("Repository mutex is shutting down"); + expect(queuedStarted).toBe(false); + + const remainingWaiters = await db.query.repositoryLockWaitersTable.findMany({ + where: { repositoryId: { eq: repoId } }, + }); + expect(remainingWaiters).toEqual([]); + expect(repoMutex.isLocked(repoId)).toBe(false); + }); + + test("shutdown should release owned lock rows after the timeout when an operation ignores abort", async () => { + const repoId = "shutdown-timeout-release"; + const started = createDeferred(); + const finish = createDeferred(); + const operation = repoMutex.runExclusive(repoId, "stuck", async ({ signal }) => { + started.resolve(signal); + await finish.promise; + }); + + const signal = await started.promise; + expect(repoMutex.isLocked(repoId)).toBe(true); + + await repoMutex.shutdown({ timeoutMs: 10 }); + + expect(signal.aborted).toBe(true); + expect(repoMutex.isLocked(repoId)).toBe(false); + + finish.resolve(); + await operation; + expect(repoMutex.isLocked(repoId)).toBe(false); + }); + + test("should reuse the same global repository mutex when the module is evaluated more than once", async () => { + const { repoMutex: reloadedMutex } = await loadRepositoryMutexModule(); + const repoId = "shutdown-other-instance"; + const started = createDeferred(); + const operation = reloadedMutex.runShared(repoId, "active", async ({ signal }) => { + started.resolve(signal); + await new Promise((resolve) => signal.addEventListener("abort", () => resolve(), { once: true })); + }); + + const signal = await started.promise; + expect(reloadedMutex).toBe(repoMutex); + expect(reloadedMutex.isLocked(repoId)).toBe(true); + + await repoMutex.shutdown({ timeoutMs: 100 }); + await operation; + + expect(signal.aborted).toBe(true); + expect(reloadedMutex.isLocked(repoId)).toBe(false); + }); + + test("shutdown should be idempotent", async () => { + await Promise.all([repoMutex.shutdown({ timeoutMs: 10 }), repoMutex.shutdown({ timeoutMs: 10 })]); + + await expect(repoMutex.runShared("post-shutdown", "op", () => "ok")).resolves.toBe("ok"); + }); }); diff --git a/app/server/core/repository-mutex.ts b/app/server/core/repository-mutex.ts index 760140e7..3ed3c879 100644 --- a/app/server/core/repository-mutex.ts +++ b/app/server/core/repository-mutex.ts @@ -24,6 +24,24 @@ interface AcquiredLock { acquiredAt: number; } +interface RepositoryLease { + signal: AbortSignal; + release: () => void; +} + +interface ActiveRepositoryOperation { + abortController: AbortController; + cleanup: () => void; + completion: Promise; + release: (() => void) | null; +} + +type RepositoryOperationContext = { + signal: AbortSignal; +}; + +type RepositoryOperation = (context: RepositoryOperationContext) => T | Promise; + type RepositoryMutexTransaction = Parameters[0]>[0]; type HeartbeatTarget = "lock" | "waiter"; type QueueAttempt = { status: "acquired"; lock: AcquiredLock } | { status: "waiting" } | { status: "missing" }; @@ -32,11 +50,27 @@ const LOCK_LEASE_MS = 30_000; const LOCK_HEARTBEAT_MS = 5_000; const LOCK_POLL_MS = 250; const LOCK_POLL_CLEANUP_MS = 5_000; +const SHUTDOWN_WAIT_MS = 5_000; +const REPOSITORY_MUTEX_INSTANCE = Symbol.for("zerobyte.repositoryMutex.instance"); -class RepositoryMutex { +function getRepositoryMutex() { + const globalObject = globalThis as typeof globalThis & Record; + const mutex = globalObject[REPOSITORY_MUTEX_INSTANCE]; + + if (mutex) return mutex; + + const newMutex = new RepositoryMutex(); + globalObject[REPOSITORY_MUTEX_INSTANCE] = newMutex; + return newMutex; +} + +export class RepositoryMutex { private ownerId = `owner_${Bun.randomUUIDv7()}`; private heartbeatTimers = new Map>(); + private activeOperations = new Map(); private nextPollCleanupAt = 0; + private shuttingDown = false; + private shutdownPromise: Promise | null = null; private generateLockId(): string { return `lock_${Bun.randomUUIDv7()}`; @@ -58,6 +92,126 @@ class RepositoryMutex { throw this.abortReason(signal); } + private throwIfShuttingDown() { + if (this.shuttingDown) { + throw new Error("Repository mutex is shutting down"); + } + } + + private createOperationController(signal?: AbortSignal) { + const abortController = new AbortController(); + const abortOperation = () => { + if (!abortController.signal.aborted) { + abortController.abort(signal ? this.abortReason(signal) : new Error("Operation aborted")); + } + }; + + if (signal?.aborted) { + abortOperation(); + } else { + signal?.addEventListener("abort", abortOperation, { once: true }); + } + + return { + abortController, + cleanup: () => signal?.removeEventListener("abort", abortOperation), + }; + } + + private createLease(releaseLock: () => void, signal?: AbortSignal): RepositoryLease { + const abortController = new AbortController(); + let released = false; + const abortLease = () => { + if (!abortController.signal.aborted) { + abortController.abort(signal ? this.abortReason(signal) : new Error("Operation aborted")); + } + }; + + if (signal?.aborted) { + abortLease(); + } else { + signal?.addEventListener("abort", abortLease, { once: true }); + } + + return { + signal: abortController.signal, + release: () => { + if (released) return; + + released = true; + signal?.removeEventListener("abort", abortLease); + releaseLock(); + }, + }; + } + + private async runWithLease(lease: RepositoryLease, operation: RepositoryOperation) { + try { + return await operation({ signal: lease.signal }); + } finally { + lease.release(); + } + } + + private runManagedOperation( + openLease: (signal: AbortSignal) => Promise, + operation: RepositoryOperation, + signal?: AbortSignal, + ) { + this.throwIfShuttingDown(); + + const operationId = `operation_${Bun.randomUUIDv7()}`; + const { abortController, cleanup } = this.createOperationController(signal); + const activeOperation: ActiveRepositoryOperation = { + abortController, + cleanup, + completion: Promise.resolve(), + release: null, + }; + + const completion = (async () => { + const lease = await openLease(abortController.signal); + activeOperation.release = lease.release; + return await this.runWithLease(lease, operation); + })(); + + activeOperation.completion = completion.finally(() => { + cleanup(); + this.activeOperations.delete(operationId); + }); + this.activeOperations.set(operationId, activeOperation); + + return activeOperation.completion as Promise; + } + + private waitForShutdownSettled(operations: ActiveRepositoryOperation[], timeoutMs: number) { + if (operations.length === 0) { + return Promise.resolve(); + } + + return new Promise((resolve) => { + const timeout = setTimeout(resolve, timeoutMs); + void Promise.allSettled(operations.map((operation) => operation.completion)).then(() => { + clearTimeout(timeout); + resolve(); + }); + }); + } + + private releaseOwnedRows() { + db.transaction((tx) => { + tx.delete(repositoryLockWaitersTable).where(eq(repositoryLockWaitersTable.ownerId, this.ownerId)).run(); + tx.delete(repositoryLocksTable).where(eq(repositoryLocksTable.ownerId, this.ownerId)).run(); + }); + } + + private stopAllHeartbeats() { + for (const timer of this.heartbeatTimers.values()) { + clearInterval(timer); + } + this.heartbeatTimers.clear(); + } + private waitForPoll(signal?: AbortSignal) { this.throwIfAborted(signal); @@ -328,40 +482,23 @@ class RepositoryMutex { } } - async acquireShared(repositoryId: string, operation: string, signal?: AbortSignal) { + private async openSingleLease(request: LockRequest, signal?: AbortSignal) { this.throwIfAborted(signal); - const request: LockRequest = { repositoryId, type: "shared", operation }; const releaseLock = this.tryAcquireImmediately(request, signal); if (releaseLock) { - return releaseLock; + return this.createLease(releaseLock, signal); } - logger.debug(`[Mutex] Waiting for shared lock on repo ${repositoryId}: ${operation}`); - return await this.waitForQueuedLock(request, signal); + logger.debug(`[Mutex] Waiting for ${request.type} lock on repo ${request.repositoryId}: ${request.operation}`); + return this.createLease(await this.waitForQueuedLock(request, signal), signal); } - async acquireExclusive(repositoryId: string, operation: string, signal?: AbortSignal) { - this.throwIfAborted(signal); - - const request: LockRequest = { repositoryId, type: "exclusive", operation }; - const releaseLock = this.tryAcquireImmediately(request, signal); - if (releaseLock) { - logger.debug(`[Mutex] Acquired exclusive lock for repo ${repositoryId}: ${operation}`); - return releaseLock; - } - - logger.debug(`[Mutex] Waiting for exclusive lock on repo ${repositoryId}: ${operation}`); - const queuedReleaseLock = await this.waitForQueuedLock(request, signal); - logger.debug(`[Mutex] Acquired exclusive lock for repo ${repositoryId}: ${operation}`); - return queuedReleaseLock; - } - - async acquireMany(requests: LockRequest[], signal?: AbortSignal) { + private async openManyLease(requests: LockRequest[], signal?: AbortSignal) { this.throwIfAborted(signal); if (requests.length === 0) { - return () => {}; + return this.createLease(() => {}, signal); } const seenRepositoryIds = new Set(); @@ -379,13 +516,87 @@ class RepositoryMutex { const releaseLocks = this.createReleaseMany(locks); this.releaseIfAborted(releaseLocks, signal); - return releaseLocks; + return this.createLease(releaseLocks, signal); } await this.waitForPoll(signal); } } + async runShared( + repositoryId: string, + operation: string, + callback: RepositoryOperation, + signal?: AbortSignal, + ) { + return this.runManagedOperation( + (operationSignal) => this.openSingleLease({ repositoryId, type: "shared", operation }, operationSignal), + callback, + signal, + ); + } + + async runExclusive( + repositoryId: string, + operation: string, + callback: RepositoryOperation, + signal?: AbortSignal, + ) { + return this.runManagedOperation( + async (operationSignal) => { + const lease = await this.openSingleLease( + { repositoryId, type: "exclusive", operation }, + operationSignal, + ); + logger.debug(`[Mutex] Acquired exclusive lock for repo ${repositoryId}: ${operation}`); + return lease; + }, + callback, + signal, + ); + } + + async runMany(requests: LockRequest[], callback: RepositoryOperation, signal?: AbortSignal) { + return this.runManagedOperation( + (operationSignal) => this.openManyLease(requests, operationSignal), + callback, + signal, + ); + } + + async shutdown(options: { timeoutMs?: number } = {}) { + if (this.shutdownPromise) { + return await this.shutdownPromise; + } + + this.shuttingDown = true; + this.shutdownPromise = (async () => { + const activeOperations = [...this.activeOperations.values()]; + const reason = new Error("Repository mutex is shutting down"); + + for (const operation of activeOperations) { + if (!operation.abortController.signal.aborted) { + operation.abortController.abort(reason); + } + } + + await this.waitForShutdownSettled(activeOperations, options.timeoutMs ?? SHUTDOWN_WAIT_MS); + + for (const operation of activeOperations) { + operation.release?.(); + operation.cleanup(); + } + + this.releaseOwnedRows(); + this.stopAllHeartbeats(); + })().finally(() => { + this.shuttingDown = false; + this.shutdownPromise = null; + }); + + return await this.shutdownPromise; + } + isLocked(repositoryId: string) { const now = Date.now(); @@ -493,4 +704,4 @@ class RepositoryMutex { } } -export const repoMutex = new RepositoryMutex(); +export const repoMutex = getRepositoryMutex(); diff --git a/app/server/modules/agents/agents-manager.ts b/app/server/modules/agents/agents-manager.ts index 98868b35..c039f3be 100644 --- a/app/server/modules/agents/agents-manager.ts +++ b/app/server/modules/agents/agents-manager.ts @@ -431,9 +431,16 @@ export const agentManager = { }); getActiveBackupScheduleIdsByJobId().set(request.payload.jobId, request.scheduleId); }); + const cancelOnAbort = () => { + void requestBackupCancellation(agentId, request.scheduleId).catch((error) => { + logger.error(`Failed to cancel backup ${request.scheduleId} after abort: ${String(error)}`); + }); + }; + request.signal.addEventListener("abort", cancelOnAbort, { once: true }); try { if (!(await Effect.runPromise(runtime.sendBackup(agentId, request.payload)))) { + request.signal.removeEventListener("abort", cancelOnAbort); clearActiveBackupRun(request.scheduleId); return { status: "unavailable", @@ -445,10 +452,13 @@ export const agentManager = { await requestBackupCancellation(agentId, request.scheduleId); } - return completion; + return await completion; } catch (error) { + request.signal.removeEventListener("abort", cancelOnAbort); clearActiveBackupRun(request.scheduleId); throw error; + } finally { + request.signal.removeEventListener("abort", cancelOnAbort); } }, cancelBackup: async (agentId: string, scheduleId: number) => { @@ -494,9 +504,16 @@ export const agentManager = { cancellationRequested: false, }); }); + const cancelOnAbort = () => { + void requestRestoreCancellation(agentId, request.payload.restoreId).catch((error) => { + logger.error(`Failed to cancel restore ${request.payload.restoreId} after abort: ${String(error)}`); + }); + }; + request.signal.addEventListener("abort", cancelOnAbort, { once: true }); try { if (!(await Effect.runPromise(runtime.sendRestore(agentId, request.payload)))) { + request.signal.removeEventListener("abort", cancelOnAbort); clearActiveRestoreRun(request.payload.restoreId); return { status: "unavailable", @@ -508,8 +525,12 @@ export const agentManager = { await requestRestoreCancellation(agentId, request.payload.restoreId); } - return { status: "started", result: completion }; + return { + status: "started", + result: completion.finally(() => request.signal.removeEventListener("abort", cancelOnAbort)), + }; } catch (error) { + request.signal.removeEventListener("abort", cancelOnAbort); clearActiveRestoreRun(request.payload.restoreId); throw error; } diff --git a/app/server/modules/backups/__tests__/backups.service.execution.test.ts b/app/server/modules/backups/__tests__/backups.service.execution.test.ts index 9edc89a8..ac809295 100644 --- a/app/server/modules/backups/__tests__/backups.service.execution.test.ts +++ b/app/server/modules/backups/__tests__/backups.service.execution.test.ts @@ -16,7 +16,6 @@ import { NotFoundError } from "http-errors-enhanced"; import { fromAny } from "@total-typescript/shoehorn"; import { scheduleQueries } from "../backups.queries"; import { repositoriesService } from "~/server/modules/repositories/repositories.service"; -import { repoMutex } from "~/server/core/repository-mutex"; import { notificationsService } from "~/server/modules/notifications/notifications.service"; import { agentManager } from "~/server/modules/agents/agents-manager"; import { createAgentBackupMocks } from "~/test/helpers/agent-mock"; @@ -26,6 +25,7 @@ import { db } from "~/server/db/db"; import { config } from "~/server/core/config"; import { Effect } from "effect"; import { taskStore } from "~/server/modules/tasks/tasks.store"; +import { holdExclusiveLock } from "~/test/helpers/repository-mutex"; const setup = () => { const resticBackupMock = vi.fn((_: SafeSpawnParams) => @@ -842,7 +842,7 @@ describe("stop backup", () => { repositoryId: repository.id, }); - const releaseLock = await repoMutex.acquireExclusive(repository.id, "test"); + const releaseLock = await holdExclusiveLock(repository.id, "test"); const executePromise = backupsService.executeBackup(schedule.id); try { @@ -857,7 +857,7 @@ describe("stop backup", () => { await backupsService.stopBackup(schedule.id); } finally { - releaseLock(); + await releaseLock(); } await executePromise; diff --git a/app/server/modules/backups/backup-executor.ts b/app/server/modules/backups/backup-executor.ts index 8016b098..28c16172 100644 --- a/app/server/modules/backups/backup-executor.ts +++ b/app/server/modules/backups/backup-executor.ts @@ -122,7 +122,7 @@ export const backupExecutor = { }, execute: async (request: BackupExecutionRequest) => { const trackedExecution = activeControllersByScheduleId.get(request.scheduleId); - if (!trackedExecution || trackedExecution.abortController.signal !== request.signal) { + if (!trackedExecution) { throw new Error(`Backup execution for schedule ${request.scheduleId} was not tracked`); } diff --git a/app/server/modules/backups/backups.service.ts b/app/server/modules/backups/backups.service.ts index 0ebd3767..d01b1ea0 100644 --- a/app/server/modules/backups/backups.service.ts +++ b/app/server/modules/backups/backups.service.ts @@ -451,70 +451,67 @@ const executeBackup = async (scheduleId: number, manual = false) => { let domainHandlerCompleted = false; try { - const releaseLock = await repoMutex.acquireShared( + await repoMutex.runShared( ctx.repository.id, `backup:${ctx.volume.name}`, + async ({ signal }) => { + taskStore.markRunning(task.id); + + const executionResult = await backupExecutor.execute({ + jobId: task.id, + scheduleId, + schedule: ctx.schedule, + volume: ctx.volume, + repository: ctx.repository, + organizationId: ctx.organizationId, + signal, + onProgress: (progress) => { + updateBackupProgress(ctx, progress); + try { + taskStore.updateProgress(task.id, { kind: "backup", progress }); + } catch (error) { + logger.error(`Failed to persist backup task progress for ${task.id}: ${toMessage(error)}`); + } + }, + }); + + switch (executionResult.status) { + case "unavailable": { + await handleBackupFailure(scheduleId, ctx.organizationId, executionResult.error, manual, ctx); + domainHandlerCompleted = true; + taskStore.fail(task.id, toMessage(executionResult.error)); + return; + } + case "completed": + await finalizeSuccessfulBackup( + ctx, + executionResult.exitCode, + executionResult.result, + executionResult.warningDetails, + ); + domainHandlerCompleted = true; + taskStore.complete(task.id, { + kind: "backup", + exitCode: executionResult.exitCode, + result: executionResult.result, + warningDetails: executionResult.warningDetails, + }); + return; + case "failed": { + await handleBackupFailure(scheduleId, ctx.organizationId, executionResult.error, manual, ctx); + domainHandlerCompleted = true; + taskStore.fail(task.id, toMessage(executionResult.error)); + return; + } + case "cancelled": + await handleBackupCancellation(scheduleId, ctx.organizationId, executionResult.message); + domainHandlerCompleted = true; + taskStore.cancel(task.id, executionResult.message ?? "Backup was stopped by the user"); + return; + } + }, abortController.signal, ); - - try { - taskStore.markRunning(task.id); - - const executionResult = await backupExecutor.execute({ - jobId: task.id, - scheduleId, - schedule: ctx.schedule, - volume: ctx.volume, - repository: ctx.repository, - organizationId: ctx.organizationId, - signal: abortController.signal, - onProgress: (progress) => { - updateBackupProgress(ctx, progress); - try { - taskStore.updateProgress(task.id, { kind: "backup", progress }); - } catch (error) { - logger.error(`Failed to persist backup task progress for ${task.id}: ${toMessage(error)}`); - } - }, - }); - - switch (executionResult.status) { - case "unavailable": { - await handleBackupFailure(scheduleId, ctx.organizationId, executionResult.error, manual, ctx); - domainHandlerCompleted = true; - taskStore.fail(task.id, toMessage(executionResult.error)); - return; - } - case "completed": - await finalizeSuccessfulBackup( - ctx, - executionResult.exitCode, - executionResult.result, - executionResult.warningDetails, - ); - domainHandlerCompleted = true; - taskStore.complete(task.id, { - kind: "backup", - exitCode: executionResult.exitCode, - result: executionResult.result, - warningDetails: executionResult.warningDetails, - }); - return; - case "failed": { - await handleBackupFailure(scheduleId, ctx.organizationId, executionResult.error, manual, ctx); - domainHandlerCompleted = true; - taskStore.fail(task.id, toMessage(executionResult.error)); - return; - } - case "cancelled": - await handleBackupCancellation(scheduleId, ctx.organizationId, executionResult.message); - domainHandlerCompleted = true; - taskStore.cancel(task.id, executionResult.message ?? "Backup was stopped by the user"); - return; - } - } finally { - releaseLock(); - } } catch (error) { if (abortController.signal.aborted) { taskStore.cancel(task.id, "Backup was stopped by the user"); diff --git a/app/server/modules/backups/helpers/backup-maintenance.ts b/app/server/modules/backups/helpers/backup-maintenance.ts index 1416cc9b..ae668d4c 100644 --- a/app/server/modules/backups/helpers/backup-maintenance.ts +++ b/app/server/modules/backups/helpers/backup-maintenance.ts @@ -20,6 +20,7 @@ export async function runForget(scheduleId: number, repositoryId?: string, organ if (!schedule.retentionPolicy) { throw new BadRequestError("No retention policy configured for this schedule"); } + const retentionPolicy = schedule.retentionPolicy; const repository = await repositoryQueries.findById(repositoryId ?? schedule.repositoryId, organizationId); @@ -28,16 +29,12 @@ export async function runForget(scheduleId: number, repositoryId?: string, organ } logger.info(`running retention policy (forget) for schedule ${scheduleId}`); - const releaseLock = await repoMutex.acquireExclusive(repository.id, `forget:${scheduleId}`); - - try { + await repoMutex.runExclusive(repository.id, `forget:${scheduleId}`, async ({ signal }) => { await runEffectPromise( - restic.forget(repository.config, schedule.retentionPolicy, { tag: schedule.shortId, organizationId }), + restic.forget(repository.config, retentionPolicy, { tag: schedule.shortId, organizationId, signal }), ); cache.delByPrefix(cacheKeys.repository.all(repository.id)); - } finally { - releaseLock(); - } + }); logger.info(`Retention policy applied successfully for schedule ${scheduleId}`); } @@ -114,23 +111,23 @@ export async function syncSnapshotsToMirror( lastCopyError: null, }); - const releaseLocks = await repoMutex.acquireMany([ - { repositoryId: sourceRepository.id, type: "shared", operation: `mirror_sync_source:${scheduleId}` }, - { repositoryId: mirrorRepository.id, type: "exclusive", operation: `mirror_sync:${scheduleId}` }, - ]); - - try { - await runEffectPromise( - restic.copy(sourceRepository.config, mirrorRepository.config, { - tag: schedule.shortId, - organizationId, - snapshotIds, - }), - ); - cache.delByPrefix(cacheKeys.repository.all(mirrorRepository.id)); - } finally { - releaseLocks(); - } + await repoMutex.runMany( + [ + { repositoryId: sourceRepository.id, type: "shared", operation: `mirror_sync_source:${scheduleId}` }, + { repositoryId: mirrorRepository.id, type: "exclusive", operation: `mirror_sync:${scheduleId}` }, + ], + async ({ signal }) => { + await runEffectPromise( + restic.copy(sourceRepository.config, mirrorRepository.config, { + tag: schedule.shortId, + organizationId, + snapshotIds, + signal, + }), + ); + cache.delByPrefix(cacheKeys.repository.all(mirrorRepository.id)); + }, + ); if (schedule.retentionPolicy) { void runForget(scheduleId, mirrorRepository.id, organizationId).catch((error) => { @@ -204,22 +201,22 @@ async function copyToSingleMirror( lastCopyError: null, }); - const releaseLocks = await repoMutex.acquireMany([ - { repositoryId: sourceRepository.id, type: "shared", operation: `mirror_source:${scheduleId}` }, - { repositoryId: mirror.repository.id, type: "exclusive", operation: `mirror:${scheduleId}` }, - ]); - - try { - await runEffectPromise( - restic.copy(sourceRepository.config, mirror.repository.config, { - tag: schedule.shortId, - organizationId, - }), - ); - cache.delByPrefix(cacheKeys.repository.all(mirror.repository.id)); - } finally { - releaseLocks(); - } + await repoMutex.runMany( + [ + { repositoryId: sourceRepository.id, type: "shared", operation: `mirror_source:${scheduleId}` }, + { repositoryId: mirror.repository.id, type: "exclusive", operation: `mirror:${scheduleId}` }, + ], + async ({ signal }) => { + await runEffectPromise( + restic.copy(sourceRepository.config, mirror.repository.config, { + tag: schedule.shortId, + organizationId, + signal, + }), + ); + cache.delByPrefix(cacheKeys.repository.all(mirror.repository.id)); + }, + ); if (retentionPolicy) { void runForget(scheduleId, mirror.repository.id, organizationId).catch((error) => { diff --git a/app/server/modules/lifecycle/__tests__/shutdown.test.ts b/app/server/modules/lifecycle/__tests__/shutdown.test.ts index 11dbed84..05dcd75c 100644 --- a/app/server/modules/lifecycle/__tests__/shutdown.test.ts +++ b/app/server/modules/lifecycle/__tests__/shutdown.test.ts @@ -1,5 +1,6 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { Scheduler } from "../../../core/scheduler"; +import { repoMutex } from "../../../core/repository-mutex"; import * as bootstrapModule from "../bootstrap"; import { agentManager } from "../../agents/agents-manager"; import { createTestVolume } from "~/test/helpers/volume"; @@ -32,6 +33,9 @@ describe("shutdown", () => { const stopApplicationRuntime = vi.fn(async () => { events.push("agents.stop"); }); + const shutdownRepoMutex = vi.fn(async () => { + events.push("repo-mutex.shutdown"); + }); const runVolumeCommand = vi.spyOn(agentManager, "runVolumeCommand"); const volume = await createTestVolume({ @@ -45,13 +49,14 @@ describe("shutdown", () => { }); vi.spyOn(Scheduler, "stop").mockImplementation(stopScheduler); + vi.spyOn(repoMutex, "shutdown").mockImplementation(shutdownRepoMutex); vi.spyOn(bootstrapModule, "stopApplicationRuntime").mockImplementation(stopApplicationRuntime); const { shutdown } = await loadShutdownModule(); await shutdown(); - expect(events).toEqual(["scheduler.stop", "agents.stop"]); + expect(events).toEqual(["scheduler.stop", "repo-mutex.shutdown", "agents.stop"]); expect(runVolumeCommand).not.toHaveBeenCalled(); const updated = await db.query.volumesTable.findFirst({ where: { id: volume.id } }); expect(updated).toBeDefined(); @@ -64,6 +69,9 @@ describe("shutdown", () => { vi.spyOn(Scheduler, "stop").mockImplementation(async () => { events.push("scheduler.stop"); }); + vi.spyOn(repoMutex, "shutdown").mockImplementation(async () => { + events.push("repo-mutex.shutdown"); + }); vi.spyOn(bootstrapModule, "stopApplicationRuntime").mockImplementation(async () => { events.push("agents.stop"); }); @@ -82,7 +90,7 @@ describe("shutdown", () => { await shutdown(); const updated = await db.query.volumesTable.findFirst({ where: { id: volume.id } }); - expect(events).toEqual(["scheduler.stop", "agents.stop"]); + expect(events).toEqual(["scheduler.stop", "repo-mutex.shutdown", "agents.stop"]); expect(runVolumeCommand).not.toHaveBeenCalled(); expect(updated).toBeDefined(); expect(updated!.status).toBe("mounted"); diff --git a/app/server/modules/lifecycle/shutdown.ts b/app/server/modules/lifecycle/shutdown.ts index ba01018a..11e0be0f 100644 --- a/app/server/modules/lifecycle/shutdown.ts +++ b/app/server/modules/lifecycle/shutdown.ts @@ -1,4 +1,5 @@ import { Scheduler } from "../../core/scheduler"; +import { repoMutex } from "../../core/repository-mutex"; import { db } from "../../db/db"; import { logger } from "@zerobyte/core/node"; import { stopApplicationRuntime } from "./bootstrap"; @@ -8,9 +9,7 @@ import { toMessage } from "../../utils/errors"; import { config } from "../../core/config"; import { LOCAL_AGENT_ID } from "../agents/constants"; -export const shutdown = async () => { - await Scheduler.stop(); - +const unmountVolumes = async () => { if (!config.flags.enableLocalAgent) { const volumes = await db.query.volumesTable.findMany({ where: { AND: [{ status: "mounted" }, { agentId: LOCAL_AGENT_ID }] }, @@ -24,6 +23,8 @@ export const shutdown = async () => { logger.info(`Volume ${volume.name} unmount status: ${status}${error ? `, error: ${error}` : ""}`); } } - - await stopApplicationRuntime(); +}; + +export const shutdown = async () => { + await Promise.allSettled([Scheduler.stop(), repoMutex.shutdown(), unmountVolumes(), stopApplicationRuntime()]); }; diff --git a/app/server/modules/repositories/__tests__/repositories.service.test.ts b/app/server/modules/repositories/__tests__/repositories.service.test.ts index 9d50bf5b..eba63d7b 100644 --- a/app/server/modules/repositories/__tests__/repositories.service.test.ts +++ b/app/server/modules/repositories/__tests__/repositories.service.test.ts @@ -22,6 +22,7 @@ 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"; +import { holdExclusiveLock } from "~/test/helpers/repository-mutex"; const createTestRepository = async (organizationId: string, overrides: Partial = {}) => { const id = randomUUID(); @@ -220,9 +221,17 @@ describe("repositoriesService.listSnapshotFiles", () => { let releaseAll = false; let exclusiveAcquired = false; let releaseExclusive: (() => void) | undefined; - let exclusivePromise: Promise<() => void> | undefined; + let exclusivePromise: Promise | undefined; const releaseWaiters: Array<() => void> = []; const exclusiveController = new AbortController(); + let resolveExclusiveStarted: () => void = () => {}; + let resolveExclusiveReleased: () => void = () => {}; + const exclusiveStarted = new Promise((resolve) => { + resolveExclusiveStarted = resolve; + }); + const exclusiveReleased = new Promise((resolve) => { + resolveExclusiveReleased = resolve; + }); const releaseWaitingCommands = () => { const waiters = releaseWaiters.splice(0); @@ -300,22 +309,27 @@ describe("repositoriesService.listSnapshotFiles", () => { }); expect(maxActive).toBe(2); - exclusivePromise = repoMutex - .acquireExclusive(repository.id, "delete", exclusiveController.signal) - .then((release) => { + exclusivePromise = repoMutex.runExclusive( + repository.id, + "delete", + async () => { exclusiveAcquired = true; - releaseExclusive = release; - return release; - }); + releaseExclusive = resolveExclusiveReleased; + resolveExclusiveStarted(); + await exclusiveReleased; + }, + exclusiveController.signal, + ); releaseWaitingCommands(); - releaseExclusive = await resolveWithin(exclusivePromise, 2000); + await resolveWithin(exclusiveStarted, 2000); expect(exclusiveAcquired).toBe(true); expect(active).toBe(0); - releaseExclusive(); + releaseExclusive?.(); releaseExclusive = undefined; + await resolveWithin(exclusivePromise, 2000); await waitForExpect(() => { expect(releaseWaiters).toHaveLength(2); @@ -373,10 +387,11 @@ describe("repositoriesService.dumpSnapshot", () => { snapshotPaths?: string[]; }) => { const organizationId = session.organizationId; + const repositoryId = randomUUID(); const shortId = generateShortId(); await db.insert(repositoriesTable).values({ - id: randomUUID(), + id: repositoryId, shortId, name: `Repository-${randomUUID()}`, type: "local", @@ -421,10 +436,27 @@ describe("repositoriesService.dumpSnapshot", () => { return { organizationId, userId: session.user.id, + repositoryId, shortId, basePath, }; }; + const rejectWithin = async (promise: Promise, timeoutMs: number) => { + let timeout: ReturnType | undefined; + try { + return await Promise.race([ + promise, + new Promise((_resolve, reject) => { + timeout = setTimeout( + () => reject(new Error(`Expected promise to reject within ${timeoutMs}ms`)), + timeoutMs, + ); + }), + ]); + } finally { + clearTimeout(timeout); + } + }; test("returns a tar download rooted at the selected directory within the snapshot", async () => { const { organizationId, userId, shortId, basePath } = await setupDumpSnapshotScenario({ @@ -547,6 +579,29 @@ describe("repositoriesService.dumpSnapshot", () => { }), ); }); + + test("rejects when shutdown cancels the dump before it acquires the mutex", async () => { + const { organizationId, userId, repositoryId, shortId } = await setupDumpSnapshotScenario({ + snapshotId: "snapshot-shutdown", + basePath: "/var/lib/zerobyte/volumes/vol123/_data", + }); + const holderStarted = new Promise((resolve) => { + void repoMutex.runExclusive(repositoryId, "holder", async ({ signal }) => { + resolve(signal); + await new Promise((done) => signal.addEventListener("abort", () => done(), { once: true })); + }); + }); + + await holderStarted; + const dump = withContext({ organizationId, userId }, () => + repositoriesService.dumpSnapshot(shortId, "snapshot-shutdown"), + ); + + await new Promise((resolve) => setTimeout(resolve, 20)); + await repoMutex.shutdown({ timeoutMs: 100 }); + + await expect(rejectWithin(dump, 1000)).rejects.toThrow("Repository mutex is shutting down"); + }); }); describe("repositoriesService.restoreSnapshot", () => { @@ -696,7 +751,7 @@ describe("repositoriesService.restoreSnapshot", () => { }); restoreMock.mockResolvedValueOnce({ status: "started", result: restoreResult }); - const releaseExclusive = await repoMutex.acquireExclusive(repositoryId, "check"); + const releaseExclusive = await holdExclusiveLock(repositoryId, "check"); let restoreId = ""; try { const restoreStart = withContext({ organizationId, userId }, () => @@ -720,7 +775,7 @@ describe("repositoriesService.restoreSnapshot", () => { expect(task?.id).toBe(restoreId); expect(task?.status).toBe("queued"); } finally { - releaseExclusive(); + await releaseExclusive(); await fs.rm(targetPath, { recursive: true, force: true }); } diff --git a/app/server/modules/repositories/helpers/doctor.ts b/app/server/modules/repositories/helpers/doctor.ts index 3fe56d8a..36ca7aeb 100644 --- a/app/server/modules/repositories/helpers/doctor.ts +++ b/app/server/modules/repositories/helpers/doctor.ts @@ -130,29 +130,30 @@ export const executeDoctor = async ( const organizationId = getOrganizationId(); try { - const releaseLock = await repoMutex.acquireExclusive(repositoryId, "doctor", signal); + await repoMutex.runExclusive( + repositoryId, + "doctor", + async ({ signal: operationSignal }) => { + // Step 1: Unlock repository + const unlockStep = await runUnlockStep(repositoryConfig, operationSignal); + steps.push(unlockStep); + checkAbortSignal(operationSignal); - try { - // Step 1: Unlock repository - const unlockStep = await runUnlockStep(repositoryConfig, signal); - steps.push(unlockStep); - checkAbortSignal(signal); + // Step 2: Check repository + const checkStep = await runCheckStep(repositoryConfig, operationSignal); + steps.push(checkStep); + checkAbortSignal(operationSignal); - // Step 2: Check repository - const checkStep = await runCheckStep(repositoryConfig, signal); - steps.push(checkStep); - checkAbortSignal(signal); - - // Step 3: Repair index if suggested - const checkOutput = parseCheckOutput(checkStep.output); - if (checkOutput?.suggest_repair_index) { - const repairStep = await runRepairIndexStep(repositoryConfig, signal); - steps.push(repairStep); - checkAbortSignal(signal); - } - } finally { - releaseLock(); - } + // Step 3: Repair index if suggested + const checkOutput = parseCheckOutput(checkStep.output); + if (checkOutput?.suggest_repair_index) { + const repairStep = await runRepairIndexStep(repositoryConfig, operationSignal); + steps.push(repairStep); + checkAbortSignal(operationSignal); + } + }, + signal, + ); const finalStatus = determineRepositoryStatus(steps); await saveDoctorResults(repositoryId, steps, finalStatus); diff --git a/app/server/modules/repositories/repositories.service.ts b/app/server/modules/repositories/repositories.service.ts index c0343ede..278117f0 100644 --- a/app/server/modules/repositories/repositories.service.ts +++ b/app/server/modules/repositories/repositories.service.ts @@ -334,10 +334,9 @@ const getRepository = async (shortId: ShortId) => { }; const runAndStoreRepositoryStats = async (repository: Repository): Promise => { - const releaseLock = await repoMutex.acquireShared(repository.id, "stats"); - try { + return repoMutex.runShared(repository.id, "stats", async ({ signal }) => { const stats = await runEffectPromise( - restic.stats(repository.config, { organizationId: repository.organizationId }), + restic.stats(repository.config, { organizationId: repository.organizationId, signal }), ); await db @@ -351,9 +350,7 @@ const runAndStoreRepositoryStats = async (repository: Repository): Promise { @@ -419,24 +416,21 @@ const listSnapshots = async (shortId: ShortId, backupId?: ShortId) => { return cached; } - const releaseLock = await repoMutex.acquireShared(repository.id, "snapshots"); - try { + return repoMutex.runShared(repository.id, "snapshots", async ({ signal }) => { let snapshots = []; if (backupId) { snapshots = await runEffectPromise( - restic.snapshots(repository.config, { tags: [backupId], organizationId }), + restic.snapshots(repository.config, { tags: [backupId], organizationId, signal }), ); } else { - snapshots = await runEffectPromise(restic.snapshots(repository.config, { organizationId })); + snapshots = await runEffectPromise(restic.snapshots(repository.config, { organizationId, signal })); } cache.set(cacheKey, snapshots); return snapshots; - } finally { - releaseLock(); - } + }); }; const listSnapshotFiles = async ( @@ -483,10 +477,9 @@ const listSnapshotFiles = async ( await runEffectPromise(limiter.take(1)); try { - const releaseLock = await repoMutex.acquireShared(repository.id, `ls:${snapshotId}`); - try { + return await repoMutex.runShared(repository.id, `ls:${snapshotId}`, async ({ signal }) => { const result = await runEffectPromise( - restic.ls(repository.config, snapshotId, path, { organizationId, offset, limit }), + restic.ls(repository.config, snapshotId, path, { organizationId, offset, limit, signal }), ); if (!result.snapshot) { @@ -511,9 +504,7 @@ const listSnapshotFiles = async ( cache.set(cacheKey, result); return response; - } finally { - releaseLock(); - } + }); } finally { await runEffectPromise(limiter.release(1)); } @@ -616,60 +607,71 @@ const dumpSnapshot = async (shortId: ShortId, snapshotId: string, path?: string, throw new NotFoundError("Repository not found"); } - const releaseLock = await repoMutex.acquireShared(repository.id, `dump:${snapshotId}`); - let dumpStream: ResticDumpStream | null = null; + type DumpSnapshotResult = ResticDumpStream & { filename: string; contentType: string }; + let resolveStarted: (result: DumpSnapshotResult) => void = () => {}; + let rejectStarted: (error: unknown) => void = () => {}; + const started = new Promise((resolve, reject) => { + resolveStarted = resolve; + rejectStarted = reject; + }); + let lockCompletion: Promise; - try { - const snapshot = await getSnapshotDetails(repository.shortId, snapshotId); - const preparedDump = prepareSnapshotDump({ - snapshotId, - snapshotPaths: snapshot.paths, - requestedPath: path, - }); - const dumpOptions: Parameters[2] = { - organizationId, - path: preparedDump.path, - }; + // The response returns once the stream opens, but the lock must stay held until the stream finishes. + lockCompletion = repoMutex.runShared(repository.id, `dump:${snapshotId}`, async ({ signal }) => { + let dumpStream: ResticDumpStream | null = null; + try { + const snapshot = await getSnapshotDetails(repository.shortId, snapshotId); + const preparedDump = prepareSnapshotDump({ + snapshotId, + snapshotPaths: snapshot.paths, + requestedPath: path, + }); + const dumpOptions: Parameters[2] = { + organizationId, + path: preparedDump.path, + signal, + }; - let filename = preparedDump.filename; - let contentType = "application/x-tar"; + let filename = preparedDump.filename; + let contentType = "application/x-tar"; - if (path && preparedDump.path !== "/") { - if (!kind) { - throw new BadRequestError("Path kind is required when downloading a specific snapshot path"); - } + if (path && preparedDump.path !== "/") { + if (!kind) { + throw new BadRequestError("Path kind is required when downloading a specific snapshot path"); + } - if (kind === "file") { - dumpOptions.archive = false; - contentType = "application/octet-stream"; - const fileName = nodePath.posix.basename(preparedDump.path); - if (fileName) { - filename = fileName; + if (kind === "file") { + dumpOptions.archive = false; + contentType = "application/octet-stream"; + const fileName = nodePath.posix.basename(preparedDump.path); + if (fileName) { + filename = fileName; + } } } + + dumpStream = await runEffectPromise(restic.dump(repository.config, preparedDump.snapshotRef, dumpOptions)); + + serverEvents.emit("dump:started", { + organizationId, + repositoryId: repository.shortId, + snapshotId, + path: preparedDump.path, + filename, + }); + + resolveStarted({ ...dumpStream, completion: lockCompletion, filename, contentType }); + + await dumpStream.completion; + } catch (error) { + rejectStarted(error); + dumpStream?.abort(); + throw error; } + }); + void lockCompletion.catch(rejectStarted); - dumpStream = await runEffectPromise(restic.dump(repository.config, preparedDump.snapshotRef, dumpOptions)); - - serverEvents.emit("dump:started", { - organizationId, - repositoryId: repository.shortId, - snapshotId, - path: preparedDump.path, - filename, - }); - - const completion = dumpStream.completion.finally(releaseLock); - void completion.catch(() => {}); - - return { ...dumpStream, completion, filename, contentType }; - } catch (error) { - if (dumpStream) { - dumpStream.abort(); - } - releaseLock(); - throw error; - } + return await started; }; const getSnapshotDetails = async (shortId: ShortId, snapshotId: string) => { @@ -684,13 +686,10 @@ const getSnapshotDetails = async (shortId: ShortId, snapshotId: string) => { let snapshots = cache.get>>(cacheKey); if (!snapshots) { - const releaseLock = await repoMutex.acquireShared(repository.id, `snapshot_details:${snapshotId}`); - try { - snapshots = await runEffectPromise(restic.snapshots(repository.config, { organizationId })); - cache.set(cacheKey, snapshots); - } finally { - releaseLock(); - } + snapshots = await repoMutex.runShared(repository.id, `snapshot_details:${snapshotId}`, async ({ signal }) => { + return await runEffectPromise(restic.snapshots(repository.config, { organizationId, signal })); + }); + cache.set(cacheKey, snapshots); } const snapshot = snapshots.find((snap) => snap.id === snapshotId || snap.short_id === snapshotId); @@ -712,9 +711,10 @@ const checkHealth = async (shortId: ShortId) => { throw new NotFoundError("Repository not found"); } - const releaseLock = await repoMutex.acquireExclusive(repository.id, "check"); - try { - const { hasErrors, error } = await runEffectPromise(restic.check(repository.config, { organizationId })); + return repoMutex.runExclusive(repository.id, "check", async ({ signal }) => { + const { hasErrors, error } = await runEffectPromise( + restic.check(repository.config, { organizationId, signal }), + ); await db .update(repositoriesTable) @@ -731,9 +731,7 @@ const checkHealth = async (shortId: ShortId) => { ); return { lastError: error }; - } finally { - releaseLock(); - } + }); }; const startDoctor = async (shortId: ShortId) => { @@ -810,18 +808,15 @@ const deleteSnapshot = async (shortId: ShortId, snapshotId: string) => { throw new NotFoundError("Repository not found"); } - const releaseLock = await repoMutex.acquireExclusive(repository.id, `delete:${snapshotId}`); - try { - await runEffectPromise(restic.deleteSnapshot(repository.config, snapshotId, { organizationId })); + await repoMutex.runExclusive(repository.id, `delete:${snapshotId}`, async ({ signal }) => { + await runEffectPromise(restic.deleteSnapshot(repository.config, snapshotId, { organizationId, signal })); cache.delByPrefix(cacheKeys.repository.all(repository.id)); void runAndStoreRepositoryStats(repository).catch((error) => { logger.error( `Failed to refresh repository stats after snapshot deletion for ${repository.shortId}: ${toMessage(error)}`, ); }); - } finally { - releaseLock(); - } + }); }; const deleteSnapshots = async (shortId: ShortId, snapshotIds: string[]) => { @@ -833,14 +828,11 @@ const deleteSnapshots = async (shortId: ShortId, snapshotIds: string[]) => { } let shouldRefreshStats = false; - const releaseLock = await repoMutex.acquireExclusive(repository.id, `delete:bulk`); - try { - await runEffectPromise(restic.deleteSnapshots(repository.config, snapshotIds, { organizationId })); + await repoMutex.runExclusive(repository.id, `delete:bulk`, async ({ signal }) => { + await runEffectPromise(restic.deleteSnapshots(repository.config, snapshotIds, { organizationId, signal })); cache.delByPrefix(cacheKeys.repository.all(repository.id)); shouldRefreshStats = true; - } finally { - releaseLock(); - } + }); if (!shouldRefreshStats) { return; @@ -865,13 +857,10 @@ const tagSnapshots = async ( throw new NotFoundError("Repository not found"); } - const releaseLock = await repoMutex.acquireExclusive(repository.id, `tag:bulk`); - try { - await runEffectPromise(restic.tagSnapshots(repository.config, snapshotIds, tags, { organizationId })); + await repoMutex.runExclusive(repository.id, `tag:bulk`, async ({ signal }) => { + await runEffectPromise(restic.tagSnapshots(repository.config, snapshotIds, tags, { organizationId, signal })); cache.delByPrefix(cacheKeys.repository.all(repository.id)); - } finally { - releaseLock(); - } + }); }; const refreshSnapshots = async (shortId: ShortId) => { @@ -884,9 +873,8 @@ const refreshSnapshots = async (shortId: ShortId) => { cache.delByPrefix(cacheKeys.repository.all(repository.id)); - const releaseLock = await repoMutex.acquireShared(repository.id, "refresh"); - try { - const snapshots = await runEffectPromise(restic.snapshots(repository.config, { organizationId })); + return repoMutex.runShared(repository.id, "refresh", async ({ signal }) => { + const snapshots = await runEffectPromise(restic.snapshots(repository.config, { organizationId, signal })); const cacheKey = cacheKeys.repository.snapshots(repository.id); cache.set(cacheKey, snapshots); @@ -894,9 +882,7 @@ const refreshSnapshots = async (shortId: ShortId) => { message: "Snapshot cache cleared and refreshed", count: snapshots.length, }; - } finally { - releaseLock(); - } + }); }; const updateRepository = async (shortId: ShortId, updates: UpdateRepositoryBody) => { @@ -982,13 +968,9 @@ const unlockRepository = async (shortId: ShortId) => { throw new NotFoundError("Repository not found"); } - const releaseLock = await repoMutex.acquireExclusive(repository.id, "unlock"); - try { - const result = await runEffectPromise(restic.unlock(repository.config, { organizationId })); - return result; - } finally { - releaseLock(); - } + return repoMutex.runExclusive(repository.id, "unlock", async ({ signal }) => { + return await runEffectPromise(restic.unlock(repository.config, { organizationId, signal })); + }); }; const execResticCommand = async ( diff --git a/app/server/modules/repositories/restore-executor.ts b/app/server/modules/repositories/restore-executor.ts index d631938f..f89b4a66 100644 --- a/app/server/modules/repositories/restore-executor.ts +++ b/app/server/modules/repositories/restore-executor.ts @@ -113,24 +113,25 @@ const executeRestoreWithRepositoryLock = async ( request: RestoreExecutionRequest, signal: AbortSignal, ): Promise => { - let releaseLock: (() => void) | null = null; - try { - releaseLock = await repoMutex.acquireShared(request.repositoryId, `restore:${request.restoreId}`, signal); + return await repoMutex.runShared( + request.repositoryId, + `restore:${request.restoreId}`, + async ({ signal: operationSignal }) => { + if (shouldRunInController(request.executionAgentId)) { + return await executeControllerRestore(request, operationSignal); + } - if (shouldRunInController(request.executionAgentId)) { - return await executeControllerRestore(request, signal); - } - - return await executeAgentRestore(request, signal); + return await executeAgentRestore(request, operationSignal); + }, + signal, + ); } catch (error) { if (signal.aborted) { return { status: "cancelled", message: "Restore was cancelled" }; } return { status: "failed", error: toMessage(error) }; - } finally { - releaseLock?.(); } }; diff --git a/app/server/plugins/bootstrap.ts b/app/server/plugins/bootstrap.ts index ede70123..2c689c80 100644 --- a/app/server/plugins/bootstrap.ts +++ b/app/server/plugins/bootstrap.ts @@ -1,10 +1,36 @@ import { definePlugin } from "nitro"; -import { bootstrapApplication, stopApplicationRuntime } from "../modules/lifecycle/bootstrap"; +import { bootstrapApplication } from "../modules/lifecycle/bootstrap"; +import { shutdown } from "../modules/lifecycle/shutdown"; import { logger } from "@zerobyte/core/node"; import { toMessage } from "../utils/errors"; +let shutdownPromise: Promise | undefined; + +const runGracefulShutdown = async (reason: string) => { + if (!shutdownPromise) { + logger.info(`${reason}, starting graceful shutdown...`); + shutdownPromise = shutdown().catch((err) => { + logger.error("Error during shutdown", err); + }); + } + + await shutdownPromise; +}; + +const runSignalShutdown = (reason: string) => { + void runGracefulShutdown(reason).finally(() => { + process.exit(0); + }); +}; + export default definePlugin(async (nitroApp) => { - nitroApp.hooks.hook("close", stopApplicationRuntime); + process.on("SIGTERM", () => { + runSignalShutdown("SIGTERM received"); + }); + process.on("SIGINT", () => { + runSignalShutdown("SIGINT received"); + }); + nitroApp.hooks.hook("close", () => runGracefulShutdown("Server closing")); await bootstrapApplication().catch((err) => { logger.error(`Bootstrap failed: ${toMessage(err)}`); diff --git a/app/test/helpers/repository-mutex.ts b/app/test/helpers/repository-mutex.ts new file mode 100644 index 00000000..b8e5e22b --- /dev/null +++ b/app/test/helpers/repository-mutex.ts @@ -0,0 +1,43 @@ +import { repoMutex } from "~/server/core/repository-mutex"; + +export const createDeferred = () => { + let resolve: (value: T | PromiseLike) => void = () => {}; + let reject: (reason?: unknown) => void = () => {}; + const promise = new Promise((promiseResolve, promiseReject) => { + resolve = promiseResolve; + reject = promiseReject; + }); + + return { promise, resolve, reject }; +}; + +const holdWithRun = async (run: (callback: () => Promise) => Promise) => { + const started = createDeferred(); + const finished = createDeferred(); + const operation = run(async () => { + started.resolve(); + await finished.promise; + }); + operation.catch((error) => started.reject(error)); + + await started.promise; + + let released = false; + return async () => { + if (!released) { + released = true; + finished.resolve(); + } + + await operation; + }; +}; + +export const holdSharedLock = (repositoryId: string, operation: string, signal?: AbortSignal) => + holdWithRun((callback) => repoMutex.runShared(repositoryId, operation, callback, signal)); + +export const holdExclusiveLock = (repositoryId: string, operation: string, signal?: AbortSignal) => + holdWithRun((callback) => repoMutex.runExclusive(repositoryId, operation, callback, signal)); + +export const holdManyLocks = (requests: Parameters[0], signal?: AbortSignal) => + holdWithRun((callback) => repoMutex.runMany(requests, callback, signal)); diff --git a/packages/core/src/restic/commands/copy.ts b/packages/core/src/restic/commands/copy.ts index 88a21f3c..a7b1d94e 100644 --- a/packages/core/src/restic/commands/copy.ts +++ b/packages/core/src/restic/commands/copy.ts @@ -18,7 +18,7 @@ class ResticCopyCommandError extends Data.TaggedError("ResticCopyCommandError")< export const copy = ( sourceConfig: RepositoryConfig, destConfig: RepositoryConfig, - options: { organizationId: string; tag?: string; snapshotIds?: string[] }, + options: { organizationId: string; tag?: string; snapshotIds?: string[]; signal?: AbortSignal }, deps: ResticDeps, ) => { return Effect.tryPromise({ @@ -63,11 +63,16 @@ export const copy = ( logger.info(`Copying snapshots from ${sourceRepoUrl} to ${destRepoUrl}...`); logger.debug(`Executing: restic ${args.join(" ")}`); - const res = await safeExec({ command: "restic", args, env }); + const res = await safeExec({ command: "restic", args, env, signal: options.signal }); await cleanupTemporaryKeys(sourceEnv, deps); await cleanupTemporaryKeys(destEnv, deps); + if (options.signal?.aborted) { + logger.warn("Restic copy was aborted by signal."); + throw new Error("Operation aborted"); + } + const { stdout, stderr } = res; if (res.exitCode !== 0) { diff --git a/packages/core/src/restic/commands/delete-snapshots.ts b/packages/core/src/restic/commands/delete-snapshots.ts index 523524e9..39dad0fc 100644 --- a/packages/core/src/restic/commands/delete-snapshots.ts +++ b/packages/core/src/restic/commands/delete-snapshots.ts @@ -17,7 +17,7 @@ class ResticDeleteSnapshotCommandError extends Data.TaggedError("ResticDeleteSna export const deleteSnapshots = ( config: RepositoryConfig, snapshotIds: string[], - options: { organizationId: string }, + options: { organizationId: string; signal?: AbortSignal }, deps: ResticDeps, ) => { return Effect.tryPromise({ @@ -33,9 +33,14 @@ export const deleteSnapshots = ( addCommonArgs(args, env, config); args.push("--", ...snapshotIds); - const res = await safeExec({ command: "restic", args, env }); + const res = await safeExec({ command: "restic", args, env, signal: options.signal }); await cleanupTemporaryKeys(env, deps); + if (options.signal?.aborted) { + logger.warn("Restic snapshot deletion was aborted by signal."); + throw new Error("Operation aborted"); + } + if (res.exitCode !== 0) { logger.error(`Restic snapshot deletion failed: ${res.stderr}`); throw createResticError(res.exitCode, res.stderr); @@ -59,7 +64,7 @@ export const deleteSnapshots = ( export const deleteSnapshot = ( config: RepositoryConfig, snapshotId: string, - options: { organizationId: string }, + options: { organizationId: string; signal?: AbortSignal }, deps: ResticDeps, ) => { return deleteSnapshots(config, [snapshotId], options, deps); diff --git a/packages/core/src/restic/commands/dump.ts b/packages/core/src/restic/commands/dump.ts index 0dd0c0b6..01a70c20 100644 --- a/packages/core/src/restic/commands/dump.ts +++ b/packages/core/src/restic/commands/dump.ts @@ -32,6 +32,7 @@ export const dump = ( organizationId: string; path?: string; archive?: false; + signal?: AbortSignal; }, deps: ResticDeps, ) => { @@ -64,6 +65,12 @@ export const dump = ( let stream: Readable | null = null; let abortController: AbortController | null = new AbortController(); + const abortDump = () => abortController?.abort(options.signal?.reason); + if (options.signal?.aborted) { + abortDump(); + } else { + options.signal?.addEventListener("abort", abortDump, { once: true }); + } const maxStderrChars = 64 * 1024; let stderrTail = ""; @@ -97,6 +104,7 @@ export const dump = ( throw createResticError(result.exitCode, stderr); }) .finally(async () => { + options.signal?.removeEventListener("abort", abortDump); abortController = null; await cleanup(); }); diff --git a/packages/core/src/restic/commands/forget.ts b/packages/core/src/restic/commands/forget.ts index f2fef598..148fc940 100644 --- a/packages/core/src/restic/commands/forget.ts +++ b/packages/core/src/restic/commands/forget.ts @@ -18,7 +18,7 @@ class ResticForgetCommandError extends Data.TaggedError("ResticForgetCommandErro export const forget = ( config: RepositoryConfig, options: RetentionPolicy, - extra: { tag: string; organizationId: string; dryRun?: boolean }, + extra: { tag: string; organizationId: string; dryRun?: boolean; signal?: AbortSignal }, deps: ResticDeps, ) => { return Effect.tryPromise({ @@ -60,9 +60,14 @@ export const forget = ( addCommonArgs(args, env, config); - const res = await safeExec({ command: "restic", args, env }); + const res = await safeExec({ command: "restic", args, env, signal: extra.signal }); await cleanupTemporaryKeys(env, deps); + if (extra.signal?.aborted) { + logger.warn("Restic forget was aborted by signal."); + throw new Error("Operation aborted"); + } + if (res.exitCode !== 0) { logger.error(`Restic forget failed: ${res.stderr}`); throw createResticError(res.exitCode, res.stderr); diff --git a/packages/core/src/restic/commands/ls.ts b/packages/core/src/restic/commands/ls.ts index f8db1202..ee3d2c96 100644 --- a/packages/core/src/restic/commands/ls.ts +++ b/packages/core/src/restic/commands/ls.ts @@ -60,7 +60,7 @@ export const ls = ( config: RepositoryConfig, snapshotId: string, path: string | undefined, - options: { organizationId: string; offset?: number; limit?: number }, + options: { organizationId: string; offset?: number; limit?: number; signal?: AbortSignal }, deps: ResticDeps, ) => { return Effect.tryPromise({ @@ -92,6 +92,7 @@ export const ls = ( command: "restic", args, env, + signal: options.signal, onStdout: (line) => { const trimmedLine = line.trim(); if (!trimmedLine) { @@ -132,6 +133,11 @@ export const ls = ( await cleanupTemporaryKeys(env, deps); + if (options.signal?.aborted) { + logger.warn("Restic ls was aborted by signal."); + throw new Error("Operation aborted"); + } + if (res.exitCode !== 0) { logger.error(`Restic ls failed: ${res.error}`); throw createResticError(res.exitCode, res.stderr || res.error); diff --git a/packages/core/src/restic/commands/snapshots.ts b/packages/core/src/restic/commands/snapshots.ts index 1c6a8d97..41089145 100644 --- a/packages/core/src/restic/commands/snapshots.ts +++ b/packages/core/src/restic/commands/snapshots.ts @@ -33,7 +33,7 @@ const snapshotInfoSchema = z.object({ export const snapshots = ( config: RepositoryConfig, - options: { tags?: string[]; organizationId: string }, + options: { tags?: string[]; organizationId: string; signal?: AbortSignal }, deps: ResticDeps, ) => { return Effect.tryPromise({ @@ -58,12 +58,18 @@ export const snapshots = ( command: "restic", args, env, + signal: options.signal, onStdout: (line) => { stdoutLines.push(line); }, }); await cleanupTemporaryKeys(env, deps); + if (options.signal?.aborted) { + logger.warn("Restic snapshots was aborted by signal."); + throw new Error("Operation aborted"); + } + if (res.exitCode !== 0) { const errorMessage = res.stderr || res.error; logger.error(`Restic snapshots retrieval failed: ${errorMessage}`); diff --git a/packages/core/src/restic/commands/stats.ts b/packages/core/src/restic/commands/stats.ts index b50bb125..84d6fa24 100644 --- a/packages/core/src/restic/commands/stats.ts +++ b/packages/core/src/restic/commands/stats.ts @@ -16,7 +16,11 @@ class ResticStatsCommandError extends Data.TaggedError("ResticStatsCommandError" message: string; }> {} -export const stats = (config: RepositoryConfig, options: { organizationId: string }, deps: ResticDeps) => { +export const stats = ( + config: RepositoryConfig, + options: { organizationId: string; signal?: AbortSignal }, + deps: ResticDeps, +) => { return Effect.tryPromise({ try: async () => { const repoUrl = buildRepoUrl(config); @@ -25,9 +29,14 @@ export const stats = (config: RepositoryConfig, options: { organizationId: strin const args = ["--repo", repoUrl, "stats", "--mode", "raw-data"]; addCommonArgs(args, env, config); - const res = await safeExec({ command: "restic", args, env }); + const res = await safeExec({ command: "restic", args, env, signal: options.signal }); await cleanupTemporaryKeys(env, deps); + if (options.signal?.aborted) { + logger.warn("Restic stats was aborted by signal."); + throw new Error("Operation aborted"); + } + if (res.exitCode !== 0) { logger.error(`Restic stats retrieval failed: ${res.stderr}`); throw createResticError(res.exitCode, res.stderr); diff --git a/packages/core/src/restic/commands/tag-snapshots.ts b/packages/core/src/restic/commands/tag-snapshots.ts index 9c2422b4..8a24503c 100644 --- a/packages/core/src/restic/commands/tag-snapshots.ts +++ b/packages/core/src/restic/commands/tag-snapshots.ts @@ -18,7 +18,7 @@ export const tagSnapshots = ( config: RepositoryConfig, snapshotIds: string[], tags: { add?: string[]; remove?: string[]; set?: string[] }, - options: { organizationId: string }, + options: { organizationId: string; signal?: AbortSignal }, deps: ResticDeps, ) => { return Effect.tryPromise({ @@ -53,9 +53,14 @@ export const tagSnapshots = ( addCommonArgs(args, env, config); args.push("--", ...snapshotIds); - const res = await safeExec({ command: "restic", args, env }); + const res = await safeExec({ command: "restic", args, env, signal: options.signal }); await cleanupTemporaryKeys(env, deps); + if (options.signal?.aborted) { + logger.warn("Restic snapshot tagging was aborted by signal."); + throw new Error("Operation aborted"); + } + if (res.exitCode !== 0) { logger.error(`Restic snapshot tagging failed: ${res.stderr}`); throw createResticError(res.exitCode, res.stderr);