diff --git a/cloud/apps/relay/src/app.ts b/cloud/apps/relay/src/app.ts index d7a114d680b..bdbd3f03d39 100644 --- a/cloud/apps/relay/src/app.ts +++ b/cloud/apps/relay/src/app.ts @@ -481,7 +481,11 @@ export function createRelayApp( }) } catch (error) { if (!(error instanceof ReservePaceError)) throw error - plan = { kind: 'retry', retryAfterSeconds: error.retryAfterSeconds, reason: 'paced' } + plan = { + kind: 'retry', + retryAfterSeconds: error.retryAfterSeconds, + reason: error.drained ? 'paced-drained' : 'paced' + } } if (plan.kind === 'answer') { operations.recordAssignmentAdmission?.(plan.lane === 'sticky' ? 'sticky' : 'placement') @@ -495,13 +499,7 @@ export function createRelayApp( } if (plan.kind === 'retry') { operations.recordRegionSelection?.({ targetRegion, fallback: false }) - operations.recordAssignmentUnavailable?.( - plan.reason === 'map-incomplete' - ? 'reserve-map-incomplete' - : plan.reason === 'database' - ? 'reserve-database' - : 'reserve-paced' - ) + operations.recordAssignmentUnavailable?.(RESERVE_RETRY_CAUSE[plan.reason]) context.header('Retry-After', String(plan.retryAfterSeconds)) return context.json({ error: 'assignment_temporarily_unavailable' }, 503) } @@ -643,7 +641,7 @@ export function createRelayApp( } catch (error) { if (error instanceof ReservePaceError) { operations.recordRegionSelection?.({ targetRegion, fallback: false }) - operations.recordAssignmentUnavailable?.('reserve-paced') + operations.recordAssignmentUnavailable?.(error.drained ? 'reserve-paced-drained' : 'reserve-paced') context.header('Retry-After', String(error.retryAfterSeconds)) return context.json({ error: 'assignment_temporarily_unavailable' }, 503) } @@ -1025,7 +1023,8 @@ export function createRelayApp( return context.json({ error: 'regional_rehome_trust_not_configured' }, 409) } try { - const source = await operations.assignments.cellDeploymentStatus( + // Runtime rows only, which a reserve-mode cell heartbeats as a database one does. + const { status: source } = await operations.assignments.reserveAwareCellDeploymentStatus( body.data.sourceCellId ) // Any cell that can be drained can be a rehome source, in either @@ -1792,8 +1791,23 @@ export function createRelayApp( const body = AdminCellStatusSchema.safeParse(await context.req.json().catch(() => null)) if (!body.success) return context.json({ error: 'invalid_request' }, 400) try { - const status = await operations.assignments.cellDeploymentStatus(body.data.cellId) - return context.json({ v: 1, status }) + // Without the opt-in a reserve-mode cell stays a 404, so an older script never reads its + // mirror rows as host counts. + if (body.data.view !== 'reserve-aware') { + const status = await operations.assignments.cellDeploymentStatus(body.data.cellId) + return context.json({ v: 1, status }) + } + const result = await operations.assignments.reserveAwareCellDeploymentStatus(body.data.cellId) + if (result.admitMode === 'db') return context.json({ v: 1, status: { ...result.status, reserve: null } }) + const view = operations.reservePlacement?.cellStatus( + body.data.cellId, + result.status.runtime?.cellIncarnation + ) ?? { ok: false, reason: 'placement-off' } + if (!view.ok) { + context.header('Retry-After', String(config.publicAssignmentRetryAfterSeconds)) + return context.json({ error: 'reserve_status_unavailable', reason: view.reason }, 503) + } + return context.json({ v: 1, status: { ...result.status, perHost: null, reserve: view.reserve } }) } catch (error) { return rejectAdminOperation(context, error, 404) } @@ -1801,6 +1815,13 @@ export function createRelayApp( return app } +const RESERVE_RETRY_CAUSE = { + 'map-incomplete': 'reserve-map-incomplete', + database: 'reserve-database', + paced: 'reserve-paced', + 'paced-drained': 'reserve-paced-drained' +} as const satisfies Record<(ReservePlan & { kind: 'retry' })['reason'], AssignmentUnavailableCause> + const AdminHostWhereaboutsSchema = z .object({ v: z.literal(1), @@ -2357,7 +2378,12 @@ const RegionalHostDrainSchema = z .strict() const AdminCellStatusSchema = z - .object({ v: z.literal(1), cellId: z.string().min(1).max(128) }) + .object({ + v: z.literal(1), + cellId: z.string().min(1).max(128), + // Opt-in: a director without it answers 400, which tells a script to keep today's refusal. + view: z.literal('reserve-aware').optional() + }) .strict() function cellFenceAttemptEvidence( diff --git a/cloud/apps/relay/src/assignment-deployment-status-postgres.test.ts b/cloud/apps/relay/src/assignment-deployment-status-postgres.test.ts index 9e8e3b85bd9..8f045a0930d 100644 --- a/cloud/apps/relay/src/assignment-deployment-status-postgres.test.ts +++ b/cloud/apps/relay/src/assignment-deployment-status-postgres.test.ts @@ -84,4 +84,27 @@ describePostgres('PostgreSQL deployment status', () => { restartBlockingReservedRequests: -1 }) }) + + it('reads the same counts through the reserve-aware status, and none for a reserve-mode cell', async () => { + const store = new RelayAssignmentStore(database!, () => 100) + const reserveCell = { ...cell, id: `${cell.id}-reserve`, url: 'https://deployment-status-reserve.example.com' } + await store.reconcileCells([cell, reserveCell]) + await store.assign({ userId: 'aware-user', relayHostId: 'd1e2f3a4b5c6' }) + expect(await store.reserveAwareCellDeploymentStatus(cell.id)).toEqual({ + admitMode: 'db', + status: await store.cellDeploymentStatus(cell.id) + }) + + await store.setCellAdmitMode(reserveCell.id, 'reserve') + await expect(store.cellDeploymentStatus(reserveCell.id)).rejects.toThrow('cell_in_reserve_mode') + const reserve = await store.reserveAwareCellDeploymentStatus(reserveCell.id) + expect(reserve).toMatchObject({ + admitMode: 'reserve', + status: { cellId: reserveCell.id, cellUrl: reserveCell.url, enabled: true, capacityRequests: 20 } + }) + for (const field of ['assignments', 'activityLeases', 'reservedRequests', 'restartBlockingReservedRequests']) { + expect(reserve.status).not.toHaveProperty(field) + } + await expect(store.reserveAwareCellDeploymentStatus('missing-cell')).rejects.toThrow('cell_not_found') + }) }) diff --git a/cloud/apps/relay/src/assignment-store.ts b/cloud/apps/relay/src/assignment-store.ts index 3eab00da52c..8427b0a3e72 100644 --- a/cloud/apps/relay/src/assignment-store.ts +++ b/cloud/apps/relay/src/assignment-store.ts @@ -296,22 +296,14 @@ export type CellEvacuationCapacity = { availableTargetUnits: number } -export type CellDeploymentStatus = { +// Rows the cell heartbeats itself: true in either admit mode. +export type CellDeploymentBaseStatus = { cellId: string cellUrl: string region: RelayRegion enabled: boolean admissionState: CellAdmissionState capacityRequests: number - reservedRequests: number - assignments: number - activityLeases: number - activityRequestUnits: number - restartBlockingActivityLeases: number - restartBlockingActivityRequestUnits: number - restartBlockingReservedRequests: number - outgoingMigrations: number - incomingMigrations: number connectionCapacity: null | { hardCap: number controlRebindReserve: number @@ -337,6 +329,25 @@ export type CellDeploymentStatus = { } } +// Per-host rows: a reserve-mode cell's are a lagging mirror, so they are never reported for one. +export type CellPerHostCounts = { + reservedRequests: number + assignments: number + activityLeases: number + activityRequestUnits: number + restartBlockingActivityLeases: number + restartBlockingActivityRequestUnits: number + restartBlockingReservedRequests: number + outgoingMigrations: number + incomingMigrations: number +} + +export type CellDeploymentStatus = CellDeploymentBaseStatus & CellPerHostCounts + +export type ReserveAwareCellDeploymentStatus = + | { admitMode: 'db'; status: CellDeploymentStatus } + | { admitMode: 'reserve'; status: CellDeploymentBaseStatus } + export type CellFenceAttemptEvidence = { attemptId: string environment: 'staging' | 'production' @@ -3935,63 +3946,41 @@ export class RelayAssignmentStore { // Its host counts here are stale for a reserve-mode cell, so tooling must not read it as empty. async cellDeploymentStatus(cellId: string): Promise { await this.refuseReserveCell(cellId) + const base = await this.cellDeploymentBaseStatus(cellId) + return { ...base, ...(await this.cellPerHostCounts(cellId)) } + } + + // For tooling that knows a reserve-mode cell: it gets no per-host counts, rather than stale ones. + async reserveAwareCellDeploymentStatus(cellId: string): Promise { + const reserveCells = await this.reserveModeCells() + if (reserveCells === null) throw new RelayCellInReserveModeError(null) + const base = await this.cellDeploymentBaseStatus(cellId) + if (reserveCells.has(cellId)) return { admitMode: 'reserve', status: base } + return { admitMode: 'db', status: { ...base, ...(await this.cellPerHostCounts(cellId)) } } + } + + private async cellDeploymentBaseStatus(cellId: string): Promise { const cellRow = ( await this.database.query( - `WITH activity AS ( - SELECT COUNT(*) AS activity_lease_count, - COALESCE(SUM(request_units), 0) AS activity_request_units, - COALESCE(SUM(CASE WHEN activity_kind = 'control' - AND activity_id LIKE 'control-pending:%' - AND request_units = 1 - THEN 1 ELSE 0 END), 0) AS pending_control_units - FROM relay_assignment_activity_leases WHERE cell_id = ? - ) - SELECT cell.*, runtime.cell_url AS runtime_cell_url, + `SELECT cell.*, runtime.cell_url AS runtime_cell_url, admission.admission_state, runtime.cell_incarnation AS runtime_cell_incarnation, runtime.started_at AS runtime_started_at, runtime.ready AS runtime_ready, runtime.observed_requests AS runtime_observed_requests, runtime.last_heartbeat_at AS runtime_last_heartbeat_at, COALESCE(capabilities.regional_rehome_protocol, 0) - AS runtime_regional_rehome_protocol, - activity.activity_lease_count, activity.activity_request_units, - activity.activity_lease_count - activity.pending_control_units - AS restart_blocking_activity_leases, - activity.activity_request_units - activity.pending_control_units - AS restart_blocking_activity_request_units, - cell.reserved_requests - activity.pending_control_units - AS restart_blocking_reserved_requests + AS runtime_regional_rehome_protocol FROM relay_cells cell - CROSS JOIN activity LEFT JOIN relay_cell_admission admission ON admission.cell_id = cell.cell_id LEFT JOIN relay_cell_runtime runtime ON runtime.cell_id = cell.cell_id LEFT JOIN relay_cell_capabilities capabilities ON capabilities.cell_id = runtime.cell_id AND capabilities.cell_incarnation = runtime.cell_incarnation WHERE cell.cell_id = ?`, - [cellId, cellId] + [cellId] ) )[0] if (!cellRow) throw new Error('cell_not_found') - const assignmentRow = ( - await this.database.query( - `SELECT COUNT(*) AS assignment_count FROM relay_assignments WHERE cell_id = ?`, - [cellId] - ) - )[0]! - const migrationRow = ( - await this.database.query( - `SELECT - COALESCE(SUM(CASE WHEN source_cell_id = ? THEN 1 ELSE 0 END), 0) - AS outgoing_migrations, - COALESCE(SUM(CASE WHEN target_cell_id = ? THEN 1 ELSE 0 END), 0) - AS incoming_migrations - FROM relay_assignment_migrations - WHERE completed_at IS NULL AND aborted_at IS NULL - AND (source_cell_id = ? OR target_cell_id = ?)`, - [cellId, cellId, cellId, cellId] - ) - )[0]! const connectionRow = ( await this.database.query( `SELECT limits.hard_cap, limits.unobserved_bound, @@ -4025,24 +4014,6 @@ export class RelayAssignmentStore { enabled: integer(cellRow, 'enabled') === 1, admissionState: parseCellAdmissionState(text(cellRow, 'admission_state')), capacityRequests: integer(cellRow, 'capacity_requests'), - reservedRequests: integer(cellRow, 'reserved_requests'), - assignments: integer(assignmentRow, 'assignment_count'), - activityLeases: integer(cellRow, 'activity_lease_count'), - activityRequestUnits: integer(cellRow, 'activity_request_units'), - restartBlockingActivityLeases: integer( - cellRow, - 'restart_blocking_activity_leases' - ), - restartBlockingActivityRequestUnits: integer( - cellRow, - 'restart_blocking_activity_request_units' - ), - restartBlockingReservedRequests: integer( - cellRow, - 'restart_blocking_reserved_requests' - ), - outgoingMigrations: integer(migrationRow, 'outgoing_migrations'), - incomingMigrations: integer(migrationRow, 'incoming_migrations'), connectionCapacity: connectionRow ? { hardCap: integer(connectionRow, 'hard_cap'), @@ -4093,6 +4064,74 @@ export class RelayAssignmentStore { } } + private async cellPerHostCounts(cellId: string): Promise { + const activityRow = ( + await this.database.query( + `WITH activity AS ( + SELECT COUNT(*) AS activity_lease_count, + COALESCE(SUM(request_units), 0) AS activity_request_units, + COALESCE(SUM(CASE WHEN activity_kind = 'control' + AND activity_id LIKE 'control-pending:%' + AND request_units = 1 + THEN 1 ELSE 0 END), 0) AS pending_control_units + FROM relay_assignment_activity_leases WHERE cell_id = ? + ) + SELECT cell.reserved_requests, + activity.activity_lease_count, activity.activity_request_units, + activity.activity_lease_count - activity.pending_control_units + AS restart_blocking_activity_leases, + activity.activity_request_units - activity.pending_control_units + AS restart_blocking_activity_request_units, + cell.reserved_requests - activity.pending_control_units + AS restart_blocking_reserved_requests + FROM relay_cells cell + CROSS JOIN activity + WHERE cell.cell_id = ?`, + [cellId, cellId] + ) + )[0] + if (!activityRow) throw new Error('cell_not_found') + const assignmentRow = ( + await this.database.query( + `SELECT COUNT(*) AS assignment_count FROM relay_assignments WHERE cell_id = ?`, + [cellId] + ) + )[0]! + const migrationRow = ( + await this.database.query( + `SELECT + COALESCE(SUM(CASE WHEN source_cell_id = ? THEN 1 ELSE 0 END), 0) + AS outgoing_migrations, + COALESCE(SUM(CASE WHEN target_cell_id = ? THEN 1 ELSE 0 END), 0) + AS incoming_migrations + FROM relay_assignment_migrations + WHERE completed_at IS NULL AND aborted_at IS NULL + AND (source_cell_id = ? OR target_cell_id = ?)`, + [cellId, cellId, cellId, cellId] + ) + )[0]! + return { + reservedRequests: integer(activityRow, 'reserved_requests'), + assignments: integer(assignmentRow, 'assignment_count'), + activityLeases: integer(activityRow, 'activity_lease_count'), + activityRequestUnits: integer(activityRow, 'activity_request_units'), + restartBlockingActivityLeases: integer( + activityRow, + 'restart_blocking_activity_leases' + ), + restartBlockingActivityRequestUnits: integer( + activityRow, + 'restart_blocking_activity_request_units' + ), + restartBlockingReservedRequests: integer( + activityRow, + 'restart_blocking_reserved_requests' + ), + outgoingMigrations: integer(migrationRow, 'outgoing_migrations'), + incomingMigrations: integer(migrationRow, 'incoming_migrations') + } + } + async verifyCellAssignment(input: AssignmentIdentity & { cellId: string assignmentEpoch: number diff --git a/cloud/apps/relay/src/cell-status-reserve-view.test.ts b/cloud/apps/relay/src/cell-status-reserve-view.test.ts new file mode 100644 index 00000000000..f51cd0e0d46 --- /dev/null +++ b/cloud/apps/relay/src/cell-status-reserve-view.test.ts @@ -0,0 +1,330 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import type { RelayConfig } from './config.js' + +vi.mock('./admin-token-verifier.js', () => ({ + createAdminTokenVerifier: () => async (token: string) => token === 'deploy-token', + createReadOnlyAdminTokenVerifier: () => async () => false, + createRegionalRehomeControlApplyTokenVerifier: () => async () => false, + createRegionalRehomeRuntimeTokenVerifier: () => async () => false, + createRegionalRehomeTokenVerifier: () => async () => false, + createRuntimeTokenVerifier: () => async () => false +})) + +vi.mock('./relay-token-verifier.js', () => ({ + createRelayTokenVerifier: () => async () => null, + readBearer: (value: string | undefined) => value?.replace(/^Bearer /, '') ?? null +})) + +import { createRelayApp } from './app.js' +import { RelayAssignmentStore } from './assignment-store.js' +import type { CellReserveClient } from './cell-reserve-client.js' +import { openInMemoryRelayDatabase, type RelayDatabase } from './database.js' +import { ReserveAssignment, type ReservePlacementMode } from './reserve-assignment.js' +import { ReservePlacer } from './reserve-placement.js' +import { ShadowSeatDirectory, type SeatFeedCell, type SeatFeedResponse } from './shadow-seat-directory.js' + +const INCARNATION = '11111111-1111-4111-8111-111111111111' +const CELLS = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } +] +const PER_HOST_FIELDS = [ + 'reservedRequests', + 'assignments', + 'activityLeases', + 'activityRequestUnits', + 'restartBlockingActivityLeases', + 'restartBlockingActivityRequestUnits', + 'restartBlockingReservedRequests', + 'outgoingMigrations', + 'incomingMigrations' +] + +function feed(cellId: string, overrides: Partial = {}): SeatFeedResponse { + return { + v: 1, + cellId, + incarnation: INCARNATION, + seq: 1, + at: 0, + draining: false, + counts: { seats: 40, controls: 42, bookings: 3, units: 45, ceiling: 500 }, + intake: { perSec: 5, burst: 20, tokens: 20 }, + flagsApplied: { generation: 3, flags: { admitMode: 'reserve' } }, + admitModeEffective: 'reserve', + admitModeRaw: 'reserve', + full: [], + ...overrides + } +} + +describe('reserve-aware cell status', () => { + let database: RelayDatabase | undefined + + afterEach(async () => { + await database?.close() + database = undefined + }) + + async function setup(options: { mode?: ReservePlacementMode; startedAt?: number; placement?: boolean } = {}) { + const now = { value: 100_000 } + database = await openInMemoryRelayDatabase() + const store = new RelayAssignmentStore(database, () => now.value, { heartbeatTtlMs: 45_000 }) + await store.reconcileCells(CELLS) + for (const cell of CELLS) { + await store.recordCellHeartbeat({ + cellId: cell.id, + cellUrl: cell.url, + cellIncarnation: INCARNATION, + startedAt: 50, + ready: true, + observedRequests: 0 + }) + } + await store.assign({ userId: 'user-a', relayHostId: 'host000000000001' }) + await store.setCellAdmitMode('cell-a', 'reserve') + const cells: SeatFeedCell[] = CELLS.map((cell) => ({ + cellId: cell.id, + cellUrl: cell.url, + region: 'us-central1', + heartbeatExpiresAt: Number.MAX_SAFE_INTEGER, + requiredForComplete: true, + general: true + })) + const directory = new ShadowSeatDirectory() + directory.setCells(cells.map((cell) => cell.cellId)) + const reservePlacement = new ReserveAssignment({ + mode: options.mode ?? 'on', + directory, + cells: () => cells, + startedAt: options.startedAt ?? 0, + placer: new ReservePlacer(() => now.value, () => 0.25), + // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: the status route never books or demotes. + client: {} as CellReserveClient, + readRow: async () => null, + now: () => now.value + }) + const app = createRelayApp(config(), { + store: {} as never, + assignments: store, + ...(options.placement === false ? {} : { reservePlacement }), + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + const status = async (body: Record) => + await app.request('/v1/admin/cell-status', { + method: 'POST', + headers: { authorization: 'Bearer deploy-token', 'content-type': 'application/json' }, + body: JSON.stringify({ v: 1, ...body }) + }) + return { store, directory, now, status } + } + + it('reads a database cell exactly as today, and leaves a reserve cell without per-host counts', async () => { + const { store } = await setup() + expect(await store.reserveAwareCellDeploymentStatus('cell-b')).toEqual({ + admitMode: 'db', + status: await store.cellDeploymentStatus('cell-b') + }) + await expect(store.cellDeploymentStatus('cell-a')).rejects.toThrow('cell_in_reserve_mode') + const reserve = await store.reserveAwareCellDeploymentStatus('cell-a') + expect(reserve).toMatchObject({ + admitMode: 'reserve', + status: { cellId: 'cell-a', runtime: { cellIncarnation: INCARNATION, ready: true, heartbeatFresh: true } } + }) + expect(Object.keys(reserve.status).filter((key) => PER_HOST_FIELDS.includes(key))).toEqual([]) + await database!.query('ALTER TABLE relay_cell_admit_modes RENAME TO relay_cell_admit_modes_gone') + const unknown = new RelayAssignmentStore(database!, () => 100_000) + await expect(unknown.reserveAwareCellDeploymentStatus('cell-b')).rejects.toThrow('reserve_mode_unknown') + }) + + it('keeps the 404 for a reserve cell without the opt-in, and 400s an unknown view', async () => { + const { status } = await setup() + const refused = await status({ cellId: 'cell-a' }) + expect(refused.status).toBe(404) + expect(await refused.json()).toEqual({ error: 'cell_in_reserve_mode' }) + expect((await status({ cellId: 'cell-a', view: 'all' })).status).toBe(400) + }) + + it('answers a database cell with the opt-in as today, plus reserve: null', async () => { + const { status } = await setup() + const today = await statusBody(await status({ cellId: 'cell-b' })) + const response = await status({ cellId: 'cell-b', view: 'reserve-aware' }) + expect(response.status).toBe(200) + expect(await response.json()).toEqual({ ...today, status: { ...today.status, reserve: null } }) + }) + + it('answers a reserve cell from the map, with perHost null and never zeros', async () => { + const { directory, now, status } = await setup() + directory.apply('cell-a', feed('cell-a', { draining: true }), now.value, now.value - 1_500) + directory.apply('cell-b', feed('cell-b', { admitModeEffective: 'db', admitModeRaw: 'db' }), now.value, now.value) + now.value += 20_000 + const response = await status({ cellId: 'cell-a', view: 'reserve-aware' }) + expect(response.status).toBe(200) + const body = await statusBody(response) + expect(body.status).toMatchObject({ + cellId: 'cell-a', + admissionState: 'general', + runtime: { cellIncarnation: INCARNATION }, + perHost: null, + reserve: { + admitMode: 'reserve', + feedStatus: 'live', + // Reported, not enforced: the caller rejects a stale map. + polledAgeMs: 21_500, + incarnation: INCARNATION, + draining: true, + controls: 42, + seats: 40, + bookings: 3, + units: 45, + ceiling: 500, + intakePerSec: 5, + free: 500 - (45 - 3) - 3 + } + }) + expect(Object.keys(body.status).filter((key) => PER_HOST_FIELDS.includes(key))).toEqual([]) + }) + + it('reports nulls, not zeros, for a restarted cell that has not applied its reserve switch', async () => { + const { directory, now, status } = await setup() + directory.apply( + 'cell-a', + feed('cell-a', { + counts: { seats: 0, controls: 0 }, + intake: undefined, + flagsApplied: { generation: 0, flags: {} }, + admitModeEffective: 'db', + admitModeRaw: 'db' + }), + now.value, + now.value + ) + const body = await statusBody(await status({ cellId: 'cell-a', view: 'reserve-aware' })) + expect(body.status.reserve).toMatchObject({ + admitMode: 'db', + bookings: null, + units: null, + ceiling: null, + intakePerSec: null, + free: null + }) + }) + + it('answers a retryable 503 while this director cannot vouch for the cell', async () => { + const unavailable = async (response: Response) => ({ + status: response.status, + retryAfter: response.headers.get('Retry-After'), + body: await response.json() + }) + const expected = (reason: string) => ({ + status: 503, + retryAfter: '5', + body: { error: 'reserve_status_unavailable', reason } + }) + for (const options of [{ mode: 'dry-run' as const }, { placement: false }]) { + const { directory, now, status } = await setup(options) + directory.apply('cell-a', feed('cell-a'), now.value, now.value) + expect(await unavailable(await status({ cellId: 'cell-a', view: 'reserve-aware' }))).toEqual( + expected('placement-off') + ) + await database?.close() + } + // A starting director with one cell unheard. + const starting = await setup({ startedAt: 99_000 }) + starting.directory.apply('cell-a', feed('cell-a'), starting.now.value, starting.now.value) + expect(await unavailable(await starting.status({ cellId: 'cell-a', view: 'reserve-aware' }))).toEqual( + expected('map-incomplete') + ) + await database?.close() + const { directory, now, status } = await setup() + directory.apply('cell-b', feed('cell-b'), now.value, now.value) + directory.markNoFeed('cell-a', now.value) + expect(await unavailable(await status({ cellId: 'cell-a', view: 'reserve-aware' }))).toEqual(expected('no-feed')) + // A feed without the drain flag or counts is an older image: no feed, never an empty cell. + directory.apply('cell-a', feed('cell-a', { draining: undefined }), now.value, now.value) + expect(await unavailable(await status({ cellId: 'cell-a', view: 'reserve-aware' }))).toEqual(expected('no-feed')) + directory.apply('cell-a', feed('cell-a', { counts: undefined }), now.value, now.value) + expect(await unavailable(await status({ cellId: 'cell-a', view: 'reserve-aware' }))).toEqual(expected('no-feed')) + // The process that answered the feed is not the one heartbeating the runtime row. + directory.apply( + 'cell-a', + feed('cell-a', { incarnation: '22222222-2222-4222-8222-222222222222' }), + now.value, + now.value + ) + expect(await unavailable(await status({ cellId: 'cell-a', view: 'reserve-aware' }))).toEqual( + expected('incarnation-mismatch') + ) + }) + + it('probes rehome trust on a reserve cell from its runtime rows', async () => { + const { store } = await setup() + await database!.query( + `INSERT INTO relay_cell_capabilities (cell_id, cell_incarnation, regional_rehome_protocol, last_heartbeat_at) + VALUES ('cell-a', ?, 1, 100000)`, + [INCARNATION] + ) + const sourceFetch = vi.fn(async () => + Response.json({ v: 1, outcome: 'host-not-connected', sharedRuntimeIdentityRejected: true }) + ) + const app = createRelayApp( + config({ + rehomeDirectorServiceAccount: 'relay-director@example.test', + rehomeAudience: 'https://relay.example.test/v1/admin/host-drain' + }), + { + store: {} as never, + assignments: store, + drain: vi.fn(), + regionalRehomeIdentityToken: vi.fn(async () => 'rehome-token'), + // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: the probe only awaits a Response. + regionalRehomeFetch: sourceFetch as unknown as typeof fetch, + ready: vi.fn(async () => true) + } + ) + const response = await app.request('/v1/admin/regional-rehome-trust-probe', { + method: 'POST', + headers: { authorization: 'Bearer deploy-token', 'content-type': 'application/json' }, + body: JSON.stringify({ v: 1, sourceCellId: 'cell-a', sourceCellIncarnation: INCARNATION }) + }) + expect(response.status).toBe(200) + expect(await response.json()).toMatchObject({ proven: true }) + expect(sourceFetch).toHaveBeenCalledTimes(2) + }) +}) + +async function statusBody(response: Response) { + // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: the route answers { v, status } on 200. + return (await response.json()) as { v: number; status: Record } +} + +function config(overrides: Partial = {}): RelayConfig { + return { + port: 8080, + publicUrl: 'https://relay.example.test', + cellUrl: 'https://relay.example.test', + region: 'us-central1', + authIssuer: 'https://auth.example.test', + authAudience: 'orca-relay', + jwksUrl: 'https://auth.example.test/jwks', + assignmentSigningKey: new Uint8Array(32), + role: 'director', + cellId: 'director', + cells: [], + adminAudience: 'https://relay.example.test/v1/admin/drain', + deployServiceAccount: 'deploy@example.test', + runtimeServiceAccount: 'relay-cell@example.test', + adminJwksUrl: 'https://auth.example.test/jwks', + databasePoolMax: 10, + publicAssignmentsEnabled: true, + publicAssignmentConcurrency: 2, + publicAssignmentQueueMax: 128, + publicAssignmentWaitMs: 4_000, + publicResolveConcurrency: 1, + publicResolveWaitMs: 5_000, + publicAssignmentRetryAfterSeconds: 5, + dataDir: './data', + ...overrides + } +} diff --git a/cloud/apps/relay/src/regional-host-drain-app.test.ts b/cloud/apps/relay/src/regional-host-drain-app.test.ts index ebd853dca31..95886981c6c 100644 --- a/cloud/apps/relay/src/regional-host-drain-app.test.ts +++ b/cloud/apps/relay/src/regional-host-drain-app.test.ts @@ -474,20 +474,23 @@ describe('regional rehome director controls', () => { it('probes dedicated trust twice and returns only aggregate proof', async () => { const requests: Array<{ url: string; init?: RequestInit }> = [] - const cellDeploymentStatus = vi.fn().mockResolvedValue({ - cellId: 'production-gce-c7', - cellUrl: 'https://c7.relay.example.test', - region: 'us-central1', - runtime: { - cellIncarnation, - ready: true, - heartbeatFresh: true, - regionalRehomeProtocol: 1 + const reserveAwareCellDeploymentStatus = vi.fn().mockResolvedValue({ + admitMode: 'db', + status: { + cellId: 'production-gce-c7', + cellUrl: 'https://c7.relay.example.test', + region: 'us-central1', + runtime: { + cellIncarnation, + ready: true, + heartbeatFresh: true, + regionalRehomeProtocol: 1 + } } }) const app = createRelayApp(config({ role: 'director', cellId: 'director' }), { store: {} as never, - assignments: { cellDeploymentStatus } as never, + assignments: { reserveAwareCellDeploymentStatus } as never, drain: vi.fn(), regionalRehomeIdentityToken: vi.fn(async () => 'rehome-token'), regionalRehomeFetch: (async (url, init) => { @@ -540,20 +543,23 @@ describe('regional rehome director controls', () => { it('probes a source cell in any region, not only the default one', async () => { // Rehoming moves hosts in both directions, so an asia-east2 cell is a // source too and its trust has to be provable the same way. - const cellDeploymentStatus = vi.fn().mockResolvedValue({ - cellId: 'production-gce-c27', - cellUrl: 'https://c27.relay.example.test', - region: 'asia-east2', - runtime: { - cellIncarnation, - ready: true, - heartbeatFresh: true, - regionalRehomeProtocol: 1 + const reserveAwareCellDeploymentStatus = vi.fn().mockResolvedValue({ + admitMode: 'db', + status: { + cellId: 'production-gce-c27', + cellUrl: 'https://c27.relay.example.test', + region: 'asia-east2', + runtime: { + cellIncarnation, + ready: true, + heartbeatFresh: true, + regionalRehomeProtocol: 1 + } } }) const app = createRelayApp(config({ role: 'director', cellId: 'director' }), { store: {} as never, - assignments: { cellDeploymentStatus } as never, + assignments: { reserveAwareCellDeploymentStatus } as never, drain: vi.fn(), regionalRehomeIdentityToken: vi.fn(async () => 'rehome-token'), regionalRehomeFetch: (async () => @@ -576,21 +582,24 @@ describe('regional rehome director controls', () => { }) it('still refuses a trust probe against a cell without the drain protocol', async () => { - const cellDeploymentStatus = vi.fn().mockResolvedValue({ - cellId: 'production-gce-c27', - cellUrl: 'https://c27.relay.example.test', - region: 'asia-east2', - runtime: { - cellIncarnation, - ready: true, - heartbeatFresh: true, - regionalRehomeProtocol: 0 + const reserveAwareCellDeploymentStatus = vi.fn().mockResolvedValue({ + admitMode: 'db', + status: { + cellId: 'production-gce-c27', + cellUrl: 'https://c27.relay.example.test', + region: 'asia-east2', + runtime: { + cellIncarnation, + ready: true, + heartbeatFresh: true, + regionalRehomeProtocol: 0 + } } }) const sourceFetch = vi.fn() const app = createRelayApp(config({ role: 'director', cellId: 'director' }), { store: {} as never, - assignments: { cellDeploymentStatus } as never, + assignments: { reserveAwareCellDeploymentStatus } as never, drain: vi.fn(), regionalRehomeIdentityToken: vi.fn(async () => 'rehome-token'), regionalRehomeFetch: sourceFetch, @@ -633,14 +642,17 @@ describe('regional rehome director controls', () => { }) it('fails closed when the source rejects the dedicated identity', async () => { - const cellDeploymentStatus = vi.fn().mockResolvedValue({ - cellUrl: 'https://c7.relay.example.test', - region: 'us-central1', - runtime: { - cellIncarnation, - ready: true, - heartbeatFresh: true, - regionalRehomeProtocol: 1 + const reserveAwareCellDeploymentStatus = vi.fn().mockResolvedValue({ + admitMode: 'db', + status: { + cellUrl: 'https://c7.relay.example.test', + region: 'us-central1', + runtime: { + cellIncarnation, + ready: true, + heartbeatFresh: true, + regionalRehomeProtocol: 1 + } } }) const sourceFetch = vi @@ -648,7 +660,7 @@ describe('regional rehome director controls', () => { .mockResolvedValue(Response.json({ error: 'invalid_token' }, { status: 401 })) const app = createRelayApp(config({ role: 'director', cellId: 'director' }), { store: {} as never, - assignments: { cellDeploymentStatus } as never, + assignments: { reserveAwareCellDeploymentStatus } as never, drain: vi.fn(), regionalRehomeIdentityToken: vi.fn(async () => 'rehome-token'), regionalRehomeFetch: sourceFetch, diff --git a/cloud/apps/relay/src/relay-observability.test.ts b/cloud/apps/relay/src/relay-observability.test.ts index f677cd13e7b..98b50fe33b0 100644 --- a/cloud/apps/relay/src/relay-observability.test.ts +++ b/cloud/apps/relay/src/relay-observability.test.ts @@ -243,6 +243,8 @@ describe('relay observability', () => { for (let index = 0; index < 40; index++) { observability.recordAssignmentUnavailable('drain-return-deferred') } + // A drained reserve host's intake pace is scheduled by the drain too. + observability.recordAssignmentUnavailable('reserve-paced-drained') observability.recordAssignmentUnavailable('sticky-lane') observability.recordAssignmentUnavailable('relay_capacity_exhausted') observability.recordAssignmentUnavailable('relay_capacity_exhausted') @@ -252,6 +254,7 @@ describe('relay observability', () => { expect(entries[0]).toMatchObject({ assign503sByCauseDelta: { 'drain-return-deferred': 40, + 'reserve-paced-drained': 1, 'sticky-lane': 1, relay_capacity_exhausted: 2 }, diff --git a/cloud/apps/relay/src/relay-observability.ts b/cloud/apps/relay/src/relay-observability.ts index 27ab9178dc8..40c7ff45415 100644 --- a/cloud/apps/relay/src/relay-observability.ts +++ b/cloud/apps/relay/src/relay-observability.ts @@ -70,6 +70,8 @@ export type AssignmentUnavailableCause = | 'reserve-map-incomplete' | 'reserve-database' | 'reserve-paced' + // Intake pacing for a host re-asking off a draining cell: scheduled by the drain, like a deferral. + | 'reserve-paced-drained' // Row-lock waiters seen in one pg_stat_activity sample, by who waits, on which // table, behind whom. Roles come from each pool's application_name. @@ -124,6 +126,12 @@ export type HostPreAuthPath = 'beforeCheck' | 'db' | 'reserve' // Abandoned: the host hung up before a decision, or before the cell could act on one. export type HostPreAuthOutcome = 'admitted' | 'refused' | 'abandoned' +// Left out of assignNonDrain503sDelta: the drain schedules these 503s. +const DRAIN_503_CAUSES: ReadonlySet = new Set([ + 'drain-return-deferred', + 'reserve-paced-drained' +]) + // Upper bounds in ms; the last bucket is unbounded. 450 is the step-5 t_preauth gate. const PRE_AUTH_BUCKET_BOUNDS_MS = [50, 100, 200, 300, 450, 700, 1_000, 2_000, 5_000] as const @@ -704,7 +712,7 @@ export class RelayObservability implements RelayRuntimeObserver { ...stickySlotFields(deltas), assign503sByCauseDelta: deltas.assign503sByCause, assignNonDrain503sDelta: Object.entries(deltas.assign503sByCause) - .filter(([cause]) => cause !== 'drain-return-deferred') + .filter(([cause]) => !DRAIN_503_CAUSES.has(cause)) .reduce((sum, [, count]) => sum + count, 0), // Waiters summed over samples; divided by samples it is the mean number of // backends waiting, i.e. lock-wait seconds per second. diff --git a/cloud/apps/relay/src/relay-region-app.test.ts b/cloud/apps/relay/src/relay-region-app.test.ts index a990fffe529..6161bcf5a1c 100644 --- a/cloud/apps/relay/src/relay-region-app.test.ts +++ b/cloud/apps/relay/src/relay-region-app.test.ts @@ -11,6 +11,7 @@ vi.mock('./relay-token-verifier.js', () => ({ })) import { createRelayApp } from './app.js' +import { ReservePaceError } from './reserve-assignment.js' describe('Relay region API', () => { it('passes a valid preference to placement and records coarse outcomes', async () => { @@ -204,6 +205,37 @@ describe('Relay region API', () => { expect(plan).not.toHaveBeenCalled() }) + it('names a drained host’s pace as the drain’s on the reserve lane and in the database lane', async () => { + const causes: string[] = [] + const drainedPace = new ReservePaceError(3, true) + const plan = vi.fn(async () => { + throw drainedPace + }) + const assign = vi.fn(async () => { + throw new ReservePaceError(2, true) + }) + const app = createRelayApp(config(), { + store: {} as never, + assignments: { assign } as never, + reservePlacement: { plan } as never, + drain: vi.fn(), + ready: async () => true, + recordAssignmentUnavailable: (cause) => causes.push(cause) + }) + const fromPlan = await app.request('/v1/assign', assignmentRequest('abcdefghijklmnop', {})) + expect(fromPlan.status).toBe(503) + expect(fromPlan.headers.get('Retry-After')).toBe('3') + plan.mockResolvedValue({ kind: 'database' } as never) + const fromAssign = await app.request('/v1/assign', assignmentRequest('abcdefghijklmnop', {})) + expect(fromAssign.status).toBe(503) + expect(fromAssign.headers.get('Retry-After')).toBe('2') + expect(causes).toEqual(['reserve-paced-drained', 'reserve-paced-drained']) + // Any other host's pace keeps today's cause. + plan.mockRejectedValue(new ReservePaceError(1)) + await app.request('/v1/assign', assignmentRequest('abcdefghijklmnop', {})) + expect(causes.at(-1)).toBe('reserve-paced') + }) + it('returns successful placement when optional window storage is unavailable', async () => { const app = createRelayApp(config(), { store: {} as never, diff --git a/cloud/apps/relay/src/reserve-assignment.test.ts b/cloud/apps/relay/src/reserve-assignment.test.ts index 498ff4a500e..6f1f0e13794 100644 --- a/cloud/apps/relay/src/reserve-assignment.test.ts +++ b/cloud/apps/relay/src/reserve-assignment.test.ts @@ -893,4 +893,172 @@ describe('reserve assignment on a director', () => { expect(outcomes(lines)).toMatchObject([{ outcome: 'unreachable' }]) }) }) + + describe('for a host re-asking off a draining cell', () => { + const events = (lines: string[], event: string) => + lines.map((line) => JSON.parse(line)).filter((entry) => entry.event === event) + const flushedSummary = async (assignment: ReserveAssignment, now: { value: number }, lines: string[]) => { + now.value += 60_001 + await assignment.plan({ ...HOST, relayHostId: 'zzzzzzzzzzzzzzzz' }, { reconnect: false, region: US }) + return events(lines, 'orca_relay_reserve_placement_summary')[0] + } + + // graceMs 0: the host re-asks while its seat still shows, but the cell's feed already says draining. + it('counts the placement by source cell before any leave lands', async () => { + const lines: string[] = [] + const { assignment, now } = setup({ + feeds: [feed('c1', { draining: true, full: [seated(7)] }), feed('c2')], + log: (line) => lines.push(line) + }) + expect(await assignment.plan(HOST, { reconnect: true, region: US })).toMatchObject({ + kind: 'answer', + lane: 'placement', + assignment: { cellId: 'c2', assignmentEpoch: 8 } + }) + expect(events(lines, 'orca_relay_reserve_drain_replaced')).toEqual([ + { event: 'orca_relay_reserve_drain_replaced', fromCellId: 'c1', toCellId: 'c2', region: US } + ]) + expect((await flushedSummary(assignment, now, lines)).drainReplaced).toEqual({ c1: 1 }) + }) + + it('counts a drain-only seat, and a newest leave on a draining cell', async () => { + const lines: string[] = [] + const drainOnly = setup({ + feeds: [feed('c1', { full: [{ ...seated(7), state: 'drain-only' }] }), feed('c2')], + log: (line) => lines.push(line) + }) + await drainOnly.assignment.plan(HOST, { reconnect: true, region: US }) + const left = setup({ feeds: [feed('c1', { draining: true }), feed('c2')], log: (line) => lines.push(line) }) + left.directory.apply( + 'c1', + feed('c1', { + draining: true, + full: undefined, + seq: 3, + changes: [ + { seq: 2, kind: 'join', ...HOST, epoch: 7, generation: 1, at: 1 }, + { seq: 3, kind: 'leave', ...HOST, epoch: 7, generation: 1, closeCode: 4503, at: 2 } + ] + }), + left.now.value, + left.now.value + ) + await left.assignment.plan(HOST, { reconnect: true, region: US }) + expect(events(lines, 'orca_relay_reserve_drain_replaced').map((entry) => entry.fromCellId)).toEqual([ + 'c1', + 'c1' + ]) + }) + + it('counts a host once: after its re-booking, a second ask is not a drain re-ask', async () => { + const lines: string[] = [] + const { assignment, directory, now } = setup({ + feeds: [feed('c1', { draining: true, full: [{ ...seated(7), state: 'drain-only' }] }), feed('c2')], + log: (line) => lines.push(line) + }) + await assignment.plan(HOST, { reconnect: true, region: US }) + // c2 reports the booking before the host's join: the map still holds only the drain-only seat. + directory.apply( + 'c2', + feed('c2', { + full: undefined, + seq: 2, + changes: [{ seq: 2, kind: 'reserve', ...HOST, epoch: 8, generation: 0, reservedBy: 'director', at: now.value }] + }), + now.value, + now.value + ) + expect(directory.bookingOf(HOST.userId, HOST.relayHostId)).toMatchObject({ cellId: 'c2', epoch: 8 }) + await assignment.plan(HOST, { reconnect: false, region: US }) + expect(events(lines, 'orca_relay_reserve_drain_replaced')).toHaveLength(1) + }) + + it('counts a host once when it asks again before the new cell reports the booking', async () => { + const lines: string[] = [] + const { assignment } = setup({ + feeds: [feed('c1', { draining: true, full: [seated(7)] }), feed('c2')], + log: (line) => lines.push(line) + }) + await assignment.plan(HOST, { reconnect: true, region: US }) + await assignment.plan(HOST, { reconnect: false, region: US }) + expect(events(lines, 'orca_relay_reserve_drain_replaced')).toHaveLength(1) + }) + + // Placement memories are clock-free: a host that never reached c2 is not counted again later. + it('counts a host once per drain, also when it asks again after 30 s', async () => { + const lines: string[] = [] + const { assignment, directory, now } = setup({ + feeds: [feed('c1', { draining: true, full: [{ ...seated(7), state: 'drain-only' }] }), feed('c2')], + log: (line) => lines.push(line) + }) + await assignment.plan(HOST, { reconnect: true, region: US }) + now.value += 31_000 + for (const cellId of ['c1', 'c2']) { + directory.apply(cellId, feed(cellId, { draining: cellId === 'c1', full: undefined, seq: 1, changes: [] }), now.value, now.value) + } + await assignment.plan(HOST, { reconnect: true, region: US }) + expect(events(lines, 'orca_relay_reserve_drain_replaced')).toHaveLength(1) + }) + + it('counts afresh once the drained cell restarts', async () => { + const lines: string[] = [] + const { assignment, directory, now } = setup({ + feeds: [feed('c1', { draining: true, full: [seated(7)] }), feed('c2')], + log: (line) => lines.push(line) + }) + await assignment.plan(HOST, { reconnect: true, region: US }) + // The host came back to the restarted c1, which drains again. + directory.apply( + 'c1', + feed('c1', { incarnation: 'c1-incarnation-2', draining: true, full: [seated(20)] }), + now.value, + now.value + ) + await assignment.plan(HOST, { reconnect: true, region: US }) + expect(events(lines, 'orca_relay_reserve_drain_replaced')).toHaveLength(2) + }) + + it('marks its pace as the drain’s, on the reserve lane, in the region fallback and in the spill', async () => { + const { assignment } = setup({ + feeds: [feed('c1', { draining: true, full: [seated(7)] }), feed('c2')], + reserve: () => ({ outcome: 'intake' }) + }) + await expect(assignment.plan(HOST, { reconnect: true, region: US })).rejects.toMatchObject({ drained: true }) + const plan = await assignment.plan(HOST, { reconnect: true, region: US, regionCorrection: true }) + await expect(plan.kind === 'database' ? plan.regionReserve?.() : null).rejects.toMatchObject({ drained: true }) + await expect( + plan.kind === 'database' ? plan.reserveSpill?.({ skipOwnRegion: false, rowEpoch: 0 }) : null + ).rejects.toMatchObject({ drained: true }) + }) + + it('leaves every other host as today: no drain count, and a plain pace', async () => { + const lines: string[] = [] + const { assignment } = setup({ feeds: [feed('c1', { full: [seated(7)] }), feed('c2')], log: (line) => lines.push(line) }) + // A seat on another draining cell than the newest one does not count either. + const paced = setup({ + feeds: [feed('c1', { full: [seated(7)] }), feed('c2', { draining: true, full: [seated(6)] })], + reserve: () => ({ outcome: 'intake' }) + }) + expect(await assignment.plan({ ...HOST, relayHostId: 'qrstuvwxyzabcdef' }, { reconnect: false, region: US })).toMatchObject({ + kind: 'answer', + lane: 'placement' + }) + await expect( + paced.assignment.plan(HOST, { reconnect: false, region: US }) + ).rejects.toMatchObject({ drained: false }) + expect(events(lines, 'orca_relay_reserve_drain_replaced')).toEqual([]) + }) + + // The database lane logs orca_relay_sticky_replaced_off_isolated_cell for these; counting both would double the canary's sum. + it('does not count a placement the database lane makes through its region fallback', async () => { + const lines: string[] = [] + const { assignment } = setup({ + feeds: [feed('c1', { draining: true, full: [seated(7)] }), feed('c2')], + log: (line) => lines.push(line) + }) + const plan = await assignment.plan(HOST, { reconnect: true, region: US, regionCorrection: true }) + await expect(plan.kind === 'database' ? plan.regionReserve?.() : null).resolves.toMatchObject({ cellId: 'c2' }) + expect(events(lines, 'orca_relay_reserve_drain_replaced')).toEqual([]) + }) + }) }) diff --git a/cloud/apps/relay/src/reserve-assignment.ts b/cloud/apps/relay/src/reserve-assignment.ts index fd92744aa5f..1d7f8605e7e 100644 --- a/cloud/apps/relay/src/reserve-assignment.ts +++ b/cloud/apps/relay/src/reserve-assignment.ts @@ -37,26 +37,59 @@ const NO_REJOIN_CLOSE_CODES: ReadonlySet = new Set([ export type ReservePlacementMode = 'off' | 'dry-run' | 'on' export class ReservePaceError extends Error { - constructor(readonly retryAfterSeconds: number) { + // drained: the host was re-asking off a draining cell, so the 503 is the drain's, not a fault. + constructor( + readonly retryAfterSeconds: number, + readonly drained = false + ) { super('reserve_cells_paced') } } type Identity = { userId: string; relayHostId: string } +function hostKey(identity: Identity): string { + return `${identity.userId}\u0000${identity.relayHostId}` +} + export type ReservePlan = | { kind: 'answer'; assignment: RelayAssignment; lane: 'sticky' | 'placement' | 'rehome' } - | { kind: 'retry'; retryAfterSeconds: number; reason: 'map-incomplete' | 'database' | 'paced' } + | { + kind: 'retry' + retryAfterSeconds: number + reason: 'map-incomplete' | 'database' | 'paced' | 'paced-drained' + } | ({ kind: 'database' } & Pick) // A host whose seat the ledger reconcile demoted mints above its row for this long, so the // re-assign does not collide with the row again (the mirror refuses an equal epoch). export const RESERVE_ROW_FLOOR_TTL_MS = 2 * 60 * 60_000 const RESERVE_ROW_FLOOR_MAX = 100_000 +const OWN_BOOKINGS_MAX = 100_000 // The WRONG_CELL row read's bound: well under a database stall (5-7 s). export const RESERVE_WRONG_CELL_ROW_READ_MS = 1_000 +// A reserve-mode cell as this director's map holds it, for the reserve-aware cell-status route. +export type ReserveCellStatus = { + admitMode: 'db' | 'reserve' + feedStatus: string + // Reported, not enforced: callers set their own freshness bound. + polledAgeMs: number + incarnation: string + draining: boolean + controls: number + seats: number + // Null: the cell reports no step-5 numbers (for example before its switch file applies). + bookings: number | null + units: number | null + ceiling: number | null + intakePerSec: number | null + free: number | null +} + +export type ReserveCellStatusUnavailable = 'placement-off' | 'map-incomplete' | 'no-feed' | 'incarnation-mismatch' + // Rejects after `ms`; the read itself still settles later, and its failure is swallowed. async function withinMs(read: Promise, ms: number): Promise { let timer: ReturnType | undefined @@ -75,6 +108,10 @@ export class ReserveAssignment { private summary: PlacementSummary = emptySummary() private summaryAt: number private readonly rowFloors = new Map() + // This director's own bookings, which the cell's feed reports only a poll or more later. + private readonly ownBookings = new Map() + // Hosts already counted off each drained cell, for that cell incarnation's drain. + private readonly drainCounted = new Map }>() constructor( private readonly input: { @@ -157,7 +194,7 @@ export class ReserveAssignment { const now = this.now() const { directory } = this.input // Until every live cell has answered, only a host that may be on an unheard cell waits. - const incomplete = !directory.isComplete() && now - this.input.startedAt < RESERVE_STARTUP_GATE_MS + const incomplete = this.mapIncomplete(now) const seats = directory.seatsOf(identity.userId, identity.relayHostId) const left = directory .recentlyLeftOf(identity.userId, identity.relayHostId, now) @@ -165,6 +202,7 @@ export class ReserveAssignment { .at(-1) const booking = directory.bookingOf(identity.userId, identity.relayHostId) const rowFloor = this.rowFloor(identity, now) + const drainedFrom = this.drainSource(identity, seats, left, booking) const known = [ ...seats.map((seat) => seat.epoch), ...(left ? [left.epoch] : []), @@ -210,12 +248,12 @@ export class ReserveAssignment { return { kind: 'database' } } // A never-seen host's spill mints above the row the database lane read. - const reserveSpill = this.reserveSpill(identity, request.region, epochFloor, seats) + const reserveSpill = this.reserveSpill(identity, request.region, epochFloor, seats, drainedFrom) return known.length === 0 ? { kind: 'database', reserveSpill } : { kind: 'database', - regionReserve: this.regionReserve(identity, request.region, epochFloor, seats, now), + regionReserve: this.regionReserve(identity, request.region, epochFloor, seats, now, drainedFrom), reserveSpill } } @@ -236,9 +274,11 @@ export class ReserveAssignment { } } const placeFresh = async (): Promise => - await this.placeFresh(identity, request.region, epochFloor, seats) - const regionReserve = this.regionReserve(identity, request.region, epochFloor, seats, now) - const reserveSpill = incomplete ? undefined : this.reserveSpill(identity, request.region, epochFloor, seats) + await this.placeFresh(identity, request.region, epochFloor, seats, false, false, drainedFrom) + const regionReserve = this.regionReserve(identity, request.region, epochFloor, seats, now, drainedFrom) + const reserveSpill = incomplete + ? undefined + : this.reserveSpill(identity, request.region, epochFloor, seats, drainedFrom) // Where the host was last seen decides who owns it: a database-mode cell keeps today's // path, sticky included; only a fresh placement may be booked on a reserve-mode cell. const lastOnReserve = lastCellId !== undefined && directory.placementModeOf(lastCellId) === 'reserve' @@ -246,10 +286,105 @@ export class ReserveAssignment { return { kind: 'database', epochFloor, placeFresh, regionReserve, reserveSpill } } const placed = await placeFresh() - if (placed) return { kind: 'answer', assignment: placed, lane: 'placement' } + if (placed) { + if (drainedFrom !== undefined) this.recordDrainReplaced(identity, drainedFrom, placed) + return { kind: 'answer', assignment: placed, lane: 'placement' } + } return { kind: 'database', epochFloor, regionReserve, reserveSpill } } + private mapIncomplete(now: number): boolean { + return !this.input.directory.isComplete() && now - this.input.startedAt < RESERVE_STARTUP_GATE_MS + } + + // Keyed on the cell, not the leave: a same-cap drain (graceMs 0) has the host re-ask before + // its leave reaches the map, but the cell's feed says draining from the first poll after. + // Only the newest grant counts, so a host already re-booked elsewhere is not a drain re-ask. + private drainSource( + identity: Identity, + seats: readonly ShadowSeat[], + left: { cellId: string; epoch: number } | undefined, + booking: { cellId: string; epoch: number } | undefined + ): string | undefined { + const own = this.ownBookings.get(hostKey(identity)) + let newest: { cellId: string; epoch: number; state?: string } | undefined + for (const entry of [...seats, ...(left ? [left] : []), ...(booking ? [booking] : []), ...(own ? [own] : [])]) { + if (!newest || entry.epoch > newest.epoch) newest = entry + } + if (!newest) return undefined + const drained = + newest.state === 'drain-only' || this.input.directory.cellState(newest.cellId)?.draining === true + return drained ? newest.cellId : undefined + } + + // Only the reserve lane's own answers: a database-lane placement off the drained (isolated) + // cell already logs orca_relay_sticky_replaced_off_isolated_cell, and the canary sums both. + // Each host counts once per drain: the canary bar must never pass on a host counted twice. + private recordDrainReplaced(identity: Identity, fromCellId: string, placed: RelayAssignment): void { + const incarnation = this.input.directory.cellState(fromCellId)?.incarnation + let counted = this.drainCounted.get(fromCellId) + if (!counted || counted.incarnation !== incarnation) { + counted = { incarnation, hosts: new Set() } + this.drainCounted.set(fromCellId, counted) + } + const key = hostKey(identity) + if (counted.hosts.has(key)) return + counted.hosts.add(key) + this.summary.drainReplaced[fromCellId] = (this.summary.drainReplaced[fromCellId] ?? 0) + 1 + this.log( + JSON.stringify({ + event: 'orca_relay_reserve_drain_replaced', + fromCellId, + toCellId: placed.cellId, + region: placed.region + }) + ) + } + + // The reserve-aware cell-status view, or why this director cannot vouch for the cell now. + cellStatus( + cellId: string, + runtimeIncarnation: string | undefined + ): { ok: true; reserve: ReserveCellStatus } | { ok: false; reason: ReserveCellStatusUnavailable } { + const now = this.now() + if (this.input.mode !== 'on') return { ok: false, reason: 'placement-off' } + if (this.mapIncomplete(now)) return { ok: false, reason: 'map-incomplete' } + const { directory } = this.input + const state = directory.cellState(cellId) + const placement = this.placementCells().find((cell) => cell.cellId === cellId) + // An older image without these reads as no feed, never as an empty cell. + if ( + !state || + !placement || + state.status === 'no-feed' || + state.incarnation === undefined || + state.polledAt === undefined || + state.draining === undefined || + state.reportedSeats === undefined || + state.reportedControls === undefined + ) { + return { ok: false, reason: 'no-feed' } + } + if (state.incarnation !== runtimeIncarnation) return { ok: false, reason: 'incarnation-mismatch' } + return { + ok: true, + reserve: { + admitMode: directory.placementModeOf(cellId), + feedStatus: state.status, + polledAgeMs: now - state.polledAt, + incarnation: state.incarnation, + draining: state.draining, + controls: state.reportedControls, + seats: state.reportedSeats, + bookings: state.bookings ?? null, + units: state.units ?? null, + ceiling: state.ceiling ?? null, + intakePerSec: state.intake?.perSec ?? null, + free: state.ceiling === undefined ? null : this.input.placer.free(placement) + } + } + } + // For the share draw's losers: the database lane tries the region first, then this, then spills. // Undefined while no in-region reserve cell has a seat, so the database lane runs as today. private regionReserve( @@ -257,14 +392,15 @@ export class ReserveAssignment { region: RelayRegion, epochFloor: number, seats: readonly ShadowSeat[], - now: number + now: number, + drainedFrom: string | undefined ): AssignOptions['regionReserve'] { const hasRoom = this.placementCells().some( (cell) => this.input.placer.isCandidate(cell, region, now) && this.input.placer.hasSeat(cell) ) if (!hasRoom) return undefined return async () => { - const placed = await this.placeFresh(identity, region, epochFloor, seats, true) + const placed = await this.placeFresh(identity, region, epochFloor, seats, true, false, drainedFrom) if (placed) this.summary.regionKept += 1 return placed } @@ -276,7 +412,8 @@ export class ReserveAssignment { identity: Identity, region: RelayRegion, epochFloor: number, - seats: readonly ShadowSeat[] + seats: readonly ShadowSeat[], + drainedFrom: string | undefined ): NonNullable { return async ({ skipOwnRegion, rowEpoch }) => { const floor = Math.max(epochFloor, rowEpoch) @@ -292,7 +429,7 @@ export class ReserveAssignment { continue } try { - const placed = await this.placeFresh(identity, target, floor, seats, true, true) + const placed = await this.placeFresh(identity, target, floor, seats, true, true, drainedFrom) if (placed) { this.recordSpill(region, target) return placed @@ -478,7 +615,8 @@ export class ReserveAssignment { regionFull = false, // The caller counts its own outcome (spill, rehome): no database lane follows it, and a pace // may not be the answer. - spill = false + spill = false, + drainedFrom?: string ): Promise { const placements = this.placementCells() if (!regionFull && this.random() >= this.input.placer.reservePoolShare(placements, region)) return null @@ -504,6 +642,7 @@ export class ReserveAssignment { this.summary.calls += result.calls if (result.kind === 'placed') { this.summary.placed += 1 + this.rememberOwnBooking(identity, result.cellId, result.epoch) this.supersede(identity, seats, result.cellId, result.epoch) const cell = this.input.cells().find((entry) => entry.cellId === result.cellId)! return { @@ -519,7 +658,7 @@ export class ReserveAssignment { if (!spill) { this.summary.retries.paced += 1 } - throw new ReservePaceError(result.retryAfterSeconds) + throw new ReservePaceError(result.retryAfterSeconds, drainedFrom !== undefined) } // Database cells in this region take it; a region with none spills as today. const hasDatabaseCells = placements.some( @@ -563,7 +702,7 @@ export class ReserveAssignment { // The ledger reconcile demoted this host's seat behind (or level with) a row at `epoch`. raiseEpochFloor(identity: Identity, epoch: number): void { - const key = `${identity.userId}\u0000${identity.relayHostId}` + const key = hostKey(identity) const current = this.rowFloors.get(key) this.rowFloors.delete(key) this.rowFloors.set(key, { @@ -575,8 +714,17 @@ export class ReserveAssignment { } } + private rememberOwnBooking(identity: Identity, cellId: string, epoch: number): void { + const key = hostKey(identity) + this.ownBookings.delete(key) + this.ownBookings.set(key, { cellId, epoch }) + if (this.ownBookings.size > OWN_BOOKINGS_MAX) { + this.ownBookings.delete(this.ownBookings.keys().next().value!) + } + } + private rowFloor(identity: Identity, now: number): number | undefined { - const key = `${identity.userId}\u0000${identity.relayHostId}` + const key = hostKey(identity) const floor = this.rowFloors.get(key) if (!floor) return undefined if (floor.expiresAt > now) return floor.epoch @@ -653,6 +801,13 @@ export class ReserveAssignment { const now = this.now() if (now - this.summaryAt < SUMMARY_INTERVAL_MS) return this.summaryAt = now + // A drain ends when the cell stops reporting it or restarts; its next drain counts afresh. + for (const [cellId, counted] of this.drainCounted) { + const state = this.input.directory.cellState(cellId) + if (!state || state.draining === false || state.incarnation !== counted.incarnation) { + this.drainCounted.delete(cellId) + } + } const { latenciesMs, outcomes, cells } = this.summary.dryRun latenciesMs.sort((left, right) => left - right) const percentile = (share: number) => @@ -701,6 +856,8 @@ type PlacementSummary = { calls: number database: number supersedes: number + // Reserve-lane placements of hosts re-asking off a draining cell, keyed by that cell. + drainReplaced: Record outcomes: Record retries: { mapIncomplete: number; database: number; paced: number } dryRun: { @@ -732,6 +889,7 @@ function emptySummary(): PlacementSummary { calls: 0, database: 0, supersedes: 0, + drainReplaced: Object.create(null), outcomes: {}, retries: { mapIncomplete: 0, database: 0, paced: 0 }, dryRun: { outcomes: {}, latenciesMs: [], cells: Object.create(null) }