From 7b108abf710d21069691554da8bb4e7d3e2c805a Mon Sep 17 00:00:00 2001 From: Jinwoo Hong <73622457+Jinwoo-H@users.noreply.github.com> Date: Fri, 4 Sep 2026 04:34:22 -0400 Subject: [PATCH] fix(relay): stop taking the fleet-wide cell inventory lock on per-connection paths (#18606) * fix(relay): stop taking the fleet-wide cell inventory lock on per-connection paths activateControl, acquireActivity, changeActivity and removeSupersededSameCellControls each adjust exactly one cell's reservation, yet took SELECT * FROM relay_cells FOR UPDATE, so every desktop rebind and phone reconnect in the fleet queued behind every other one and behind placement. They now use the single-row atomic update (or lock only their own cell row), leaving the inventory lock to placement and sweeps. Fleet-wide 55P03 retries ran p50 430 / p99 1320 per five minutes on 2026-09-03, every cell pinned sqlLatencyMsMax at the lock timeout, and the old cell image crashed on the resulting pool timeouts ~every 15 minutes. A real-Postgres test holds another cell's row and asserts a rebind proceeds; re-adding the inventory lock fails it. * fix(relay): lock the touched cell rows in order on cross-cell activity moves Review found that acquireActivity's existing-lease branch could lock the old lease's cell row (via removeActivityLease) before the new cell's row, which cycles with placement's ascending inventory lock; reproduced on real Postgres as paired 55P03 retries. lockCellRows now takes the one or two rows a per-connection path touches in cell_id order with the 500 ms request bound, and the census fails on any inline relay_cells FOR UPDATE outside the named lock helpers. A three-cell Postgres test moves an activity from the highest cell to a lower one while the target row is held and asserts the mover holds nothing else; five revert-mutants (inventory lock on each path, dropped ordering, dropped ORDER BY) fail it. * test(relay): make the inline relay_cells lock census scan whole statements Review showed two evasions: a FOR UPDATE inside query() and a queryLocked whose FROM relay_cells sat past a fixed line window. The guard now matches every query()/queryLocked() template statement in full; both evasions fail it. Also clears relay_cell_connection_snapshots in the connection- headroom Postgres suite so an aborted run does not poison the next. --- ...nment-connection-headroom-postgres.test.ts | 6 + ...ment-control-supersession-postgres.test.ts | 4 + cloud/apps/relay/src/assignment-store.ts | 32 ++- .../src/cell-inventory-lock-census.test.ts | 73 ++++- ...rol-rebind-inventory-lock-postgres.test.ts | 260 ++++++++++++++++++ 5 files changed, 361 insertions(+), 14 deletions(-) create mode 100644 cloud/apps/relay/src/control-rebind-inventory-lock-postgres.test.ts diff --git a/cloud/apps/relay/src/assignment-connection-headroom-postgres.test.ts b/cloud/apps/relay/src/assignment-connection-headroom-postgres.test.ts index 6ac9521c3d6..80a74a47eeb 100644 --- a/cloud/apps/relay/src/assignment-connection-headroom-postgres.test.ts +++ b/cloud/apps/relay/src/assignment-connection-headroom-postgres.test.ts @@ -44,6 +44,12 @@ describePostgres('PostgreSQL assignment connection headroom', () => { `DELETE FROM relay_assignments WHERE user_id LIKE 'connection-headroom-postgres-%'` ) + // A snapshot left by an aborted run rejects the replayed watermark + // with stale_connection_snapshot. + await database.query( + `DELETE FROM relay_cell_connection_snapshots WHERE cell_id = ?`, + [cell.id] + ) await database.query( `DELETE FROM relay_cell_connection_runtime WHERE cell_id = ?`, [cell.id] diff --git a/cloud/apps/relay/src/assignment-control-supersession-postgres.test.ts b/cloud/apps/relay/src/assignment-control-supersession-postgres.test.ts index 10193b78cc6..cf8819686b5 100644 --- a/cloud/apps/relay/src/assignment-control-supersession-postgres.test.ts +++ b/cloud/apps/relay/src/assignment-control-supersession-postgres.test.ts @@ -38,6 +38,10 @@ describePostgres('PostgreSQL control supersession', () => { [identity.userId] ) await database.query(`DELETE FROM relay_assignments WHERE user_id = ?`, [identity.userId]) + // A snapshot left by an aborted run rejects the replayed watermark with stale_connection_snapshot. + await database.query(`DELETE FROM relay_cell_connection_snapshots WHERE cell_id = ?`, [ + cell.id + ]) await database.query(`DELETE FROM relay_cell_connection_runtime WHERE cell_id = ?`, [cell.id]) await database.query(`DELETE FROM relay_cell_connection_limits WHERE cell_id = ?`, [cell.id]) await database.query(`DELETE FROM relay_cell_runtime WHERE cell_id = ?`, [cell.id]) diff --git a/cloud/apps/relay/src/assignment-store.ts b/cloud/apps/relay/src/assignment-store.ts index d0517d46746..9d240e304a7 100644 --- a/cloud/apps/relay/src/assignment-store.ts +++ b/cloud/apps/relay/src/assignment-store.ts @@ -3202,8 +3202,7 @@ export class RelayAssignmentStore { ) const requestDelta = ACTIVITY_REQUEST_UNITS[kind] * (after - before) if (requestDelta !== 0) { - await this.lockCellInventory(transaction, 'request') - await this.adjustCellReservation(transaction, text(row, 'cell_id'), requestDelta) + await this.adjustCellReservationAtomically(transaction, text(row, 'cell_id'), requestDelta) } }) }) @@ -3263,9 +3262,12 @@ export class RelayAssignmentStore { } const units = ACTIVITY_REQUEST_UNITS[input.kind] if (existing) { - await this.lockCellInventory(transaction, 'request') + // Why: a client-chosen activity id can move between cells, so lock the + // one or two rows this path touches in cell_id order, the same order + // placement takes the inventory in, and no cycle can form. + await this.lockCellRows(transaction, [text(existing, 'cell_id'), input.cellId]) await this.removeActivityLease(transaction, identity, existing, now) - await this.adjustCellReservation(transaction, input.cellId, units) + await this.adjustCellReservationAtomically(transaction, input.cellId, units) } await this.adjustActivityCount(transaction, identity, input.kind, 1, expiresAt, now) await transaction.query( @@ -3580,8 +3582,7 @@ export class RelayAssignmentStore { ) await this.touchAssignment(transaction, identity, expiresAt, now) } else { - await this.lockCellInventory(transaction, 'request') - await this.adjustCellReservation(transaction, input.cellId, 1) + await this.adjustCellReservationAtomically(transaction, input.cellId, 1) await this.adjustActivityCount(transaction, identity, 'control', 1, expiresAt, now) await transaction.query( `INSERT INTO relay_assignment_activity_leases @@ -6954,6 +6955,19 @@ export class RelayAssignmentStore { return rows } + // Per-connection paths touch one or two cells. Locking exactly those rows, + // in the same ascending order the inventory lock uses (ORDER BY fixes the + // row-lock order), keeps them off the fleet-wide lock without a cycle. + private async lockCellRows(database: RelayDatabase, cellIds: string[]): Promise { + const distinct = [...new Set(cellIds)] + return await database.queryLocked( + `SELECT * FROM relay_cells WHERE cell_id IN (${distinct.map(() => '?').join(', ')}) + ORDER BY cell_id ASC`, + distinct, + { lockTimeoutMs: CELL_INVENTORY_LOCK_TIMEOUT_MS } + ) + } + private async lockGeneralCellInventory( database: RelayDatabase, mode: CellInventoryLockMode @@ -7590,7 +7604,10 @@ export class RelayAssignmentStore { ) { throw new Error('activity_lease_shape_mismatch') } - const cells = await this.lockCellInventory(database, 'request') + // Why: this recomputes one cell's reservation from its leases, so only that + // row needs to be held; the 23-row inventory lock here serialised every + // desktop control rebind in the fleet behind every other one. + const cellRow = (await this.lockCellRows(database, [cellId]))[0] await database.query( `DELETE FROM relay_assignment_activity_leases WHERE user_id = ? AND relay_host_id = ? AND activity_kind = 'control' @@ -7611,7 +7628,6 @@ export class RelayAssignmentStore { [cellId] ) )[0]! - const cellRow = cells.find((cell) => text(cell, 'cell_id') === cellId) const cellUnits = integer(cellUnitsRow, 'request_units') if (!cellRow) throw new Error('assigned_cell_missing') if (cellUnits > integer(cellRow, 'capacity_requests')) { diff --git a/cloud/apps/relay/src/cell-inventory-lock-census.test.ts b/cloud/apps/relay/src/cell-inventory-lock-census.test.ts index 8ca7cee55f5..a26e15f8e1d 100644 --- a/cloud/apps/relay/src/cell-inventory-lock-census.test.ts +++ b/cloud/apps/relay/src/cell-inventory-lock-census.test.ts @@ -25,10 +25,12 @@ const CENSUS: CensusEntry[] = [ { method: 'assignOnce', mode: 'nowait', reach: 'both' }, { method: 'assignOnce', mode: 'nowait', reach: 'both' }, { method: 'refreshDrainMigrationLeasesOnce', mode: 'request', reach: 'request' }, - // Reachable from neither: changeActivity has no production callers, only tests. - { method: 'changeActivity', mode: 'request', reach: 'orphan' }, - { method: 'acquireActivity', mode: 'request', reach: 'request' }, - { method: 'activateControl', mode: 'request', reach: 'request' }, + // changeActivity, acquireActivity, activateControl and + // removeSupersededSameCellControls no longer take the inventory: they lock + // only the one or two cell rows they touch, in cell_id order (lockCellRows), + // so they cannot cycle with placement's ordered inventory lock, and the + // 23-row lock there had serialised every reconnect in the fleet behind every + // other one. { method: 'startEvacuation', mode: 'request', reach: 'request' }, { method: 'completeEvacuationFromDeadSourceOnce', mode: 'request', reach: 'request' }, { method: 'completeEvacuationFromDeadSourceOnce', mode: 'nowait', reach: 'request' }, @@ -48,8 +50,31 @@ const CENSUS: CensusEntry[] = [ { method: 'releaseExpiredActivityLeases', mode: 'nowait', reach: 'sweep' }, { method: 'releaseExpiredActivity', mode: 'nowait', reach: 'sweep' }, { method: 'reconcileReservationAccounting', mode: 'pool-default', reach: 'both' }, - { method: 'leastLoadedCell', mode: 'pool-default', reach: 'both' }, - { method: 'removeSupersededSameCellControls', mode: 'request', reach: 'request' } + { method: 'leastLoadedCell', mode: 'pool-default', reach: 'both' } +] + +// Every inline `FROM relay_cells ... FOR UPDATE` outside the named lock helpers, +// in source order: whole-table locks in reconciliation and sticky placement, +// and single-row locks for a cell the method is already scoped to (heartbeat, +// fence, drain generation, configuration, or a reservation adjust that runs +// under a lock its caller already holds). A new inline lock fails the census +// below until it is listed here; per-connection paths that touch more than one +// cell go through lockCellRows so the order is fixed. +const NAMED_LOCK_HELPERS = ['lockCellInventory', 'lockGeneralCellInventory', 'lockCellRows'] + +const INLINE_CELL_LOCK_SITES = [ + 'reconcileCellsWithOptions', + 'assignStickyOnce', + 'recordCellHeartbeat', + 'attestCellFence', + 'adoptLegacyCellFence', + 'commitLegacyCellFenceAdoption', + 'prepareCellFenceAttempt', + 'attestCellFenceAttempt', + 'attestCellFenceAttempt', + 'configureCell', + 'assertDrainCellGeneration', + 'adjustCellReservation' ] // The background sweeps, and nothing else. A method reachable from one of these @@ -151,6 +176,42 @@ describe('cell inventory lock call-site census', () => { ) }) + // Why: the census only sees lockCellInventory calls, so a hand-written + // `relay_cells ... FOR UPDATE` would escape classification entirely. + it('routes every relay_cells row lock through a named lock helper', () => { + const lines = storeSource() + const rawSites: string[] = [] + // Whole statements, not a fixed window: a wide column list or a raw + // FOR UPDATE inside query() must not slip past. + const source = lines.join('\n') + const bounds: { name: string; start: number }[] = [] + lines.forEach((line, index) => { + const declaration = DECLARATION.exec(line) + if (declaration) bounds.push({ name: declaration[1]!, start: index }) + }) + const methodAt = (offset: number): string => { + const lineIndex = source.slice(0, offset).split('\n').length - 1 + let name = '' + for (const bound of bounds) if (bound.start <= lineIndex) name = bound.name + return name + } + const tick = String.fromCharCode(96) + const statementCall = new RegExp( + '\\.(queryLocked|query)\\(\\s*' + tick + '([^' + tick + ']*)' + tick, + 'g' + ) + for (const call of source.matchAll(statementCall)) { + const statement = call[2]! + if (!/\bFROM\s+relay_cells\b/.test(statement)) continue + const locks = call[1] === 'queryLocked' || /\bFOR\s+UPDATE\b/.test(statement) + if (!locks) continue + const method = methodAt(call.index) + if (NAMED_LOCK_HELPERS.includes(method)) continue + rawSites.push(method) + } + expect(rawSites).toEqual(INLINE_CELL_LOCK_SITES) + }) + it('leaves no call site taking the inventory without naming a mode', () => { const source = readFileSync(new URL('./assignment-store.ts', import.meta.url), 'utf8') const unclassified = source diff --git a/cloud/apps/relay/src/control-rebind-inventory-lock-postgres.test.ts b/cloud/apps/relay/src/control-rebind-inventory-lock-postgres.test.ts new file mode 100644 index 00000000000..e990ac1ed1a --- /dev/null +++ b/cloud/apps/relay/src/control-rebind-inventory-lock-postgres.test.ts @@ -0,0 +1,260 @@ +import { afterAll, beforeAll, describe, expect, it } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.js' +import { openRelayDatabase, type RelayDatabase } from './database.js' + +const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL +const describePostgres = databaseUrl ? describe : describe.skip + +// Three cells: the inventory lock covers more than the rows a move touches, and +// a high-to-low move exposes any lock taken out of cell_id order. +const cells = [ + { + id: 'rebind-inventory-postgres-a', + url: 'https://rebind-inventory-postgres-a.example.com', + capacityRequests: 1_000, + connectionHardCap: 600 as const, + connectionUnobservedBound: 50 + }, + { + id: 'rebind-inventory-postgres-b', + url: 'https://rebind-inventory-postgres-b.example.com', + capacityRequests: 1_000, + connectionHardCap: 600 as const, + connectionUnobservedBound: 50 + }, + { + id: 'rebind-inventory-postgres-c', + url: 'https://rebind-inventory-postgres-c.example.com', + capacityRequests: 1_000, + connectionHardCap: 600 as const, + connectionUnobservedBound: 50 + } +] +const identity = { userId: 'rebind-inventory-postgres-user', relayHostId: 'rebindinvhost001' } + +function heartbeat(cell: (typeof cells)[number]) { + return { + cellId: cell.id, + cellUrl: cell.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0, + connectionInclusionWatermark: 1, + connectionHardCap: 600 as const, + connectionUnobservedBound: 50 + } +} + +// Why: every desktop control rebind used to take the fleet-wide relay_cells +// FOR UPDATE lock, so a rebind on one cell queued behind whatever held any +// other cell's row, until COMMIT (55P03 at the request bound). A rebind only +// touches its own cell row, so it must proceed while another cell's row is +// held elsewhere. +describePostgres('PostgreSQL control rebind under a held cell row', () => { + const databases: RelayDatabase[] = [] + + beforeAll(async () => { + databases.push( + await openRelayDatabase({ databaseUrl, dataDir: '' }), + await openRelayDatabase({ databaseUrl, dataDir: '' }) + ) + }) + + async function removeTestRows(database: RelayDatabase): Promise { + await database.query( + `DELETE FROM relay_control_connection_reservations WHERE user_id = ?`, + [identity.userId] + ) + for (const table of [ + 'relay_assignment_activity_leases', + 'relay_post_drain_migration_pins', + 'relay_assignment_migration_incarnations', + 'relay_assignment_migrations', + 'relay_assignments' + ]) { + await database.query(`DELETE FROM ${table} WHERE user_id = ?`, [identity.userId]) + } + for (const cell of cells) { + for (const table of [ + 'relay_cell_connection_snapshots', + 'relay_cell_connection_runtime', + 'relay_cell_connection_limits', + 'relay_cell_runtime', + 'relay_cells' + ]) { + await database.query(`DELETE FROM ${table} WHERE cell_id = ?`, [cell.id]) + } + } + } + + afterAll(async () => { + if (databases[0]) await removeTestRows(databases[0]) + for (const connection of databases) await connection.close() + }) + + it("rebinds and supersedes a control while another cell's row is held", async () => { + // A prior aborted run leaves connection snapshots that reject a replayed watermark. + await removeTestRows(databases[0]!) + const store = new RelayAssignmentStore(databases[0]!, () => 100) + await store.reconcileCells(cells) + for (const cell of cells) await store.recordCellHeartbeat(heartbeat(cell)) + // Pin the host to cell A so placement is deterministic. + await store.setCellEnabled(cells[1]!.id, false) + await store.setCellEnabled(cells[2]!.id, false) + const assignment = await store.assign(identity) + expect(assignment.cellId).toBe(cells[0]!.id) + await store.setCellEnabled(cells[1]!.id, true) + await store.setCellEnabled(cells[2]!.id, true) + await store.activateControl(identity, { + cellId: cells[0]!.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1, + connectionInclusionWatermark: 10 + }) + + // Hold only cell B's row on a second connection, the way a rebind on B + // does, for longer than the request-path lock bound. + let releaseInventory!: () => void + const inventoryReleased = new Promise((resolve) => { + releaseInventory = resolve + }) + let inventoryHeld!: () => void + const inventoryHeldPromise = new Promise((resolve) => { + inventoryHeld = resolve + }) + const holder = databases[1]!.transaction(async (transaction) => { + await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [cells[1]!.id]) + inventoryHeld() + await inventoryReleased + }) + await inventoryHeldPromise + + // A generation-2 rebind on cell A supersedes generation 1. It must not + // wait on cell B's row. + const startedAt = Date.now() + const blockedStatement = async (): Promise => { + const rows = await databases[1]!.query( + `SELECT left(query, 160) AS q FROM pg_stat_activity + WHERE datname = current_database() AND wait_event_type = 'Lock'` + ) + return rows.map((row) => String(row.q)).join(' | ') + } + const timeout = new Promise((_, reject) => + setTimeout( + () => + void blockedStatement().then((statement) => + reject(new Error(`rebind on cell A blocked behind cell B's row: ${statement}`)) + ), + 2_000 + ) + ) + const rebound = await Promise.race([ + store.activateControl(identity, { + cellId: cells[0]!.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 2, + connectionInclusionWatermark: 11 + }), + timeout + ]) + const elapsedMs = Date.now() - startedAt + releaseInventory() + await holder + + expect(rebound).toBe(`control:${cells[0]!.id}:2`) + expect(elapsedMs).toBeLessThan(2_000) + const controls = await databases[0]!.query( + `SELECT activity_id FROM relay_assignment_activity_leases + WHERE user_id = ? AND activity_kind = 'control' ORDER BY activity_id`, + [identity.userId] + ) + expect(controls).toEqual([{ activity_id: `control:${cells[0]!.id}:2` }]) + const reserved = await databases[0]!.query( + `SELECT reserved_requests FROM relay_cells WHERE cell_id = ?`, + [cells[0]!.id] + ) + expect(Number(reserved[0]!.reserved_requests)).toBe(1) + }, 15_000) + + // Why: a phone's activity id is client-chosen and can follow the host across + // a migration, so acquireActivity may touch two cell rows. Moving from the + // higher cell to the lower one is where an unordered lock cycles with + // placement's ascending inventory lock (reproduced live before this fix). + it('moves an activity from a higher cell to a lower one in cell_id order', async () => { + await removeTestRows(databases[0]!) + const [cellA, cellB, cellC] = cells as [typeof cells[0], typeof cells[0], typeof cells[0]] + const store = new RelayAssignmentStore(databases[0]!, () => 100) + await store.reconcileCells(cells) + for (const cell of cells) await store.recordCellHeartbeat(heartbeat(cell)) + await store.setCellEnabled(cellA.id, false) + await store.setCellEnabled(cellB.id, false) + const assignment = await store.assign(identity) + expect(assignment.cellId).toBe(cellC.id) + await store.setCellEnabled(cellA.id, true) + await store.setCellEnabled(cellB.id, true) + const activityId = 'splice:rebind-inventory-postgres' + await store.acquireActivity(identity, { activityId, kind: 'splice', cellId: cellC.id }) + // The migration makes B authoritative; the lease still sits on C. + const migration = await store.startEvacuation(identity, cellB.id) + expect(migration.targetCellId).toBe(cellB.id) + + // Hold B elsewhere. An ordered move locks B first and queues here holding + // nothing else. Locking C first (the old lease's row, as an unordered move + // does) or the whole inventory (which takes A) shows up as a held row. + let releaseRow!: () => void + const rowReleased = new Promise((resolve) => { + releaseRow = resolve + }) + let rowHeld!: () => void + const rowHeldPromise = new Promise((resolve) => { + rowHeld = resolve + }) + const heldWhileMoverWaits: string[] = [] + const holder = databases[1]!.transaction(async (transaction) => { + await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [cellB.id]) + rowHeld() + await rowReleased + for (const cell of [cellA, cellC]) { + try { + await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [cell.id], { + failIfUnavailable: true + }) + } catch { + heldWhileMoverWaits.push(cell.id) + } + } + }) + await rowHeldPromise + const move = store.acquireActivity(identity, { activityId, kind: 'splice', cellId: cellB.id }) + let moved = false + void move.then(() => { + moved = true + }) + await new Promise((resolve) => setTimeout(resolve, 250)) + expect(moved).toBe(false) + releaseRow() + await holder + await move + expect(heldWhileMoverWaits).toEqual([]) + + const reservations = await databases[0]!.query( + `SELECT cell_id, reserved_requests FROM relay_cells + WHERE cell_id IN (?, ?, ?) ORDER BY cell_id ASC`, + [cellA.id, cellB.id, cellC.id] + ) + const reserved = reservations.map((row) => [String(row.cell_id), Number(row.reserved_requests)]) + expect(reserved).toEqual([ + [cellA.id, 0], + // Migration grant plus the moved splice, as in the SQLite origin-scoped + // reservation case: the lock change did not alter accounting. + [cellB.id, 6], + // The sticky grant stays on the source until the migration completes. + [cellC.id, 1] + ]) + }, 15_000) +})