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 cf8819686b5..756706a182b 100644 --- a/cloud/apps/relay/src/assignment-control-supersession-postgres.test.ts +++ b/cloud/apps/relay/src/assignment-control-supersession-postgres.test.ts @@ -101,6 +101,12 @@ describePostgres('PostgreSQL control supersession', () => { cell.id ] ) + // The store reserves a unit per control lease, so a hand-written pair has to + // carry its own reservation or the fixture starts out of balance. + await databases[0]!.query( + `UPDATE relay_cells SET reserved_requests = reserved_requests + 2 WHERE cell_id = ?`, + [cell.id] + ) await stores[0]!.activateControl(identity, { cellId: cell.id, assignmentEpoch: assignment.assignmentEpoch, diff --git a/cloud/apps/relay/src/assignment-row-lock-order-postgres.test.ts b/cloud/apps/relay/src/assignment-row-lock-order-postgres.test.ts new file mode 100644 index 00000000000..4b3893654f3 --- /dev/null +++ b/cloud/apps/relay/src/assignment-row-lock-order-postgres.test.ts @@ -0,0 +1,138 @@ +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 + +const cells = [ + { + id: 'row-lock-order-a', + url: 'https://row-lock-order-a.example.com', + capacityRequests: 200, + connectionHardCap: 600 as const, + connectionUnobservedBound: 50 + }, + { + id: 'row-lock-order-b', + url: 'https://row-lock-order-b.example.com', + capacityRequests: 200, + connectionHardCap: 600 as const, + connectionUnobservedBound: 50 + } +] +const userId = 'row-lock-order-user' +const hosts = ['rowlockhost00001', 'rowlockhost00002', 'rowlockhost00003', 'rowlockhost00004'].map( + (relayHostId) => ({ userId, relayHostId }) +) + +// Why a deadlock counter and not "it eventually succeeded": the transaction +// wrapper retries 40P01 three times, so a cycle that fires on every wave still +// reports success to the caller while burning the retry budget that turns into +// a 503 under load. PostgreSQL counts every detected cycle in pg_stat_database, +// which sees through the retry. +describePostgres('PostgreSQL row lock order', () => { + const databases: RelayDatabase[] = [] + + beforeAll(async () => { + for (let index = 0; index < 4; index++) { + databases.push(await openRelayDatabase({ databaseUrl, dataDir: '' })) + } + }) + + async function removeTestRows(database: RelayDatabase): Promise { + await database.query( + `DELETE FROM relay_control_connection_reservations WHERE user_id = ?`, + [userId] + ) + for (const table of [ + 'relay_control_capabilities', + '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 = ?`, [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() + }) + + async function deadlockCount(): Promise { + const rows = await databases[0]!.query( + `SELECT deadlocks FROM pg_stat_database WHERE datname = current_database()` + ) + return Number(rows[0]!.deadlocks) + } + + it('runs the cell accept and the placement retry concurrently without a cycle', async () => { + await removeTestRows(databases[0]!) + const stores = databases.map((database) => new RelayAssignmentStore(database, () => 100)) + await stores[0]!.reconcileCells(cells) + for (const cell of cells) { + await stores[0]!.recordCellHeartbeat({ + 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, + connectionUnobservedBound: 50 + }) + } + const placements = new Map() + for (const identity of hosts) { + const assignment = await stores[0]!.assign(identity) + placements.set(identity.relayHostId, assignment) + } + + const before = await deadlockCount() + // The accept path (host rows, then the cell row last) against the paths + // that must read the inventory first: placement and evacuation. + for (let round = 0; round < 12; round++) { + await Promise.allSettled( + hosts.flatMap((identity, index) => { + const placement = placements.get(identity.relayHostId)! + const store = stores[index % stores.length]! + const other = stores[(index + 1) % stores.length]! + return [ + store.activateControl(identity, { + cellId: placement.cellId, + assignmentEpoch: placement.assignmentEpoch, + generation: round + 2 + }), + other.assign(identity), + other.startEvacuation( + identity, + placement.cellId === cells[0]!.id ? cells[1]!.id : cells[0]!.id + ) + ] + }) + ) + } + const after = await deadlockCount() + + expect(after - before).toBe(0) + }, 120_000) +}) diff --git a/cloud/apps/relay/src/assignment-store.test.ts b/cloud/apps/relay/src/assignment-store.test.ts index e9d27a83db4..e4a8ce54ebf 100644 --- a/cloud/apps/relay/src/assignment-store.test.ts +++ b/cloud/apps/relay/src/assignment-store.test.ts @@ -1886,6 +1886,11 @@ describe('RelayAssignmentStore', () => { 'control:cell-b:3' ] ) + // The store reserves a unit per control lease, so a hand-written pair has to + // carry its own reservation or the fixture starts out of balance. + await database!.query( + `UPDATE relay_cells SET reserved_requests = reserved_requests + 2 WHERE cell_id = 'cell-b'` + ) const latest = await store.activateControl(identity, { cellId: 'cell-b', assignmentEpoch: migration.assignmentEpoch, diff --git a/cloud/apps/relay/src/assignment-store.ts b/cloud/apps/relay/src/assignment-store.ts index 871a2be1c54..b34e25a69a0 100644 --- a/cloud/apps/relay/src/assignment-store.ts +++ b/cloud/apps/relay/src/assignment-store.ts @@ -342,6 +342,26 @@ const ACTIVITY_REQUEST_UNITS: Record = { } const ASSIGNMENT_LOCK_RETRY_DEADLINE_MS = 15_000 + +// THE ROW LOCK ORDER. Every transaction that takes more than one of these +// takes them in this order, whichever role it runs on: +// +// 1. relay_assignments (the host's row) +// 2. relay_assignment_migrations +// relay_assignment_activity_leases +// 3. relay_control_connection_reservations (lockControlConnectionReservations) +// 4. relay_cells (lockCellInventory / lockCellRows / the +// conditional reservation UPDATE) +// +// relay_cells is last because it is the only row shared by every host on a +// cell: a transaction that takes it early holds it across every round trip +// that follows, and on a cell far from PostgreSQL that is what turns accepts +// into a queue. Everything above it is per-host, so holding it longer costs +// only that host. Paths that read the inventory to make a placement decision +// cannot defer relay_cells, so they lock the host's rows from tiers 1-3 up +// front instead, before the inventory, and re-check what they read afterwards. +// Tier 1 is what serialises two transactions on the same host; the tiers below +// it keep transactions on *different* hosts from cycling through relay_cells. // Why: one global FOR UPDATE over a 23-row table serialises every director and // cell. At the 1s pool lock_timeout each blocked waiter also holds a pooled // client for a full second, so the queue converts contention into pool @@ -854,6 +874,9 @@ export class RelayAssignmentStore { let retryScope: RetriedAssignmentInventoryScope = inventoryScope === 'all' ? 'all' : 'general' return await this.database.transaction(async (transaction) => { + // The retry paths below open with the inventory, so this path takes its + // host rows before any of them rather than where the others do. + await this.lockControlConnectionReservations(transaction, identity, lockMode) let lockedCells = inventoryScope === 'all' ? await this.lockCellInventory(transaction, lockMode) @@ -3821,14 +3844,16 @@ export class RelayAssignmentStore { ? activityId : pendingId : activityId - await this.removeSupersededSameCellControls( + // Accumulated, not applied: every cell-row write on this path is folded + // into one conditional statement issued last, below. + let reservationDelta = -(await this.removeSupersededSameCellControls( transaction, identity, activityLeases, input.cellId, retainedActivityId, now - ) + )) if (existing) { await transaction.query( `UPDATE relay_assignment_activity_leases SET expires_at = ?, updated_at = ? @@ -3846,7 +3871,7 @@ export class RelayAssignmentStore { ) await this.touchAssignment(transaction, identity, expiresAt, now) } else { - await this.adjustCellReservationAtomically(transaction, input.cellId, 1) + reservationDelta += ACTIVITY_REQUEST_UNITS.control await this.adjustActivityCount(transaction, identity, 'control', 1, expiresAt, now) await transaction.query( `INSERT INTO relay_assignment_activity_leases @@ -3902,6 +3927,17 @@ export class RelayAssignmentStore { input.idleRegionalRehome && input.cellIncarnation ? 1 : 0 ] ) + // Last, and only if the count actually moved: the cell row is shared by + // every host on the cell, and this transaction spans a dozen round + // trips. Holding its write lock from the first of them capped a + // far-from-Postgres cell at a couple of accepts a second. + if (reservationDelta !== 0) { + await this.adjustCellReservationAtomically( + transaction, + input.cellId, + reservationDelta + ) + } return activityId }) }) @@ -3946,6 +3982,7 @@ export class RelayAssignmentStore { } if (sourceCellId === targetCellId) throw new Error('target_matches_source') await this.lockAssignmentActivities(transaction, identity) + await this.lockControlConnectionReservations(transaction, identity) const cells = await this.lockCellInventory(transaction, 'request') const target = cells.find((row) => text(row, 'cell_id') === targetCellId) if (!target || integer(target, 'enabled') !== 1) throw new Error('target_cell_unavailable') @@ -4151,6 +4188,7 @@ export class RelayAssignmentStore { ): Promise { const now = this.now() return await this.database.transaction(async (transaction) => { + await this.lockControlConnectionReservations(transaction, identity) let lockedCells: SqlRow[] | undefined if (inventoryFirst) { try { @@ -4298,6 +4336,7 @@ export class RelayAssignmentStore { ): Promise { const now = this.now() return await this.database.transaction(async (transaction) => { + await this.lockControlConnectionReservations(transaction, identity) const lockedCells = inventoryFirst ? await this.lockCellInventory(transaction, 'request') : undefined @@ -4790,7 +4829,10 @@ export class RelayAssignmentStore { throw new Error('migration_activity_topology_mismatch') } } - if (obsoleteLeases.length > 0) await this.lockCellInventory(transaction, 'request') + if (obsoleteLeases.length > 0) { + await this.lockControlConnectionReservations(transaction, identity) + await this.lockCellInventory(transaction, 'request') + } for (const lease of obsoleteLeases) { await this.removeActivityLease(transaction, identity, lease, now) } @@ -5133,6 +5175,7 @@ export class RelayAssignmentStore { const sourceCellId = text(assignment, 'cell_id') if (sourceCellId === targetCellId) throw new Error('target_matches_source') await this.lockAssignmentActivities(transaction, identity) + await this.lockControlConnectionReservations(transaction, identity) const cells = await this.lockCellInventory(transaction, 'request') const admission = await cellAdmissionStates(transaction) const targetRow = cells.find( @@ -5464,6 +5507,7 @@ export class RelayAssignmentStore { } const activityLeases = await this.lockAssignmentActivities(transaction, input.identity) assertAssignmentActivityCounts(assignment, activityLeases, 0) + await this.lockControlConnectionReservations(transaction, input.identity, 'nowait') const cells = await this.lockCellInventory(transaction, 'nowait') const admission = await cellAdmissionStates(transaction) const regions = new Map( @@ -6383,6 +6427,7 @@ export class RelayAssignmentStore { integer(lease, 'expires_at') > now ) if (targetActive) return false + await this.lockControlConnectionReservations(transaction, identity, 'nowait') const cells = await this.lockCellInventory(transaction, 'nowait') const source = cells.find((cell) => text(cell, 'cell_id') === sourceCellId) const admission = await cellAdmissionStates(transaction) @@ -6549,7 +6594,10 @@ export class RelayAssignmentStore { ] .map((activityId) => activityLeaseById(activityLeases, activityId)) .filter((lease): lease is SqlRow => lease !== undefined) - if (obsoleteLeases.length > 0) await this.lockCellInventory(transaction, 'nowait') + if (obsoleteLeases.length > 0) { + await this.lockControlConnectionReservations(transaction, identity, 'nowait') + await this.lockCellInventory(transaction, 'nowait') + } for (const lease of obsoleteLeases) { await this.removeActivityLease(transaction, identity, lease, now) } @@ -6567,6 +6615,7 @@ export class RelayAssignmentStore { ) return true } + await this.lockControlConnectionReservations(transaction, identity, 'nowait') const cells = await this.lockCellInventory(transaction, 'nowait') const sourceCellId = text(row, 'source_cell_id') const admissionRows = await transaction.query( @@ -6986,6 +7035,24 @@ export class RelayAssignmentStore { ) } + // Tier 3 of the row lock order: a path that will take relay_cells and also + // touch this host's reservations takes them here, before the cell rows. The + // set is the host's own rows, so it is small and known before any placement + // decision is read. + private async lockControlConnectionReservations( + database: RelayDatabase, + identity: AssignmentIdentity, + mode: CellInventoryLockMode = 'request' + ): Promise { + const { measureHoldMs: _sampled, ...wait } = cellInventoryLockOptions(mode) + await database.queryLocked( + `SELECT reservation_id FROM relay_control_connection_reservations + WHERE user_id = ? AND relay_host_id = ? ORDER BY reservation_id ASC`, + [identity.userId, identity.relayHostId], + wait + ) + } + // Unlocked on purpose: this only names the row to lock next, and the caller // re-checks the pin once the assignment row is held. private async pinnedCellId( @@ -7641,6 +7708,9 @@ export class RelayAssignmentStore { await this.adjustActivityCount(database, identity, kind, -1, now, now) } + // Returns the request units this removal frees on `cellId`. The caller folds + // them into the one conditional cell-row write it makes just before COMMIT, + // so no statement here touches the fleet's most contended row. private async removeSupersededSameCellControls( database: RelayDatabase, identity: AssignmentIdentity, @@ -7648,14 +7718,14 @@ export class RelayAssignmentStore { cellId: string, retainedActivityId: string, now: number - ): Promise { + ): Promise { const superseded = leases.filter( (lease) => activityKind(lease) === 'control' && text(lease, 'cell_id') === cellId && text(lease, 'activity_id') !== retainedActivityId ) - if (superseded.length === 0) return + if (superseded.length === 0) return 0 if ( superseded.some( (lease) => integer(lease, 'request_units') !== ACTIVITY_REQUEST_UNITS.control @@ -7663,10 +7733,6 @@ export class RelayAssignmentStore { ) { throw new Error('activity_lease_shape_mismatch') } - // 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' @@ -7680,22 +7746,9 @@ export class RelayAssignmentStore { WHERE user_id = ? AND relay_host_id = ?`, [remainingControls, now, identity.userId, identity.relayHostId] ) - const cellUnitsRow = ( - await database.query( - `SELECT COALESCE(SUM(request_units), 0) AS request_units - FROM relay_assignment_activity_leases WHERE cell_id = ?`, - [cellId] - ) - )[0]! - const cellUnits = integer(cellUnitsRow, 'request_units') - if (!cellRow) throw new Error('assigned_cell_missing') - if (cellUnits > integer(cellRow, 'capacity_requests')) { - throw new Error('relay_capacity_exhausted') - } - await database.query( - `UPDATE relay_cells SET reserved_requests = ?, updated_at = ? WHERE cell_id = ?`, - [cellUnits, now, cellId] - ) + // The shape check above proved every superseded lease holds exactly the + // control unit, so the freed units are exact without re-summing the cell. + return superseded.length * ACTIVITY_REQUEST_UNITS.control } private async adjustActivityCount( 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 929173eb9da..6f564844ba6 100644 --- a/cloud/apps/relay/src/cell-inventory-lock-census.test.ts +++ b/cloud/apps/relay/src/cell-inventory-lock-census.test.ts @@ -32,7 +32,9 @@ const CENSUS: CensusEntry[] = [ // 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. + // other one. The control accept path went one step further and takes no cell + // read lock at all: its single conditional write is the last statement before + // COMMIT. { method: 'startEvacuation', mode: 'request', reach: 'request' }, { method: 'completeEvacuationFromDeadSourceOnce', mode: 'request', reach: 'request' }, { method: 'completeEvacuationFromDeadSourceOnce', mode: 'nowait', reach: 'request' }, @@ -172,6 +174,63 @@ function readCallSites(): { method: string; mode: CensusMode }[] { return sites } + +// Tier 3 and tier 4 of the row lock order documented in assignment-store.ts. A +// transaction that takes relay_cells before this host's reservation rows can +// cycle with one that takes them the other way round, and PostgreSQL resolves +// that as a 40P01 during exactly the drain and rehome waves these paths exist +// to run. The cell row is the one every host on a cell shares, so it is the +// lock that must be taken last, which fixes the direction for everyone else. +const CELL_LOCK_CALL = + /this\.(?:lockCellInventory|lockGeneralCellInventory|lockCellRows|adjustCellReservationAtomically|adjustCellReservation)\(|UPDATE relay_cells/ +const RESERVATION_LOCK_CALL = + /this\.(?:lockControlConnectionReservations|insertControlConnectionReservation|claimControlConnectionReservation|releaseSupersededControlConnectionReservations)\(|(?:UPDATE|INTO|DELETE FROM)\s+relay_control_connection_reservations/ + +// The lock helpers themselves, plus the one reporting query that reads both +// tables without locking either. +const ROW_LOCK_ORDER_EXEMPT = [ + 'lockCellInventory', + 'lockGeneralCellInventory', + 'lockCellRows', + 'lockControlConnectionReservations', + 'adjustCellReservation', + 'adjustCellReservationAtomically', + 'insertControlConnectionReservation', + 'claimControlConnectionReservation', + 'releaseSupersededControlConnectionReservations', + 'cellDeploymentStatus' +] + +function methodSpans(lines: string[]): { name: string; start: number; end: number }[] { + const starts: { name: string; start: number }[] = [] + lines.forEach((line, index) => { + const declaration = DECLARATION.exec(line) + if (declaration) starts.push({ name: declaration[1]!, start: index }) + }) + return starts.map((entry, index) => ({ + ...entry, + end: starts[index + 1]?.start ?? lines.length + })) +} + +function pathsTakingCellsBeforeReservations(lines: string[]): string[] { + const offending: string[] = [] + for (const span of methodSpans(lines)) { + if (ROW_LOCK_ORDER_EXEMPT.includes(span.name)) continue + let cell = Number.POSITIVE_INFINITY + let reservation = Number.POSITIVE_INFINITY + for (let index = span.start; index < span.end; index++) { + const line = lines[index]! + if (CELL_LOCK_CALL.test(line)) cell = Math.min(cell, index) + if (RESERVATION_LOCK_CALL.test(line)) reservation = Math.min(reservation, index) + } + if (cell < reservation && reservation !== Number.POSITIVE_INFINITY) { + offending.push(span.name) + } + } + return offending +} + describe('cell inventory lock call-site census', () => { it('classifies every call site exactly as recorded', () => { expect(readCallSites()).toEqual(CENSUS.map(({ method, mode }) => ({ method, mode }))) @@ -213,6 +272,10 @@ describe('cell inventory lock call-site census', () => { expect(rawSites).toEqual(INLINE_CELL_LOCK_SITES) }) + it('takes the host reservation rows before the shared cell row everywhere', () => { + expect(pathsTakingCellsBeforeReservations(storeSource())).toEqual([]) + }) + 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-accept-cell-row-lock-postgres.test.ts b/cloud/apps/relay/src/control-accept-cell-row-lock-postgres.test.ts new file mode 100644 index 00000000000..11a7b4d1b0e --- /dev/null +++ b/cloud/apps/relay/src/control-accept-cell-row-lock-postgres.test.ts @@ -0,0 +1,199 @@ +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 + +const cell = { + id: 'accept-lock-postgres', + url: 'https://accept-lock-postgres.example.com', + capacityRequests: 2, + connectionHardCap: 600 as const, + connectionUnobservedBound: 50 +} +const userId = 'accept-lock-postgres-user' +const first = { userId, relayHostId: 'acceptlockhost01' } +const second = { userId, relayHostId: 'acceptlockhost02' } + +// Why: the accept path now enforces the capacity ceiling inside its single +// conditional cell-row write instead of behind a SELECT ... FOR UPDATE it held +// for the rest of the transaction. Two accepts reaching for the same last slot +// are what would expose a lost update if that check were no longer atomic. +describePostgres('PostgreSQL control accept without 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 = ?`, + [userId] + ) + for (const table of [ + 'relay_control_capabilities', + '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 = ?`, [userId]) + } + 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('lets exactly one of two racing accepts take the last capacity slot', async () => { + await removeTestRows(databases[0]!) + const stores = databases.map((database) => new RelayAssignmentStore(database, () => 100)) + await prepareCell(stores[0]!) + + // Both hosts hold a grant, then drop the control the grant reserved, so the + // cell has exactly one free slot and two accepts that each want it. + const epochs = new Map() + for (const identity of [first, second]) { + const assignment = await stores[0]!.assign(identity) + epochs.set(identity.relayHostId, assignment.assignmentEpoch) + const control = await stores[0]!.activateControl(identity, { + cellId: cell.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + await stores[0]!.releaseActivity(identity, control) + } + await databases[0]!.query( + `UPDATE relay_cells SET reserved_requests = ? WHERE cell_id = ?`, + [cell.capacityRequests - 1, cell.id] + ) + await databases[0]!.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, 'splice:ballast', 'splice', ?, 1, 90100, 100)`, + [userId, 'acceptlockhost03', cell.id] + ) + + const outcomes = await Promise.allSettled([ + stores[0]!.activateControl(first, { + cellId: cell.id, + assignmentEpoch: epochs.get(first.relayHostId)!, + generation: 2 + }), + stores[1]!.activateControl(second, { + cellId: cell.id, + assignmentEpoch: epochs.get(second.relayHostId)!, + generation: 2 + }) + ]) + + expect(outcomes.filter((outcome) => outcome.status === 'fulfilled')).toHaveLength(1) + const rejection = outcomes.find((outcome) => outcome.status === 'rejected') + expect(String((rejection as PromiseRejectedResult).reason)).toContain( + 'relay_capacity_exhausted' + ) + const cells = await databases[0]!.query( + `SELECT reserved_requests FROM relay_cells WHERE cell_id = ?`, + [cell.id] + ) + expect(Number(cells[0]!.reserved_requests)).toBe(cell.capacityRequests) + const units = await databases[0]!.query( + `SELECT COALESCE(SUM(request_units), 0) AS units + FROM relay_assignment_activity_leases WHERE cell_id = ?`, + [cell.id] + ) + expect(Number(units[0]!.units)).toBe(cell.capacityRequests) + }, 20_000) + + it('rebinds a control while another connection holds the cell row', async () => { + await removeTestRows(databases[0]!) + const store = new RelayAssignmentStore(databases[0]!, () => 100) + await prepareCell(store) + const assignment = await store.assign(first) + await store.activateControl(first, { + cellId: cell.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + + let release!: () => void + const released = new Promise((resolve) => { + release = resolve + }) + let held!: () => void + const heldPromise = new Promise((resolve) => { + held = resolve + }) + const holder = databases[1]!.transaction(async (transaction) => { + await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [cell.id]) + held() + await released + }) + await heldPromise + + // Retiring generation 1 and installing generation 2 leaves the cell's + // reservation where it was, so the accept has no reason to wait on the row + // at all. Reading it up front is what used to make it wait, and then fail + // at the request-path lock bound. + try { + await expect( + store.activateControl(first, { + cellId: cell.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 2 + }) + ).resolves.toBe(`control:${cell.id}:2`) + } finally { + release() + await holder + } + + const cells = await databases[0]!.query( + `SELECT reserved_requests FROM relay_cells WHERE cell_id = ?`, + [cell.id] + ) + const units = await databases[0]!.query( + `SELECT COALESCE(SUM(request_units), 0) AS units + FROM relay_assignment_activity_leases WHERE cell_id = ?`, + [cell.id] + ) + expect(Number(cells[0]!.reserved_requests)).toBe(Number(units[0]!.units)) + }, 20_000) + + async function prepareCell(store: RelayAssignmentStore): Promise { + await store.reconcileCells([cell]) + await store.recordCellHeartbeat({ + 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, + connectionUnobservedBound: 50 + }) + } + +}) diff --git a/cloud/apps/relay/src/control-accept-cell-row-lock.test.ts b/cloud/apps/relay/src/control-accept-cell-row-lock.test.ts new file mode 100644 index 00000000000..e6136b18e5e --- /dev/null +++ b/cloud/apps/relay/src/control-accept-cell-row-lock.test.ts @@ -0,0 +1,148 @@ +import { beforeEach, describe, expect, it } from 'vitest' +import type { RelayCellConfig } from './config.js' +import { RelayAssignmentStore } from './assignment-store.js' +import { + openInMemoryRelayDatabase, + type RelayDatabase, + type RelayLockOptions, + type RelayTransactionOptions, + type SqlRow +} from './database.js' + +const CELL: RelayCellConfig = { + id: 'accept-cell-a', + url: 'https://accept-a.example.com', + capacityRequests: 2 +} +const host = { userId: 'accept-user', relayHostId: 'acceptho00000001' } +const second = { userId: 'accept-user', relayHostId: 'acceptho00000002' } +const third = { userId: 'accept-user', relayHostId: 'acceptho00000003' } + +type Statement = { sql: string; locked: boolean } + +// Records the statements a transaction issues, in order, so what the accept +// path does with the shared cell row can be asserted rather than described. +class RecordingDatabase implements RelayDatabase { + constructor( + private readonly inner: RelayDatabase, + readonly statements: Statement[] = [] + ) {} + + async query(sql: string, params: unknown[] = []): Promise { + this.statements.push({ sql, locked: false }) + return await this.inner.query(sql, params) + } + + async queryLocked( + sql: string, + params: unknown[] = [], + options: RelayLockOptions = {} + ): Promise { + this.statements.push({ sql, locked: true }) + return await this.inner.queryLocked(sql, params, options) + } + + async transaction( + operation: (transaction: RelayDatabase) => Promise, + options: RelayTransactionOptions = {} + ): Promise { + return await this.inner.transaction( + async (transaction) => + await operation(new RecordingDatabase(transaction, this.statements)), + options + ) + } + + async close(): Promise { + await this.inner.close() + } +} + +describe('control accept cell-row lock span', () => { + let recorder: RecordingDatabase + let store: RelayAssignmentStore + let assignmentEpoch: number + + beforeEach(async () => { + recorder = new RecordingDatabase(await openInMemoryRelayDatabase()) + store = new RelayAssignmentStore(recorder, () => 100) + await store.reconcileCells([CELL]) + assignmentEpoch = (await store.assign(host)).assignmentEpoch + }) + + async function accept(generation: number): Promise { + recorder.statements.length = 0 + return await store.activateControl(host, { + cellId: CELL.id, + assignmentEpoch, + generation + }) + } + + function cellStatements(): Statement[] { + return recorder.statements.filter((statement) => /relay_cells/.test(statement.sql)) + } + + async function reservedRequests(): Promise { + const row = ( + await recorder.query(`SELECT reserved_requests FROM relay_cells WHERE cell_id = ?`, [ + CELL.id + ]) + )[0]! + return Number(row.reserved_requests) + } + + async function cellLeaseUnits(): Promise { + const row = ( + await recorder.query( + `SELECT COALESCE(SUM(request_units), 0) AS units + FROM relay_assignment_activity_leases WHERE cell_id = ?`, + [CELL.id] + ) + )[0]! + return Number(row.units) + } + + it('writes the shared cell row once, last, and never takes it as a read lock', async () => { + const control = await accept(1) + await store.releaseActivity(host, control) + + await accept(2) + + const cells = cellStatements() + expect(cells.map((statement) => statement.locked)).toEqual([false]) + expect(cells[0]!.sql).toContain('RETURNING cell_id') + // The contended row is written by the last statement of the transaction, so + // its write lock is held across the commit alone, not the whole accept. + expect(recorder.statements.at(-1)).toBe(cells[0]) + }) + + it('leaves the cell row untouched when a rebind retires and installs one control', async () => { + await accept(1) + + await accept(2) + + expect(cellStatements()).toEqual([]) + expect(await reservedRequests()).toBe(1) + expect(await cellLeaseUnits()).toBe(1) + }) + + it('keeps the reservation equal to the cell lease units across repeated rebinds', async () => { + for (const generation of [1, 2, 3, 4, 5]) await accept(generation) + + expect(await reservedRequests()).toBe(await cellLeaseUnits()) + expect(await reservedRequests()).toBe(1) + }) + + it('still refuses an accept that would exceed the cell capacity', async () => { + const control = await accept(1) + await store.assign(second) + await store.releaseActivity(host, control) + await store.assign(third) + expect(await reservedRequests()).toBe(CELL.capacityRequests) + + await expect(accept(2)).rejects.toThrow('relay_capacity_exhausted') + expect(await reservedRequests()).toBe(CELL.capacityRequests) + expect(await cellLeaseUnits()).toBe(CELL.capacityRequests) + }) +}) diff --git a/tests/e2e/markdown-fence-links.spec.ts b/tests/e2e/markdown-fence-links.spec.ts new file mode 100644 index 00000000000..80c6bfbb553 --- /dev/null +++ b/tests/e2e/markdown-fence-links.spec.ts @@ -0,0 +1,42 @@ +import { test, expect } from './helpers/orca-app' +import { waitForActiveWorktree, waitForSessionReady } from './helpers/store' +import { + cleanupMarkdownFixture, + createMarkdownFixture, + getActiveWorktreeContext, + openMarkdownFixture +} from './helpers/markdown-editor-fixture' + +test('source links stay outside mixed code fences', async ({ orcaPage }, testInfo) => { + await waitForSessionReady(orcaPage) + await waitForActiveWorktree(orcaPage) + const context = await getActiveWorktreeContext(orcaPage) + const file = await createMarkdownFixture( + context, + 'markdown-fences', + 'links', + testInfo.workerIndex, + '# Fence boundaries\n\n~~~text\n```\n[[inside-code]]\n~~~\n\n[[outside-code]]\n' + ) + try { + await openMarkdownFixture(orcaPage, context, file) + await orcaPage.evaluate(() => { + const state = window.__store!.getState() + if (!state.activeFileId) { + throw new Error('missing active file') + } + state.setMarkdownViewMode(state.activeFileId, 'source') + }) + const links = orcaPage.locator('.monaco-editor .view-line').filter({ + has: orcaPage.locator('.monaco-markdown-doc-link') + }) + await expect(links).toHaveCount(1) + await expect(links).toHaveText('[[outside-code]]') + await testInfo.attach('fence-links', { + body: await orcaPage.screenshot({ path: testInfo.outputPath('fence-links.png') }), + contentType: 'image/png' + }) + } finally { + await cleanupMarkdownFixture(file) + } +})