diff --git a/server/src/services/skill-version-cache-lock.test.ts b/server/src/services/skill-version-cache-lock.test.ts new file mode 100644 index 00000000000..fbf63ff141e --- /dev/null +++ b/server/src/services/skill-version-cache-lock.test.ts @@ -0,0 +1,75 @@ +import { randomUUID } from "node:crypto"; +import { setTimeout as delay } from "node:timers/promises"; +import { sql } from "drizzle-orm"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { createDb } from "@paperclipai/db"; +import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../__tests__/helpers/embedded-postgres.js"; +import { SkillCacheBusy, withSkillVersionCacheLock } from "./skill-version-cache.js"; + +const support = await getEmbeddedPostgresTestSupport(); +const describePostgres = support.supported ? describe : describe.skip; +if (!support.supported) console.warn(`Skill version cache lock tests require Postgres: ${support.reason}`); + +describePostgres("skill version cache lock during slow preparation", () => { + let database: Awaited>; + let db: ReturnType; + + beforeAll(async () => { + database = await startEmbeddedPostgresTestDatabase("skill-version-cache-lock-"); + db = createDb(database.connectionString, { maxConnections: 1 }); + 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'`); + }); + + afterAll(async () => { + try { + await db?.$client.end({ timeout: 1 }); + } finally { + await database?.cleanup(); + } + }); + + it("keeps the lock beyond the idle timeout while preparation uses the ordinary pool", async () => { + const companyId = randomUUID(); + let entered!: () => void, release!: () => void; + const started = new Promise(resolve => { entered = resolve; }); + const pending = new Promise(resolve => { release = resolve; }); + const preparation = withSkillVersionCacheLock(db, companyId, async () => { + await db.execute(sql`select 1`); + entered(); + await pending; + return "prepared"; + }).then(value => ({ value }), error => ({ error })); + + try { + // Surface a preparation that settles before entering instead of timing out. + let hasEntered = false; + await Promise.race([started.then(() => { hasEntered = true; }), preparation.then(result => { + if (hasEntered) return; + throw "error" in result ? result.error : new Error("Preparation finished before entering the lock"); + })]); + await delay(1400); + await expect(db.execute(sql`select 1 as usable`)).resolves.toMatchObject([{ usable: 1 }]); + await expect(withSkillVersionCacheLock(db, companyId, async () => "overlap", true)) + .rejects.toBeInstanceOf(SkillCacheBusy); + await expect(withSkillVersionCacheLock(db, randomUUID(), async () => "independent", true)) + .resolves.toBe("independent"); + } finally { + release(); + await preparation; + } + expect(await preparation).toEqual({ value: "prepared" }); + await expect(withSkillVersionCacheLock(db, companyId, async () => "next", true)).resolves.toBe("next"); + }); + + it("releases a failed preparation and preserves the original failure", async () => { + const companyId = randomUUID(); + const failure = new Error("skill preparation failed"); + await expect(withSkillVersionCacheLock(db, companyId, async () => { + await delay(1400); + throw failure; + })).rejects.toBe(failure); + await expect(withSkillVersionCacheLock(db, companyId, async () => "recovered", true)).resolves.toBe("recovered"); + }); +}); diff --git a/server/src/services/skill-version-cache.ts b/server/src/services/skill-version-cache.ts index ccb232eb100..f2f08ff774f 100644 --- a/server/src/services/skill-version-cache.ts +++ b/server/src/services/skill-version-cache.ts @@ -13,6 +13,7 @@ import type { CompanySkillVersionFileInventoryEntry } from "@paperclipai/shared" import { noteSkillVersionCacheChange } from "./skill-version-retention-notifications.js"; import { parseNativeExecutionInput } from "../vendor/paperclip-runner/index.js"; import { beginRunPhase } from "./execution-preparation-timing.js"; +import { withWorkspaceLockKeepalive } from "./native-runtime/native-workspace-lock-keepalive.js"; export const SKILL_VERSION_CACHE_IDLE_MS = 60 * 60 * 1000; const locks = new AsyncLocalStorage>(); @@ -32,7 +33,13 @@ export async function withSkillVersionCacheLock(db: Db, companyId: string, ac if (!rows[0]?.acquired) throw new SkillCacheBusy(); } else await tx.execute(sql`select pg_advisory_xact_lock(hashtextextended(${key}, 0))`); acquired(); - return locks.run([...held, { db, companyId }], action); + // Preparation can outlast idle_in_transaction_session_timeout while this + // lock-only transaction waits; losing the session would release the lock. + const settings = await tx.execute<{ idle_timeout_ms: string }>( + sql`select setting::bigint as idle_timeout_ms from pg_settings where name = 'idle_in_transaction_session_timeout'`, + ); + return withWorkspaceLockKeepalive(tx, Number(settings[0]?.idle_timeout_ms), + () => locks.run([...held, { db, companyId }], action)); })).catch(error => { acquired(true); throw error; @@ -207,4 +214,3 @@ export async function allocatedBytes(directory: string, shouldContinue: () => bo } return bytes; } -