diff --git a/src/quota-pool/identity.ts b/src/quota-pool/identity.ts index 32b397cf2a5..ea838ea652d 100644 --- a/src/quota-pool/identity.ts +++ b/src/quota-pool/identity.ts @@ -14,8 +14,8 @@ * either. Deleting the file rotates `homeId` and every `accountRef` together. */ import { createHmac, randomBytes, randomUUID } from "node:crypto"; -import { closeSync, constants as fsConstants, linkSync, lstatSync, openSync, readFileSync, unlinkSync, writeSync } from "node:fs"; -import { join } from "node:path"; +import { closeSync, constants as fsConstants, fsyncSync, linkSync, lstatSync, openSync, readFileSync, unlinkSync, writeSync } from "node:fs"; +import { dirname, join } from "node:path"; import { getConfigDir } from "../config/paths"; export const QUOTA_POOL_IDENTITY_FILE = "quota-pool-identity.json"; @@ -61,34 +61,101 @@ export function readQuotaPoolIdentity(configDir = getConfigDir()): QuotaPoolIden } const NOFOLLOW = process.platform === "win32" ? 0 : fsConstants.O_NOFOLLOW; +/** How long a loser waits for a concurrent winner's file to become readable before giving up. */ +const PUBLISHED_READ_ATTEMPTS = 50; +const PUBLISHED_READ_INTERVAL_MS = 10; -/** Write a whole new owner-only file; fails if anything (a symlink included) already exists. */ +/** Test seams for the publication windows; production never sets them. */ +export interface QuotaPoolIdentityTestHooks { + /** After the first read found no identity, before the existence check. */ + afterAbsentRead?: (path: string) => void; + /** After a file was exclusively created, before its bytes are written. */ + afterCreate?: (path: string) => void; + /** Replace `linkSync` (to force the in-place fallback). */ + link?: (existing: string, target: string) => void; + /** Override the bounded wait for a concurrent winner. */ + publishedReadAttempts?: number; +} +let hooks: QuotaPoolIdentityTestHooks = {}; + +export function setQuotaPoolIdentityHooksForTests(next: QuotaPoolIdentityTestHooks = {}): void { + hooks = next; +} + +/** Write a whole new owner-only file, flushed to disk; fails if anything (a symlink included) already exists. */ function createExclusive(path: string, content: string): void { const fd = openSync(path, fsConstants.O_WRONLY | fsConstants.O_CREAT | fsConstants.O_EXCL | NOFOLLOW, 0o600); - try { writeSync(fd, content); } finally { closeSync(fd); } + try { + hooks.afterCreate?.(path); + writeSync(fd, content); + fsyncSync(fd); + } finally { closeSync(fd); } } function errorCode(error: unknown): string | undefined { return (error as NodeJS.ErrnoException | null)?.code; } +/** + * Make the new directory entry durable too: fsyncing the file alone can lose the entry on power + * loss, and a lost entry would mint a new identity and rotate every accountRef on the next start. + * Best effort; Windows cannot open a directory for fsync. + */ +function syncParentDirectory(path: string): void { + if (process.platform === "win32") return; + try { + const fd = openSync(dirname(path), fsConstants.O_RDONLY); + try { fsyncSync(fd); } finally { closeSync(fd); } + } catch { /* the file itself is flushed; a directory fsync is an extra guarantee */ } +} + +function pause(ms: number): void { + Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, ms); +} + +/** + * Something already sits at the path: adopt it once it reads as a valid identity. + * + * A concurrent winner's file is complete the moment it appears when it was hard-linked, but the + * in-place fallback creates the file before writing it, so a loser can briefly see it empty. The + * wait is bounded; an invalid result is never returned or memoized. A symlink, or a file that stays + * invalid past the bound, is refused and left alone. + */ +function adoptPublishedIdentity(configDir: string, path: string): QuotaPoolIdentity { + const attempts = hooks.publishedReadAttempts ?? PUBLISHED_READ_ATTEMPTS; + for (let attempt = 0; attempt < attempts; attempt++) { + let stat; + try { stat = lstatSync(path); } catch (error) { + if (errorCode(error) !== "ENOENT") throw error; + stat = null; + } + if (stat && (stat.isSymbolicLink() || !stat.isFile())) break; + const identity = readQuotaPoolIdentity(configDir); + if (identity) return identity; + if (attempt + 1 < attempts) pause(PUBLISHED_READ_INTERVAL_MS); + } + throw new Error("quota-pool identity file is not a valid regular file; remove it to rotate the identity"); +} + /** * Create the identity once. Called off the request path (server startup); idempotent. * - * Publication is exclusive and never follows a link: the secret is written to a private temp - * file and hard-linked into place, which fails if anything already sits at the path. Two servers - * starting on one fresh home therefore cannot each keep a different identity: the loser re-reads - * and adopts the winner's. A symlink or a malformed file at the path is refused and left alone - * (the read keeps reporting the identity as unavailable until an operator removes it); it is never - * overwritten. + * Publication is exclusive and never follows a link: the secret is written and flushed to a + * private temp file, then hard-linked into place, which fails if anything already sits at the path. + * Where hard links are unavailable it is created in place exclusively. Two servers starting on one + * fresh home therefore cannot each keep a different identity: a loser, including one that finds + * the winner's file appear between its own checks, re-reads within a bound and adopts the winner's. + * A symlink or a persistently malformed file at the path is refused and left alone (the read keeps + * reporting the identity as unavailable until an operator removes it); it is never overwritten. */ export function ensureQuotaPoolIdentity(configDir = getConfigDir()): QuotaPoolIdentity { const existing = readQuotaPoolIdentity(configDir); if (existing) return existing; const path = quotaPoolIdentityPath(configDir); + hooks.afterAbsentRead?.(path); let present = true; try { lstatSync(path); } catch (error) { if (errorCode(error) !== "ENOENT") throw error; present = false; } - if (present) throw new Error("quota-pool identity file is not a valid regular file; remove it to rotate the identity"); + if (present) return adoptPublishedIdentity(configDir, path); const identity: QuotaPoolIdentity = { homeId: `home_${randomUUID().replaceAll("-", "")}`, secret: randomBytes(32).toString("hex"), @@ -99,15 +166,11 @@ export function ensureQuotaPoolIdentity(configDir = getConfigDir()): QuotaPoolId try { createExclusive(temp, content); try { - linkSync(temp, path); + (hooks.link ?? linkSync)(temp, path); published = true; } catch (error) { const code = errorCode(error); - if (code === "EEXIST") { - const winner = readQuotaPoolIdentity(configDir); - if (winner) return winner; - throw new Error("quota-pool identity was created concurrently but is not readable yet"); - } + if (code === "EEXIST") return adoptPublishedIdentity(configDir, path); // No hard links on this filesystem: exclusive create in place is still never a replacement. if (code !== "EPERM" && code !== "ENOTSUP" && code !== "EXDEV" && code !== "ENOSYS") throw error; } @@ -119,11 +182,10 @@ export function ensureQuotaPoolIdentity(configDir = getConfigDir()): QuotaPoolId createExclusive(path, content); } catch (error) { if (errorCode(error) !== "EEXIST") throw error; - const winner = readQuotaPoolIdentity(configDir); - if (winner) return winner; - throw error; + return adoptPublishedIdentity(configDir, path); } } + syncParentDirectory(path); memo = { path, identity }; return identity; } diff --git a/structure/quota-pool.md b/structure/quota-pool.md index 3f39547e1f1..58ab4e6a206 100644 --- a/structure/quota-pool.md +++ b/structure/quota-pool.md @@ -65,10 +65,12 @@ the pooled windows. `homeId` (`home_` + 32 hex, the Home's stable instance id, not a hostname) and a 256-bit secret. The server creates it in `acquireServerBackgroundLifecycle` before it serves; the read only reads it and answers 503 `quota_pool_identity_unavailable` instead of creating it. Creation is exclusive and -never follows a link: the secret is written to a private temp file and hard-linked into place (an -in-place exclusive create where hard links are unavailable), so two servers starting on one fresh -home adopt the same identity, and a symlink or malformed file at the path is refused and left -untouched rather than overwritten. +never follows a link: the secret is written and flushed to a private temp file and hard-linked into +place (an in-place exclusive create where hard links are unavailable), so two servers starting on +one fresh home adopt the same identity. A loser that finds a file at the path, including one that +appears between its own absence check and existence check, re-reads it for up to 500 ms and adopts +it once it is valid; an empty or partial file is never returned or memoized. A symlink, or a file +that stays malformed past that bound, is refused and left untouched rather than overwritten. `accountRef` is `acct_` + base64url of the first 128 bits of HMAC-SHA256(secret, `ocx.quota-pool/1` NUL provider NUL accountId). It is stable across restarts, distinct per provider, diff --git a/tests/usage/quota-pool.test.ts b/tests/usage/quota-pool.test.ts index 7d1f4333691..29a13b02baf 100644 --- a/tests/usage/quota-pool.test.ts +++ b/tests/usage/quota-pool.test.ts @@ -4,12 +4,18 @@ * Codex header parser, and the opaque, stable `accountRef`. */ import { afterEach, beforeEach, describe, expect, test } from "bun:test"; -import { lstatSync, mkdtempSync, readFileSync, symlinkSync, unlinkSync, writeFileSync } from "node:fs"; +import { lstatSync, mkdtempSync, readdirSync, readFileSync, symlinkSync, unlinkSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import type { OcxConfig } from "../../src/types"; import { codexQuotaWindows, providerQuotaWindows, readQuotaPool } from "../../src/quota-pool/assemble"; -import { ensureQuotaPoolIdentity, quotaPoolIdentityPath, readQuotaPoolIdentity, resetQuotaPoolIdentityMemoForTests } from "../../src/quota-pool/identity"; +import { + ensureQuotaPoolIdentity, + quotaPoolIdentityPath, + readQuotaPoolIdentity, + resetQuotaPoolIdentityMemoForTests, + setQuotaPoolIdentityHooksForTests, +} from "../../src/quota-pool/identity"; import { parseUpstreamQuotaHeaders } from "../../src/codex/quota"; import { getAccountSet, saveCredential } from "../../src/oauth/store"; import { clearAccountQuotaCache, setCachedProviderAccountQuotaForTests } from "../../src/providers/quota"; @@ -65,6 +71,7 @@ beforeEach(async () => { }); afterEach(() => { + setQuotaPoolIdentityHooksForTests(); setProviderQuotaReportCache(null); clearAccountQuotaCache("xai"); resetQuotaPoolIdentityMemoForTests(); @@ -176,6 +183,57 @@ describe("quota-pool identity creation", () => { expect(readQuotaPoolIdentity()).toBeNull(); }); + test("a winner published between the absence read and the existence check is adopted, not refused", () => { + const winner = { homeId: `home_${"a".repeat(32)}`, secret: "b".repeat(64) }; + setQuotaPoolIdentityHooksForTests({ + // Another server links its complete file in exactly inside this process's check window. + afterAbsentRead: path => writeFileSync(path, `${JSON.stringify({ version: 1, ...winner })}\n`, { mode: 0o600 }), + }); + expect(ensureQuotaPoolIdentity()).toEqual(winner); + expect(readFileSync(quotaPoolIdentityPath(), "utf8")).toContain(winner.homeId); + }); + + test("losing the hard-link race (EEXIST) adopts the winner's complete file", () => { + const winner = { homeId: `home_${"c".repeat(32)}`, secret: "d".repeat(64) }; + setQuotaPoolIdentityHooksForTests({ + // The winner links its complete file first; this process's own link then reports EEXIST. + link: (_temp, target) => { + writeFileSync(target, `${JSON.stringify({ version: 1, ...winner })}\n`, { mode: 0o600 }); + throw Object.assign(new Error("file exists"), { code: "EEXIST" }); + }, + }); + expect(ensureQuotaPoolIdentity()).toEqual(winner); + expect(readQuotaPoolIdentity()).toEqual(winner); + // The loser's private temp file is cleaned up. + expect(readdirSync(home).filter(name => name.startsWith("quota-pool-identity.json.") && name.endsWith(".tmp"))).toEqual([]); + }); + + test("an in-place create seen before its bytes land is never adopted or memoized", () => { + const seen: { read?: unknown; concurrent?: unknown; concurrentError?: string } = {}; + let writerPaused = false; + setQuotaPoolIdentityHooksForTests({ + // No hard links on this filesystem: publication falls back to exclusive create in place. + link: () => { throw Object.assign(new Error("hard links unsupported"), { code: "EPERM" }); }, + publishedReadAttempts: 2, + afterCreate: path => { + if (path !== quotaPoolIdentityPath() || writerPaused) return; + writerPaused = true; + // The file exists and is still empty: a reader and a concurrent loser look now. + seen.read = readQuotaPoolIdentity(); + try { seen.concurrent = ensureQuotaPoolIdentity(); } catch (error) { seen.concurrentError = (error as Error).message; } + }, + }); + const identity = ensureQuotaPoolIdentity(); + expect(writerPaused).toBe(true); + expect(seen.read).toBeNull(); + expect(seen.concurrent).toBeUndefined(); + expect(seen.concurrentError).toContain("not a valid regular file"); + expect(identity.homeId).toMatch(/^home_[0-9a-f]{32}$/); + // Nothing invalid was memoized: a fresh read returns the writer's complete identity. + resetQuotaPoolIdentityMemoForTests(); + expect(readQuotaPoolIdentity()).toEqual(identity); + }); + test("servers racing on a fresh home agree on one identity", async () => { const script = `const { ensureQuotaPoolIdentity } = await import(${JSON.stringify(join(import.meta.dir, "../../src/quota-pool/identity.ts"))}); await Bun.sleep(Math.max(0, Number(process.env.QP_START_AT) - Date.now())); @@ -186,10 +244,14 @@ console.log(ensureQuotaPoolIdentity().homeId);`; const runs = Array.from({ length: 4 }, () => Bun.spawn([process.execPath, "-e", script], { env: { ...process.env, OPENCODEX_HOME: dir, QP_START_AT: startAt }, stdout: "pipe", stderr: "pipe", })); - const ids = await Promise.all(runs.map(async run => (await new Response(run.stdout).text()).trim())); + // Read each child to completion (stdout, stderr and exit) so a loser that fails names its error. + const results = await Promise.all(runs.map(async run => { + const [out, err, code] = await Promise.all([new Response(run.stdout).text(), new Response(run.stderr).text(), run.exited]); + return { id: out.trim(), code, error: err.trim() }; + })); resetQuotaPoolIdentityMemoForTests(); const persisted = readQuotaPoolIdentity(dir)!.homeId; - expect(ids).toEqual([persisted, persisted, persisted, persisted]); + expect(results).toEqual(Array.from({ length: 4 }, () => ({ id: persisted, code: 0, error: "" }))); } }); });