feat(relay): per-cell switches for the rejection fence and the lost-reply margin (#27057)

* feat(relay): per-cell switches for the rejection fence and the lost-reply margin, and report ignored keys

rejectionFence (default on) turns the unhandled-rejection fence off without a roll.
readTimeoutMarginMs (1-60 s) replaces the boot margin per query through pg's query_timeout.
The applied log line and flagsApplied name keys this image does not know.

* feat(relay): runtime-status names the switches and values this image accepts

The flag tool checks a write against it, so it can refuse ticketCheck=enforce (or any key)
on an image that would void or drop it.
This commit is contained in:
Jinwoo Hong
2026-10-10 01:54:00 -04:00
committed by GitHub
parent df54bbb5ed
commit 9c41d9e8d8
10 changed files with 209 additions and 34 deletions
+4 -2
View File
@@ -40,7 +40,7 @@ import { AssignmentRejectionLogWindow } from './assignment-rejection-log-window.
import { CELL_ADMISSION_STATES } from './cell-admission-selector.js'
import { registerCellSeatFeedRoute } from './cell-seat-feed-route.js'
import type { CellSeatFeedPage } from './cell-seat-log.js'
import type { CellFlags } from './cell-flags.js'
import { type CellFlags, supportedCellFlags } from './cell-flags.js'
import type { AppliedControlFlags } from './relay-control-flag-channel.js'
import { RELAY_MAX_CELL_CAPACITY_REQUESTS, type RelayConfig } from './config.js'
import type { RelayCredentialStore } from './credential-store.js'
@@ -783,7 +783,9 @@ export function createRelayApp(
draining: operations.isDraining?.() ?? false,
regionalRehomeProtocol: config.rehomeAudience && config.rehomeDirectorServiceAccount ? 3 : 0,
// The flag workflow's read-back: applied switches, never the desired object.
...(operations.cellFlags ? { flagsApplied: operations.cellFlags() } : {}),
...(operations.cellFlags
? { flagsApplied: operations.cellFlags(), supportedFlags: supportedCellFlags() }
: {}),
connectionCapacity:
config.connectionHardCap === undefined
? null
+56 -3
View File
@@ -6,16 +6,63 @@ import type { ControlFlagParse } from './relay-control-flag-channel.js'
export type CellFlags = {
readinessLocal: boolean
ticketCheck: 'off' | 'shadow'
// Off restores Node's default: any unhandled rejection ends the process (#26817 reverted).
rejectionFence: boolean
// Past statement_timeout before a lost reply ends its connection; absent keeps the boot value.
readTimeoutMarginMs?: number
}
export const CELL_FLAG_DEFAULTS: CellFlags = { readinessLocal: false, ticketCheck: 'off' }
export const CELL_FLAG_DEFAULTS: CellFlags = {
readinessLocal: false,
ticketCheck: 'off',
rejectionFence: true
}
export const READ_TIMEOUT_MARGIN_MIN_MS = 1_000
export const READ_TIMEOUT_MARGIN_MAX_MS = 60_000
export const CELL_FLAGS_APPLIED_EVENT = 'orca_relay_cell_flags_applied'
const CellFlagsSchema = z.object({
readinessLocal: z.boolean().default(CELL_FLAG_DEFAULTS.readinessLocal),
ticketCheck: z.enum(['off', 'shadow']).default(CELL_FLAG_DEFAULTS.ticketCheck)
ticketCheck: z.enum(['off', 'shadow']).default(CELL_FLAG_DEFAULTS.ticketCheck),
rejectionFence: z.boolean().default(CELL_FLAG_DEFAULTS.rejectionFence),
readTimeoutMarginMs: z
.number()
.int()
.min(READ_TIMEOUT_MARGIN_MIN_MS)
.max(READ_TIMEOUT_MARGIN_MAX_MS)
.optional()
})
const KNOWN_FLAG_KEYS = new Set(Object.keys(CellFlagsSchema.shape))
export type SupportedCellFlag =
| { type: 'boolean' }
| { type: 'enum'; values: string[] }
| { type: 'number'; min?: number; max?: number; integer?: boolean }
// What this image accepts, read from the schema itself: the flag tool refuses a key or value
// the cell would drop or void before it writes.
export function supportedCellFlags(): Record<string, SupportedCellFlag> {
const supported: Record<string, SupportedCellFlag> = {}
for (const [key, field] of Object.entries(CellFlagsSchema.shape)) {
let inner: z.ZodTypeAny = field
while (inner instanceof z.ZodDefault || inner instanceof z.ZodOptional) {
inner = inner instanceof z.ZodDefault ? inner.removeDefault() : inner.unwrap()
}
if (inner instanceof z.ZodBoolean) supported[key] = { type: 'boolean' }
else if (inner instanceof z.ZodEnum) supported[key] = { type: 'enum', values: [...inner.options] }
else if (inner instanceof z.ZodNumber) {
supported[key] = {
type: 'number',
...(inner.minValue === null ? {} : { min: inner.minValue }),
...(inner.maxValue === null ? {} : { max: inner.maxValue }),
...(inner.isInt ? { integer: true } : {})
}
}
}
return supported
}
export function cellFlagObjectName(cellId: string): string {
return `cells/${cellId}.json`
@@ -29,6 +76,12 @@ export function cellFlagParser(cellId: string): ControlFlagParse<CellFlags> {
})
return (body) => {
const parsed = ObjectSchema.safeParse(body)
return parsed.success ? parsed.data.flags : null
if (!parsed.success) return null
const named = body !== null && typeof body === 'object' && 'flags' in body ? body.flags : {}
const ignoredKeys =
named !== null && typeof named === 'object'
? Object.keys(named).filter((key) => !KNOWN_FLAG_KEYS.has(key)).sort()
: []
return { flags: parsed.data.flags, ignoredKeys }
}
}
+17 -3
View File
@@ -67,7 +67,12 @@ const SeatFeedReplySchema = z.object({
flagsApplied: z
.object({
generation: z.number(),
flags: z.object({ readinessLocal: z.boolean(), ticketCheck: z.enum(['off', 'shadow']) })
flags: z.object({
readinessLocal: z.boolean(),
ticketCheck: z.enum(['off', 'shadow']),
rejectionFence: z.boolean()
}),
ignoredKeys: z.array(z.string()).optional()
})
.optional(),
seq: z.number(),
@@ -333,14 +338,23 @@ describe('cell seat feed', () => {
let applied: AppliedControlFlags<CellFlags> = { generation: 0, flags: CELL_FLAG_DEFAULTS }
const { read, app } = createCell(config(), () => applied)
expect((await read()).flagsApplied).toEqual({ generation: 0, flags: CELL_FLAG_DEFAULTS })
applied = { generation: 12, flags: { readinessLocal: true, ticketCheck: 'shadow' } }
applied = { generation: 12, flags: { ...CELL_FLAG_DEFAULTS, readinessLocal: true, ticketCheck: 'shadow' } }
expect((await read()).flagsApplied).toEqual(applied)
const runtime = await app.request('/v1/admin/runtime-status', {
method: 'POST',
headers: { authorization: 'Bearer deploy-token', 'content-type': 'application/json' },
body: JSON.stringify({ v: 1 })
})
expect(await runtime.json()).toMatchObject({ flagsApplied: applied })
expect(await runtime.json()).toMatchObject({
flagsApplied: applied,
// What the flag tool checks a write against before it writes.
supportedFlags: {
readinessLocal: { type: 'boolean' },
ticketCheck: { type: 'enum', values: ['off', 'shadow'] },
rejectionFence: { type: 'boolean' },
readTimeoutMarginMs: { type: 'number', min: 1_000, max: 60_000, integer: true }
}
})
})
it('accepts only the directors rehome identity, verified once per poll', async () => {
+47 -9
View File
@@ -888,6 +888,27 @@ function inlinePostgresParameters(sql: string, params: unknown[], client: pg.Poo
return inlined
}
// A per-query read timeout, when the cell's switch overrides the pool's (pg reads
// `query_timeout` from a query config object first). Undefined keeps the pool's.
type PostgresQueryTimeout = () => number | undefined
async function timedQuery(
client: pg.PoolClient,
text: string,
values: unknown[] | undefined,
queryTimeout: PostgresQueryTimeout | undefined
): Promise<pg.QueryResult> {
const timeout = queryTimeout?.()
if (timeout === undefined) return values === undefined ? await client.query(text) : await client.query(text, values)
// pg@8.22 reads it at client.js:660; @types/pg does not declare it on QueryConfig.
const config: pg.QueryConfig & { query_timeout: number } = {
text,
...(values === undefined ? {} : { values }),
query_timeout: timeout
}
return await client.query(config)
}
class PostgresTransaction implements RelayDatabase {
readonly dialect = 'postgres' as const
private held: { fromMs: number; site: CellLockHoldSite } | undefined
@@ -899,7 +920,10 @@ class PostgresTransaction implements RelayDatabase {
// later COMMIT would commit a statement its caller saw fail. The connection is finished.
private lostReply: Error | undefined
constructor(protected readonly client: pg.PoolClient) {}
constructor(
protected readonly client: pg.PoolClient,
private readonly queryTimeout?: PostgresQueryTimeout
) {}
get poisonedBy(): Error | undefined {
return this.lostReply
@@ -926,7 +950,7 @@ class PostgresTransaction implements RelayDatabase {
`WITH final_write AS (${inlinePostgresParameters(sql, params, this.client)}) ` +
'SELECT 1 / (SELECT count(*)::int FROM final_write); COMMIT'
try {
await this.client.query(message)
await timedQuery(this.client, message, undefined, this.queryTimeout)
} catch (error) {
// Simple query stops at the first error, so any server error means COMMIT never ran:
// retryable codes take the caller's normal rollback-and-retry path. A lost connection
@@ -977,7 +1001,7 @@ class PostgresTransaction implements RelayDatabase {
this.assertNotCommitted()
this.failIfPoisoned()
try {
const result = await this.client.query(postgresSql(sql), params)
const result = await timedQuery(this.client, postgresSql(sql), params, this.queryTimeout)
return returnsRows(sql) ? (result.rows as SqlRow[]) : [{ changes: result.rowCount ?? 0 }]
} catch (error) {
this.noteFailure(error)
@@ -1142,7 +1166,10 @@ export class PostgresDatabase implements RelayDatabase {
return this.holds.consumeCounts()
}
constructor(private readonly pool: pg.Pool) {
constructor(
private readonly pool: pg.Pool,
private readonly queryTimeout?: PostgresQueryTimeout
) {
this.pressure = new PostgresPoolPressure(pool)
}
@@ -1166,7 +1193,7 @@ export class PostgresDatabase implements RelayDatabase {
try {
client = await this.pressure.connect(lane)
phase = 'execute'
const result = await client.query(postgresSql(sql), params)
const result = await timedQuery(client, postgresSql(sql), params, this.queryTimeout)
return returnsRows(sql) ? (result.rows as SqlRow[]) : [{ changes: result.rowCount ?? 0 }]
} catch (error) {
markRelayDatabaseError(error)
@@ -1215,14 +1242,14 @@ export class PostgresDatabase implements RelayDatabase {
): Promise<T> {
for (let attempt = 1; attempt <= POSTGRES_TRANSACTION_ATTEMPTS; attempt++) {
const client = await this.pressure.connect()
const transaction = new PostgresTransaction(client)
const transaction = new PostgresTransaction(client, this.queryTimeout)
let lostReply: Error | undefined
try {
await client.query('BEGIN')
await timedQuery(client, 'BEGIN', undefined, this.queryTimeout)
const result = await operation(transaction)
// Even when the operation caught it: COMMIT must not run after a lost reply.
if (transaction.poisonedBy) throw transaction.poisonedBy
if (transaction.open) await client.query('COMMIT')
if (transaction.open) await timedQuery(client, 'COMMIT', undefined, this.queryTimeout)
recordMeasuredHold(this.holds, transaction)
this.holds.recordUnavailable(transaction.consumeLockUnavailable())
this.holds.recordLockTimeout(transaction.consumeLockTimeouts())
@@ -1427,6 +1454,8 @@ export type RelayDatabaseOpenInput = {
applicationName?: string
statementTimeoutMs?: number
readTimeoutMarginMs?: number
// Read per query, so a cell's switch moves the margin without a restart.
readTimeoutMarginOverrideMs?: () => number | undefined
// Directors own the PostgreSQL schema. A cell skips it and never touches the database
// at boot, so it starts listening while the database is down and stays unready until
// its first successful query.
@@ -1460,7 +1489,16 @@ export async function openRelayDatabase(input: RelayDatabaseOpenInput): Promise<
absorbPostgresIdleClientErrors(pool)
keepPostgresClientErrorsHandled(pool)
reportPostgresIdleSessionTimeoutOnce(pool)
database = new PostgresDatabase(pool)
const override = input.readTimeoutMarginOverrideMs
database = new PostgresDatabase(
pool,
override
? () => {
const marginMs = override()
return marginMs === undefined ? undefined : statementTimeoutMs + marginMs
}
: undefined
)
} else {
mkdirSync(input.dataDir, { recursive: true })
const sqlite = new DatabaseSync(join(input.dataDir, 'orca-relay.sqlite'))
+9 -2
View File
@@ -9,6 +9,7 @@ import { startCellHeartbeat } from './cell-heartbeat-client.js'
import {
CELL_FLAG_DEFAULTS,
CELL_FLAGS_APPLIED_EVENT,
type CellFlags,
cellFlagObjectName,
cellFlagParser
} from './cell-flags.js'
@@ -36,15 +37,20 @@ import {
readRegisteredMigrationInventory
} from './registered-migration-inventory.js'
// Both read the cell's switch file once it exists; until then, the boot values.
let cellFlagsApplied: (() => CellFlags) | null = null
// Before anything can start a database promise.
process.on('unhandledRejection', handleRelayUnhandledRejection)
process.on('unhandledRejection', (reason) =>
handleRelayUnhandledRejection(reason, cellFlagsApplied?.().rejectionFence ?? true)
)
const config = loadRelayConfig()
const database = await openRelayDatabaseAtBoot({
databaseUrl: config.databaseUrl,
dataDir: config.dataDir,
poolMax: config.databasePoolMax,
applicationName: `orca-relay/${config.role}/${config.cellId}`,
appliesPostgresSchema: config.role !== 'cell'
appliesPostgresSchema: config.role !== 'cell',
readTimeoutMarginOverrideMs: () => cellFlagsApplied?.().readTimeoutMarginMs
})
await reconcileCellAdmissionAtStartup(config, new RelayAssignmentStore(database))
// Never awaited: the cell listens on defaults (all off) until the first read lands.
@@ -57,6 +63,7 @@ const cellFlagChannel =
appliedEvent: CELL_FLAGS_APPLIED_EVENT
})
: null
cellFlagsApplied = cellFlagChannel ? () => cellFlagChannel.applied().flags : null
const {
server,
sessions,
@@ -68,7 +68,11 @@ describePostgres('a lost PostgreSQL reply against a real server', () => {
async function open(
url: string,
options: { readTimeoutMarginMs?: number; poolMax?: number } = {}
options: {
readTimeoutMarginMs?: number
poolMax?: number
readTimeoutMarginOverrideMs?: () => number | undefined
} = {}
): Promise<RelayDatabase> {
const database = await openRelayDatabase({
databaseUrl: url,
@@ -76,7 +80,8 @@ describePostgres('a lost PostgreSQL reply against a real server', () => {
appliesPostgresSchema: false,
statementTimeoutMs: STATEMENT_TIMEOUT_MS,
readTimeoutMarginMs: options.readTimeoutMarginMs ?? READ_TIMEOUT_MARGIN_MS,
poolMax: options.poolMax
poolMax: options.poolMax,
readTimeoutMarginOverrideMs: options.readTimeoutMarginOverrideMs
})
cleanups.push(async () => await database.close())
return database
@@ -125,6 +130,32 @@ describePostgres('a lost PostgreSQL reply against a real server', () => {
).resolves.toEqual([{ id: 'c1' }])
}, 20_000)
// The per-cell switch moves the margin per query, with no restart and no new pool.
it('applies the switch file margin per query, over the pool margin', async () => {
const direct = await open(databaseUrl!)
await direct.query(`CREATE TABLE IF NOT EXISTS ${PROBE_TABLE} (id TEXT PRIMARY KEY)`)
cleanups.push(async () => {
await direct.query(`DROP TABLE IF EXISTS ${PROBE_TABLE}`)
})
const proxy = await blackholeProxy(new URL(databaseUrl!), `FROM ${PROBE_TABLE}`)
cleanups.push(proxy.close)
let override: number | undefined = 1_000
// The pool's own margin would hold the lost reply for 30 s.
const database = await open(proxy.url, {
readTimeoutMarginMs: 30_000,
readTimeoutMarginOverrideMs: () => override
})
await expect(database.query('SELECT 1 AS one')).resolves.toEqual([{ one: 1 }])
const startedAt = performance.now()
const failure = await database
.query(`SELECT id FROM ${PROBE_TABLE}`)
.catch((error: unknown) => error)
expect(isRelayDatabaseTransientError(failure)).toBe(true)
expect(performance.now() - startedAt).toBeLessThan(STATEMENT_TIMEOUT_MS + 1_000 + 1_000)
override = undefined
await expect(database.query('SELECT 1 AS one')).resolves.toEqual([{ one: 1 }])
}, 20_000)
// A reply the server sent after statement_timeout, held up by a whole-VM stall, is still a
// reply: the default margin must outlast the routine 5-7 s stalls.
it('keeps a reply delayed 6.5 s by a stall, on the same connection', async () => {
@@ -88,7 +88,7 @@ describe('control flag channel', () => {
expect(google.state.authorizations[0]).toBe(`Bearer ${ACCESS_TOKEN}`)
expect(channel.applied()).toEqual({
generation: 5,
flags: { readinessLocal: true, ticketCheck: 'off' }
flags: { ...CELL_FLAG_DEFAULTS, readinessLocal: true }
})
})
@@ -106,7 +106,7 @@ describe('control flag channel', () => {
{
event: CELL_FLAGS_APPLIED_EVENT,
generation: 5,
flags: { readinessLocal: true, ticketCheck: 'shadow' }
flags: { ...CELL_FLAG_DEFAULTS, readinessLocal: true, ticketCheck: 'shadow' }
},
{ event: CELL_FLAGS_APPLIED_EVENT, generation: 7, flags: CELL_FLAG_DEFAULTS }
])
@@ -134,6 +134,16 @@ describe('control flag channel', () => {
expect(output).not.toContain(ACCESS_TOKEN)
})
it('applies the fence and read-timeout switches within their bounds', async () => {
const google = fakeGoogle(cellObject(5, { rejectionFence: false, readTimeoutMarginMs: 4_000 }))
const channel = cellChannel(google.fetchImpl)
await channel.poll()
expect(channel.applied()).toEqual({
generation: 5,
flags: { ...CELL_FLAG_DEFAULTS, rejectionFence: false, readTimeoutMarginMs: 4_000 }
})
})
it('ignores a generation below the applied one', async () => {
const google = fakeGoogle(cellObject(9, { readinessLocal: true }))
const channel = cellChannel(google.fetchImpl)
@@ -142,7 +152,7 @@ describe('control flag channel', () => {
await channel.poll()
expect(channel.applied()).toEqual({
generation: 9,
flags: { readinessLocal: true, ticketCheck: 'off' }
flags: { ...CELL_FLAG_DEFAULTS, readinessLocal: true }
})
})
@@ -150,11 +160,19 @@ describe('control flag channel', () => {
const google = fakeGoogle(cellObject(5, { readinessLocal: true, admitMode: 'memory' }))
const channel = cellChannel(google.fetchImpl)
await channel.poll()
expect(channel.applied().flags).toEqual({ readinessLocal: true, ticketCheck: 'off' })
expect(channel.applied()).toEqual({
generation: 5,
flags: { ...CELL_FLAG_DEFAULTS, readinessLocal: true },
ignoredKeys: ['admitMode']
})
expect(appliedLines().at(-1)).toMatchObject({ ignoredKeys: ['admitMode'] })
for (const [generation, object] of [
[6, cellObject(6, { readinessLocal: false, ticketCheck: 'enforce' })],
[7, cellObject(7, { readinessLocal: 'no' })],
[8, cellObject(8, { readinessLocal: false }, 'production-gce-c8')]
[8, cellObject(8, { readTimeoutMarginMs: 500 })],
[9, cellObject(9, { readTimeoutMarginMs: 61_000 })],
[10, cellObject(10, { rejectionFence: 'off' })],
[11, cellObject(11, { readinessLocal: false }, 'production-gce-c8')]
] as const) {
google.state.object = object
await channel.poll()
@@ -14,9 +14,10 @@ export const CONTROL_FLAG_POLL_MS = 5_000
// Refresh a cached access token this long before Google says it expires.
const ACCESS_TOKEN_EARLY_REFRESH_MS = 60_000
export type ControlFlagParse<Flags> = (body: unknown) => Flags | null
// `ignoredKeys` are names this image does not know: a newer writer's, or a typo.
export type ControlFlagParse<Flags> = (body: unknown) => { flags: Flags; ignoredKeys: string[] } | null
export type AppliedControlFlags<Flags> = { generation: number; flags: Flags }
export type AppliedControlFlags<Flags> = { generation: number; flags: Flags; ignoredKeys?: string[] }
export type ControlFlagChannel<Flags> = {
applied: () => AppliedControlFlags<Flags>
@@ -136,10 +137,11 @@ export function startControlFlagChannel<Flags>(input: {
// Only a body read whole is "seen": a read cut off mid-body must be asked for again, or
// every later poll would answer 304 and strand that write.
seenGeneration = generation
const flags = input.parse(body)
if (flags === null) throw new ControlFlagPollError('void')
applied = { generation, flags }
console.log(JSON.stringify({ event: input.appliedEvent, generation, flags }))
const parsed = input.parse(body)
if (parsed === null) throw new ControlFlagPollError('void')
const { flags, ignoredKeys } = parsed
applied = ignoredKeys.length > 0 ? { generation, flags, ignoredKeys } : { generation, flags }
console.log(JSON.stringify({ event: input.appliedEvent, ...applied }))
}
const poll = (): Promise<void> => {
@@ -87,6 +87,15 @@ describe('database rejection fence', () => {
expect(error).not.toHaveBeenCalled()
})
it('keeps every rejection fatal while the cell switch has the fence off', async () => {
const duplicate = await rejectionFrom(() =>
databaseFailingWith(serverError('23505', 'duplicate key')).query('SELECT 1')
)
expect(() => handleRelayUnhandledRejection(duplicate, false)).toThrow()
expect(error).toHaveBeenCalledTimes(1)
expect(() => handleRelayUnhandledRejection(duplicate, true)).not.toThrow()
})
// A whole-VM stall of 6-7 s outlasts the 5 s idle-in-transaction limit for any backend caught
// idle mid-transaction, and the server ends that session with FATAL 25P03 on resume.
it('fences an idle-in-transaction timeout but keeps the rest of class 25 fatal', async () => {
@@ -67,8 +67,9 @@ function rejectionFields(reason: unknown): { code: string; message: string } {
// Installed once by the process entry point. Fatal rejections are rethrown, which ends the
// process exactly as Node's default would, after one line that names them.
export function handleRelayUnhandledRejection(reason: unknown): void {
if (isFencedRelayDatabaseRejection(reason)) {
// `fenced` false is the per-cell switch turned off: every rejection is fatal, as before #26817.
export function handleRelayUnhandledRejection(reason: unknown, fenced = true): void {
if (fenced && isFencedRelayDatabaseRejection(reason)) {
console.warn(
JSON.stringify({ event: 'orca_relay_database_rejection_fenced', ...rejectionFields(reason) })
)