mirror of
https://github.com/paperclipai/paperclip.git
synced 2026-11-11 05:11:49 +01:00
fix(runtime): keep the copyback lock transaction active (#16340)
Keep the dedicated workspace-copyback transaction active with bounded SELECT 1 probes while filesystem work runs. Preserve existing connection, owner-token, takeover and publication fences.
Verified head 7fb0ad45bf: full CI, E2E, Superagent, legacy statuses and Apex 5/5; no unresolved review findings. Real PostgreSQL idle-timeout reproduction and 71 affected cases pass. Strengthened disabled-timeout controls fail when their guard is removed.
Co-Authored-By: Paperclip <noreply@paperclip.ing>
This commit is contained in:
DottaandPaperclip authored and GitHub committed
1 parent
cc96c9f832
commit
aa04c5c2a7
5 files changed
+264
-139
No files matched your search
@@ -7,6 +7,7 @@ import os from "node:os";
|
||||
import { and, eq, inArray, sql } from "drizzle-orm";
|
||||
import { heartbeatRuns, issues, issueRecoveryActions, nativeRunFinalizations, workspaceOperations, withDedicatedDbConnection, type Db } from "@paperclipai/db";
|
||||
import { currentNativeWorkspaceRecoveryResult } from "./native-workspace-recovery-binding.js";
|
||||
import { withWorkspaceLockKeepalive } from "./native-workspace-lock-keepalive.js";
|
||||
import {
|
||||
currentNativeControllerIdentity,
|
||||
evaluateNativeControllerTakeover,
|
||||
@@ -71,149 +72,152 @@ export async function withNativeWorkspaceFinalizationOwnership<T>(
|
||||
let actionPromise: Promise<T> | undefined;
|
||||
try {
|
||||
return await withDedicatedDbConnection(input.db, (dedicated) => dedicated.transaction(async (lock) => {
|
||||
const rows = await lock.execute(sql`select pg_try_advisory_xact_lock(hashtextextended(${`native-workspace-finalization:${input.companyId}:${input.runId}`}, 0)) as acquired`);
|
||||
const rows = await lock.execute(sql`select pg_try_advisory_xact_lock(hashtextextended(${`native-workspace-finalization:${input.companyId}:${input.runId}`}, 0)) as acquired,
|
||||
(select setting::bigint from pg_settings where name = 'idle_in_transaction_session_timeout') as idle_timeout_ms`);
|
||||
if (rows[0]?.acquired !== true) return { acquired: false } as const;
|
||||
// Reads and durable progress use the ordinary pool, not the lock transaction.
|
||||
const [run] = await input.db.select({ profile: heartbeatRuns.runnerProfileJson, issueId: heartbeatRuns.nativeIssueId, agentId: heartbeatRuns.agentId }).from(heartbeatRuns).where(scope).limit(1);
|
||||
if (!run) throw new Error("native_workspace_finalization_binding_missing");
|
||||
const rawPrior = run.profile?.[OWNER_KEY];
|
||||
const prior = readOwner(rawPrior);
|
||||
let automaticResultId: string | undefined;
|
||||
const currentRecoveryResult = async (db: Db) => {
|
||||
if (!run.issueId) return null;
|
||||
// Use the arbiter's lock order and release row locks before copyback I/O.
|
||||
const [coordinator] = await db.select().from(nativeRunFinalizations).where(and(
|
||||
eq(nativeRunFinalizations.companyId, input.companyId), eq(nativeRunFinalizations.runId, input.runId),
|
||||
eq(nativeRunFinalizations.issueId, run.issueId),
|
||||
)).for("update");
|
||||
const [issue] = await db.select().from(issues).where(and(
|
||||
eq(issues.companyId, input.companyId), eq(issues.id, run.issueId),
|
||||
)).for("update");
|
||||
const [currentRun] = await db.select().from(heartbeatRuns).where(scope).for("update");
|
||||
return coordinator && issue && currentRun ? currentNativeWorkspaceRecoveryResult(db, {
|
||||
companyId: input.companyId, coordinator, issue, run: currentRun,
|
||||
}) : null;
|
||||
};
|
||||
let platformProof: WorkspaceWorkerStopProof | null = null;
|
||||
if (rawPrior != null) {
|
||||
const requireStopVerification = async () => {
|
||||
const fingerprint = nativeSha256({ runId: input.runId, owner: rawPrior });
|
||||
const nextAction = "Verify the previous controller and its workspace-copyback processes have stopped, then release only this exact workspace owner receipt using doc/native-workspace-finalization-recovery.md. Resume workspace finalization without another provider turn.";
|
||||
const recovery = issueRecoveryActionService(input.db);
|
||||
const existing = run.issueId ? await recovery.getActiveForIssue(input.companyId, run.issueId) : null;
|
||||
// Rechecking the same physical owner is not a new recovery attempt.
|
||||
if (existing?.cause === NATIVE_WORKSPACE_OWNER_UNVERIFIED && existing.fingerprint === fingerprint) {
|
||||
if (!existing.evidence?.workspaceOwnerStop) {
|
||||
await input.db.update(issueRecoveryActions).set({ nextAction, updatedAt: new Date(),
|
||||
evidence: sql`jsonb_set(coalesce(${issueRecoveryActions.evidence}, '{}'::jsonb), '{workspaceOwnerRecovery}', '{"state":"operator_required"}'::jsonb)`,
|
||||
}).where(and(
|
||||
eq(issueRecoveryActions.companyId, input.companyId), eq(issueRecoveryActions.id, existing.id),
|
||||
eq(issueRecoveryActions.fingerprint, fingerprint), eq(issueRecoveryActions.cause, NATIVE_WORKSPACE_OWNER_UNVERIFIED),
|
||||
inArray(issueRecoveryActions.status, ["active", "escalated"]),
|
||||
sql`${issueRecoveryActions.evidence}->'workspaceOwnerStop' is null`,
|
||||
sql`(${issueRecoveryActions.nextAction} is distinct from ${nextAction} or ${issueRecoveryActions.evidence}#>>'{workspaceOwnerRecovery,state}' is distinct from 'operator_required')`,
|
||||
));
|
||||
}
|
||||
return { acquired: false } as const;
|
||||
}
|
||||
if (run.issueId) await recovery.upsertSourceScoped({
|
||||
companyId: input.companyId, sourceIssueId: run.issueId,
|
||||
kind: "active_run_watchdog", ownerType: "board", returnOwnerAgentId: run.agentId,
|
||||
cause: NATIVE_WORKSPACE_OWNER_UNVERIFIED,
|
||||
fingerprint,
|
||||
evidence: { runId: input.runId, owner: rawPrior, workspaceOwnerRecovery: { state: "operator_required" } },
|
||||
nextAction,
|
||||
wakePolicy: null, maxAttempts: 1, supersedeOnIdentityChange: true,
|
||||
});
|
||||
return { acquired: false } as const;
|
||||
return withWorkspaceLockKeepalive(lock as unknown as Db, Number(rows[0].idle_timeout_ms), async () => {
|
||||
// Reads and durable progress use the ordinary pool, not the lock transaction.
|
||||
const [run] = await input.db.select({ profile: heartbeatRuns.runnerProfileJson, issueId: heartbeatRuns.nativeIssueId, agentId: heartbeatRuns.agentId }).from(heartbeatRuns).where(scope).limit(1);
|
||||
if (!run) throw new Error("native_workspace_finalization_binding_missing");
|
||||
const rawPrior = run.profile?.[OWNER_KEY];
|
||||
const prior = readOwner(rawPrior);
|
||||
let automaticResultId: string | undefined;
|
||||
const currentRecoveryResult = async (db: Db) => {
|
||||
if (!run.issueId) return null;
|
||||
// Use the arbiter's lock order and release row locks before copyback I/O.
|
||||
const [coordinator] = await db.select().from(nativeRunFinalizations).where(and(
|
||||
eq(nativeRunFinalizations.companyId, input.companyId), eq(nativeRunFinalizations.runId, input.runId),
|
||||
eq(nativeRunFinalizations.issueId, run.issueId),
|
||||
)).for("update");
|
||||
const [issue] = await db.select().from(issues).where(and(
|
||||
eq(issues.companyId, input.companyId), eq(issues.id, run.issueId),
|
||||
)).for("update");
|
||||
const [currentRun] = await db.select().from(heartbeatRuns).where(scope).for("update");
|
||||
return coordinator && issue && currentRun ? currentNativeWorkspaceRecoveryResult(db, {
|
||||
companyId: input.companyId, coordinator, issue, run: currentRun,
|
||||
}) : null;
|
||||
};
|
||||
// The harness checks the exact owned deployment, including orphaned children.
|
||||
// A live deployment, unsigned answer or older receipt stays fenced.
|
||||
if (prior?.worker) platformProof = await requestWorkspaceWorkerStopProof({ runId: input.runId, ownerToken: prior.token, worker: prior.worker });
|
||||
if (!prior) return requireStopVerification();
|
||||
if (!platformProof && prior.hostname !== owner.hostname) return requireStopVerification();
|
||||
const knownJoinedHere = prior.controllerBootId === identity.bootId
|
||||
&& prior.pid === identity.pid && joinedOwnerTokens.has(prior.token)
|
||||
&& !activeOwnerTokens.has(prior.token);
|
||||
if (!platformProof && !knownJoinedHere) {
|
||||
const takeover = await evaluateNativeControllerTakeover({
|
||||
owner: { leaseOwner: prior.token, leaseExpiresAt: new Date(0), controllerPid: prior.pid,
|
||||
controllerProcessStartedAt: new Date(prior.processStartedAt) }, now: new Date(),
|
||||
});
|
||||
if (!takeover.allowed) return { acquired: false } as const;
|
||||
// A dead parent may have orphaned tar/Git children. Only a durable
|
||||
// completed-copyback barrier proves they joined; otherwise board proof
|
||||
// is required even on the same host. Both callers reuse this barrier.
|
||||
const [completedCopyback] = await input.db.select({ id: workspaceOperations.id }).from(workspaceOperations).where(and(
|
||||
eq(workspaceOperations.companyId, input.companyId),
|
||||
eq(workspaceOperations.heartbeatRunId, input.runId),
|
||||
eq(workspaceOperations.issueId, run.issueId!),
|
||||
eq(workspaceOperations.phase, "workspace_finalize"),
|
||||
eq(workspaceOperations.status, "succeeded"),
|
||||
)).limit(1);
|
||||
if (!completedCopyback) return requireStopVerification();
|
||||
}
|
||||
}
|
||||
// Claim before the potentially ambiguous write so finally can remove
|
||||
// only our token even if the server committed but its response was lost.
|
||||
const claimOwner = async (db: Db) => {
|
||||
activeOwnerTokens.add(owner.token);
|
||||
claimed = true;
|
||||
const [claimedRun] = await db.update(heartbeatRuns).set({
|
||||
runnerProfileJson: sql`jsonb_set(coalesce(${heartbeatRuns.runnerProfileJson}, '{}'::jsonb), array[${OWNER_KEY}], ${JSON.stringify(owner)}::jsonb) || ${JSON.stringify(platformProof ? { nativeWorkspaceOwnerStopProof: platformProof } : {})}::jsonb`,
|
||||
}).where(and(scope, prior
|
||||
? sql`${heartbeatRuns.runnerProfileJson}->${OWNER_KEY} = ${JSON.stringify(rawPrior)}::jsonb`
|
||||
: sql`${heartbeatRuns.runnerProfileJson}->>${OWNER_KEY} is null`
|
||||
)).returning({ id: heartbeatRuns.id });
|
||||
return claimedRun;
|
||||
};
|
||||
// A stop receipt proves the old worker exited, not that its task is still
|
||||
// current. Revalidate after the provider read and atomically with the CAS.
|
||||
const claimedRun = platformProof ? await input.db.transaction(async tx => {
|
||||
const result = await currentRecoveryResult(tx as unknown as Db);
|
||||
if (!result) return null;
|
||||
const claimedRun = await claimOwner(tx as unknown as Db);
|
||||
if (claimedRun) automaticResultId = result.id;
|
||||
return claimedRun;
|
||||
}) : await claimOwner(input.db);
|
||||
if (!claimedRun) return { acquired: false } as const;
|
||||
if (platformProof && run.issueId) await logActivity(input.db, {
|
||||
companyId: input.companyId, actorType: "system", actorId: "workspace-recovery",
|
||||
action: "issue.workspace_owner_stop_verified", entityType: "issue", entityId: run.issueId,
|
||||
runId: input.runId, details: { proof: platformProof },
|
||||
});
|
||||
if (prior) joinedOwnerTokens.delete(prior.token);
|
||||
const ownership: NativeWorkspaceFinalizationOwnership = {
|
||||
token: owner.token,
|
||||
assertHeld: async () => {
|
||||
try {
|
||||
// A failed transaction is never transparently reconnected by the DB client.
|
||||
await lock.execute(sql`select 1`);
|
||||
const [current] = await input.db.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(owned).limit(1);
|
||||
if (!current) throw new Error("owner_replaced");
|
||||
if (automaticResultId) {
|
||||
const result = await input.db.transaction(tx => currentRecoveryResult(tx as unknown as Db));
|
||||
if (result?.id !== automaticResultId) throw new Error("task_replaced");
|
||||
let platformProof: WorkspaceWorkerStopProof | null = null;
|
||||
if (rawPrior != null) {
|
||||
const requireStopVerification = async () => {
|
||||
const fingerprint = nativeSha256({ runId: input.runId, owner: rawPrior });
|
||||
const nextAction = "Verify the previous controller and its workspace-copyback processes have stopped, then release only this exact workspace owner receipt using doc/native-workspace-finalization-recovery.md. Resume workspace finalization without another provider turn.";
|
||||
const recovery = issueRecoveryActionService(input.db);
|
||||
const existing = run.issueId ? await recovery.getActiveForIssue(input.companyId, run.issueId) : null;
|
||||
// Rechecking the same physical owner is not a new recovery attempt.
|
||||
if (existing?.cause === NATIVE_WORKSPACE_OWNER_UNVERIFIED && existing.fingerprint === fingerprint) {
|
||||
if (!existing.evidence?.workspaceOwnerStop) {
|
||||
await input.db.update(issueRecoveryActions).set({ nextAction, updatedAt: new Date(),
|
||||
evidence: sql`jsonb_set(coalesce(${issueRecoveryActions.evidence}, '{}'::jsonb), '{workspaceOwnerRecovery}', '{"state":"operator_required"}'::jsonb)`,
|
||||
}).where(and(
|
||||
eq(issueRecoveryActions.companyId, input.companyId), eq(issueRecoveryActions.id, existing.id),
|
||||
eq(issueRecoveryActions.fingerprint, fingerprint), eq(issueRecoveryActions.cause, NATIVE_WORKSPACE_OWNER_UNVERIFIED),
|
||||
inArray(issueRecoveryActions.status, ["active", "escalated"]),
|
||||
sql`${issueRecoveryActions.evidence}->'workspaceOwnerStop' is null`,
|
||||
sql`(${issueRecoveryActions.nextAction} is distinct from ${nextAction} or ${issueRecoveryActions.evidence}#>>'{workspaceOwnerRecovery,state}' is distinct from 'operator_required')`,
|
||||
));
|
||||
}
|
||||
return { acquired: false } as const;
|
||||
}
|
||||
} catch (cause) {
|
||||
throw new NativeWorkspaceFinalizationOwnershipLostError({ cause });
|
||||
if (run.issueId) await recovery.upsertSourceScoped({
|
||||
companyId: input.companyId, sourceIssueId: run.issueId,
|
||||
kind: "active_run_watchdog", ownerType: "board", returnOwnerAgentId: run.agentId,
|
||||
cause: NATIVE_WORKSPACE_OWNER_UNVERIFIED,
|
||||
fingerprint,
|
||||
evidence: { runId: input.runId, owner: rawPrior, workspaceOwnerRecovery: { state: "operator_required" } },
|
||||
nextAction,
|
||||
wakePolicy: null, maxAttempts: 1, supersedeOnIdentityChange: true,
|
||||
});
|
||||
return { acquired: false } as const;
|
||||
};
|
||||
// The harness checks the exact owned deployment, including orphaned children.
|
||||
// A live deployment, unsigned answer or older receipt stays fenced.
|
||||
if (prior?.worker) platformProof = await requestWorkspaceWorkerStopProof({ runId: input.runId, ownerToken: prior.token, worker: prior.worker });
|
||||
if (!prior) return requireStopVerification();
|
||||
if (!platformProof && prior.hostname !== owner.hostname) return requireStopVerification();
|
||||
const knownJoinedHere = prior.controllerBootId === identity.bootId
|
||||
&& prior.pid === identity.pid && joinedOwnerTokens.has(prior.token)
|
||||
&& !activeOwnerTokens.has(prior.token);
|
||||
if (!platformProof && !knownJoinedHere) {
|
||||
const takeover = await evaluateNativeControllerTakeover({
|
||||
owner: { leaseOwner: prior.token, leaseExpiresAt: new Date(0), controllerPid: prior.pid,
|
||||
controllerProcessStartedAt: new Date(prior.processStartedAt) }, now: new Date(),
|
||||
});
|
||||
if (!takeover.allowed) return { acquired: false } as const;
|
||||
// A dead parent may have orphaned tar/Git children. Only a durable
|
||||
// completed-copyback barrier proves they joined; otherwise board proof
|
||||
// is required even on the same host. Both callers reuse this barrier.
|
||||
const [completedCopyback] = await input.db.select({ id: workspaceOperations.id }).from(workspaceOperations).where(and(
|
||||
eq(workspaceOperations.companyId, input.companyId),
|
||||
eq(workspaceOperations.heartbeatRunId, input.runId),
|
||||
eq(workspaceOperations.issueId, run.issueId!),
|
||||
eq(workspaceOperations.phase, "workspace_finalize"),
|
||||
eq(workspaceOperations.status, "succeeded"),
|
||||
)).limit(1);
|
||||
if (!completedCopyback) return requireStopVerification();
|
||||
}
|
||||
},
|
||||
};
|
||||
await ownership.assertHeld();
|
||||
if (platformProof && run.issueId) await input.db.update(issueRecoveryActions).set({
|
||||
evidence: sql`jsonb_set(coalesce(${issueRecoveryActions.evidence}, '{}'::jsonb), '{workspaceOwnerRecovery}', '{"state":"resuming"}'::jsonb)`,
|
||||
updatedAt: new Date(),
|
||||
}).where(and(
|
||||
eq(issueRecoveryActions.companyId, input.companyId), eq(issueRecoveryActions.sourceIssueId, run.issueId),
|
||||
eq(issueRecoveryActions.cause, NATIVE_WORKSPACE_OWNER_UNVERIFIED),
|
||||
eq(issueRecoveryActions.fingerprint, nativeSha256({ runId: input.runId, owner: rawPrior })),
|
||||
inArray(issueRecoveryActions.status, ["active", "escalated"]),
|
||||
));
|
||||
actionPromise = action(ownership);
|
||||
const value = await actionPromise;
|
||||
await ownership.assertHeld();
|
||||
return { acquired: true, value } as const;
|
||||
}
|
||||
// Claim before the potentially ambiguous write so finally can remove
|
||||
// only our token even if the server committed but its response was lost.
|
||||
const claimOwner = async (db: Db) => {
|
||||
activeOwnerTokens.add(owner.token);
|
||||
claimed = true;
|
||||
const [claimedRun] = await db.update(heartbeatRuns).set({
|
||||
runnerProfileJson: sql`jsonb_set(coalesce(${heartbeatRuns.runnerProfileJson}, '{}'::jsonb), array[${OWNER_KEY}], ${JSON.stringify(owner)}::jsonb) || ${JSON.stringify(platformProof ? { nativeWorkspaceOwnerStopProof: platformProof } : {})}::jsonb`,
|
||||
}).where(and(scope, prior
|
||||
? sql`${heartbeatRuns.runnerProfileJson}->${OWNER_KEY} = ${JSON.stringify(rawPrior)}::jsonb`
|
||||
: sql`${heartbeatRuns.runnerProfileJson}->>${OWNER_KEY} is null`
|
||||
)).returning({ id: heartbeatRuns.id });
|
||||
return claimedRun;
|
||||
};
|
||||
// A stop receipt proves the old worker exited, not that its task is still
|
||||
// current. Revalidate after the provider read and atomically with the CAS.
|
||||
const claimedRun = platformProof ? await input.db.transaction(async tx => {
|
||||
const result = await currentRecoveryResult(tx as unknown as Db);
|
||||
if (!result) return null;
|
||||
const claimedRun = await claimOwner(tx as unknown as Db);
|
||||
if (claimedRun) automaticResultId = result.id;
|
||||
return claimedRun;
|
||||
}) : await claimOwner(input.db);
|
||||
if (!claimedRun) return { acquired: false } as const;
|
||||
if (platformProof && run.issueId) await logActivity(input.db, {
|
||||
companyId: input.companyId, actorType: "system", actorId: "workspace-recovery",
|
||||
action: "issue.workspace_owner_stop_verified", entityType: "issue", entityId: run.issueId,
|
||||
runId: input.runId, details: { proof: platformProof },
|
||||
});
|
||||
if (prior) joinedOwnerTokens.delete(prior.token);
|
||||
const ownership: NativeWorkspaceFinalizationOwnership = {
|
||||
token: owner.token,
|
||||
assertHeld: async () => {
|
||||
try {
|
||||
// A failed transaction is never transparently reconnected by the DB client.
|
||||
await lock.execute(sql`select 1`);
|
||||
const [current] = await input.db.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(owned).limit(1);
|
||||
if (!current) throw new Error("owner_replaced");
|
||||
if (automaticResultId) {
|
||||
const result = await input.db.transaction(tx => currentRecoveryResult(tx as unknown as Db));
|
||||
if (result?.id !== automaticResultId) throw new Error("task_replaced");
|
||||
}
|
||||
} catch (cause) {
|
||||
throw new NativeWorkspaceFinalizationOwnershipLostError({ cause });
|
||||
}
|
||||
},
|
||||
};
|
||||
await ownership.assertHeld();
|
||||
if (platformProof && run.issueId) await input.db.update(issueRecoveryActions).set({
|
||||
evidence: sql`jsonb_set(coalesce(${issueRecoveryActions.evidence}, '{}'::jsonb), '{workspaceOwnerRecovery}', '{"state":"resuming"}'::jsonb)`,
|
||||
updatedAt: new Date(),
|
||||
}).where(and(
|
||||
eq(issueRecoveryActions.companyId, input.companyId), eq(issueRecoveryActions.sourceIssueId, run.issueId),
|
||||
eq(issueRecoveryActions.cause, NATIVE_WORKSPACE_OWNER_UNVERIFIED),
|
||||
eq(issueRecoveryActions.fingerprint, nativeSha256({ runId: input.runId, owner: rawPrior })),
|
||||
inArray(issueRecoveryActions.status, ["active", "escalated"]),
|
||||
));
|
||||
actionPromise = action(ownership);
|
||||
const value = await actionPromise;
|
||||
await ownership.assertHeld();
|
||||
return { acquired: true, value } as const;
|
||||
});
|
||||
}));
|
||||
} finally {
|
||||
// postgres.js may reject the transaction when the socket dies before the
|
||||
|
||||
@@ -892,4 +892,21 @@ describe("native workspace finalization recovery", () => {
|
||||
}),
|
||||
]);
|
||||
});
|
||||
it("preserves a live copyback owner across the configured idle transaction timeout", async () => {
|
||||
const seeded = await seedRun({ executionWorkspaceId: randomUUID(), title: "Idle copyback lock" });
|
||||
const rows = await db.execute(sql`select current_user as name`);
|
||||
const role = String(rows[0].name);
|
||||
await db.execute(sql`alter role ${sql.identifier(role)} set idle_in_transaction_session_timeout = '1s'`);
|
||||
try {
|
||||
const result = await withNativeWorkspaceFinalizationOwnership({ db, companyId, runId: seeded.runId }, async ownership => {
|
||||
await new Promise(resolve => setTimeout(resolve, 3200));
|
||||
await ownership.assertHeld();
|
||||
return "copied";
|
||||
});
|
||||
expect(result).toEqual({ acquired: true, value: "copied" });
|
||||
} finally {
|
||||
await db.execute(sql`alter role ${sql.identifier(role)} reset idle_in_transaction_session_timeout`);
|
||||
}
|
||||
});
|
||||
|
||||
});
|
||||
@@ -0,0 +1,71 @@
|
||||
import { setImmediate } from "node:timers/promises";
|
||||
import * as timers from "node:timers/promises";
|
||||
import { expect, it, vi } from "vitest";
|
||||
import type { Db } from "@paperclipai/db";
|
||||
import { withWorkspaceLockKeepalive } from "./native-workspace-lock-keepalive.js";
|
||||
|
||||
vi.mock("node:timers/promises", async importOriginal => {
|
||||
const actual = await importOriginal<typeof import("node:timers/promises")>();
|
||||
return { ...actual, setTimeout: vi.fn(actual.setTimeout) };
|
||||
});
|
||||
|
||||
it.each([0, -1, NaN, Infinity])("does not probe when the idle timeout is %s", async timeout => {
|
||||
const execute = vi.fn();
|
||||
const value = {};
|
||||
let finish!: (value: object) => void;
|
||||
const pending = new Promise<object>(resolve => { finish = resolve; });
|
||||
vi.useFakeTimers();
|
||||
const sleep = vi.mocked(timers.setTimeout);
|
||||
const originalSleep = sleep.getMockImplementation()!;
|
||||
sleep.mockClear().mockImplementation((_delay, result, options) => new Promise((resolve, reject) => {
|
||||
const timer = setTimeout(() => resolve(result), 30_000);
|
||||
options?.signal?.addEventListener("abort", () => {
|
||||
clearTimeout(timer);
|
||||
reject(new Error("aborted"));
|
||||
}, { once: true });
|
||||
}));
|
||||
const result = withWorkspaceLockKeepalive({ execute } as unknown as Db, timeout, () => pending);
|
||||
try {
|
||||
await vi.advanceTimersByTimeAsync(60_001);
|
||||
expect(execute).not.toHaveBeenCalled();
|
||||
expect(sleep).not.toHaveBeenCalled();
|
||||
} finally {
|
||||
finish(value);
|
||||
await expect(result).resolves.toBe(value);
|
||||
sleep.mockImplementation(originalSleep);
|
||||
vi.useRealTimers();
|
||||
}
|
||||
});
|
||||
|
||||
it("joins an in-flight probe before releasing the lock scope", async () => {
|
||||
let entered!: () => void, finish!: () => void;
|
||||
const started = new Promise<void>(resolve => { entered = resolve; });
|
||||
const pending = new Promise<void>(resolve => { finish = resolve; });
|
||||
const execute = vi.fn(async () => { entered(); await pending; });
|
||||
let settled = false;
|
||||
const result = withWorkspaceLockKeepalive({ execute } as unknown as Db, 3, async () => {
|
||||
await started;
|
||||
return 42;
|
||||
}).then(value => { settled = true; return value; });
|
||||
try {
|
||||
await started;
|
||||
await setImmediate();
|
||||
expect(settled).toBe(false);
|
||||
expect(execute).toHaveBeenCalledOnce();
|
||||
} finally { finish(); }
|
||||
await expect(result).resolves.toBe(42);
|
||||
await setImmediate();
|
||||
expect(execute).toHaveBeenCalledOnce();
|
||||
});
|
||||
|
||||
it("preserves the work error when its keepalive query fails", async () => {
|
||||
let entered!: () => void;
|
||||
const started = new Promise<void>(resolve => { entered = resolve; });
|
||||
const execute = vi.fn(async () => { entered(); throw new Error("connection lost"); });
|
||||
const original = new Error("copyback failed");
|
||||
const work = vi.fn(async () => { await started; throw original; });
|
||||
await expect(withWorkspaceLockKeepalive({ execute } as unknown as Db, 3, work)).rejects.toBe(original);
|
||||
expect(work).toHaveBeenCalledOnce();
|
||||
await setImmediate();
|
||||
expect(execute).toHaveBeenCalledOnce();
|
||||
});
|
||||
@@ -0,0 +1,25 @@
|
||||
import { setTimeout } from "node:timers/promises";
|
||||
import { sql } from "drizzle-orm";
|
||||
import type { Db } from "@paperclipai/db";
|
||||
|
||||
export async function withWorkspaceLockKeepalive<T>(
|
||||
lock: Pick<Db, "execute">,
|
||||
idleTimeoutMs: number,
|
||||
work: () => Promise<T>,
|
||||
): Promise<T> {
|
||||
if (!Number.isFinite(idleTimeoutMs) || idleTimeoutMs <= 0) return work();
|
||||
const stop = new AbortController();
|
||||
const interval = Math.max(1, Math.min(30_000, Math.floor(idleTimeoutMs / 3)));
|
||||
const keeper = (async () => {
|
||||
while (!stop.signal.aborted) {
|
||||
await setTimeout(interval, undefined, { signal: stop.signal, ref: false });
|
||||
await lock.execute(sql`select 1`);
|
||||
}
|
||||
})().catch(() => undefined);
|
||||
try {
|
||||
return await work();
|
||||
} finally {
|
||||
stop.abort();
|
||||
await keeper;
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user