From 57840c6f5230d0faea386491e228f34cc0a5a7b4 Mon Sep 17 00:00:00 2001 From: phamngocquy Date: Fri, 31 Jul 2026 14:09:22 +0800 Subject: [PATCH] queue(locks): register the actuation and contributor-cap locks in the shutdown held-lock registry Fixes #10021 --- src/queue/transient-locks.ts | 31 +++--- test/unit/transient-locks.test.ts | 157 +++++++++++++++++++++++++++++- 2 files changed, 174 insertions(+), 14 deletions(-) diff --git a/src/queue/transient-locks.ts b/src/queue/transient-locks.ts index f29678e468..7b709ba85e 100644 --- a/src/queue/transient-locks.ts +++ b/src/queue/transient-locks.ts @@ -20,6 +20,7 @@ // delete a later claimer's live lock (#2129/#2135). import { randomUUID } from "node:crypto"; +import { registerHeldLock, unregisterHeldLock } from "./held-lock-registry"; import { RetryableJobError } from "./retryable"; /** Result of a transient-lock claim attempt. `ownerToken` is the random value THIS call wrote when it actually @@ -212,11 +213,12 @@ export async function claimPrActuationLock( repoFullName: string, prNumber: number, ): Promise { - return claimTransientLock( - env, - prActuationLockKey(repoFullName, prNumber), - PR_ACTUATION_LOCK_TTL_SECONDS, - ); + const key = prActuationLockKey(repoFullName, prNumber); + const claim = await claimTransientLock(env, key, PR_ACTUATION_LOCK_TTL_SECONDS); + if (claim.acquired && claim.ownerToken !== null) { + registerHeldLock(key, claim.ownerToken, () => releaseTransientLockIfOwner(env, key, claim.ownerToken)); + } + return claim; } export async function releasePrActuationLock( env: Env, @@ -224,7 +226,9 @@ export async function releasePrActuationLock( prNumber: number, ownerToken: string | null, ): Promise { - await releaseTransientLockIfOwner(env, prActuationLockKey(repoFullName, prNumber), ownerToken); + const key = prActuationLockKey(repoFullName, prNumber); + await releaseTransientLockIfOwner(env, key, ownerToken); + if (ownerToken !== null) unregisterHeldLock(key, ownerToken); } // A plain thrown Error still reaches the queue's retry path (this call site is deliberately uncaught, same as @@ -266,11 +270,12 @@ export async function claimContributorCapLock( repoFullName: string, authorLogin: string, ): Promise { - return claimTransientLock( - env, - contributorCapLockKey(repoFullName, authorLogin), - CONTRIBUTOR_CAP_LOCK_TTL_SECONDS, - ); + const key = contributorCapLockKey(repoFullName, authorLogin); + const claim = await claimTransientLock(env, key, CONTRIBUTOR_CAP_LOCK_TTL_SECONDS); + if (claim.acquired && claim.ownerToken !== null) { + registerHeldLock(key, claim.ownerToken, () => releaseTransientLockIfOwner(env, key, claim.ownerToken)); + } + return claim; } export async function releaseContributorCapLock( env: Env, @@ -278,5 +283,7 @@ export async function releaseContributorCapLock( authorLogin: string, ownerToken: string | null, ): Promise { - await releaseTransientLockIfOwner(env, contributorCapLockKey(repoFullName, authorLogin), ownerToken); + const key = contributorCapLockKey(repoFullName, authorLogin); + await releaseTransientLockIfOwner(env, key, ownerToken); + if (ownerToken !== null) unregisterHeldLock(key, ownerToken); } diff --git a/test/unit/transient-locks.test.ts b/test/unit/transient-locks.test.ts index 6a11d73865..1735f73357 100644 --- a/test/unit/transient-locks.test.ts +++ b/test/unit/transient-locks.test.ts @@ -1,4 +1,5 @@ import { afterEach, describe, expect, it, vi } from "vitest"; +import { heldLockCountForTest, releaseAllHeldLocksAtShutdown } from "../../src/queue/held-lock-registry"; import { SubmissionLock } from "../../src/queue/submission-lock"; import { claimContributorCapLock, @@ -501,15 +502,18 @@ describe("domain wrappers + PrActuationLockContendedError (#8896)", () => { }); delete env.SUBMISSION_LOCK; - await claimPrActuationLock(env, "acme/widgets", 7); + const prClaim = await claimPrActuationLock(env, "acme/widgets", 7); const [actuationTtl] = ttlCalls; ttlCalls.length = 0; - await claimContributorCapLock(env, "acme/widgets", "alice"); + const capClaim = await claimContributorCapLock(env, "acme/widgets", "alice"); const [capTtl] = ttlCalls; expect(capTtl).toBe(actuationTtl); expect(capTtl).toBe(600); + + await releasePrActuationLock(env, "acme/widgets", 7, prClaim.ownerToken); + await releaseContributorCapLock(env, "acme/widgets", "alice", capClaim.ownerToken); }); it("builds a fast-retry contended error with a distinct retryKind", () => { @@ -841,3 +845,152 @@ describe("lock heartbeat end-to-end (#9467)", () => { expect((await claimTransientLock(env, "lock:pr", 60)).acquired).toBe(true); }); }); + +// #10021: register actuation and contributor-cap locks in the shutdown held-lock registry (same shape as +// claimAiReviewLock / releaseAiReviewLock in ai-review-orchestration.ts). +describe("claimPrActuationLock / releasePrActuationLock register with the held-lock registry (#10021)", () => { + afterEach(async () => { + await releaseAllHeldLocksAtShutdown(); + }); + + function cacheWithReleaseTracking() { + const releases: Array<{ key: string; value: string }> = []; + const held = new Map(); + const cache = { + get: async (key: string) => held.get(key) ?? null, + set: async () => undefined, + claim: async (key: string, value: string) => { + if (held.has(key)) return false; + held.set(key, value); + return true; + }, + releaseIfValue: async (key: string, value: string) => { + releases.push({ key, value }); + if (held.get(key) !== value) return false; + held.delete(key); + return true; + }, + }; + return { cache, releases }; + } + + it("registers a real claim, and shutdown release issues releaseIfValue for the pr-actuation-lock key", async () => { + const { cache, releases } = cacheWithReleaseTracking(); + const env = createTestEnv({ SELFHOST_TRANSIENT_CACHE: cache }); + delete env.SUBMISSION_LOCK; + + const before = heldLockCountForTest(); + const claim = await claimPrActuationLock(env, "acme/widgets", 7); + expect(claim.acquired).toBe(true); + expect(heldLockCountForTest()).toBe(before + 1); + expect(releases).toHaveLength(0); + + expect(await releaseAllHeldLocksAtShutdown()).toBe(before + 1); + expect(releases).toEqual([ + { key: "pr-actuation-lock:acme/widgets#7", value: claim.ownerToken }, + ]); + }); + + it("unregisters on the normal release path with a matching owner token", async () => { + const env = createTestEnv(); + const before = heldLockCountForTest(); + const claim = await claimPrActuationLock(env, "acme/widgets", 8); + await releasePrActuationLock(env, "acme/widgets", 8, claim.ownerToken); + expect(heldLockCountForTest()).toBe(before); + }); + + it("does not register a fail-open claim — there is nothing real to release", async () => { + const env = createTestEnv(); + delete (env as { SELFHOST_TRANSIENT_CACHE?: unknown }).SELFHOST_TRANSIENT_CACHE; + const before = heldLockCountForTest(); + const claim = await claimPrActuationLock(env, "acme/widgets", 10); + expect(claim.ownerToken).toBeNull(); + expect(heldLockCountForTest()).toBe(before); + }); + + it("a null-token release does not unregister a real held entry", async () => { + const env = createTestEnv(); + const before = heldLockCountForTest(); + await claimPrActuationLock(env, "acme/widgets", 11); + await releasePrActuationLock(env, "acme/widgets", 11, null); + expect(heldLockCountForTest()).toBe(before + 1); + }); + + it("#10021 releasePrActuationLock with a non-matching ownerToken does not evict the registry entry (#9468)", async () => { + const env = createTestEnv(); + const before = heldLockCountForTest(); + await claimPrActuationLock(env, "acme/widgets", 9); + expect(heldLockCountForTest()).toBe(before + 1); + await releasePrActuationLock(env, "acme/widgets", 9, "tok-someone-else"); + expect(heldLockCountForTest()).toBe(before + 1); + }); +}); + +describe("claimContributorCapLock / releaseContributorCapLock register with the held-lock registry (#10021)", () => { + afterEach(async () => { + await releaseAllHeldLocksAtShutdown(); + }); + + function cacheWithReleaseTracking() { + const releases: Array<{ key: string; value: string }> = []; + const held = new Map(); + const cache = { + get: async (key: string) => held.get(key) ?? null, + set: async () => undefined, + claim: async (key: string, value: string) => { + if (held.has(key)) return false; + held.set(key, value); + return true; + }, + releaseIfValue: async (key: string, value: string) => { + releases.push({ key, value }); + if (held.get(key) !== value) return false; + held.delete(key); + return true; + }, + }; + return { cache, releases }; + } + + it("registers a real claim, and shutdown release issues releaseIfValue for the contributor-cap-lock key", async () => { + const { cache, releases } = cacheWithReleaseTracking(); + const env = createTestEnv({ SELFHOST_TRANSIENT_CACHE: cache }); + delete env.SUBMISSION_LOCK; + + const before = heldLockCountForTest(); + const claim = await claimContributorCapLock(env, "Acme/Widgets", "Alice"); + expect(claim.acquired).toBe(true); + expect(heldLockCountForTest()).toBe(before + 1); + expect(releases).toHaveLength(0); + + expect(await releaseAllHeldLocksAtShutdown()).toBe(before + 1); + expect(releases).toEqual([ + { key: "contributor-cap-lock:acme/widgets:alice", value: claim.ownerToken }, + ]); + }); + + it("unregisters on the normal release path with a matching owner token", async () => { + const env = createTestEnv(); + const before = heldLockCountForTest(); + const claim = await claimContributorCapLock(env, "acme/widgets", "bob"); + await releaseContributorCapLock(env, "acme/widgets", "bob", claim.ownerToken); + expect(heldLockCountForTest()).toBe(before); + }); + + it("does not register a fail-open claim — there is nothing real to release", async () => { + const env = createTestEnv(); + delete (env as { SELFHOST_TRANSIENT_CACHE?: unknown }).SELFHOST_TRANSIENT_CACHE; + const before = heldLockCountForTest(); + const claim = await claimContributorCapLock(env, "acme/widgets", "dave"); + expect(claim.ownerToken).toBeNull(); + expect(heldLockCountForTest()).toBe(before); + }); + + it("a null-token release does not unregister a real held entry", async () => { + const env = createTestEnv(); + const before = heldLockCountForTest(); + await claimContributorCapLock(env, "acme/widgets", "carol"); + await releaseContributorCapLock(env, "acme/widgets", "carol", null); + expect(heldLockCountForTest()).toBe(before + 1); + }); +});