diff --git a/.github/workflows/cloud-deploy-relay-production-same-cap-job.yml b/.github/workflows/cloud-deploy-relay-production-same-cap-job.yml index fec048bf4b2..bf31b0b2f55 100644 --- a/.github/workflows/cloud-deploy-relay-production-same-cap-job.yml +++ b/.github/workflows/cloud-deploy-relay-production-same-cap-job.yml @@ -557,6 +557,10 @@ jobs: if test "${ENTRY_ADMISSION}" = migration-only; then jq -e '.changed == false' <<< "${ISOLATE_RESULT}" >/dev/null fi + # The director re-places hosts off a cell only when the isolate stamped it, so a + # script too old to ask for the stamp produces today's behaviour and the canary + # reads as "the fix did nothing" with nothing to tell that from a wrong premise. + jq -e '.rollIsolated == true' <<< "${ISOLATE_RESULT}" >/dev/null ISOLATE_GENERATION="$(jq -er '.generation' <<< "${ISOLATE_RESULT}")" echo "SELECTOR_GENERATION_AFTER_ISOLATE=${ISOLATE_GENERATION}" >> "${GITHUB_ENV}" node dev/scripts/prepare-relay-production-capacity-canary.mjs \ @@ -907,6 +911,9 @@ jobs: --director-origin "${DIRECTOR_ORIGIN}" --cell-origin "${CELL_ORIGIN}" \ --cell-id "${TARGET_CELL_ID}" --approved-cells same-cap --mode isolate)" echo "${ISOLATE_RESULT}" + # Same reason as the isolate step above: a failed wave leaves this cell isolated + # deliberately, and the stamp is what lets its hosts leave. + jq -e '.rollIsolated == true' <<< "${ISOLATE_RESULT}" >/dev/null # The isolate result carries the authoritative post-isolate generation; # fixed offsets are wrong whenever an earlier isolate was a no-op. FAILSAFE_GENERATION="$(jq -er '.generation' <<< "${ISOLATE_RESULT}")" diff --git a/cloud/apps/relay/src/app.ts b/cloud/apps/relay/src/app.ts index b5c76a57f86..cb3a47f7f01 100644 --- a/cloud/apps/relay/src/app.ts +++ b/cloud/apps/relay/src/app.ts @@ -1614,7 +1614,15 @@ 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 reaching an updated director is unchanged: + // the cell goes unmarked and its hosts stay pinned, today's behaviour. The + // other direction is NOT ignored — the schema below is .strict(), so an + // updated caller reaching an older director is a 400. That fails closed, + // before the isolate step sets MUTATION_STARTED and before anything is + // written, but it is a deploy ordering constraint: the director ships + // first, then any workflow run that uses the updated script. + 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 19f5d46aa33..21e1d6aa110 100644 --- a/cloud/apps/relay/src/assignment-connection-headroom.test.ts +++ b/cloud/apps/relay/src/assignment-connection-headroom.test.ts @@ -448,34 +448,48 @@ describe('relay assignment connection headroom', () => { }) }) - it.each(['existing-only', 'migration-only'] as const)( - 'keeps a zero-activity assignment pinned on a %s cell without connection headroom', - async (admission) => { - const { database, store, source } = await setupHeadroomReassignment() - const identity = { - userId: `pinned-${admission}-user`, - relayHostId: `pinned${admission.replace('-', '')}` - } - await store.setCellAdmissionState(source.id, admission) - await database.query( - `INSERT INTO relay_assignments - (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, - last_activity_at, reserved_controls, reserved_splices, reserved_invites, - pending_installs, pending_confirmations, migration_leases) - VALUES (?, ?, ?, ?, ?, ?, 0, 0, 0, 0, 0, 0)`, - [identity.userId, identity.relayHostId, source.id, 7, 10_000, 100] - ) + it('keeps a zero-activity assignment pinned on an existing-only cell', async () => { + const { database, store, source } = await setupHeadroomReassignment() + const identity = { userId: 'pinned-existing-only-user', relayHostId: 'pinnedexistingonl' } + await store.setCellAdmissionState(source.id, 'existing-only') + await database.query( + `INSERT INTO relay_assignments + (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, + last_activity_at, reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases) + VALUES (?, ?, ?, ?, ?, ?, 0, 0, 0, 0, 0, 0)`, + [identity.userId, identity.relayHostId, source.id, 7, 10_000, 100] + ) - await expect(store.assign(identity)).rejects.toThrow( - 'relay_connection_headroom_exhausted' - ) - expect(await store.resolve(identity)).toMatchObject({ - cellId: source.id, - assignmentEpoch: 7 - }) - expect(await database.query(`SELECT * FROM relay_assignment_activity_leases`)).toEqual([]) - } - ) + await expect(store.assign(identity)).rejects.toThrow('relay_connection_headroom_exhausted') + expect(await store.resolve(identity)).toMatchObject({ + cellId: source.id, + assignmentEpoch: 7 + }) + expect(await database.query(`SELECT * FROM relay_assignment_activity_leases`)).toEqual([]) + }) + + 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( + `INSERT INTO relay_assignments + (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, + last_activity_at, reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases) + VALUES (?, ?, ?, ?, ?, ?, 0, 0, 0, 0, 0, 0)`, + [identity.userId, identity.relayHostId, source.id, 7, 10_000, 100] + ) + + await expect(store.assign(identity)).rejects.toThrow('relay_connection_headroom_exhausted') + expect(await store.resolve(identity)).toMatchObject({ + cellId: source.id, + assignmentEpoch: 7 + }) + }) it('renames a pending reservation exactly once on control activation', async () => { const { database, store } = await setup(449) 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 new file mode 100644 index 00000000000..678c6350303 --- /dev/null +++ b/cloud/apps/relay/src/assignment-isolated-cell-replacement-postgres.test.ts @@ -0,0 +1,364 @@ +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' + +const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL +const describePostgres = databaseUrl ? describe : describe.skip + +const USER_PREFIX = 'isolated-replacement-postgres' +const NOW = 100 +const CAPPED = { + capacityRequests: 1_000, + connectionHardCap: 600, + connectionUnobservedBound: 50 +} as const +const ISOLATED: RelayCellConfig = { + id: 'isolated-replacement-source', + url: 'https://isolated-replacement-source.example.com', + region: 'us-central1', + ...CAPPED +} +const TARGETS: RelayCellConfig[] = [ + { + id: 'isolated-replacement-target-a', + url: 'https://isolated-replacement-target-a.example.com', + region: 'us-central1', + ...CAPPED + }, + { + id: 'isolated-replacement-target-b', + url: 'https://isolated-replacement-target-b.example.com', + region: 'us-central1', + ...CAPPED + } +] +const CELLS = [ISOLATED, ...TARGETS] +const HOST_COUNT = 50 + +function hostIdentity(index: number): { userId: string; relayHostId: string } { + return { + userId: `${USER_PREFIX}-${index}`, + // relay host ids are fixed-width opaque ids. + relayHostId: `isolatedhost${String(index).padStart(4, '0')}` + } +} + +describePostgres('PostgreSQL re-placement off a cell isolated for a roll', () => { + const databases: RelayDatabase[] = [] + let stores: RelayAssignmentStore[] = [] + + async function reservedRequests(cellId: string): Promise { + const rows = await databases[0]!.query( + `SELECT reserved_requests FROM relay_cells WHERE cell_id = ?`, + [cellId] + ) + return Number(rows[0]!['reserved_requests']) + } + + async function heartbeatAll(): Promise { + for (const [index, cell] of CELLS.entries()) { + await stores[0]!.recordCellHeartbeat({ + cellId: cell.id, + cellUrl: cell.url, + cellIncarnation: `1111111${index}-1111-4111-8111-111111111111`, + startedAt: 50, + ready: true, + observedRequests: 0, + region: cell.region, + totalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }) + } + } + + // 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() + 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 { + for (const table of [ + 'relay_control_connection_reservations', + 'relay_assignment_activity_leases', + 'relay_assignment_migrations', + 'relay_assignments' + ]) { + await databases[0]!.query(`DELETE FROM ${table} WHERE user_id LIKE '${USER_PREFIX}-%'`) + } + } + + beforeAll(async () => { + // Four connections so the concurrent dials below really contend on the + // fleet-wide relay_cells lock rather than queueing in one client. + for (let index = 0; index < 4; index += 1) { + databases.push(await openRelayDatabase({ databaseUrl, dataDir: '' })) + } + stores = databases.map( + (database) => + new RelayAssignmentStore(database, () => NOW, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + ) + await deleteHostRows() + // 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', + 'relay_cell_connection_runtime', + 'relay_cell_connection_limits', + 'relay_cell_runtime', + 'relay_cell_admission', + 'relay_cell_regions', + 'relay_cells' + ]) { + 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('ignores a stamp older than the roll it is supposed to describe', async () => { + // A failed wave keeps its stamp on purpose and can sit for hours; past the + // bound the cell stops shedding hosts one dial at a time. + await resetFleet() + const identity = hostIdentity(902) + const first = await stores[0]!.assign(identity, 'us-central1') + await applySelector({ [ISOLATED.id]: 'migration-only' }, [ISOLATED.id]) + await databases[0]!.query( + `UPDATE relay_cell_admission SET roll_isolated_at = ? WHERE cell_id = ?`, + [NOW - (2 * 60 * 60_000 + 1), ISOLATED.id] + ) + + expect(await stores[0]!.assign(identity, 'us-central1')).toMatchObject({ + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch + }) + }, 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)) + const sourceBaseline = await reservedRequests(ISOLATED.id) + const targetBaseline = + (await reservedRequests(TARGETS[0]!.id)) + (await reservedRequests(TARGETS[1]!.id)) + + // Everyone lands on the cell about to be isolated. + 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 applySelector({ [TARGETS[0]!.id]: 'general', [TARGETS[1]!.id]: 'general' }) + + // 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( + async (identity, index) => + await stores[1 + (index % (stores.length - 1))]!.assign(identity, 'us-central1') + ) + ) + + for (const [index, grant] of grants.entries()) { + const identity = identities[index]! + expect(grant.cellId).not.toBe(ISOLATED.id) + expect(TARGETS.map(({ id }) => id)).toContain(grant.cellId) + // Exactly once: a double bump would mean two transactions both moved it. + expect(grant.assignmentEpoch).toBe(first.get(identity.relayHostId)!.assignmentEpoch + 1) + } + + expect(await reservedRequests(ISOLATED.id)).toBe(sourceBaseline) + expect( + (await reservedRequests(TARGETS[0]!.id)) + (await reservedRequests(TARGETS[1]!.id)) + ).toBe(targetBaseline + HOST_COUNT) + + const rows = await databases[0]!.query( + `SELECT COUNT(*) AS count FROM relay_assignments + WHERE user_id LIKE '${USER_PREFIX}-%' AND cell_id = ?`, + [ISOLATED.id] + ) + expect(Number(rows[0]!['count'])).toBe(0) + }, 60_000) + + it('lets exactly one of a host’s racing dials win the re-placement', async () => { + await resetFleet() + const identity = hostIdentity(900) + 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 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 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. + const raced = await Promise.allSettled([ + stores[1]!.assign(identity, 'us-central1'), + stores[2]!.assign(identity, 'us-central1') + ]) + const granted = raced.flatMap((result) => + result.status === 'fulfilled' ? [result.value] : [] + ) + expect(granted.length).toBeGreaterThan(0) + + // However many dials were granted, only one re-placement may have + // committed: the epoch advances by exactly one and the reservation moves + // exactly one unit. + const settled = await stores[0]!.resolve(identity) + expect(settled?.assignmentEpoch).toBe(first.assignmentEpoch + 1) + expect(settled?.cellId).not.toBe(ISOLATED.id) + for (const grant of granted) expect(grant.cellId).toBe(settled?.cellId) + expect(await reservedRequests(ISOLATED.id)).toBe(sourceBefore - 1) + expect( + (await reservedRequests(TARGETS[0]!.id)) + (await reservedRequests(TARGETS[1]!.id)) + ).toBe(targetBaseline + 1) + }, 30_000) + + it('grants no host the isolated cell after the admission flip commits', async () => { + await resetFleet() + const identity = hostIdentity(901) + const first = await stores[0]!.assign(identity, 'us-central1') + 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( + identity, + 'us-central1' + ) + expect(grant.cellId).not.toBe(ISOLATED.id) + } + // Only the first dial re-places; the rest are ordinary sticky re-grants. + expect((await stores[0]!.resolve(identity))?.assignmentEpoch).toBe( + first.assignmentEpoch + 1 + ) + }, 30_000) +}) diff --git a/cloud/apps/relay/src/assignment-isolated-cell-replacement.test.ts b/cloud/apps/relay/src/assignment-isolated-cell-replacement.test.ts new file mode 100644 index 00000000000..cc1ae6a26bb --- /dev/null +++ b/cloud/apps/relay/src/assignment-isolated-cell-replacement.test.ts @@ -0,0 +1,609 @@ +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, + type RelayDatabase, + type RelayLockOptions, + type RelayTransactionOptions, + type SqlRow +} from './database.js' + +// `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 answers 503 relay_home_cell_unavailable. +const CAPPED = { + capacityRequests: 1_000, + connectionHardCap: 600, + connectionUnobservedBound: 50 +} as const +const CELLS: RelayCellConfig[] = [ + { id: 'us-c1', url: 'https://us-c1.example.com', region: 'us-central1', ...CAPPED }, + { id: 'us-c2', url: 'https://us-c2.example.com', region: 'us-central1', ...CAPPED }, + { id: 'asia-c1', url: 'https://asia-c1.example.com', region: 'asia-east2', ...CAPPED } +] + +const databases: RelayDatabase[] = [] + +afterEach(async () => { + vi.restoreAllMocks() + for (const database of databases.splice(0)) await database.close() +}) + +// Counts the reads this change adds, so the general hot path can be held to a +// single admission lookup and no migration lookup at all. +class QueryCountingDatabase implements RelayDatabase { + readonly sql: string[] = [] + // Fails the first statement containing this fragment, so a test can roll a + // transaction back at a chosen point. + failOnce: string | undefined + + constructor(private readonly delegate: RelayDatabase) {} + + get dialect(): 'sqlite' | 'postgres' | undefined { + return this.delegate.dialect + } + + count(fragment: string): number { + return this.sql.filter((statement) => statement.includes(fragment)).length + } + + record(sql: string): void { + this.sql.push(sql) + if (this.failOnce !== undefined && sql.includes(this.failOnce)) { + this.failOnce = undefined + throw new Error('injected_placement_write_failure') + } + } + + async query(sql: string, params?: unknown[]): Promise { + this.record(sql) + return await this.delegate.query(sql, params) + } + + async queryLocked( + sql: string, + params?: unknown[], + options?: RelayLockOptions + ): Promise { + this.record(sql) + return await this.delegate.queryLocked(sql, params, options) + } + + async transaction( + operation: (transaction: RelayDatabase) => Promise, + options?: RelayTransactionOptions + ): Promise { + return await this.delegate.transaction( + async (inner) => await operation(this.recording(inner)), + options + ) + } + + async close(): Promise { + await this.delegate.close() + } + + // A nested transaction shares this one's tape, so a whole assign() is + // measured as a unit. + private recording(inner: RelayDatabase): RelayDatabase { + return { + dialect: inner.dialect, + query: async (sql, params) => { + this.record(sql) + return await inner.query(sql, params) + }, + queryLocked: async (sql, params, options) => { + this.record(sql) + return await inner.queryLocked(sql, params, options) + }, + transaction: async (operation, options) => await inner.transaction(operation, options), + close: async () => await inner.close() + } + } +} + +interface Harness { + store: RelayAssignmentStore + database: RelayDatabase + counter: QueryCountingDatabase + 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 { + const inner = await openInMemoryRelayDatabase() + databases.push(inner) + const counter = new QueryCountingDatabase(inner) + let now = START_MS + const store = new RelayAssignmentStore(counter, () => now, { + requireLiveCells: true, + heartbeatTtlMs: HEARTBEAT_TTL_MS + }) + await store.reconcileCells(cells, true) + 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) + + 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 insertMigration( + database: RelayDatabase, + 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, ?, ?, ?, ?)`, + [ + IDENTITY.userId, + IDENTITY.relayHostId, + input.sourceCellId, + input.targetCellId, + 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 + ] + ) + await database.query( + `UPDATE relay_assignments SET migration_leases = ? + WHERE user_id = ? AND relay_host_id = ?`, + [input.leases, IDENTITY.userId, IDENTITY.relayHostId] + ) +} + +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() + const first = await store.assign(IDENTITY, 'us-central1') + + counter.sql.length = 0 + const second = await store.assign(IDENTITY, 'us-central1') + + // 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) + 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('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 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') + + expect(moved.cellId).not.toBe(first.cellId) + expect(moved.assignmentEpoch).toBe(first.assignmentEpoch + 1) + expect(counter.count('relay_assignment_migrations')).toBeGreaterThan(0) + + // Durable: the next dial does not bounce back to the isolated cell. + expect(await store.assign(IDENTITY, 'us-central1')).toMatchObject({ + cellId: moved.cellId, + assignmentEpoch: moved.assignmentEpoch + }) + }) + + it.each([ + { label: 'just inside the bound', age: 2 * 60 * 60_000 - 1, moves: true }, + { label: 'just outside the bound', age: 2 * 60 * 60_000 + 1, moves: false } + ])('treats a stamp $label as $moves', async ({ age, moves }) => { + // Why the bound exists: a roll isolates and restores inside ~15 minutes, so + // an older stamp is a failed wave waiting on an operator, or an orphan left + // by a director rollback whose restore predates the clearing clause. Both + // mean a possibly healthy cell, and the safe answer is the pre-existing one. + const { store, heartbeat, isolateForRoll, setNow } = await setup() + const first = await store.assign(IDENTITY, 'us-central1') + await isolateForRoll(first.cellId) + + const later = START_MS + age + setNow(later) + for (const cell of CELLS) await heartbeat(cell, later) + const grant = await store.assign(IDENTITY, 'us-central1') + + if (moves) { + expect(grant.cellId).not.toBe(first.cellId) + expect(grant.assignmentEpoch).toBe(first.assignmentEpoch + 1) + } else { + expect(grant.cellId).toBe(first.cellId) + expect(grant.assignmentEpoch).toBe(first.assignmentEpoch) + } + }) + + 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, park, setNow } = await setup() + const first = await store.assign(IDENTITY, 'us-central1') + 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, + request_units, expires_at, updated_at) + VALUES (?, ?, 'control:live-1', 'control', ?, 1, ?, ?)`, + [IDENTITY.userId, IDENTITY.relayHostId, first.cellId, START_MS + 600_000, START_MS] + ) + + // 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, START_MS + 61_000) + expect(await store.assign(IDENTITY, 'us-central1')).toMatchObject({ + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch + }) + }) + + 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 park(other.id, 'migration-only') + await isolateForRoll(first.cellId) + + warn.mockClear() + expect(await store.assign(IDENTITY, 'us-central1')).toMatchObject({ + 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', async () => { + const { store, isolateForRoll } = await setup() + const first = await store.assign(IDENTITY, 'asia-east2') + expect(first.region).toBe('asia-east2') + 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: first.cellId, + assignmentEpoch: first.assignmentEpoch + }) + }) + + 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 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) + expect(moved).toMatchObject({ assignmentEpoch: first.assignmentEpoch + 1 }) + }) + + 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, isolateForRoll } = await setup() + const first = await store.assign(IDENTITY, 'us-central1') + await database.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:keep-1', 'splice', ?, 1, ?, ?)`, + [IDENTITY.userId, IDENTITY.relayHostId, first.cellId, START_MS + 600_000, START_MS] + ) + await isolateForRoll(first.cellId) + + await store.assign(IDENTITY, 'us-central1') + expect( + await database.query( + `SELECT activity_id FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? AND activity_id = 'splice:keep-1'`, + [IDENTITY.userId, IDENTITY.relayHostId] + ) + ).toHaveLength(1) + }) + + it('emits nothing when the placement rolls back after the decision', async () => { + // Why: the decision and the writes share one transaction. A line already on + // stdout cannot be rolled back with it, so an event written where it is + // decided would have the canary counting moves that never happened. + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + const { store, counter, isolateForRoll } = await setup() + const first = await store.assign(IDENTITY, 'us-central1') + await isolateForRoll(first.cellId) + + warn.mockClear() + // The first write after the event is decided. + counter.failOnce = 'INSERT INTO relay_assignments' + await expect(store.assign(IDENTITY, 'us-central1')).rejects.toThrow( + 'injected_placement_write_failure' + ) + + expect(jsonEvents(warn, 'orca_relay_sticky_replaced_off_isolated_cell')).toEqual([]) + // The precondition for reading that absence: the move really was rolled + // back, so the event would have been a lie rather than merely early. + expect(await store.resolve(IDENTITY)).toMatchObject({ + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch + }) + + // Control: the same dial with nothing injected does emit, so the assertion + // above is measuring the rollback and not a broken harness. + warn.mockClear() + const moved = await store.assign(IDENTITY, 'us-central1') + expect(moved.cellId).not.toBe(first.cellId) + expect(jsonEvents(warn, 'orca_relay_sticky_replaced_off_isolated_cell')).toHaveLength(1) + }) + + it('logs one event naming both cells, the admission state and the region', async () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + const { store, isolateForRoll } = await setup() + const first = await store.assign(IDENTITY, 'us-central1') + await isolateForRoll(first.cellId) + + warn.mockClear() + const moved = await store.assign(IDENTITY, 'us-central1') + + 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 e4a8ce54ebf..57e3b9d25e1 100644 --- a/cloud/apps/relay/src/assignment-store.test.ts +++ b/cloud/apps/relay/src/assignment-store.test.ts @@ -1556,6 +1556,11 @@ describe('RelayAssignmentStore', () => { }) 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' } diff --git a/cloud/apps/relay/src/assignment-store.ts b/cloud/apps/relay/src/assignment-store.ts index fd1193e0d2f..d050fe59182 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, @@ -392,6 +393,20 @@ export type CellInventoryLockMode = // ordinary dormancy — which stays governed by the 24h rule. const STRANDED_MIN_GRANT_AGE_MS = 60_000 const STRANDED_RECENT_ACTIVITY_MS = 15 * 60_000 +// Why: a roll's isolate step writes exactly this state +// (cloud/dev/scripts/prepare-relay-production-capacity-canary.mjs:136) and its +// restore writes 'general' (same line). 'existing-only' is C3's decommission +// posture, where the cell still serves the hosts it already has, so it is +// deliberately not an isolation signal here. +const ROLL_ISOLATED_ADMISSION: CellAdmissionState = 'migration-only' +// Why the stamp expires: a roll isolates and restores one cell inside ~15 +// minutes, so a stamp older than this is not a roll in progress. It is a failed +// or stalled wave whose failsafe re-isolated a possibly healthy cell and is +// waiting on an operator, or a director rollback whose restore wrote 'general' +// without the clause that clears the stamp, leaving an orphan that the next +// park would reactivate. Both want the same answer, and it is the pre-existing +// one: keep the pin and let the host retry its own cell. +const ROLL_ISOLATION_STAMP_MAX_AGE_MS = 2 * 60 * 60_000 const REGION_PREFERENCE_RETENTION_MS = 30 * 24 * 60 * 60_000 const REGIONAL_REHOME_UNREGISTERED_REFRESH_MS = 5 * 60_000 const REGIONAL_REHOME_MAX_REFRESH_MS = 24 * 60 * 60_000 @@ -736,10 +751,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 } @@ -776,6 +805,22 @@ export class RelayAssignmentStore { ) { return null } + // Why: a null here means "fall through to placement", which is exactly + // 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, + now, + undefined, + pinnedAdmission + ) + ) { + return null + } if ( !hadControl && @@ -826,7 +871,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 ( @@ -835,14 +882,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 = ( @@ -878,7 +922,15 @@ export class RelayAssignmentStore { const now = this.now() let retryScope: RetriedAssignmentInventoryScope = inventoryScope === 'all' ? 'all' : 'general' - return await this.database.transaction(async (transaction) => { + // Why the events ride back out rather than being written where they are + // decided: everything below runs in one transaction, and a reservation or + // lease write that fails after the decision rolls the placement back. A + // line already on stdout cannot be rolled back with it, so the canary would + // count re-placements that never happened. Returning them means only a + // committed attempt emits, and a transaction retry cannot leave a stale + // line behind either. + const outcome = await this.database.transaction(async (transaction) => { + const events: string[] = [] // 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) @@ -902,6 +954,8 @@ export class RelayAssignmentStore { let forcedDeadReassignment = false 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) @@ -922,10 +976,46 @@ 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. 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 && - (!this.requireLiveCells || (await this.cellIsLive(transaction, current.cellId, now))) + currentIsLive && + (await this.incumbentCellIsolatedForRoll( + transaction, + identity, + existing, + now, + admission + )) ) { + 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. + events.push( + JSON.stringify({ + event: 'orca_relay_sticky_replacement_deferred', + reason: 'no_same_region_headroom', + cellId: current.cellId, + region: current.region + }) + ) + } + } + if (!strandedReassignment && !isolatedIncumbent && currentIsLive) { const hadControl = holdsControlLease( activityLeases, current.cellId, @@ -956,7 +1046,10 @@ export class RelayAssignmentStore { now ) } - return this.result(identity, existing, current, leaseExpiresAt) + return { + assignment: this.result(identity, existing, current, leaseExpiresAt), + events + } } if (requestUnits(existing) > 0) { throw new Error('relay_connection_headroom_exhausted') @@ -966,7 +1059,14 @@ export class RelayAssignmentStore { } connectionHeadroomReassignment = true } - if (!strandedReassignment && !connectionHeadroomReassignment) { + // 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. `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( transaction, @@ -988,12 +1088,27 @@ 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 + // drained cell's host count. + events.push( + JSON.stringify({ + event: 'orca_relay_sticky_replaced_off_isolated_cell', + fromCellId: isolatedIncumbent.cellId, + fromRegion: isolatedIncumbent.region, + admissionState: ROLL_ISOLATED_ADMISSION, + toCellId: target.cellId, + region: target.region + }) + ) + } const previousUnits = existing ? requestUnits(existing) : 0 if (existing) { await this.adjustCellReservation(transaction, text(existing, 'cell_id'), -previousUnits) @@ -1060,13 +1175,18 @@ export class RelayAssignmentStore { assignmentEpoch, now ) - return { ...identity, ...target, assignmentEpoch, leaseExpiresAt } + return { + assignment: { ...identity, ...target, assignmentEpoch, leaseExpiresAt }, + events + } }).catch((error: unknown) => { if (isDatabaseLockUnavailable(error)) { throw new AssignmentInventoryLockUnavailable(retryScope) } throw error }) + for (const event of outcome.events) console.warn(event) + return outcome.assignment } async resolve(identity: AssignmentIdentity): Promise { @@ -1114,6 +1234,7 @@ export class RelayAssignmentStore { expectedGeneration: number expectedMembershipSha256?: string membership: CellAdmissionMembership + rollIsolatedCells?: string[] }): Promise<{ changed: boolean selector: { @@ -7147,7 +7268,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) => [ @@ -7193,7 +7319,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 @@ -7314,6 +7440,76 @@ export class RelayAssignmentStore { return rows.length === 1 } + // Why: a cell isolated for a roll keeps heartbeating ready=1 for the whole + // drain while its registry refuses every attach with 4503, so `cellIsLive` + // 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, + // The caller's clock, not a second this.now(): one assign reasons about one + // instant, and the stamp's age decides whether a host moves. + now: number, + // 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, + // Sticky has already read the pinned row for the stranded rule. + pinnedAdmission?: PinnedCellAdmission + ): Promise { + const cellId = text(existing, 'cell_id') + // 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 (now - pinned.rollIsolatedAt >= ROLL_ISOLATION_STAMP_MAX_AGE_MS) 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 + // counter alone would re-place a host out from under a stalled migration. + const open = ( + await database.query( + `SELECT COUNT(*) AS open FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ? + AND completed_at IS NULL AND aborted_at IS NULL`, + [identity.userId, identity.relayHostId] + ) + )[0] + return integer(open!, 'open') === 0 + } + + // 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 { + const row = ( + await database.query( + `SELECT admission_state, roll_isolated_at FROM relay_cell_admission WHERE cell_id = ?`, + [cellId] + ) + )[0] + // 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 // separates an expected drain or boot from a cell whose readiness went out // from under its hosts. @@ -7898,6 +8094,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) {