From abff7ab86a567ce5757bb6c2880c1a34fc7e15d4 Mon Sep 17 00:00:00 2001 From: Jinwoo Hong <73622457+Jinwoo-H@users.noreply.github.com> Date: Sat, 10 Oct 2026 02:18:37 -0400 Subject: [PATCH] feat(relay): flip a reserve-mode cell back without disconnecting anyone (E-mix) (#27043) Rebased onto main as one commit after its parent squash-merged. --- cloud/apps/relay/src/app.ts | 14 +- cloud/apps/relay/src/assignment-store.ts | 52 +++- cloud/apps/relay/src/cell-flags.ts | 5 +- cloud/apps/relay/src/cell-reserve-dead-man.ts | 64 +++++ .../cell-reserve-flip-back-postgres.test.ts | 116 ++++++++ .../cell-reserve-rollback.blackbox.test.ts | 268 ++++++++++++++++++ cloud/apps/relay/src/cell-seat-feed-route.ts | 5 + cloud/apps/relay/src/cell-seat-feed.test.ts | 4 +- cloud/apps/relay/src/host-session-registry.ts | 249 +++++++++++++++- cloud/apps/relay/src/relay-server.ts | 39 ++- .../src/test-fixtures/reserve-mode-cell.ts | 219 ++++++++++++++ 11 files changed, 1017 insertions(+), 18 deletions(-) create mode 100644 cloud/apps/relay/src/cell-reserve-dead-man.ts create mode 100644 cloud/apps/relay/src/cell-reserve-flip-back-postgres.test.ts create mode 100644 cloud/apps/relay/src/cell-reserve-rollback.blackbox.test.ts create mode 100644 cloud/apps/relay/src/test-fixtures/reserve-mode-cell.ts diff --git a/cloud/apps/relay/src/app.ts b/cloud/apps/relay/src/app.ts index 91fe5fcf3f6..7786432d585 100644 --- a/cloud/apps/relay/src/app.ts +++ b/cloud/apps/relay/src/app.ts @@ -130,6 +130,8 @@ export function createRelayApp( cellFlags?: () => AppliedControlFlags cellReserve?: (request: ReserveRequest) => ReserveOutcome[] cellDemote?: (request: DemoteRequest) => string + cellReserverPoll?: () => void + cellAdmitModeEffective?: () => 'db' | 'reserve' cellReserveCounts?: () => { bookings: number intake: { perSec: number; burst: number; tokens: number } @@ -343,6 +345,8 @@ export function createRelayApp( isDraining: operations.isDraining, runtimeCounts: operations.runtimeCounts, cellFlags: operations.cellFlags, + onReserverPoll: operations.cellReserverPoll, + admitModeEffective: operations.cellAdmitModeEffective, reserveCounts: operations.cellReserveCounts, placementCeiling: operations.cellPlacementCeiling }) @@ -800,7 +804,15 @@ export function createRelayApp( regionalRehomeProtocol: config.rehomeAudience && config.rehomeDirectorServiceAccount ? 3 : 0, // The flag workflow's read-back: applied switches, never the desired object. ...(operations.cellFlags - ? { flagsApplied: operations.cellFlags(), supportedFlags: supportedCellFlags() } + ? { + flagsApplied: operations.cellFlags(), + supportedFlags: supportedCellFlags(), + // The flag tool paces a flip back's wait from it. + databasePoolMax: config.databasePoolMax + } + : {}), + ...(operations.cellAdmitModeEffective + ? { admitModeEffective: operations.cellAdmitModeEffective() } : {}), connectionCapacity: config.connectionHardCap === undefined diff --git a/cloud/apps/relay/src/assignment-store.ts b/cloud/apps/relay/src/assignment-store.ts index f0812babc6d..2541348d128 100644 --- a/cloud/apps/relay/src/assignment-store.ts +++ b/cloud/apps/relay/src/assignment-store.ts @@ -140,6 +140,15 @@ type RelayAssignmentStoreOptions = { export type { ControlRenewalOutcome, ControlRenewalRequest } +type ControlActivationInput = { + cellId: string + assignmentEpoch: number + generation: number + idleRegionalRehome?: boolean + cellIncarnation?: string + connectionInclusionWatermark?: number +} + export type RelayAssignment = AssignmentIdentity & { cellId: string cellUrl: string @@ -4279,15 +4288,38 @@ export class RelayAssignmentStore { async activateControl( identity: AssignmentIdentity, - input: { - cellId: string - assignmentEpoch: number - generation: number - idleRegionalRehome?: boolean - cellIncarnation?: string - connectionInclusionWatermark?: number - } + input: ControlActivationInput ): Promise { + return (await this.activateControlOnce(identity, input, false)).activityId + } + + // Step 5 flip back: the same activation, but the cell-row delta is returned for the caller + // to fold into one relay_cells write per batch (commitCellReservationDelta), not committed + // here. A dozen re-registrations from Asia each holding that row's lock convoy on it. + async activateControlDeferringCell( + identity: AssignmentIdentity, + input: ControlActivationInput + ): Promise<{ activityId: string; reservationDelta: number }> { + return await this.activateControlOnce(identity, input, true) + } + + // Unguarded by capacity: these hosts are already connected to the cell. + async commitCellReservationDelta(cellId: string, delta: number): Promise { + if (delta === 0) return + await this.database.query( + `UPDATE relay_cells SET reserved_requests = + CASE WHEN reserved_requests + ? < 0 THEN 0 ELSE reserved_requests + ? END, + updated_at = ? + WHERE cell_id = ?`, + [delta, delta, this.now(), cellId] + ) + } + + private async activateControlOnce( + identity: AssignmentIdentity, + input: ControlActivationInput, + deferCell: boolean + ): Promise<{ activityId: string; reservationDelta: number }> { const activityId = `control:${input.cellId}:${input.generation}` validateActivityId(activityId) return await this.activityQueue.run(identity, async () => { @@ -4399,14 +4431,14 @@ export class RelayAssignmentStore { // every host on the cell, and this transaction spans a dozen round // trips. Holding its write lock from the first of them capped a // far-from-Postgres cell at a couple of accepts a second. - if (reservationDelta !== 0) { + if (reservationDelta !== 0 && !deferCell) { await this.commitCellReservationAtomically( transaction, input.cellId, reservationDelta ) } - return activityId + return { activityId, reservationDelta: deferCell ? reservationDelta : 0 } }) }) } diff --git a/cloud/apps/relay/src/cell-flags.ts b/cloud/apps/relay/src/cell-flags.ts index 6543a826681..ccf22204d59 100644 --- a/cloud/apps/relay/src/cell-flags.ts +++ b/cloud/apps/relay/src/cell-flags.ts @@ -17,6 +17,8 @@ export type CellFlags = { intakePerSec?: number // Answers a director's dry-run booking check; off, the reserve endpoint is inert in db mode. reserveDryRun: boolean + // Flip-back leases in flight at once; absent is a third of the pool. Speeds up one cell. + reregisterInFlight?: number } export const CELL_FLAG_DEFAULTS: CellFlags = { @@ -49,7 +51,8 @@ const CellFlagsSchema = z.object({ .optional(), admitMode: z.enum(['db', 'reserve']).default(CELL_FLAG_DEFAULTS.admitMode), intakePerSec: z.number().positive().max(1_000).optional(), - reserveDryRun: z.boolean().default(CELL_FLAG_DEFAULTS.reserveDryRun) + reserveDryRun: z.boolean().default(CELL_FLAG_DEFAULTS.reserveDryRun), + reregisterInFlight: z.number().int().min(1).max(16).optional() }) const KNOWN_FLAG_KEYS = new Set(Object.keys(CellFlagsSchema.shape)) diff --git a/cloud/apps/relay/src/cell-reserve-dead-man.ts b/cloud/apps/relay/src/cell-reserve-dead-man.ts new file mode 100644 index 00000000000..38bdc879d00 --- /dev/null +++ b/cloud/apps/relay/src/cell-reserve-dead-man.ts @@ -0,0 +1,64 @@ +import { createHash } from 'node:crypto' +import type { CellFlags } from './cell-flags.js' +import type { AppliedControlFlags } from './relay-control-flag-channel.js' + +// A reserve-mode cell keeps its hosts in memory only while a placing director is there to +// book and supersede them. With none for this long (a director rolled back, or placement +// turned off) it flips itself back to db and re-registers, as a switch flip would. +export const RESERVE_DEAD_MAN_MS = 60_000 +const RESERVE_DEAD_MAN_SPREAD_MS = 60_000 + +// 60-120 s, fixed per cell: when every director goes at once, the cells trip (and start +// re-registering against the database) spread over a minute, not together. +export function reserveDeadManWindowMs(cellId: string): number { + const digest = createHash('sha256').update(cellId).digest() + return RESERVE_DEAD_MAN_MS + (digest.readUInt32BE(0) % (RESERVE_DEAD_MAN_SPREAD_MS + 1)) +} + +export class ReserveDeadMan { + private lastContactAt: number + private reserveGeneration: number | null = null + private trippedGeneration: number | null = null + + constructor( + private readonly cellId: string, + private readonly now: () => number = Date.now, + private readonly windowMs = RESERVE_DEAD_MAN_MS, + private readonly log: (line: string) => void = (line) => console.error(line) + ) { + this.lastContactAt = now() + } + + // A feed poll marked reserver=1, or a real (not dry-run) booking. + contact(): void { + this.lastContactAt = this.now() + } + + // The admitMode the cell acts on. Latched per switch-file generation: once tripped, only a + // new write brings reserve back, so a director that returns cannot flap the cell. + mode(applied: AppliedControlFlags): 'db' | 'reserve' { + if (applied.flags.admitMode !== 'reserve') { + this.reserveGeneration = null + return 'db' + } + const now = this.now() + if (this.reserveGeneration !== applied.generation) { + // A new reserve write starts its own window. + this.reserveGeneration = applied.generation + this.trippedGeneration = null + this.lastContactAt = Math.max(this.lastContactAt, now) + } + if (this.trippedGeneration === applied.generation) return 'db' + if (now - this.lastContactAt <= this.windowMs) return 'reserve' + this.trippedGeneration = applied.generation + this.log( + JSON.stringify({ + event: 'orca_relay_cell_reserve_dead_man_tripped', + cellId: this.cellId, + generation: applied.generation, + silentMs: now - this.lastContactAt + }) + ) + return 'db' + } +} diff --git a/cloud/apps/relay/src/cell-reserve-flip-back-postgres.test.ts b/cloud/apps/relay/src/cell-reserve-flip-back-postgres.test.ts new file mode 100644 index 00000000000..833c366b126 --- /dev/null +++ b/cloud/apps/relay/src/cell-reserve-flip-back-postgres.test.ts @@ -0,0 +1,116 @@ +import { mkdtempSync, rmSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterEach, describe, expect, it, vi } from 'vitest' +import WebSocket from 'ws' +import { openRelayDatabase, type RelayDatabase } from './database.js' +import { openDelayedPostgresDatabase, type StatementDelay } from './test-fixtures/delayed-postgres-database.js' +import { RESERVE_MODE_CELL_ID, startReserveModeCell } from './test-fixtures/reserve-mode-cell.js' + +// E-mix on PostgreSQL: a reserve-mode cell flipped back to db leases every control it admitted +// from memory, disconnecting nobody, with the pool mostly left to DB-path work and one +// relay_cells write per tick, even with a cross-region round trip on every statement. + +const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL +const describePostgres = databaseUrl ? describe : describe.skip + +async function until(check: () => Promise | boolean, timeoutMs: number): Promise { + const deadline = Date.now() + timeoutMs + while (!(await check())) { + if (Date.now() > deadline) throw new Error('condition not reached') + await new Promise((resolve) => setTimeout(resolve, 50)) + } +} + +describePostgres('flipping a reserve-mode cell back on PostgreSQL', () => { + const cleanup: Array<() => Promise | void> = [] + + afterEach(async () => { + for (const close of cleanup.splice(0).reverse()) await close() + vi.restoreAllMocks() + // The database is shared with the other PostgreSQL suites, some of which read every cell. + const database = await openRelayDatabase({ databaseUrl, dataDir: tmpdir() }) + try { + const tables = await database.query( + `SELECT table_name FROM information_schema.columns + WHERE table_schema = current_schema() AND column_name = 'cell_id' AND table_name LIKE 'relay_%'` + ) + for (const row of tables) { + await database.query(`DELETE FROM ${String(row.table_name)} WHERE cell_id = ?`, [RESERVE_MODE_CELL_ID]) + } + } finally { + await database.close() + } + }) + + async function cellOn(database: RelayDatabase, databasePoolMax = 10) { + vi.spyOn(console, 'log').mockImplementation(() => undefined) + vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const dataDir = mkdtempSync(join(tmpdir(), 'orca-relay-flip-back-pg-')) + cleanup.push(() => rmSync(dataDir, { recursive: true, force: true })) + return await startReserveModeCell({ database, dataDir, cleanup, databasePoolMax }) + } + + it('re-registers every memory-admitted control, ten flips in a row', async () => { + const database = await openRelayDatabase({ databaseUrl, dataDir: tmpdir() }) + cleanup.push(() => database.close()) + for (let flip = 0; flip < 10; flip += 1) { + const cell = await cellOn(database) + const hosts = [] + for (let index = 0; index < 3; index += 1) hosts.push(await cell.bookedHost(5 + index)) + const ids = new Set(hosts.map((entry) => entry.identity.hostId)) + const leased = async () => + (await cell.controlLeases()).filter((row) => ids.has(String(row.relay_host_id))) + expect(await leased()).toEqual([]) + cell.setAdmitMode('db') + await until(async () => (await leased()).length === 3, 10_000) + await until(() => cell.relay.sessions.reregistrationPending() === 0, 10_000) + expect(hosts.every((entry) => entry.socket.readyState === WebSocket.OPEN)).toBe(true) + for (const entry of hosts) entry.socket.close() + } + }, 120_000) + + it('keeps re-registration to a third of the pool and one cell-row write per tick, at 170 ms a trip', async () => { + const delay: StatementDelay = { enabled: false, delayMs: 170 } + let cellRowWrites = 0 + let leaseWritesInFlight = 0 + let peakInFlight = 0 + delay.beforeTrip = async (sql) => { + if (/UPDATE relay_cells\b/.test(sql)) cellRowWrites += 1 + } + const database = openDelayedPostgresDatabase(databaseUrl!, delay, 10) + cleanup.push(() => database.close()) + // The schema and the cell row, without the delay. + const setup = await openRelayDatabase({ databaseUrl, dataDir: tmpdir() }) + await setup.close() + const cell = await cellOn(database, 10) + const deferred = cell.relay.assignments.activateControlDeferringCell.bind(cell.relay.assignments) + vi.spyOn(cell.relay.assignments, 'activateControlDeferringCell').mockImplementation(async (...args) => { + leaseWritesInFlight += 1 + peakInFlight = Math.max(peakInFlight, leaseWritesInFlight) + try { + return await deferred(...args) + } finally { + leaseWritesInFlight -= 1 + } + }) + const hosts = [] + for (let index = 0; index < 9; index += 1) hosts.push(await cell.bookedHost(5 + index)) + const ids = new Set(hosts.map((entry) => entry.identity.hostId)) + delay.enabled = true + const startedAt = Date.now() + cell.setAdmitMode('db') + await until(() => cell.relay.sessions.reregistrationPending() > 0, 5_000) + await until(() => cell.relay.sessions.reregistrationPending() === 0, 45_000) + const elapsedMs = Date.now() - startedAt + delay.enabled = false + const leased = (await cell.controlLeases()).filter((row) => ids.has(String(row.relay_host_id))) + expect(leased).toHaveLength(9) + expect(hosts.every((entry) => entry.socket.readyState === WebSocket.OPEN)).toBe(true) + // Pool of 10: at most 3 activations at once, leaving 7 for hellos and renewals. + expect(peakInFlight).toBeLessThanOrEqual(3) + // Far fewer cell-row writes than activations: deltas fold into one write per tick. + expect(cellRowWrites).toBeLessThan(9) + process.stderr.write(`flip_back_latency ${JSON.stringify({ elapsedMs, peakInFlight, cellRowWrites })}\n`) + }, 60_000) +}) diff --git a/cloud/apps/relay/src/cell-reserve-rollback.blackbox.test.ts b/cloud/apps/relay/src/cell-reserve-rollback.blackbox.test.ts new file mode 100644 index 00000000000..2b039af92f2 --- /dev/null +++ b/cloud/apps/relay/src/cell-reserve-rollback.blackbox.test.ts @@ -0,0 +1,268 @@ +import { mkdtempSync, rmSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterEach, describe, expect, it, vi } from 'vitest' +import WebSocket from 'ws' +import { RESERVE_DEAD_MAN_MS, reserveDeadManWindowMs } from './cell-reserve-dead-man.js' +import { openRelayDatabase } from './database.js' +import { + REREGISTER_MAX_ATTEMPTS, + REREGISTER_STALL_LOG_EVERY, + reregisterInFlightLimit +} from './host-session-registry.js' +import { startReserveModeCell } from './test-fixtures/reserve-mode-cell.js' + +// Flip back (E-mix): a cell leaves reserve mode with hosts it admitted from memory. Nobody is +// disconnected, and the database ends up holding exactly the leases its seat report names. + +const byHost = (left: Record, right: Record) => + String(left.relay_host_id) < String(right.relay_host_id) ? -1 : 1 + +async function until(check: () => Promise | boolean, timeoutMs = 10_000): Promise { + const deadline = Date.now() + timeoutMs + while (!(await check())) { + if (Date.now() > deadline) throw new Error('condition not reached') + await new Promise((resolve) => setTimeout(resolve, 50)) + } +} + +describe('flipping a reserve-mode cell back to the database', () => { + const cleanup: Array<() => Promise | void> = [] + + afterEach(async () => { + for (const close of cleanup.splice(0).reverse()) await close() + vi.restoreAllMocks() + }) + + async function cell(now?: () => number) { + vi.spyOn(console, 'log').mockImplementation(() => undefined) + vi.spyOn(console, 'warn').mockImplementation(() => undefined) + vi.spyOn(console, 'error').mockImplementation(() => undefined) + const dataDir = mkdtempSync(join(tmpdir(), 'orca-relay-reserve-rollback-')) + cleanup.push(() => rmSync(dataDir, { recursive: true, force: true })) + const database = await openRelayDatabase({ dataDir }) + cleanup.push(() => database.close()) + return await startReserveModeCell({ database, dataDir, cleanup, ...(now ? { now } : {}) }) + } + + it('registers every memory-admitted control without disconnecting any, matching the seat report', async () => { + const reserve = await cell() + const first = await reserve.bookedHost(5) + const second = await reserve.bookedHost(6) + // The ledger lags by up to a second: this row lands only after the flip. + const late = await reserve.bookedHost(7, { row: 'never' }) + expect(await reserve.controlLeases()).toEqual([]) + const marked = vi.spyOn(reserve.relay.assignments, 'markMigrationTargetRegistered') + + reserve.setAdmitMode('db') + await until(async () => (await reserve.controlLeases()).length === 2) + // Still reporting reserve: sweeps must not resume while a control lacks its lease. + expect(reserve.relay.sessions.reregistrationPending()).toBeGreaterThan(0) + await reserve.insertRow(late.identity.hostId, 7) + await until(async () => (await reserve.controlLeases()).length === 3) + await until(() => reserve.relay.sessions.reregistrationPending() === 0) + + expect([first, second, late].every((entry) => entry.socket.readyState === WebSocket.OPEN)).toBe(true) + const page = reserve.relay.sessions.seatFeed(null) + const seats = 'full' in page ? page.full : [] + expect( + seats + .map((seat) => ({ + relay_host_id: seat.relayHostId, + cell_id: 'production-gce-c3', + activity_id: `control:production-gce-c3:${seat.generation}` + })) + .sort(byHost) + ).toEqual((await reserve.controlLeases()).sort(byHost)) + // As a database-path join does, each one also tells an open migration its host arrived. + expect(marked.mock.calls.map(([, input]) => input.assignmentEpoch).sort()).toEqual([5, 6, 7]) + // A second flip back has nothing left to register. + expect(reserve.relay.sessions.reregisterMemoryControls()).toBe(0) + }, 30_000) + + it('closes with 4409 a control whose row never names this cell, after every retry', async () => { + let now = Date.now() + const reserve = await cell(() => now) + // The database says another cell owns it: the lease can never be taken here. + const stranded = await reserve.bookedHost(5, { row: 'production-gce-c9' }) + const fine = await reserve.bookedHost(6) + const tries = vi.spyOn(reserve.relay.assignments, 'activateControlDeferringCell') + reserve.setAdmitMode('db') + for (let step = 0; step < 3 * REREGISTER_MAX_ATTEMPTS && stranded.socket.readyState === WebSocket.OPEN; step += 1) { + now += 6_000 + // Both desktops stay alive on the test clock: a pong resets the control silence timer. + for (const entry of [stranded, fine]) { + if (entry.socket.readyState === WebSocket.OPEN) entry.socket.send(JSON.stringify({ type: 'pong', t: now })) + } + await new Promise((resolve) => setTimeout(resolve, 120)) + } + expect(await stranded.closeCode).toBe(4409) + // Every retry was spent first: the row may still have been on its way. + expect(tries.mock.calls.length).toBeGreaterThanOrEqual(REREGISTER_MAX_ATTEMPTS + 1) + await until(() => reserve.relay.sessions.reregistrationPending() === 0) + expect(fine.socket.readyState).toBe(WebSocket.OPEN) + expect((await reserve.controlLeases()).map((row) => row.relay_host_id)).toEqual([fine.identity.hostId]) + }, 30_000) + + it('gives up on a refusal from a database that answered, as on a row naming another seat', async () => { + let now = Date.now() + const reserve = await cell(() => now) + const host = await reserve.bookedHost(5) + vi.spyOn(reserve.relay.assignments, 'activateControlDeferringCell').mockRejectedValue( + new Error('activity_lease_shape_mismatch') + ) + reserve.setAdmitMode('db') + for (let step = 0; step < 3 * REREGISTER_MAX_ATTEMPTS && host.socket.readyState === WebSocket.OPEN; step += 1) { + now += 6_000 + if (host.socket.readyState === WebSocket.OPEN) host.socket.send(JSON.stringify({ type: 'pong', t: now })) + await new Promise((resolve) => setTimeout(resolve, 120)) + } + expect(await host.closeCode).toBe(4409) + // The cell stops reporting reserve once nothing is left to re-register. + await until(() => reserve.relay.sessions.reregistrationPending() === 0) + expect(reserve.relay.sessions.inReserveMode()).toBe(false) + }, 30_000) + + it('retries a stalled database for as long as the control is open, and closes nobody for it', async () => { + let now = Date.now() + const reserve = await cell(() => now) + const host = await reserve.bookedHost(5) + const real = reserve.relay.assignments.activateControlDeferringCell.bind(reserve.relay.assignments) + let stalls = 0 + // Far more failures than the give-up count, as in a long fleet-wide database stall. + vi.spyOn(reserve.relay.assignments, 'activateControlDeferringCell').mockImplementation(async (...args) => { + if (stalls < 3 * REREGISTER_MAX_ATTEMPTS) { + stalls += 1 + // A dropped connection: the database never answered. + throw Object.assign(new Error('terminating connection due to administrator command'), { code: '57P01' }) + } + return await real(...args) + }) + reserve.setAdmitMode('db') + for (let step = 0; step < 6 * REREGISTER_MAX_ATTEMPTS && (await reserve.controlLeases()).length === 0; step += 1) { + now += 6_000 + if (host.socket.readyState === WebSocket.OPEN) host.socket.send(JSON.stringify({ type: 'pong', t: now })) + await new Promise((resolve) => setTimeout(resolve, 120)) + } + expect(stalls).toBe(3 * REREGISTER_MAX_ATTEMPTS) + expect(await reserve.controlLeases()).toHaveLength(1) + expect(host.socket.readyState).toBe(WebSocket.OPEN) + // The endless retry is visible: one line per run of consecutive stalls. + const stalledLines = vi + .mocked(console.warn) + .mock.calls.filter(([line]) => String(line).includes('orca_relay_cell_reregistration_stalled')) + expect(stalledLines).toHaveLength(Math.floor(stalls / REREGISTER_STALL_LOG_EVERY)) + }, 60_000) + + it('retries a lease insert that lost a race (23505) instead of counting it as a refusal', async () => { + let now = Date.now() + const reserve = await cell(() => now) + const host = await reserve.bookedHost(5) + const real = reserve.relay.assignments.activateControlDeferringCell.bind(reserve.relay.assignments) + let races = 0 + vi.spyOn(reserve.relay.assignments, 'activateControlDeferringCell').mockImplementation(async (...args) => { + if (races < 2 * REREGISTER_MAX_ATTEMPTS) { + races += 1 + throw Object.assign(new Error('duplicate key value violates unique constraint'), { code: '23505' }) + } + return await real(...args) + }) + reserve.setAdmitMode('db') + for (let step = 0; step < 5 * REREGISTER_MAX_ATTEMPTS && (await reserve.controlLeases()).length === 0; step += 1) { + now += 6_000 + if (host.socket.readyState === WebSocket.OPEN) host.socket.send(JSON.stringify({ type: 'pong', t: now })) + await new Promise((resolve) => setTimeout(resolve, 120)) + } + expect(races).toBe(2 * REREGISTER_MAX_ATTEMPTS) + expect(await reserve.controlLeases()).toHaveLength(1) + expect(host.socket.readyState).toBe(WebSocket.OPEN) + }, 60_000) + + it('never lets the switch file take the connections hellos and renewals need', () => { + expect(reregisterInFlightLimit(undefined, 10)).toBe(3) + expect(reregisterInFlightLimit(undefined, 16)).toBe(5) + expect(reregisterInFlightLimit(16, 10)).toBe(8) + expect(reregisterInFlightLimit(4, 10)).toBe(4) + expect(reregisterInFlightLimit(16, 2)).toBe(1) + }) + + it('takes its flip-back pace from the switch file', async () => { + const reserve = await cell() + for (let index = 0; index < 4; index += 1) await reserve.bookedHost(5 + index) + const real = reserve.relay.assignments.activateControlDeferringCell.bind(reserve.relay.assignments) + let inFlight = 0 + let peak = 0 + vi.spyOn(reserve.relay.assignments, 'activateControlDeferringCell').mockImplementation(async (...args) => { + inFlight += 1 + peak = Math.max(peak, inFlight) + await new Promise((resolve) => setTimeout(resolve, 50)) + try { + return await real(...args) + } finally { + inFlight -= 1 + } + }) + reserve.setAdmitMode('db', { reregisterInFlight: 1 }) + await until(async () => (await reserve.controlLeases()).length === 4) + expect(peak).toBe(1) + }, 30_000) + + it('leases a control whose proof lands after the flip, though memory admitted its hello', async () => { + const reserve = await cell() + const identity = await reserve.host() + await reserve.insertRow(identity.hostId, 4) + reserve.relay.sessions.reserve({ + v: 1, + directorId: 'director-a', + items: [{ userId: 'user-1', relayHostId: identity.hostId, epoch: 4, ttlMs: 30_000 }] + }) + // The flip lands between the hello (admitted from the booking) and the proof. + const activate = vi.spyOn(reserve.relay.assignments, 'activateControl') + const pending = reserve.connect(identity, 4) + reserve.setAdmitMode('db') + const reply = await pending + expect(reply.ack).toMatchObject({ type: 'host-hello-ack' }) + await until(async () => (await reserve.controlLeases()).length === 1) + // The proof took the lease itself, or the rescan did: either way exactly one. + expect(activate.mock.calls.length + reserve.relay.sessions.reregistrationPending()).toBeLessThanOrEqual(1) + expect(reply.socket.readyState).toBe(WebSocket.OPEN) + }, 30_000) + + it('flips itself back to db when no placing director has been in touch for the dead-man window', async () => { + let now = Date.now() + const reserve = await cell(() => now) + const seated = await reserve.bookedHost(3) + expect(reserve.relay.sessions.inReserveMode()).toBe(true) + const windowMs = reserveDeadManWindowMs(reserve.config.cellId) + expect(windowMs).toBeGreaterThanOrEqual(RESERVE_DEAD_MAN_MS) + expect(windowMs).toBeLessThanOrEqual(2 * RESERVE_DEAD_MAN_MS) + // Cells trip at different times, so a fleet-wide trip is staggered. + const windows = new Set(['c1', 'c2', 'c3', 'c4', 'c5'].map((cell) => reserveDeadManWindowMs(`production-gce-${cell}`))) + expect(windows.size).toBeGreaterThan(1) + // A booking is contact: the window restarts. + now += windowMs - 1_000 + await reserve.bookedHost(4) + now += windowMs - 1_000 + expect(reserve.relay.sessions.inReserveMode()).toBe(true) + now += 2_000 + expect(reserve.relay.sessions.inReserveMode()).toBe(false) + // The alert keys on this line and its cellId (relay-observability.tf). + const tripped = vi + .mocked(console.error) + .mock.calls.map(([line]) => String(line)) + .filter((line) => line.includes('orca_relay_cell_reserve_dead_man_tripped')) + expect(tripped.map((line) => JSON.parse(line).cellId)).toEqual([reserve.config.cellId]) + await until(async () => (await reserve.controlLeases()).length === 2) + expect(seated.socket.readyState).toBe(WebSocket.OPEN) + // Latched for this switch generation: contact again does not bring reserve back. + reserve.relay.sessions.reserve({ + v: 1, + directorId: 'director-a', + items: [{ userId: 'user-1', relayHostId: 'abcdefghijklmnop', epoch: 9, ttlMs: 30_000 }] + }) + expect(reserve.relay.sessions.inReserveMode()).toBe(false) + // A new write of the switch does. + reserve.setAdmitMode('reserve') + expect(reserve.relay.sessions.inReserveMode()).toBe(true) + }, 30_000) +}) diff --git a/cloud/apps/relay/src/cell-seat-feed-route.ts b/cloud/apps/relay/src/cell-seat-feed-route.ts index 3385353118c..422437bbb21 100644 --- a/cloud/apps/relay/src/cell-seat-feed-route.ts +++ b/cloud/apps/relay/src/cell-seat-feed-route.ts @@ -19,6 +19,9 @@ export function registerCellSeatFeedRoute( isDraining?: () => boolean runtimeCounts?: () => RelayRuntimeCounts cellFlags?: () => AppliedControlFlags + // A placing director's poll (reserver=1) feeds the cell's dead-man. + onReserverPoll?: () => void + admitModeEffective?: () => 'db' | 'reserve' // Step 5: what a director needs to estimate a seat and a booking here. reserveCounts?: () => { bookings: number @@ -39,6 +42,7 @@ export function registerCellSeatFeedRoute( if (!bearer || !(await input.verifyRegionalRehomeToken(bearer))) { return context.json({ error: 'invalid_token' }, 401) } + if (context.req.query('reserver') === '1') input.onReserverPoll?.() const cursor = parseCellSeatCursor(context.req.query('since')) if (cursor === 'invalid') return context.json({ error: 'invalid_request' }, 400) const sinceSeq = @@ -69,6 +73,7 @@ export function registerCellSeatFeedRoute( ...(reserve ? { intake: reserve.intake } : {}), // Applied, never desired: generation 0 means no object has been read since boot. ...(input.cellFlags ? { flagsApplied: input.cellFlags() } : {}), + ...(input.admitModeEffective ? { admitModeEffective: input.admitModeEffective() } : {}), ...page }) }) diff --git a/cloud/apps/relay/src/cell-seat-feed.test.ts b/cloud/apps/relay/src/cell-seat-feed.test.ts index ecfbbd0c48c..487f6356709 100644 --- a/cloud/apps/relay/src/cell-seat-feed.test.ts +++ b/cloud/apps/relay/src/cell-seat-feed.test.ts @@ -349,6 +349,7 @@ describe('cell seat feed', () => { }) expect(await runtime.json()).toMatchObject({ flagsApplied: applied, + databasePoolMax: expect.any(Number), // What the flag tool checks a write against before it writes. supportedFlags: { readinessLocal: { type: 'boolean' }, @@ -357,7 +358,8 @@ describe('cell seat feed', () => { readTimeoutMarginMs: { type: 'number', min: 1_000, max: 60_000, integer: true }, admitMode: { type: 'enum', values: ['db', 'reserve'] }, intakePerSec: { type: 'number', max: 1_000 }, - reserveDryRun: { type: 'boolean' } + reserveDryRun: { type: 'boolean' }, + reregisterInFlight: { type: 'number', min: 1, max: 16, integer: true } } }) }) diff --git a/cloud/apps/relay/src/host-session-registry.ts b/cloud/apps/relay/src/host-session-registry.ts index 5410addee44..2da057ea9af 100644 --- a/cloud/apps/relay/src/host-session-registry.ts +++ b/cloud/apps/relay/src/host-session-registry.ts @@ -34,6 +34,7 @@ import type { RelayAssignmentStore } from './assignment-store.js' import type { CellAdmitMode, CellReserveBook } from './cell-reserve-book.js' import type { DemoteRequest, ReserveOutcome, ReserveRequest } from './cell-reserve-contract.js' import { CellSeatLog, type CellSeatChange, type CellSeatFeedPage } from './cell-seat-log.js' +import { isRelayDatabaseTransientError } from './database.js' import { ControlRenewalBatch } from './control-renewal-batch.js' import { RelayCredentialStore, type CredentialReservation } from './credential-store.js' import { HostCloseReasonMemory } from './host-close-reason-memory.js' @@ -89,7 +90,8 @@ export type HostSession = { readonly relayHostId: string readonly generation: number assignmentEpoch: number - readonly controlActivityId: string | null + // Null for a control admitted from memory until a flip back registers it. + controlActivityId: string | null readonly controlResumeSecret: string // Why: reconnect churn is only actionable once it can be pinned to a client build. // Refreshed on rebind so it describes the socket that closed, not the first one. @@ -213,6 +215,10 @@ export type CellReserveAdmission = { mode: () => CellAdmitMode ticketEnforce: () => boolean dryRunEnabled: () => boolean + // A real booking: a placing director is there (the dead-man's evidence). + directorContact?: () => void + // The switch file's flip-back pace, when set. + reregisterInFlight?: () => number | undefined // The pool-pressure shed's condition: a hello memory cannot admit is shed, not queued. databaseShedding: () => boolean book: CellReserveBook @@ -231,6 +237,23 @@ const NO_REJOIN_CLOSE_CODES: ReadonlySet = new Set([ ]) // A demoted duplicate keeps its splices this long, then closes so the desktop re-assigns. export const DEMOTION_CLOSE_MS = 60_000 +// A flip back registers memory-admitted controls with at most this share of the pool in +// flight (US 10 -> 3, Asia 16 -> 5), so DB-path hellos and renewals keep the rest. +export const REREGISTER_POOL_SHARE = 1 / 3 +const REREGISTER_TICK_MS = 100 +// A row the ledger has not written yet gets this long; then the control is closed (4409). +export const REREGISTER_MAX_ATTEMPTS = 20 +// An endless database retry stays visible: one line per this many consecutive stalls. +export const REREGISTER_STALL_LOG_EVERY = 20 + +// The switch-file override never takes the last two connections hellos and renewals need. +export function reregisterInFlightLimit(flag: number | undefined, databasePoolMax: number): number { + if (flag === undefined) return Math.max(1, Math.floor(databasePoolMax * REREGISTER_POOL_SHARE)) + return Math.max(1, Math.min(flag, databasePoolMax - 2)) +} + +// `attempts` counts refusals from a database that answered; `stalls` paces retries of one that did not. +type ReregistrationEntry = { key: string; generation: number; attempts: number; stalls?: number; dueAt: number } export class HostSessionRegistry { private readonly sessions = new Map() @@ -249,6 +272,16 @@ export class HostSessionRegistry { // Seat key -> the demoted control socket. Its own set, not drainSentHosts: it ends with // that socket (close, rebind or a new join), so a later booking here is a fresh seat. private readonly demotedSeats = new Map() + private reregistration: { + queue: ReregistrationEntry[] + // Sessions already queued, given up on, or done in this pass: the rescan skips them. + seen: Set + inFlight: number + // relay_cells delta not yet written: one write per tick folds every activation's. + cellDelta: number + cellWriteInFlight: boolean + timer: ReturnType + } | null = null private readonly idleWork = new Map() private readonly idleAttempts = new Map< @@ -390,6 +423,213 @@ export class HostSessionRegistry { return this.reserveMode() } + // Leaving reserve mode: every control admitted from memory gets the database lease it never + // took, while it stays connected. Rescans until none lacks one, so a hello whose proof lands + // after the flip is caught too. A control whose row never names this cell is closed (4409). + reregisterMemoryControls(): number { + if (this.config.role !== 'cell') return 0 + if (this.reregistration) clearInterval(this.reregistration.timer) + this.reregistration = null + const state: NonNullable = { + queue: [], + seen: new Set(), + inFlight: 0, + cellDelta: 0, + cellWriteInFlight: false, + timer: setInterval(() => this.reregisterTick(state), REREGISTER_TICK_MS) + } + state.timer.unref?.() + const queued = this.rescanUnleasedControls(state) + this.reregistration = state + console.warn( + JSON.stringify({ + event: 'orca_relay_cell_reregistration_started', + ...this.logIdentity(), + controls: queued + }) + ) + return queued + } + + reregistrationPending(): number { + const state = this.reregistration + return state ? state.queue.length + state.inFlight + (state.cellDelta === 0 ? 0 : 1) : 0 + } + + private rescanUnleasedControls(state: NonNullable): number { + const now = this.now() + let queued = 0 + for (const [key, session] of this.sessions) { + if (session.controlActivityId !== null || session.state === 'closed' || !session.socket) continue + const seenKey = `${key}\u0000${session.generation}` + if (state.seen.has(seenKey)) continue + state.seen.add(seenKey) + state.queue.push({ key, generation: session.generation, attempts: 0, dueAt: now }) + queued += 1 + } + return queued + } + + private reregisterTick(state: NonNullable): void { + // Back in reserve mode, the remaining controls need no lease again. + if (this.reserveMode() || this.reregistration !== state) { + clearInterval(state.timer) + if (this.reregistration === state) this.reregistration = null + return + } + this.flushReregistrationCellDelta(state) + const now = this.now() + const maxInFlight = reregisterInFlightLimit( + this.reserveAdmission?.reregisterInFlight?.(), + this.config.databasePoolMax + ) + while (state.inFlight < maxInFlight) { + const index = state.queue.findIndex((entry) => entry.dueAt <= now) + if (index < 0) break + const [entry] = state.queue.splice(index, 1) + this.reregisterOne(state, entry!) + } + if (state.queue.length > 0 || state.inFlight > 0) return + if (this.rescanUnleasedControls(state) > 0) return + if (state.cellDelta !== 0 || state.cellWriteInFlight) return + clearInterval(state.timer) + this.reregistration = null + console.warn( + JSON.stringify({ event: 'orca_relay_cell_reregistration_finished', ...this.logIdentity() }) + ) + } + + private flushReregistrationCellDelta(state: NonNullable): void { + if (state.cellDelta === 0 || state.cellWriteInFlight) return + const delta = state.cellDelta + state.cellDelta = 0 + state.cellWriteInFlight = true + void this.assignments + .commitCellReservationDelta(this.config.cellId, delta) + .catch(() => { + // Retried next tick: the leases exist, only the cell's count is behind. + state.cellDelta += delta + }) + .finally(() => { + state.cellWriteInFlight = false + }) + } + + private reregisterOne( + state: NonNullable, + entry: ReregistrationEntry + ): void { + const session = this.sessions.get(entry.key) + if ( + !session || + session.state === 'closed' || + session.generation !== entry.generation || + session.controlActivityId !== null + ) { + return + } + const identity = { userId: session.identity.sub, relayHostId: session.relayHostId } + const epoch = session.assignmentEpoch + state.inFlight += 1 + void this.assignments + .activateControlDeferringCell(identity, { + cellId: this.config.cellId, + assignmentEpoch: epoch, + generation: session.generation, + idleRegionalRehome: + (session.socket && + this.hostCapabilities.get(session.socket)?.has(RELAY_HOST_CAPABILITY_IDLE_REGIONAL_REHOME)) ?? + false, + cellIncarnation: this.cellIncarnation + }) + .then(async ({ activityId, reservationDelta }) => { + state.cellDelta += reservationDelta + // As a database-path join does: an open migration onto this cell sees its host arrive. + try { + await this.assignments.markMigrationTargetRegistered(identity, { + cellId: this.config.cellId, + assignmentEpoch: epoch + }) + } catch (error) { + this.releaseActivityBestEffort(identity, activityId) + throw error + } + return activityId + }) + .then((activityId) => { + const current = this.sessions.get(entry.key) + if ( + current !== session || + session.state === 'closed' || + session.generation !== entry.generation || + session.assignmentEpoch !== epoch || + this.reserveMode() + ) { + // The control moved on while the lease was taken; it must not stay leased. + this.releaseActivityBestEffort(identity, activityId) + return + } + session.controlActivityId = activityId + session.activityRenewalDueAt = this.now() + RELAY_PROTOCOL_LIMITS.controlPingIntervalMs + }) + .catch((error: unknown) => { + if (this.reregistration !== state) return + const socket = session.socket + const current = + this.sessions.get(entry.key) === session && + session.generation === entry.generation && + socket?.readyState === socket?.OPEN + // A database that did not answer (a stall, a dropped connection, a busy lock) is retried + // for as long as the control is open: a fleet-wide flip back must not close everyone over + // a slow database. An answer that refuses (a row naming another seat, which may still be + // the ledger's write on its way, or any other refusal) counts toward giving up. A lease + // insert that lost a race (23505) is retried too: the next try reads the winner's row. + const unanswered = + isRelayDatabaseTransientError(error) || + (error instanceof Error && error.message === 'database_lock_unavailable') || + (error as { code?: unknown } | null)?.code === '23505' + if (unanswered) { + if (!current) return + entry.stalls = (entry.stalls ?? 0) + 1 + if (entry.stalls % REREGISTER_STALL_LOG_EVERY === 0) { + console.warn( + JSON.stringify({ + event: 'orca_relay_cell_reregistration_stalled', + ...this.logIdentity(), + relayHostIdDigest: relayHostLogDigest(session.relayHostId), + stalls: entry.stalls, + reason: error instanceof Error ? error.message : 'unknown' + }) + ) + } + entry.dueAt = this.now() + Math.min(5_000, 250 * 2 ** Math.min(entry.stalls, 5)) + state.queue.push(entry) + return + } + entry.stalls = 0 + entry.attempts += 1 + if (entry.attempts < REREGISTER_MAX_ATTEMPTS) { + entry.dueAt = this.now() + Math.min(5_000, 250 * 2 ** entry.attempts) + state.queue.push(entry) + return + } + console.warn( + JSON.stringify({ + event: 'orca_relay_cell_reregistration_gave_up', + ...this.logIdentity(), + relayHostIdDigest: relayHostLogDigest(session.relayHostId), + assignmentEpoch: epoch, + reason: error instanceof Error ? error.message : 'unknown' + }) + ) + // Its row names another seat: re-assign instead of serving unleased. + if (current && socket) socket.close(RELAY_CLOSE_CODE.WRONG_CELL, 'control could not re-register') + }) + .finally(() => { + state.inFlight -= 1 + }) + } + reserveCounts(): { bookings: number; intake: { perSec: number; burst: number; tokens: number } } | null { const book = this.reserveAdmission?.book return book ? { bookings: book.count(), intake: book.intakeSnapshot() } : null @@ -400,6 +640,7 @@ export class HostSessionRegistry { if (this.config.role !== 'cell' || !admission || (request.dryRun && !admission.dryRunEnabled())) { return request.items.map(() => ({ outcome: 'off' })) } + if (!request.dryRun) admission.directorContact?.() const context = { mode: admission.mode(), draining: this.draining, @@ -1366,7 +1607,8 @@ export class HostSessionRegistry { reserved?: { reservedBy?: string } ): Promise { let controlActivityId: string | null = null - if (this.config.role === 'cell' && !reserved) { + // Admitted from memory, but the cell left reserve mode before the proof landed: lease now. + if (this.config.role === 'cell' && !(reserved && this.reserveMode())) { try { controlActivityId = await this.assignments.activateControl( { userId: identity.sub, relayHostId: identity.relayHostId }, @@ -1417,6 +1659,9 @@ export class HostSessionRegistry { existing.orphanTimer = null existing.authorityRevision += 1 existing.assignmentEpoch = assignmentEpoch + // A memory-admitted control rebound through the database now holds a lease (the same + // id today, since a rebind keeps its generation). + if (controlActivityId) existing.controlActivityId = controlActivityId existing.socket = socket existing.state = existing.regionalDrainAttemptId ? 'drain-only' : 'active' existing.appVersion = appVersion diff --git a/cloud/apps/relay/src/relay-server.ts b/cloud/apps/relay/src/relay-server.ts index 525cd6b977a..53151ffda87 100644 --- a/cloud/apps/relay/src/relay-server.ts +++ b/cloud/apps/relay/src/relay-server.ts @@ -19,6 +19,7 @@ import { createRemoteJWKSet } from 'jose' import { WebSocketServer } from 'ws' import type WebSocket from 'ws' import type { RawData } from 'ws' +import { ReserveDeadMan, reserveDeadManWindowMs } from './cell-reserve-dead-man.js' import { createRelayApp } from './app.js' import { classifyAssignmentLease } from './assignment-lease.js' import { defaultIntakePerSec, type CellFlags } from './cell-flags.js' @@ -64,6 +65,8 @@ function decodePathSegment(value: string): string | null { // tripping it. On 10-07 every window that met both also logged SQL failures. export const HOST_HELLO_SHED_OLDEST_WAIT_MS = 1_000 const HOST_HELLO_SHED_RETRY_AFTER_SECONDS = 2 +// Well inside the flag channel's 5 s read, so a flip back starts re-registering within ~6 s. +export const RESERVE_MODE_WATCH_MS = 1_000 // Shipped desktops ignore it; it is for the ones that learn to read it. const CELL_FULL_RETRY_AFTER_SECONDS = 2 @@ -206,6 +209,15 @@ export function createRelayServer( const databaseShedding = (): boolean => readRelayDatabasePoolPressure(database).databasePoolWaiting >= config.databasePoolMax && readRelayDatabasePoolOldestWaitMs(database) >= HOST_HELLO_SHED_OLDEST_WAIT_MS + const reserveDeadMan = reserveBook + ? new ReserveDeadMan(config.cellId, options.now, reserveDeadManWindowMs(config.cellId)) + : null + // The switch file's admitMode, unless the dead-man has flipped this cell back. + const effectiveAdmitMode = (): 'db' | 'reserve' => { + const applied = options.cellFlags?.() + if (!applied) return 'db' + return reserveDeadMan ? reserveDeadMan.mode(applied) : applied.flags.admitMode + } const queuedBytes = new ProcessQueuedByteBudget() const sessions = new HostSessionRegistry( config, @@ -224,7 +236,9 @@ export function createRelayServer( }), reserveBook ? { - mode: () => options.cellFlags?.().flags.admitMode ?? 'db', + mode: effectiveAdmitMode, + directorContact: () => reserveDeadMan?.contact(), + reregisterInFlight: () => options.cellFlags?.().flags.reregisterInFlight, ticketEnforce: () => options.cellFlags?.().flags.ticketCheck === 'enforce', dryRunEnabled: () => options.cellFlags?.().flags.reserveDryRun === true, databaseShedding: () => databaseShedding(), @@ -238,8 +252,20 @@ export function createRelayServer( } : undefined ) - // Expired bookings give their units back even while no director is reserving. - const reserveSweepTimer = reserveBook ? setInterval(() => reserveBook.sweep(), 1_000) : null + // Expired bookings give their units back even while no director is reserving. A flip out of + // reserve mode voids the bookings and registers every control admitted from memory. + let appliedAdmitMode = effectiveAdmitMode() + const reserveSweepTimer = reserveBook + ? setInterval(() => { + reserveBook.sweep() + const admitMode = effectiveAdmitMode() + if (appliedAdmitMode === 'reserve' && admitMode !== 'reserve') { + reserveBook.clear() + sessions.reregisterMemoryControls() + } + appliedAdmitMode = admitMode + }, RESERVE_MODE_WATCH_MS) + : null reserveSweepTimer?.unref() const app = createRelayApp(config, { store, @@ -266,6 +292,13 @@ export function createRelayServer( ...(reserveBook && placementCeiling !== null ? { cellReserve: (request: ReserveRequest) => sessions.reserve(request), + cellReserverPoll: () => reserveDeadMan?.contact(), + // Reserve until every control it admitted from memory holds a lease again: the flag + // workflow records db in Postgres (and sweeps resume) only after this says db. + cellAdmitModeEffective: (): 'db' | 'reserve' => + effectiveAdmitMode() === 'reserve' || sessions.reregistrationPending() > 0 + ? 'reserve' + : 'db', cellDemote: (request: DemoteRequest) => sessions.demote(request), cellReserveCounts: () => sessions.reserveCounts(), cellPlacementCeiling: placementCeiling diff --git a/cloud/apps/relay/src/test-fixtures/reserve-mode-cell.ts b/cloud/apps/relay/src/test-fixtures/reserve-mode-cell.ts new file mode 100644 index 00000000000..aff1f7e1ea3 --- /dev/null +++ b/cloud/apps/relay/src/test-fixtures/reserve-mode-cell.ts @@ -0,0 +1,219 @@ +import { createHash, createHmac } from 'node:crypto' +import { createServer, type Server } from 'node:http' +import { createServer as createNetServer } from 'node:net' +import { buildHostProofMacInput, HOST_CHALLENGE_PLAINTEXT_DOMAIN } from '@orca-cloud/relay-contract' +import { exportJWK, generateKeyPair, SignJWT } from 'jose' +import nacl from 'tweetnacl' +import WebSocket from 'ws' +import type { RawData } from 'ws' +import { CELL_FLAG_DEFAULTS, type CellFlags } from '../cell-flags.js' +import type { RelayConfig } from '../config.js' +import type { RelayDatabase } from '../database.js' +import type { AppliedControlFlags } from '../relay-control-flag-channel.js' +import { createRelayServer } from '../relay-server.js' + +// One reserve-mode cell over a real database, with desktops that complete the host proof. +// For the flip-back tests (E-mix), on SQLite and on a latency-injected PostgreSQL. + +export const RESERVE_MODE_CELL_ID = 'production-gce-c3' + +export type ReserveModeCell = Awaited> + +async function unusedPort(): Promise { + const server = createNetServer() + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)) + const address = server.address() + if (!address || typeof address === 'string') throw new Error('missing test port') + await new Promise((resolve) => server.close(() => resolve())) + return address.port +} + +function nextMessage(socket: WebSocket): Promise> { + return new Promise((resolve, reject) => { + socket.once('message', (data: RawData) => { + try { + // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: test frames are JSON objects. + resolve(JSON.parse(data.toString()) as Record) + } catch (error) { + reject(error) + } + }) + socket.once('error', reject) + }) +} + +export async function startReserveModeCell(input: { + database: RelayDatabase + dataDir: string + cleanup: Array<() => Promise | void> + databasePoolMax?: number + now?: () => number +}) { + const keys = await generateKeyPair('ES256') + const publicJwk = await exportJWK(keys.publicKey) + const jwksServer: Server = createServer((_request, response) => { + response.setHeader('content-type', 'application/json') + response.end(JSON.stringify({ keys: [{ ...publicJwk, kid: 'test-key', alg: 'ES256', use: 'sig' }] })) + }) + await new Promise((resolve) => jwksServer.listen(0, '127.0.0.1', resolve)) + input.cleanup.push(() => new Promise((resolve) => jwksServer.close(() => resolve()))) + const jwksAddress = jwksServer.address() + if (!jwksAddress || typeof jwksAddress === 'string') throw new Error('missing JWKS address') + const issuer = `http://127.0.0.1:${jwksAddress.port}` + const port = await unusedPort() + const relayUrl = `http://127.0.0.1:${port}` + const cellId = RESERVE_MODE_CELL_ID + const config = { + port, + publicUrl: relayUrl, + cellUrl: relayUrl, + authIssuer: issuer, + authAudience: 'orca-relay', + jwksUrl: issuer, + assignmentSigningKey: new Uint8Array(32), + role: 'cell', + cellId, + cells: [{ id: cellId, url: relayUrl, capacityRequests: 900, connectionHardCap: 600, connectionUnobservedBound: 60 }], + adminAudience: `${relayUrl}/admin`, + deployServiceAccount: 'deploy@example.com', + runtimeServiceAccount: 'runtime@example.com', + connectionHardCap: 600, + connectionUnobservedBound: 60, + adminJwksUrl: `${issuer}/admin-jwks`, + databasePoolMax: input.databasePoolMax ?? 10, + publicAssignmentsEnabled: true, + publicAssignmentConcurrency: 1, + publicAssignmentQueueMax: 128, + publicAssignmentWaitMs: 4_000, + publicResolveConcurrency: 1, + publicResolveWaitMs: 5_000, + publicAssignmentRetryAfterSeconds: 5, + dataDir: input.dataDir + } satisfies RelayConfig + let applied: AppliedControlFlags = { + generation: 1, + flags: { ...CELL_FLAG_DEFAULTS, admitMode: 'reserve' } + } + const relay = createRelayServer(config, input.database, { + cellFlags: () => applied, + ...(input.now ? { now: input.now } : {}) + }) + relay.server.listen(port, '127.0.0.1') + await new Promise((resolve) => relay.server.once('listening', resolve)) + input.cleanup.push(() => new Promise((resolve) => relay.server.close(() => resolve()))) + await relay.assignments.reconcileCells(config.cells) + + const host = async () => { + const keyPair = nacl.box.keyPair() + const hostId = createHash('sha256').update(keyPair.publicKey).digest('base64url').slice(0, 16) + const token = await new SignJWT({ prof: 'profile-1', org: 'org-1', purpose: 'host-control', relayHostId: hostId }) + .setProtectedHeader({ alg: 'ES256', kid: 'test-key' }) + .setIssuer(issuer) + .setAudience('orca-relay') + .setSubject('user-1') + .setIssuedAt() + .setExpirationTime('30m') + .sign(keys.privateKey) + return { keyPair, hostId, token } + } + + // The first control frame after the proof, and the socket's close code once it closes. + const connect = async (identity: Awaited>, epoch: number) => { + const socket = new WebSocket(`ws://127.0.0.1:${port}/v1/host/control`, { + headers: { authorization: `Bearer ${identity.token}` }, + perMessageDeflate: false + }) + await new Promise((resolve, reject) => { + socket.once('open', () => resolve()) + socket.once('error', reject) + }) + input.cleanup.push(() => socket.terminate()) + const closeCode = new Promise((resolve) => socket.once('close', resolve)) + // A live desktop answers the relay's pings, so a test clock may run ahead. + socket.on('message', (data: RawData) => { + const text = data.toString() + if (!text.includes('"ping"')) return + const frame: unknown = JSON.parse(text) + if (frame && typeof frame === 'object' && 'type' in frame && frame.type === 'ping' && 't' in frame) { + socket.send(JSON.stringify({ type: 'pong', t: frame.t })) + } + }) + socket.send( + JSON.stringify({ + type: 'host-hello', + v: 1, + relayHostId: identity.hostId, + assignmentEpoch: epoch, + hostPublicKeyB64: Buffer.from(identity.keyPair.publicKey).toString('base64'), + appVersion: 'test' + }) + ) + const challenge = await nextMessage(socket) + const plaintext = nacl.box.open( + Buffer.from(String(challenge.ciphertextB64), 'base64'), + Buffer.from(String(challenge.nonceB64), 'base64'), + Buffer.from(String(challenge.relayEphemeralPublicKeyB64), 'base64'), + identity.keyPair.secretKey + ) + if (!plaintext) throw new Error('challenge did not open') + const domainLength = new TextEncoder().encode(`${HOST_CHALLENGE_PLAINTEXT_DOMAIN}\0`).length + const transcriptLength = new DataView(plaintext.buffer, plaintext.byteOffset + domainLength, 4).getUint32(0, false) + const transcriptStart = domainLength + 4 + const transcript = plaintext.slice(transcriptStart, transcriptStart + transcriptLength) + const secret = plaintext.slice(transcriptStart + transcriptLength) + socket.send( + JSON.stringify({ + type: 'host-challenge-ack', + challengeId: challenge.challengeId, + proofB64: createHmac('sha256', secret).update(buildHostProofMacInput(transcript)).digest('base64') + }) + ) + return { ack: await nextMessage(socket), socket, closeCode } + } + + // What the ledger writer would have written about a second after the booked join. + const insertRow = async (relayHostId: string, epoch: number, rowCellId = cellId) => + await input.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)`, + ['user-1', relayHostId, rowCellId, epoch, Date.now(), Date.now()] + ) + + // A booked host, connected from memory at `epoch`, with its row written first unless asked. + const bookedHost = async (epoch: number, options: { row?: 'now' | 'never' | string } = {}) => { + const identity = await host() + if (options.row === undefined || options.row === 'now') await insertRow(identity.hostId, epoch) + else if (options.row !== 'never') await insertRow(identity.hostId, epoch, options.row) + const outcome = relay.sessions.reserve({ + v: 1, + directorId: 'director-a', + items: [{ userId: 'user-1', relayHostId: identity.hostId, epoch, ttlMs: 30_000 }] + }) + if (outcome[0]?.outcome !== 'ok') throw new Error(`booking refused: ${JSON.stringify(outcome)}`) + const reply = await connect(identity, epoch) + if (reply.ack.type !== 'host-hello-ack') throw new Error(`not admitted: ${JSON.stringify(reply.ack)}`) + return { identity, ...reply } + } + + const controlLeases = async () => + await input.database.query( + `SELECT relay_host_id, cell_id, activity_id FROM relay_assignment_activity_leases + WHERE activity_kind = 'control' ORDER BY relay_host_id` + ) + + return { + relay, + config, + host, + connect, + insertRow, + bookedHost, + controlLeases, + setAdmitMode: (admitMode: 'db' | 'reserve', flags: Partial = {}) => { + applied = { generation: applied.generation + 1, flags: { ...CELL_FLAG_DEFAULTS, ...flags, admitMode } } + } + } +}