diff --git a/cloud/apps/relay/src/app.ts b/cloud/apps/relay/src/app.ts index b5c76a57f86..31e53a5d852 100644 --- a/cloud/apps/relay/src/app.ts +++ b/cloud/apps/relay/src/app.ts @@ -1614,7 +1614,11 @@ const AdminAdmissionSelectorApplySchema = z attemptId: AdmissionSelectorAttemptIdSchema, expectedGeneration: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER), expectedMembershipSha256: z.string().regex(/^[a-f0-9]{64}$/).optional(), - membership: AdmissionSelectorMembershipSchema + membership: AdmissionSelectorMembershipSchema, + // Optional so an older caller reaches an updated director unchanged, and an + // updated caller reaching an older director is simply ignored: either way + // the cell goes unmarked and its hosts stay pinned, today's behaviour. + rollIsolatedCells: z.array(CellIdSchema).max(256).optional() }) .strict() .refine( diff --git a/cloud/apps/relay/src/assignment-connection-headroom.test.ts b/cloud/apps/relay/src/assignment-connection-headroom.test.ts index ed6b25bfb66..21e1d6aa110 100644 --- a/cloud/apps/relay/src/assignment-connection-headroom.test.ts +++ b/cloud/apps/relay/src/assignment-connection-headroom.test.ts @@ -469,10 +469,10 @@ describe('relay assignment connection headroom', () => { expect(await database.query(`SELECT * FROM relay_assignment_activity_leases`)).toEqual([]) }) - it('re-places a zero-activity assignment off a migration-only cell', async () => { - // Why: migration-only is what a roll's isolate step writes. The source's - // own headroom is irrelevant — the cell refuses every attach either way. - const { database, store, source, target } = await setupHeadroomReassignment() + it('keeps a zero-activity assignment pinned on an unstamped migration-only cell', async () => { + // Why: an evacuation target is deliberately migration-only and deliberately + // full. Without the roll's isolate stamp it must keep the hosts it holds. + const { database, store, source } = await setupHeadroomReassignment() const identity = { userId: 'pinned-migration-only-user', relayHostId: 'pinnedmigrationonly' } await store.setCellAdmissionState(source.id, 'migration-only') await database.query( @@ -484,9 +484,10 @@ describe('relay assignment connection headroom', () => { [identity.userId, identity.relayHostId, source.id, 7, 10_000, 100] ) - expect(await store.assign(identity)).toMatchObject({ - cellId: target.id, - assignmentEpoch: 8 + await expect(store.assign(identity)).rejects.toThrow('relay_connection_headroom_exhausted') + expect(await store.resolve(identity)).toMatchObject({ + cellId: source.id, + assignmentEpoch: 7 }) }) diff --git a/cloud/apps/relay/src/assignment-isolated-cell-replacement-postgres.test.ts b/cloud/apps/relay/src/assignment-isolated-cell-replacement-postgres.test.ts index 5839856b88f..96074e7206b 100644 --- a/cloud/apps/relay/src/assignment-isolated-cell-replacement-postgres.test.ts +++ b/cloud/apps/relay/src/assignment-isolated-cell-replacement-postgres.test.ts @@ -1,5 +1,11 @@ +import { createHash } from 'node:crypto' import { afterAll, beforeAll, describe, expect, it } from 'vitest' import { RelayAssignmentStore } from './assignment-store.js' +import { + encodeMembership, + type CellAdmissionMembership, + type CellAdmissionState +} from './cell-admission-selector.js' import type { RelayCellConfig } from './config.js' import { openRelayDatabase, type RelayDatabase } from './database.js' @@ -76,11 +82,82 @@ describePostgres('PostgreSQL re-placement off a cell isolated for a roll', () => } } + // The real isolate path: one selector apply, generation CAS and all, naming + // the cells it stamps. Nothing else in the fleet may write the stamp. + async function applySelector( + states: Record, + rollIsolatedCells?: string[] + ): Promise { + const current = await stores[0]!.inspectCellAdmissionSelector() + // Built from relay_cells, not from the selector's own membership: the apply + // requires exact coverage of the fleet, and this database is shared. + const fleet = await databases[0]!.query( + `SELECT cell_id FROM relay_cells ORDER BY cell_id ASC` + ) + const membership: CellAdmissionMembership = { + existingOnly: [], + migrationOnly: [], + general: [] + } + for (const row of fleet) { + const cellId = String(row['cell_id']) + const state = + states[cellId] ?? + (current.selector.membership.existingOnly.includes(cellId) + ? 'existing-only' + : current.selector.membership.migrationOnly.includes(cellId) + ? 'migration-only' + : 'general') + if (state === 'existing-only') membership.existingOnly.push(cellId) + else if (state === 'migration-only') membership.migrationOnly.push(cellId) + else membership.general.push(cellId) + } + await stores[0]!.applyCellAdmissionSelector({ + attemptId: `isolated-${current.selector.generation}-${Object.keys(states).join('-')}`, + expectedGeneration: current.selector.generation, + ...(current.selector.generation === 0 + ? { + expectedMembershipSha256: createHash('sha256') + .update(encodeMembership(current.selector.membership)) + .digest('hex') + } + : {}), + membership, + ...(rollIsolatedCells ? { rollIsolatedCells } : {}) + }) + } + + async function rollIsolatedAt(cellId: string): Promise { + const rows = await databases[0]!.query( + `SELECT roll_isolated_at FROM relay_cell_admission WHERE cell_id = ?`, + [cellId] + ) + const value = rows[0]?.['roll_isolated_at'] + return value === undefined || value === null ? null : Number(value) + } + // Every case starts from the whole fleet general; a case that leaves a cell // isolated would otherwise starve the next one of placement candidates. async function resetFleet(): Promise { await deleteHostRows() - for (const cell of CELLS) await stores[0]!.setCellAdmissionState(cell.id, 'general') + await applySelector(Object.fromEntries(CELLS.map((cell) => [cell.id, 'general']))) + } + + // The selector is fleet-wide and this database is shared with every other + // Postgres file in the project, all of which write admission through the + // generation-0 helpers. Advancing the generation and leaving it advanced + // would fail every one of them with admission_selector_boundary_active, so + // this file puts the boundary back exactly as it found it. + async function resetSelectorBoundary(): Promise { + await databases[0]!.query( + `UPDATE relay_admission_selectors SET generation = 0, attempt_id = NULL + WHERE selector_id = 'general'` + ) + await databases[0]!.query( + `DELETE FROM relay_admission_selector_intents WHERE attempt_id LIKE 'isolated-%'` + ) + // Rewrites membership_json from the live fleet, which generation 0 allows. + await stores[0]!.reconcileCells([], false) } async function deleteHostRows(): Promise { @@ -108,13 +185,19 @@ describePostgres('PostgreSQL re-placement off a cell isolated for a roll', () => }) ) await deleteHostRows() - for (const cell of CELLS) await stores[0]!.configureCell(cell, 'general') + // Registers the cell rows without touching admission, so this works at any + // selector generation the shared database happens to be sitting at. + await stores[0]!.reconcileCells(CELLS, false) + await resetSelectorBoundary() await heartbeatAll() }) afterAll(async () => { if (databases[0]) { await deleteHostRows() + // Order matters: drop the cells first, then rebuild the selector's + // membership from what is left, or it names rows that no longer exist and + // every later read fails admission_selector_membership_drift. for (const cell of CELLS) { for (const table of [ 'relay_cell_connection_snapshots', @@ -128,10 +211,33 @@ describePostgres('PostgreSQL re-placement off a cell isolated for a roll', () => await databases[0].query(`DELETE FROM ${table} WHERE cell_id = ?`, [cell.id]) } } + await resetSelectorBoundary() } for (const connection of databases) await connection.close() }) + it('stamps only the named cell on isolate and clears it on restore', async () => { + await resetFleet() + expect(await rollIsolatedAt(ISOLATED.id)).toBeNull() + + await applySelector({ [ISOLATED.id]: 'migration-only' }, [ISOLATED.id]) + const stamped = await rollIsolatedAt(ISOLATED.id) + expect(stamped).toBe(NOW) + // Cells the isolate did not name stay unmarked even when parked in the same + // apply: that is every non-roll flow, and its hosts must keep their pin. + await applySelector({ [TARGETS[0]!.id]: 'migration-only' }) + expect(await rollIsolatedAt(TARGETS[0]!.id)).toBeNull() + + // A failed wave re-isolates rather than restoring; the stamp survives, and + // does not move, so the hosts left behind stay eligible. + await applySelector({ [ISOLATED.id]: 'migration-only' }, [ISOLATED.id]) + expect(await rollIsolatedAt(ISOLATED.id)).toBe(stamped) + + // Restore writes 'general', and the same statement clears the stamp. + await applySelector({ [ISOLATED.id]: 'general' }) + expect(await rollIsolatedAt(ISOLATED.id)).toBeNull() + }, 30_000) + it('re-places every host off an isolated cell without leaking a reservation', async () => { await resetFleet() const identities = Array.from({ length: HOST_COUNT }, (_, index) => hostIdentity(index)) @@ -140,20 +246,19 @@ describePostgres('PostgreSQL re-placement off a cell isolated for a roll', () => (await reservedRequests(TARGETS[0]!.id)) + (await reservedRequests(TARGETS[1]!.id)) // Everyone lands on the cell about to be isolated. - await stores[0]!.configureCell(TARGETS[0]!, 'migration-only') - await stores[0]!.configureCell(TARGETS[1]!, 'migration-only') + await applySelector({ [TARGETS[0]!.id]: 'migration-only', [TARGETS[1]!.id]: 'migration-only' }) const first = new Map() for (const identity of identities) { const grant = await stores[0]!.assign(identity, 'us-central1') expect(grant.cellId).toBe(ISOLATED.id) first.set(identity.relayHostId, grant) } - await stores[0]!.configureCell(TARGETS[0]!, 'general') - await stores[0]!.configureCell(TARGETS[1]!, 'general') + await applySelector({ [TARGETS[0]!.id]: 'general', [TARGETS[1]!.id]: 'general' }) - // The isolate step: admission and enabled move together under the - // fleet-wide relay_cells lock, so there is no torn state to race against. - await stores[0]!.setCellAdmissionState(ISOLATED.id, 'migration-only') + // The isolate step. The state and the stamp are written by one UPDATE under + // the fleet-wide relay_cells lock, so there is no torn state to race against. + await applySelector({ [ISOLATED.id]: 'migration-only' }, [ISOLATED.id]) + expect(await rollIsolatedAt(ISOLATED.id)).not.toBeNull() const grants = await Promise.all( identities.map( @@ -186,17 +291,15 @@ describePostgres('PostgreSQL re-placement off a cell isolated for a roll', () => it('lets exactly one of a host’s racing dials win the re-placement', async () => { await resetFleet() const identity = hostIdentity(900) - await stores[0]!.configureCell(TARGETS[0]!, 'migration-only') - await stores[0]!.configureCell(TARGETS[1]!, 'migration-only') + await applySelector({ [TARGETS[0]!.id]: 'migration-only', [TARGETS[1]!.id]: 'migration-only' }) const first = await stores[0]!.assign(identity, 'us-central1') expect(first.cellId).toBe(ISOLATED.id) - await stores[0]!.configureCell(TARGETS[0]!, 'general') - await stores[0]!.configureCell(TARGETS[1]!, 'general') + await applySelector({ [TARGETS[0]!.id]: 'general', [TARGETS[1]!.id]: 'general' }) const targetBaseline = (await reservedRequests(TARGETS[0]!.id)) + (await reservedRequests(TARGETS[1]!.id)) const sourceBefore = await reservedRequests(ISOLATED.id) - await stores[0]!.setCellAdmissionState(ISOLATED.id, 'migration-only') + await applySelector({ [ISOLATED.id]: 'migration-only' }, [ISOLATED.id]) // Two dials from the same host, on separate connections, through the same // sticky lane. The per-assignment row lock is what must serialise them. @@ -226,7 +329,7 @@ describePostgres('PostgreSQL re-placement off a cell isolated for a roll', () => await resetFleet() const identity = hostIdentity(901) const first = await stores[0]!.assign(identity, 'us-central1') - await stores[0]!.setCellAdmissionState(ISOLATED.id, 'migration-only') + await applySelector({ [ISOLATED.id]: 'migration-only' }, [ISOLATED.id]) for (let dial = 0; dial < 5; dial += 1) { const grant = await stores[1 + (dial % (stores.length - 1))]!.assign( diff --git a/cloud/apps/relay/src/assignment-isolated-cell-replacement.test.ts b/cloud/apps/relay/src/assignment-isolated-cell-replacement.test.ts index a5350d48276..9fcb67ff56b 100644 --- a/cloud/apps/relay/src/assignment-isolated-cell-replacement.test.ts +++ b/cloud/apps/relay/src/assignment-isolated-cell-replacement.test.ts @@ -1,5 +1,11 @@ +import { createHash } from 'node:crypto' import { afterEach, describe, expect, it, vi } from 'vitest' import { RelayAssignmentStore, RelayHomeCellUnavailableError } from './assignment-store.js' +import { + encodeMembership, + type CellAdmissionMembership, + type CellAdmissionState +} from './cell-admission-selector.js' import type { RelayCellConfig } from './config.js' import { openInMemoryRelayDatabase, @@ -9,17 +15,20 @@ import { type SqlRow } from './database.js' -// A roll's isolate step writes 'migration-only' and its restore writes -// 'general' (cloud/dev/scripts/prepare-relay-production-capacity-canary.mjs:136, -// driven from cloud-deploy-relay-production-same-cap-job.yml:551-566 and -// cloud-deploy-relay-production-capacity-job.yml:768). 'existing-only' is C3's -// decommission posture and is never written by a roll. +// `migration-only` alone says nothing about a roll: it is the admission class +// (cloud/docs/orca-relay-operations.md:229-233) that evacuation targets, Asia +// `--mode rollback`, a failed wave's re-isolate and newly registered cells all +// occupy durably while holding hosts. The same-cap roll's isolate step is the +// only writer of the roll stamp, via `rollIsolatedCells` on the selector apply +// (cloud/dev/scripts/prepare-relay-production-capacity-canary.mjs isolate mode); +// restore moves the cell to 'general', which clears the stamp in the same +// statement that writes the state. const HEARTBEAT_TTL_MS = 45_000 const START_MS = 100 const IDENTITY = { userId: 'user-a', relayHostId: 'host000000000001' } // A connection-limited cell is what makes deadCellRequiresCommittedFence true, -// which is the branch that used to answer 503 relay_home_cell_unavailable. +// which is the branch that answers 503 relay_home_cell_unavailable. const CAPPED = { capacityRequests: 1_000, connectionHardCap: 600, @@ -72,7 +81,7 @@ class QueryCountingDatabase implements RelayDatabase { options?: RelayTransactionOptions ): Promise { return await this.delegate.transaction( - async (inner) => await operation(new QueryCountingDatabase(this.recording(inner))), + async (inner) => await operation(this.recording(inner)), options ) } @@ -81,8 +90,8 @@ class QueryCountingDatabase implements RelayDatabase { await this.delegate.close() } - // Nested transactions get their own wrapper; share this one's tape so a whole - // assign() is measured as a unit. + // A nested transaction shares this one's tape, so a whole assign() is + // measured as a unit. private recording(inner: RelayDatabase): RelayDatabase { const tape = this.sql return { @@ -105,8 +114,15 @@ interface Harness { store: RelayAssignmentStore database: RelayDatabase counter: QueryCountingDatabase - heartbeat: (cell: RelayCellConfig) => Promise + heartbeat: (cell: RelayCellConfig, at?: number) => Promise setNow: (value: number) => void + /** The real isolate path: one selector apply that names the cell it stamps. */ + isolateForRoll: (cellId: string) => Promise + /** The real restore path: back to 'general', which clears the stamp. */ + restore: (cellId: string) => Promise + /** An admission move with no stamp — every flow that is not a same-cap roll. */ + park: (cellId: string, state: CellAdmissionState) => Promise + rollIsolatedAt: (cellId: string) => Promise } async function setup(cells: RelayCellConfig[] = CELLS): Promise { @@ -119,37 +135,105 @@ async function setup(cells: RelayCellConfig[] = CELLS): Promise { heartbeatTtlMs: HEARTBEAT_TTL_MS }) await store.reconcileCells(cells, true) - const heartbeat = async (cell: RelayCellConfig): Promise => { - await store.recordCellHeartbeat({ - cellId: cell.id, - cellUrl: cell.url, - cellIncarnation: `1111111${cells.indexOf(cell)}-1111-4111-8111-111111111111`, - startedAt: 50, - ready: true, - observedRequests: 0, - region: cell.region, - totalConnections: 0, - inFlightConnections: 0, - reservedConnectionUnits: 0, - enforcedConnectionUnits: 0, - connectionHardCap: cell.connectionHardCap ?? 600, - connectionUnobservedBound: cell.connectionUnobservedBound ?? 50 - }) + const heartbeat = async (cell: RelayCellConfig, at?: number): Promise => { + const previous = now + if (at !== undefined) now = at + try { + await store.recordCellHeartbeat({ + cellId: cell.id, + cellUrl: cell.url, + cellIncarnation: `1111111${cells.indexOf(cell)}-1111-4111-8111-111111111111`, + startedAt: 50, + ready: true, + observedRequests: 0, + region: cell.region, + totalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0, + connectionHardCap: cell.connectionHardCap ?? 600, + connectionUnobservedBound: cell.connectionUnobservedBound ?? 50 + }) + } finally { + if (at !== undefined) now = previous + } } for (const cell of cells) await heartbeat(cell) - return { store, database: inner, counter, heartbeat, setNow: (value: number) => (now = value) } + + const applySelector = async ( + states: Record, + rollIsolatedCells?: string[] + ): Promise => { + const current = await store.inspectCellAdmissionSelector() + const membership: CellAdmissionMembership = { + existingOnly: [], + migrationOnly: [], + general: [] + } + for (const cell of cells) { + const state = + states[cell.id] ?? + (current.selector.membership.migrationOnly.includes(cell.id) + ? 'migration-only' + : current.selector.membership.existingOnly.includes(cell.id) + ? 'existing-only' + : 'general') + if (state === 'migration-only') membership.migrationOnly.push(cell.id) + else if (state === 'existing-only') membership.existingOnly.push(cell.id) + else membership.general.push(cell.id) + } + await store.applyCellAdmissionSelector({ + attemptId: `attempt-${current.selector.generation}-${Object.keys(states).join('-')}`, + expectedGeneration: current.selector.generation, + ...(current.selector.generation === 0 + ? { + expectedMembershipSha256: createHash('sha256') + .update(encodeMembership(current.selector.membership)) + .digest('hex') + } + : {}), + membership, + ...(rollIsolatedCells ? { rollIsolatedCells } : {}) + }) + } + + return { + store, + database: inner, + counter, + heartbeat, + setNow: (value: number) => (now = value), + isolateForRoll: async (cellId) => + await applySelector({ [cellId]: 'migration-only' }, [cellId]), + restore: async (cellId) => await applySelector({ [cellId]: 'general' }), + park: async (cellId, state) => await applySelector({ [cellId]: state }), + rollIsolatedAt: async (cellId) => { + const rows = await inner.query( + `SELECT roll_isolated_at FROM relay_cell_admission WHERE cell_id = ?`, + [cellId] + ) + const value = rows[0]?.['roll_isolated_at'] + return value === undefined || value === null ? null : Number(value) + } + } } -async function openMigration( +async function insertMigration( database: RelayDatabase, - input: { sourceCellId: string; targetCellId: string; assignmentEpoch: number; leases: number } + input: { + sourceCellId: string + targetCellId: string + assignmentEpoch: number + leases: number + settled?: 'completed' | 'aborted' + } ): Promise { await database.query( `INSERT INTO relay_assignment_migrations (user_id, relay_host_id, source_cell_id, target_cell_id, previous_epoch, assignment_epoch, source_request_units, target_reserved_units, expires_at, target_registered_at, completed_at, aborted_at, created_at, updated_at) - VALUES (?, ?, ?, ?, ?, ?, 1, 1, ?, NULL, NULL, NULL, ?, ?)`, + VALUES (?, ?, ?, ?, ?, ?, 1, 1, ?, NULL, ?, ?, ?, ?)`, [ IDENTITY.userId, IDENTITY.relayHostId, @@ -158,6 +242,8 @@ async function openMigration( input.assignmentEpoch - 1, input.assignmentEpoch, START_MS + 900_000, + input.settled === 'completed' ? START_MS : null, + input.settled === 'aborted' ? START_MS : null, START_MS, START_MS ] @@ -169,6 +255,15 @@ async function openMigration( ) } +type ConsoleWarnSpy = { mock: { calls: unknown[][] } } + +function jsonEvents(warn: ConsoleWarnSpy, event: string): unknown[] { + return warn.mock.calls + .map((call) => (typeof call[0] === 'string' ? call[0] : '')) + .filter((line) => line.includes(`"${event}"`)) + .map((line) => JSON.parse(line) as unknown) +} + describe('re-placing a host off a cell isolated for a roll', () => { it('leaves a host on a live general cell untouched, at one added admission read', async () => { const { store, counter } = await setup() @@ -177,21 +272,37 @@ describe('re-placing a host off a cell isolated for a roll', () => { counter.sql.length = 0 const second = await store.assign(IDENTITY, 'us-central1') - // The grant is identical, not merely same-celled: the clock has not moved, - // so every field including the lease deadline must match. + // Identical, not merely same-celled: the clock has not moved, so every + // field including the lease deadline must match. expect(second).toEqual(first) + // One row read answers both the stranded rule and the roll-stamp check. expect(counter.count('relay_cell_admission')).toBe(1) - // The isolation guard's second read never runs for a general incumbent. expect(counter.count('relay_assignment_migrations')).toBe(0) // No placement: the sticky lane never reaches the fleet-wide inventory. expect(counter.count('ORDER BY cell_id ASC')).toBe(0) }) - it('re-places a host off a migration-only cell on its next dial', async () => { - const { store, counter } = await setup() + it('keeps a host pinned to a migration-only cell with no roll stamp', async () => { + // The guard for every non-roll flow that parks a loaded cell: an Asia + // `--mode rollback`, an evacuation target awaiting promotion, a failed + // wave's re-isolate. Moving these hosts would undo the operator's intent. + const { store, park, rollIsolatedAt } = await setup() const first = await store.assign(IDENTITY, 'us-central1') - await store.setCellAdmissionState(first.cellId, 'migration-only') + await park(first.cellId, 'migration-only') + expect(await rollIsolatedAt(first.cellId)).toBeNull() + expect(await store.assign(IDENTITY, 'us-central1')).toMatchObject({ + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch + }) + }) + + it('re-places a host once the roll stamp is set', async () => { + const { store, counter, isolateForRoll, rollIsolatedAt } = await setup() + const first = await store.assign(IDENTITY, 'us-central1') + await isolateForRoll(first.cellId) + + expect(await rollIsolatedAt(first.cellId)).toBe(START_MS) counter.sql.length = 0 const moved = await store.assign(IDENTITY, 'us-central1') @@ -206,13 +317,26 @@ describe('re-placing a host off a cell isolated for a roll', () => { }) }) + it('stops re-placing once restore clears the stamp', async () => { + const { store, isolateForRoll, restore, rollIsolatedAt } = await setup() + const first = await store.assign(IDENTITY, 'us-central1') + await isolateForRoll(first.cellId) + await restore(first.cellId) + + expect(await rollIsolatedAt(first.cellId)).toBeNull() + expect(await store.assign(IDENTITY, 'us-central1')).toMatchObject({ + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch + }) + }) + it('keeps a host pinned to an existing-only cell that still serves it', async () => { // Why: existing-only cells serve the hosts they already hold (PR #194). Only // assignmentStrandedOnUnservedCell may release that pin, and only on proof // the cell stopped serving this host. - const { store, database, heartbeat, setNow } = await setup() + const { store, database, heartbeat, park, setNow } = await setup() const first = await store.assign(IDENTITY, 'us-central1') - await store.setCellAdmissionState(first.cellId, 'existing-only') + await park(first.cellId, 'existing-only') await database.query( `INSERT INTO relay_assignment_activity_leases (user_id, relay_host_id, activity_id, activity_kind, cell_id, @@ -223,69 +347,147 @@ describe('re-placing a host off a cell isolated for a roll', () => { // Past the stranded rule's minimum grant age, which outlives the heartbeat TTL. setNow(START_MS + 61_000) - for (const cell of CELLS) await heartbeat(cell) + for (const cell of CELLS) await heartbeat(cell, START_MS + 61_000) expect(await store.assign(IDENTITY, 'us-central1')).toMatchObject({ cellId: first.cellId, assignmentEpoch: first.assignmentEpoch }) }) - it('never re-places onto the isolated cell or any other non-general cell', async () => { - const { store } = await setup() + it('keeps a host pinned on an unstamped cell whose migration has completed', async () => { + // The terminal state of a successful evacuation: the host's row points at + // the target, the migration is completed, and the cell waits migration-only + // for a separate promote dispatch. Re-placing here undoes the evacuation. + const { store, database, park } = await setup() + const first = await store.assign(IDENTITY, 'us-central1') + await insertMigration(database, { + sourceCellId: 'us-c2', + targetCellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + leases: 0, + settled: 'completed' + }) + await park(first.cellId, 'migration-only') + + expect(await store.assign(IDENTITY, 'us-central1')).toMatchObject({ + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch + }) + }) + + it('re-places a stamped cell whose migration has completed', async () => { + const { store, database, isolateForRoll } = await setup() + const first = await store.assign(IDENTITY, 'us-central1') + await insertMigration(database, { + sourceCellId: 'us-c2', + targetCellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + leases: 0, + settled: 'completed' + }) + await isolateForRoll(first.cellId) + + expect((await store.assign(IDENTITY, 'us-central1')).cellId).not.toBe(first.cellId) + }) + + it('keeps the pin while a migration lease is outstanding', async () => { + const { store, database, isolateForRoll } = await setup() + const first = await store.assign(IDENTITY, 'us-central1') + await insertMigration(database, { + sourceCellId: first.cellId, + targetCellId: 'us-c2', + assignmentEpoch: first.assignmentEpoch, + leases: 1 + }) + await isolateForRoll(first.cellId) + + expect(await store.assign(IDENTITY, 'us-central1')).toMatchObject({ + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch + }) + }) + + it('keeps the pin while a migration row is open but its lease has lapsed', async () => { + // Why: the durable relay_assignment_migrations row outlives the 15-minute + // lease the counter tracks, and rollBackStalledRegionalRehomes refuses to + // unwind it while the source is not general. Re-placing would strand it. + const { store, database, isolateForRoll } = await setup() + const first = await store.assign(IDENTITY, 'us-central1') + await insertMigration(database, { + sourceCellId: first.cellId, + targetCellId: 'us-c2', + assignmentEpoch: first.assignmentEpoch, + leases: 0 + }) + await isolateForRoll(first.cellId) + + expect(await store.assign(IDENTITY, 'us-central1')).toMatchObject({ + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch + }) + }) + + it('never re-places across a region boundary, and says so', async () => { + // Both US cells parked and only Asia general: ordinary placement would spill + // to asia-east2 through `preferred[0] ?? candidates[0]`. This path refuses. + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + const { store, isolateForRoll, park } = await setup() const first = await store.assign(IDENTITY, 'us-central1') const other = CELLS.find((cell) => cell.region === 'us-central1' && cell.id !== first.cellId)! - await store.setCellAdmissionState(first.cellId, 'migration-only') - await store.setCellAdmissionState(other.id, 'migration-only') + await park(other.id, 'migration-only') + await isolateForRoll(first.cellId) - // The only general cell left is out of region, so the fallback is forced. + warn.mockClear() expect(await store.assign(IDENTITY, 'us-central1')).toMatchObject({ - cellId: 'asia-c1', - region: 'asia-east2' + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch }) + expect(jsonEvents(warn, 'orca_relay_sticky_replacement_deferred')).toEqual([ + { + event: 'orca_relay_sticky_replacement_deferred', + reason: 'no_same_region_headroom', + cellId: first.cellId, + region: 'us-central1' + } + ]) + expect(jsonEvents(warn, 'orca_relay_sticky_replaced_off_isolated_cell')).toEqual([]) }) - it('keeps a re-placed host in its own region when that region has a general cell', async () => { - const { store } = await setup() + it('keeps a re-placed host in its own region', async () => { + const { store, isolateForRoll } = await setup() const first = await store.assign(IDENTITY, 'asia-east2') expect(first.region).toBe('asia-east2') - await store.setCellAdmissionState(first.cellId, 'migration-only') - // asia-c1 is the only Asia cell, so a region-honouring placement has to fall - // back to US; give Asia a second cell and it must stay. - const asiaSpare: RelayCellConfig = { - id: 'asia-c2', - url: 'https://asia-c2.example.com', - region: 'asia-east2', - ...CAPPED - } - await store.configureCell(asiaSpare, 'general') - await store.recordCellHeartbeat({ - cellId: asiaSpare.id, - cellUrl: asiaSpare.url, - cellIncarnation: '99999999-9999-4999-8999-999999999999', - startedAt: 50, - ready: true, - observedRequests: 0, - region: 'asia-east2', - totalConnections: 0, - inFlightConnections: 0, - reservedConnectionUnits: 0, - enforcedConnectionUnits: 0, - connectionHardCap: 600, - connectionUnobservedBound: 50 - }) + await isolateForRoll(first.cellId) + // asia-c1 is the only Asia cell, so the only in-region candidate is gone. expect(await store.assign(IDENTITY, 'asia-east2')).toMatchObject({ - cellId: 'asia-c2', - region: 'asia-east2' + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch }) }) - it('does not demand a committed fence for a live isolated cell', async () => { - // §1.8's regression guard: the isolated cell is capped, so the dead-cell - // fence branch would reject with 503 relay_home_cell_unavailable. - const { store } = await setup() + it('takes the dead-cell path when a stamped cell stops heartbeating', async () => { + // Why: the fence is the only proof that a cell we cannot reach has stopped + // serving the sockets it still holds. A stamp does not make it reachable. + const { store, heartbeat, isolateForRoll, setNow } = await setup() const first = await store.assign(IDENTITY, 'us-central1') - await store.setCellAdmissionState(first.cellId, 'migration-only') + await isolateForRoll(first.cellId) + + const stale = START_MS + HEARTBEAT_TTL_MS + 1_000 + setNow(stale) + for (const cell of CELLS) { + if (cell.id !== first.cellId) await heartbeat(cell, stale) + } + + await expect(store.assign(IDENTITY, 'us-central1')).rejects.toBeInstanceOf( + RelayHomeCellUnavailableError + ) + }) + + it('does not demand a committed fence while the stamped cell is live', async () => { + const { store, isolateForRoll } = await setup() + const first = await store.assign(IDENTITY, 'us-central1') + await isolateForRoll(first.cellId) const moved = await store.assign(IDENTITY, 'us-central1').catch((error: unknown) => error) expect(moved).not.toBeInstanceOf(RelayHomeCellUnavailableError) @@ -295,7 +497,7 @@ describe('re-placing a host off a cell isolated for a roll', () => { it('leaves the source cell activity leases in place', async () => { // Contrast with the fence and stranded paths, which delete them: an isolated // cell is alive and still owns drainable work behind those leases. - const { store, database } = await setup() + const { store, database, isolateForRoll } = await setup() const first = await store.assign(IDENTITY, 'us-central1') await database.query( `INSERT INTO relay_assignment_activity_leases @@ -304,7 +506,7 @@ describe('re-placing a host off a cell isolated for a roll', () => { VALUES (?, ?, 'splice:keep-1', 'splice', ?, 1, ?, ?)`, [IDENTITY.userId, IDENTITY.relayHostId, first.cellId, START_MS + 600_000, START_MS] ) - await store.setCellAdmissionState(first.cellId, 'migration-only') + await isolateForRoll(first.cellId) await store.assign(IDENTITY, 'us-central1') expect( @@ -316,82 +518,24 @@ describe('re-placing a host off a cell isolated for a roll', () => { ).toHaveLength(1) }) - it('keeps the pin while a migration lease is outstanding', async () => { - const { store, database } = await setup() - const first = await store.assign(IDENTITY, 'us-central1') - await openMigration(database, { - sourceCellId: first.cellId, - targetCellId: 'us-c2', - assignmentEpoch: first.assignmentEpoch, - leases: 1 - }) - await store.setCellAdmissionState(first.cellId, 'migration-only') - - expect(await store.assign(IDENTITY, 'us-central1')).toMatchObject({ - cellId: first.cellId, - assignmentEpoch: first.assignmentEpoch - }) - }) - - it('keeps the pin while a migration row is open but its lease has lapsed', async () => { - // Why: the durable relay_assignment_migrations row outlives the 15-minute - // lease the counter tracks, and rollBackStalledRegionalRehomes refuses to - // unwind it while the source is not general. Re-placing here would strand it. - const { store, database } = await setup() - const first = await store.assign(IDENTITY, 'us-central1') - await openMigration(database, { - sourceCellId: first.cellId, - targetCellId: 'us-c2', - assignmentEpoch: first.assignmentEpoch, - leases: 0 - }) - await store.setCellAdmissionState(first.cellId, 'migration-only') - - expect(await store.assign(IDENTITY, 'us-central1')).toMatchObject({ - cellId: first.cellId, - assignmentEpoch: first.assignmentEpoch - }) - }) - - it('re-places once a settled migration is no longer open', async () => { - const { store, database } = await setup() - const first = await store.assign(IDENTITY, 'us-central1') - await openMigration(database, { - sourceCellId: first.cellId, - targetCellId: 'us-c2', - assignmentEpoch: first.assignmentEpoch, - leases: 0 - }) - await database.query( - `UPDATE relay_assignment_migrations SET aborted_at = ? - WHERE user_id = ? AND relay_host_id = ?`, - [START_MS, IDENTITY.userId, IDENTITY.relayHostId] - ) - await store.setCellAdmissionState(first.cellId, 'migration-only') - - expect((await store.assign(IDENTITY, 'us-central1')).cellId).not.toBe(first.cellId) - }) - it('logs one event naming both cells, the admission state and the region', async () => { const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) - const { store } = await setup() + const { store, isolateForRoll } = await setup() const first = await store.assign(IDENTITY, 'us-central1') - await store.setCellAdmissionState(first.cellId, 'migration-only') + await isolateForRoll(first.cellId) warn.mockClear() const moved = await store.assign(IDENTITY, 'us-central1') - const events = warn.mock.calls - .map(([line]) => (typeof line === 'string' ? line : '')) - .filter((line) => line.includes('orca_relay_sticky_replaced_off_isolated_cell')) - expect(events).toHaveLength(1) - expect(JSON.parse(events[0]!)).toEqual({ - event: 'orca_relay_sticky_replaced_off_isolated_cell', - fromCellId: first.cellId, - fromRegion: 'us-central1', - admissionState: 'migration-only', - toCellId: moved.cellId, - region: moved.region - }) + expect(jsonEvents(warn, 'orca_relay_sticky_replaced_off_isolated_cell')).toEqual([ + { + event: 'orca_relay_sticky_replaced_off_isolated_cell', + fromCellId: first.cellId, + fromRegion: 'us-central1', + admissionState: 'migration-only', + toCellId: moved.cellId, + region: moved.region + } + ]) }) }) diff --git a/cloud/apps/relay/src/assignment-store.test.ts b/cloud/apps/relay/src/assignment-store.test.ts index 47ae75ed684..57e3b9d25e1 100644 --- a/cloud/apps/relay/src/assignment-store.test.ts +++ b/cloud/apps/relay/src/assignment-store.test.ts @@ -1555,19 +1555,20 @@ describe('RelayAssignmentStore', () => { expect(pinned.cellId).toBe(first.cellId) }) - it('re-places a host off a migration-only cell isolated for a roll', async () => { - // Why: a roll's isolate step writes migration-only while the cell keeps - // heartbeating ready=1 and refuses every attach. The stranded rule above - // never fires here (it demands existing-only), so this is the only exit. + it('keeps migration-only cells pinned inside the stranded window', async () => { + // Why: migration-only is an admission class, not a drain signal. Evacuation + // targets, an Asia `--mode rollback` and a failed wave's re-isolate all park + // loaded cells there durably, and moving those hosts would undo the + // evacuation or scatter the region. Only the same-cap roll's isolate stamp + // releases the pin, and nothing here sets it. let now = 100 const store = await setup(() => now) const identity = { userId: 'user-a', relayHostId: 'host000000000001' } const first = await store.assign(identity) await store.setCellAdmissionState(first.cellId, 'migration-only') now += 61_000 - const moved = await store.assign(identity) - expect(moved.cellId).not.toBe(first.cellId) - expect(moved.assignmentEpoch).toBe(first.assignmentEpoch + 1) + const pinned = await store.assign(identity) + expect(pinned.cellId).toBe(first.cellId) }) it('leaves quiet existing-only assignments to the normal dormancy rule', async () => { diff --git a/cloud/apps/relay/src/assignment-store.ts b/cloud/apps/relay/src/assignment-store.ts index 9afe27ef69c..48497b11e86 100644 --- a/cloud/apps/relay/src/assignment-store.ts +++ b/cloud/apps/relay/src/assignment-store.ts @@ -37,6 +37,7 @@ import { cellAdmissionState, cellAdmissionStates, ensureCellAdmission, + isCellAdmissionState, parseCellAdmissionState, RelayCellAdmissionSelector, setCellAdmissionBeforeBoundary, @@ -742,10 +743,24 @@ export class RelayAssignmentStore { const activityLeases = await this.lockAssignmentActivities(transaction, identity, true) await this.recordRegionPreference(transaction, identity, preferredRegion, now) if (mayNormallyReassign(activity(existing), now)) return null + // One read serves the stranded rule and the roll-isolation check below; + // issuing it twice inside one sticky transaction is pure waste. + const pinnedAdmission = await this.pinnedCellAdmission( + transaction, + text(existing, 'cell_id') + ) // Why: a stranded host must fall through to placement — re-granting the // pinned cell here is what refreshes its own activity and sustains the // loop (issue #225). - if (await this.assignmentStrandedOnUnservedCell(transaction, identity, existing, now)) { + if ( + await this.assignmentStrandedOnUnservedCell( + transaction, + identity, + existing, + now, + pinnedAdmission?.state + ) + ) { return null } @@ -786,7 +801,15 @@ export class RelayAssignmentStore { // what an isolated incumbent needs — and the only way out, because // re-granting the pin refreshes the host's own activity and sustains the // loop. The placement lane re-places it on a cell that will take it. - if (await this.incumbentCellIsolatedForRoll(transaction, identity, existing)) { + if ( + await this.incumbentCellIsolatedForRoll( + transaction, + identity, + existing, + undefined, + pinnedAdmission + ) + ) { return null } @@ -839,7 +862,9 @@ export class RelayAssignmentStore { transaction: RelayDatabase, identity: AssignmentIdentity, existing: SqlRow, - now: number + now: number, + // Read by the caller when it needs the same row for another question. + admissionState?: CellAdmissionState ): Promise { const lastActivityAt = integer(existing, 'last_activity_at') if ( @@ -848,14 +873,11 @@ export class RelayAssignmentStore { ) { return false } - const admissionRow = ( - await transaction.query( - `SELECT admission_state FROM relay_cell_admission WHERE cell_id = ?`, - [text(existing, 'cell_id')] - ) - )[0] + const state = + admissionState ?? + (await this.pinnedCellAdmission(transaction, text(existing, 'cell_id')))?.state // Unknown or missing admission fails safe: the pin stays. - if (admissionRow?.['admission_state'] !== 'existing-only') { + if (state !== 'existing-only') { return false } const liveLeases = ( @@ -916,6 +938,7 @@ export class RelayAssignmentStore { let connectionHeadroomReassignment = false let strandedReassignment = false let isolatedIncumbent: CellRow | undefined + let isolatedTarget: CellRow | undefined if (existing && !mayNormallyReassign(activity(existing), now)) { lockedCells ??= await this.lockCellInventory(transaction, 'nowait') const admission = await cellAdmissionStates(transaction) @@ -936,20 +959,40 @@ export class RelayAssignmentStore { existing, now ) + const currentIsLive = + strandedReassignment || + !this.requireLiveCells || + (await this.cellIsLive(transaction, current.cellId, now)) // Why: the sticky lane sends an isolated incumbent here, and the gate // below would hand the pin straight back — the cell is live and, being - // emptied, has more headroom than anyone. + // emptied, has more headroom than anyone. A cell that is *not* live is + // left to the dead-cell path below, whose fence is the only proof that + // an unreachable cell has stopped serving the sockets it still holds. if ( !strandedReassignment && + currentIsLive && (await this.incumbentCellIsolatedForRoll(transaction, identity, existing, admission)) ) { - isolatedIncumbent = current + isolatedTarget = + (await this.leastLoadedCell(transaction, lockedCells, current.region, 'require')) ?? + undefined + if (isolatedTarget) { + isolatedIncumbent = current + } else { + // Keeping the pin is today's behaviour: the host keeps retrying its + // own cell. Scattering a region across the fleet is worse, and it + // cannot be undone without the rehome worker. + console.warn( + JSON.stringify({ + event: 'orca_relay_sticky_replacement_deferred', + reason: 'no_same_region_headroom', + cellId: current.cellId, + region: current.region + }) + ) + } } - if ( - !strandedReassignment && - !isolatedIncumbent && - (!this.requireLiveCells || (await this.cellIsLive(transaction, current.cellId, now))) - ) { + if (!strandedReassignment && !isolatedIncumbent && currentIsLive) { const hadControl = holdsControlLease( activityLeases, current.cellId, @@ -990,12 +1033,13 @@ export class RelayAssignmentStore { } connectionHeadroomReassignment = true } - // Why no fence for an isolated cell: a fence proves a cell we cannot + // Why no fence for a live isolated cell: a fence proves a cell we cannot // contact has stopped serving a host. This one is contactable and // enforces the invariant itself — verifyCellAssignment reads // relay_assignments live, so the epoch bump below makes it close any - // attach naming the old epoch with 4409. Taking the branch below would - // instead 503 the host for the whole drain. + // attach naming the old epoch with 4409. `isolatedIncumbent` is only set + // when the cell is live, so a cell that stopped heartbeating while still + // holding sockets still takes this branch. if (!strandedReassignment && !connectionHeadroomReassignment && !isolatedIncumbent) { if ( (await this.deadCellRequiresCommittedFence( @@ -1018,11 +1062,12 @@ export class RelayAssignmentStore { lockedCells ??= existing ? await this.lockCellInventory(transaction, 'nowait') : await this.lockGeneralCellInventory(transaction, 'nowait') - const target = await this.leastLoadedCell( - transaction, - lockedCells, - placementRegion - ) + // The isolated target was chosen under the same inventory lock, with + // region required rather than preferred; re-picking here would reopen the + // cross-region spill it exists to refuse. + const target = + isolatedTarget ?? + (await this.leastLoadedCell(transaction, lockedCells, placementRegion)) if (!target) throw new Error('relay_capacity_exhausted') if (isolatedIncumbent) { // The canary's proof that the fix fired: count these against the @@ -1158,6 +1203,7 @@ export class RelayAssignmentStore { expectedGeneration: number expectedMembershipSha256?: string membership: CellAdmissionMembership + rollIsolatedCells?: string[] }): Promise<{ changed: boolean selector: { @@ -7191,7 +7237,12 @@ export class RelayAssignmentStore { // Required: the one caller has already locked the inventory it selects from, // and an optional parameter left a second fleet-wide lock reachable here. rows: SqlRow[], - preferredRegion: RelayRegion + preferredRegion: RelayRegion, + // Ordinary placement treats region as a preference and spills globally + // rather than refuse a host a cell. Re-placing off an isolated cell is the + // one caller that must not: a whole region parked migration-only would send + // every one of its hosts to the fallback region, permanently. + regionMode: 'prefer' | 'require' = 'prefer' ): Promise { const regions = new Map( (await database.query(`SELECT cell_id, region FROM relay_cell_regions`)).map((row) => [ @@ -7237,7 +7288,7 @@ export class RelayAssignmentStore { (candidate) => (regions.get(text(candidate, 'cell_id')) ?? RELAY_DEFAULT_REGION) === preferredRegion ) - const selected = preferred[0] ?? candidates[0] + const selected = regionMode === 'require' ? preferred[0] : (preferred[0] ?? candidates[0]) return selected ? cell(selected, regions.get(text(selected, 'cell_id')) ?? RELAY_DEFAULT_REGION) : null @@ -7363,19 +7414,30 @@ export class RelayAssignmentStore { // cannot see that the host's home has stopped taking it. Re-granting the pin // is what turns a roll into 13-16 minutes of retry loop; returning true here // sends the host down the ordinary placement lane on its next dial instead. + // + // The stamp, not the admission state, is the signal. 'migration-only' is an + // admission class that evacuation targets, Asia `--mode rollback`, a failed + // wave's re-isolate and newly registered cells all occupy durably while + // holding hosts; moving those would undo an evacuation or scatter a region. private async incumbentCellIsolatedForRoll( database: RelayDatabase, identity: AssignmentIdentity, existing: SqlRow, // Placement has already read the whole table; sticky has not, and a single // pinned cell does not justify a second fleet-wide read. - admission?: ReadonlyMap + admission?: ReadonlyMap, + // Sticky has already read the pinned row for the stranded rule. + pinnedAdmission?: PinnedCellAdmission ): Promise { const cellId = text(existing, 'cell_id') - const state = admission - ? admission.get(cellId) - : await this.pinnedCellAdmissionState(database, cellId) - if (state !== ROLL_ISOLATED_ADMISSION) return false + // Placement's map answers the common "not parked at all" case for free; the + // stamp itself is only on the row, so a migration-only incumbent still pays + // the single-row read. + if (admission && admission.get(cellId) !== ROLL_ISOLATED_ADMISSION) return false + const pinned = pinnedAdmission ?? (await this.pinnedCellAdmission(database, cellId)) + if (pinned?.state !== ROLL_ISOLATED_ADMISSION || pinned.rollIsolatedAt === undefined) { + return false + } if (integer(existing, 'migration_leases') > 0) return false // A migration owns this assignment's epoch on both sides, and its durable // row outlives the 15-minute lease that the counter above tracks, so the @@ -7391,18 +7453,26 @@ export class RelayAssignmentStore { return integer(open!, 'open') === 0 } - private async pinnedCellAdmissionState( + // One read for both the class and the roll stamp, so the sticky hot path pays + // a single indexed row lookup rather than one per question. + private async pinnedCellAdmission( database: RelayDatabase, cellId: string - ): Promise { + ): Promise { const row = ( await database.query( - `SELECT admission_state FROM relay_cell_admission WHERE cell_id = ?`, + `SELECT admission_state, roll_isolated_at FROM relay_cell_admission WHERE cell_id = ?`, [cellId] ) )[0] - // Unknown or missing admission fails safe: the pin stays. - return row ? parseCellAdmissionState(text(row, 'admission_state')) : undefined + // Unknown, missing or unrecognised admission fails safe: the pin stays. This + // reader must not throw — it now sits on the sticky path every dial takes, + // and the rule it feeds is "move the host", so silence has to mean "don't". + if (!row || !isCellAdmissionState(row['admission_state'])) return undefined + return { + state: row['admission_state'], + rollIsolatedAt: optionalInteger(row, 'roll_isolated_at') + } } // Reports which of `cellIsLive`'s conditions failed, so the rejection log @@ -7989,6 +8059,13 @@ export class RelayAssignmentStore { } } +type PinnedCellAdmission = { + state: CellAdmissionState + // Set only by the same-cap roll's isolate step; cleared by any write that + // moves the cell out of 'migration-only'. + rollIsolatedAt: number | undefined +} + type CellRow = { cellId: string; cellUrl: string; region: RelayRegion } function cell(row: SqlRow, region: RelayRegion): CellRow { diff --git a/cloud/apps/relay/src/cell-admission-selector.ts b/cloud/apps/relay/src/cell-admission-selector.ts index 92764e53e40..dbbe5ab4505 100644 --- a/cloud/apps/relay/src/cell-admission-selector.ts +++ b/cloud/apps/relay/src/cell-admission-selector.ts @@ -38,6 +38,13 @@ type ApplySelectorInput = { expectedGeneration: number expectedMembershipSha256?: string membership: CellAdmissionMembership + // Why a marker rather than reading 'migration-only' directly: that state is an + // admission class, not a drain signal (orca-relay-operations.md:229-233). + // Evacuation targets, Asia `--mode rollback`, a failed wave's re-isolate and + // newly registered cells all sit there durably while holding hosts. Only the + // same-cap isolate step names cells here; every write out of 'migration-only' + // clears the stamp, so a restore cannot leave one behind. + rollIsolatedCells?: string[] } const SELECTOR_ID = 'general' @@ -50,6 +57,12 @@ export function enabledForState(state: CellAdmissionState): number { return state === 'existing-only' ? 0 : 1 } +// The narrowing counterpart of parseCellAdmissionState, for readers that must +// answer "unknown" rather than throw. +export function isCellAdmissionState(value: unknown): value is CellAdmissionState { + return CELL_ADMISSION_STATES.some((state) => state === value) +} + export function parseCellAdmissionState(value: string): CellAdmissionState { if (!CELL_ADMISSION_STATES.includes(value as CellAdmissionState)) { throw new Error('invalid_cell_admission_state') @@ -122,9 +135,10 @@ export async function setCellAdmissionBeforeBoundary( await transaction.query( `UPDATE relay_cell_admission SET updated_at = CASE WHEN admission_state <> ? THEN ? ELSE updated_at END, - admission_state = ? + admission_state = ?, + roll_isolated_at = CASE WHEN ? = 'migration-only' THEN roll_isolated_at ELSE NULL END WHERE cell_id = ?`, - [state, now, state, cellId] + [state, now, state, state, cellId] ) await synchronizeCellAdmissionBoundary(transaction, now) } @@ -184,14 +198,30 @@ export class RelayCellAdmissionSelector { } } const now = this.now() + const rollIsolated = new Set(input.rollIsolatedCells ?? []) for (const [state, cellIds] of membershipEntries(membership)) { for (const cellId of cellIds) { + // Three-way, in the same statement as the state so they cannot tear: + // stamp a named isolate (keeping an earlier stamp, so a failed wave's + // re-isolate stays marked), clear on any move out of 'migration-only', + // and leave an unnamed 'migration-only' cell exactly as it was. + const marker = + state === 'migration-only' + ? rollIsolated.has(cellId) + ? 'stamp' + : 'keep' + : 'clear' await transaction.query( `UPDATE relay_cell_admission SET updated_at = CASE WHEN admission_state <> ? THEN ? ELSE updated_at END, - admission_state = ? + admission_state = ?, + roll_isolated_at = CASE + WHEN ? = 'clear' THEN NULL + WHEN ? = 'stamp' THEN COALESCE(roll_isolated_at, ?) + ELSE roll_isolated_at + END WHERE cell_id = ?`, - [state, now, state, cellId] + [state, now, state, marker, marker, now, cellId] ) await transaction.query( `UPDATE relay_cells diff --git a/cloud/apps/relay/src/database.ts b/cloud/apps/relay/src/database.ts index 005716c0657..4a5cf1084aa 100644 --- a/cloud/apps/relay/src/database.ts +++ b/cloud/apps/relay/src/database.ts @@ -327,7 +327,8 @@ CREATE TABLE IF NOT EXISTS relay_cell_admission ( cell_id TEXT PRIMARY KEY, admission_state TEXT NOT NULL CHECK (admission_state IN ('existing-only', 'migration-only', 'general')), - updated_at BIGINT NOT NULL + updated_at BIGINT NOT NULL, + roll_isolated_at BIGINT ); CREATE TABLE IF NOT EXISTS relay_admission_selectors ( @@ -680,6 +681,12 @@ export const POSTGRES_SCHEMA_MIGRATIONS = [ // Nullable with no default, so the rewrite is catalog-only; every row // aborted before this column existed reads as an unattributed abort. `ALTER TABLE relay_region_rehome_attempts ADD COLUMN IF NOT EXISTS abort_reason TEXT`, + // migration-only alone cannot say why a cell is parked there: five flows park loaded cells in + // that state durably. This stamps only the same-cap roll's isolate step, so placement can tell a + // cell being rolled from an evacuation target or an Asia rollback. Nullable with no default, so + // the rewrite is catalog-only; a cell isolated before this column existed reads as unmarked and + // keeps its hosts pinned, which is the pre-existing behaviour. + `ALTER TABLE relay_cell_admission ADD COLUMN IF NOT EXISTS roll_isolated_at BIGINT`, // Dropped, not created: see the comment on relay_assignment_activity_leases. Deferrable because // this is the one boot where it has to take ACCESS EXCLUSIVE on a table under continuous write, // and all 28 directors reach it at once; a lock timeout here must not restart the instance, which diff --git a/cloud/apps/relay/src/relay-schema-lock-targets.test.ts b/cloud/apps/relay/src/relay-schema-lock-targets.test.ts index 6b46675f429..5f1a8a88514 100644 --- a/cloud/apps/relay/src/relay-schema-lock-targets.test.ts +++ b/cloud/apps/relay/src/relay-schema-lock-targets.test.ts @@ -127,6 +127,7 @@ const GOLDEN_LOCK_TAKING: SchemaLockTarget[] = [ { kind: 'column', table: 'relay_control_capabilities', name: 'idle_regional_rehome', skipWhen: 'present' }, { kind: 'column', table: 'relay_region_rehome_attempts', name: 'source_generation', skipWhen: 'present' }, { kind: 'column', table: 'relay_region_rehome_attempts', name: 'abort_reason', skipWhen: 'present' }, + { kind: 'column', table: 'relay_cell_admission', name: 'roll_isolated_at', skipWhen: 'present' }, { kind: 'index-by-name', name: 'relay_assignment_activity_expiry', skipWhen: 'absent' }, { kind: 'reloption', diff --git a/cloud/dev/scripts/prepare-relay-production-capacity-canary.mjs b/cloud/dev/scripts/prepare-relay-production-capacity-canary.mjs index 9787dbdd62a..5c3452e542c 100644 --- a/cloud/dev/scripts/prepare-relay-production-capacity-canary.mjs +++ b/cloud/dev/scripts/prepare-relay-production-capacity-canary.mjs @@ -136,12 +136,18 @@ export async function prepareProductionCapacityCell(config, overrides = {}) { const desiredState = config.mode === 'isolate' ? 'migration-only' : 'general' const membership = membershipWithStates(before.selector, { [config.cellId]: desiredState }) const result = await applyExactAdmissionSelector(post, membership, { - expectedCurrentSelector: before.selector + expectedCurrentSelector: before.selector, + // The stamp that tells the director this cell is parked for a restart rather + // than held back as capacity, so it may re-place the hosts still on it. The + // activate branch omits it, and moving to 'general' clears it in the same + // statement that writes the state. + ...(config.mode === 'isolate' ? { rollIsolatedCells: [config.cellId] } : {}) }) return { changed: result.changed, generation: result.selector.generation, - admissionState: desiredState + admissionState: desiredState, + rollIsolated: config.mode === 'isolate' } } diff --git a/cloud/dev/scripts/prepare-relay-production-capacity-canary.test.mjs b/cloud/dev/scripts/prepare-relay-production-capacity-canary.test.mjs index e721bfb8c75..1378e26416d 100644 --- a/cloud/dev/scripts/prepare-relay-production-capacity-canary.test.mjs +++ b/cloud/dev/scripts/prepare-relay-production-capacity-canary.test.mjs @@ -168,6 +168,36 @@ describe('production Relay capacity cell admission', () => { assert.doesNotMatch(fake.calls.map(({ path }) => path).join(','), /\/v1\/admin\/drain/) }) + it('stamps the roll isolation on isolate and never on activate', async () => { + // The stamp is what lets the director tell a cell parked for a restart from + // an evacuation target or an Asia rollback, both of which must keep their + // hosts. Only this call site may send it. + const isolate = canaryFetch() + await prepareProductionCapacityCell( + { ...config, mode: 'isolate' }, + { fetch: isolate.fetch, token: 'token' } + ) + const isolateApply = isolate.calls.find( + ({ path }) => path === '/v1/admin/admission-selector/apply' + ) + assert.deepEqual(isolateApply.body.rollIsolatedCells, [config.cellId]) + + // Restore has to be a real apply, not the no-op an already-general cell + // takes, or the assertion below proves nothing. + isolate.calls.length = 0 + const restored = await prepareProductionCapacityCell( + { ...config, mode: 'activate' }, + { fetch: isolate.fetch, token: 'token' } + ) + assert.equal(restored.admissionState, 'general') + const restoreApply = isolate.calls.find( + ({ path }) => path === '/v1/admin/admission-selector/apply' + ) + assert.ok(restoreApply, 'restore must issue an apply') + assert.equal(restoreApply.body.rollIsolatedCells, undefined) + assert.ok(restoreApply.body.membership.general.includes(config.cellId)) + }) + it('drains the selected cell independently after durable isolation', async () => { const fake = canaryFetch() const result = await prepareProductionCapacityCell( diff --git a/cloud/dev/scripts/relay-admission-selector.mjs b/cloud/dev/scripts/relay-admission-selector.mjs index f41528c7e7f..9b993fdc9ae 100644 --- a/cloud/dev/scripts/relay-admission-selector.mjs +++ b/cloud/dev/scripts/relay-admission-selector.mjs @@ -162,6 +162,9 @@ export async function applyExactAdmissionSelector(post, membership, options = {} ...(before.selector.generation === 0 ? { expectedMembershipSha256: membershipSha256(before.selector.membership) } : {}), + // Only a same-cap roll's isolate passes this; every other caller omits it + // and leaves the cell's hosts pinned, which is the pre-existing behaviour. + ...(options.rollIsolatedCells ? { rollIsolatedCells: options.rollIsolatedCells } : {}), membership: desired }) } catch (error) {