From 7080eb06044b8eb3b953781ef033abcc35005279 Mon Sep 17 00:00:00 2001 From: Jinwoo Hong <73622457+Jinwoo-H@users.noreply.github.com> Date: Fri, 18 Sep 2026 21:16:19 -0400 Subject: [PATCH] fix(relay): bound the idle-rehome candidate poll to a window of decisions (#21557) * fix(relay): bound the idle-rehome candidate poll to a window of decisions The director's idle-regional-rehome poll built every (eligible host x target cell in its preferred region) pair, applied the cohort predicate downstream of that fan-out, sorted the lot, and took LIMIT 100 OFFSET n. Its cost was set by the size of the fleet and the width of the cohort, so raising the cohort from 10% to 100% pushed it past the serving pool's 5 s statement_timeout and the rollout stalled at 0.37 hosts/min. The poll now resolves the cell inventory once (tens of rows), takes a bounded window of decision rows in primary-key order from a keyset cursor with the cohort, freshness and cross-region predicates applied first, verifies only that window against the host-side gates, and ranks targets in the process. Same candidates in the same priority order; the work per poll no longer depends on the cohort or the fleet. Adds a once-a-minute aggregated poll summary so an operator can tell a poll gated by the dispatch budget from one that found nobody to move. Co-Authored-By: Claude * fix(relay): pin the rehome verification to the window's exact keys The window read and the verification read take separate snapshots. The verification repeated the window's predicate with its own LIMIT, so a decision that turned eligible between the two reads shifted that LIMIT and pushed the window's last host out of it -- while the cursor still advanced past that host, skipping it for a whole sweep. The verification now names the keys the window returned. Its LIMIT stays as the optimisation fence that stops Postgres flattening the subquery, but can no longer truncate a key set that is at most one window long. Co-Authored-By: Claude --------- Co-authored-by: Claude --- cloud/apps/relay/src/assignment-store.ts | 52 ++- .../src/idle-regional-rehome-selection.ts | 372 +++++++++++++----- .../src/idle-regional-rehome-window.test.ts | 310 +++++++++++++++ .../regional-rehome-poll-telemetry.test.ts | 58 +++ .../src/regional-rehome-poll-telemetry.ts | 69 ++++ .../relay/src/regional-rehome-store.test.ts | 2 +- 6 files changed, 754 insertions(+), 109 deletions(-) create mode 100644 cloud/apps/relay/src/idle-regional-rehome-window.test.ts create mode 100644 cloud/apps/relay/src/regional-rehome-poll-telemetry.test.ts create mode 100644 cloud/apps/relay/src/regional-rehome-poll-telemetry.ts diff --git a/cloud/apps/relay/src/assignment-store.ts b/cloud/apps/relay/src/assignment-store.ts index 48d91e21b20..871a2be1c54 100644 --- a/cloud/apps/relay/src/assignment-store.ts +++ b/cloud/apps/relay/src/assignment-store.ts @@ -1,5 +1,13 @@ import { createDrainMigrationRowLookup } from './drain-migration-row-lookup.js' -import { IDLE_REHOME_PAGE_SIZE, selectIdleRegionalRehomes } from './idle-regional-rehome-selection.js' +import { + selectIdleRegionalRehomes, + type IdleRegionalRehomeCandidate, + type IdleRehomeHostCursor +} from './idle-regional-rehome-selection.js' +import { + RegionalRehomePollTelemetry, + type RegionalRehomePollGate +} from './regional-rehome-poll-telemetry.js' import { readRegionCorrectionOutcomes } from './region-correction-outcomes.js' import { previewRegionalRehomeEligibility, @@ -3357,17 +3365,26 @@ export class RelayAssignmentStore { return previewRegionCorrection(this.database, this.now()) } - private idleRegionalCandidateOffset = 0 + private idleRegionalCandidateCursor: IdleRehomeHostCursor = null + private readonly regionalRehomePollTelemetry = new RegionalRehomePollTelemetry() async selectIdleRegionalRehomeCandidates( processSafety?: RegionalRehomeSafetySnapshot - ): Promise> { + ): Promise { const now = this.now() - if (!processSafety || this.regionalRehomeCohortPercent === 0) return [] + const gated = (gate: RegionalRehomePollGate): IdleRegionalRehomeCandidate[] => { + this.regionalRehomePollTelemetry.record({ now, gate, candidates: 0 }) + return [] + } + if (!processSafety) return gated('process-safety-unavailable') + if (this.regionalRehomeCohortPercent === 0) return gated('cohort-zero') const control = (await this.database.query( - "SELECT enabled, not_before FROM relay_region_rehome_control WHERE control_id = 'global'" + `SELECT enabled, not_before, preference_max_age_ms, host_cooldown_ms + FROM relay_region_rehome_control WHERE control_id = 'global'` ))[0] - if (!control || Number(control.enabled) !== 1 || Number(control.not_before) > now) return [] + if (!control || Number(control.enabled) !== 1 || Number(control.not_before) > now) { + return gated('control-closed') + } // The dispatch budget is durable and global, but until now only // `commitIdleRegionalRehome` consulted it -- after the join had already run and // the worker had already POSTed every candidate to its source cell. An absent @@ -3377,19 +3394,28 @@ export class RelayAssignmentStore { WHERE worker_id = 'global'` ))[0] if (worker && (Number(worker.paused_until) > now || Number(worker.next_dispatch_at) > now)) { - return [] + return gated('budget-closed') } const fleetSafety = await this.readRegionalRehomeFleetSafety(this.database, now) - if (regionalRehomeFleetSafetyFailure(processSafety, fleetSafety, now)) return [] - const candidates = await selectIdleRegionalRehomes({ + if (regionalRehomeFleetSafetyFailure(processSafety, fleetSafety, now)) return gated('fleet-safety') + const startedAt = performance.now() + const selection = await selectIdleRegionalRehomes({ database: this.database, now, heartbeatTtlMs: this.heartbeatTtlMs, - cohortPercent: this.regionalRehomeCohortPercent, offset: this.idleRegionalCandidateOffset, + cohortPercent: this.regionalRehomeCohortPercent, + preferenceMaxAgeMs: Number(control.preference_max_age_ms), + hostCooldownMs: Number(control.host_cooldown_ms), + cursor: this.idleRegionalCandidateCursor, connectionHeadroom: await this.connectionHeadroomByCell(this.database), cellIsClean: regionalRehomeCellSafetyIsClean }) - this.idleRegionalCandidateOffset = candidates.length < IDLE_REHOME_PAGE_SIZE - ? 0 : this.idleRegionalCandidateOffset + candidates.length - return candidates + this.idleRegionalCandidateCursor = selection.cursor + this.regionalRehomePollTelemetry.record({ + now, + gate: 'open', + candidates: selection.candidates.length, + selectionMs: performance.now() - startedAt + }) + return selection.candidates } async commitIdleRegionalRehome( diff --git a/cloud/apps/relay/src/idle-regional-rehome-selection.ts b/cloud/apps/relay/src/idle-regional-rehome-selection.ts index f8790459c0b..f24819abf11 100644 --- a/cloud/apps/relay/src/idle-regional-rehome-selection.ts +++ b/cloud/apps/relay/src/idle-regional-rehome-selection.ts @@ -4,114 +4,296 @@ import type { RelayDatabase, SqlRow } from './database.js' export const IDLE_REHOME_PAGE_SIZE = 100 -export async function selectIdleRegionalRehomes(input: { +// How many decision rows one poll is allowed to look at. The poll runs about +// fifty times a minute across the directors, so its cost has to be set by this +// number and not by the size of the fleet or the width of the cohort. +export const IDLE_REHOME_DECISION_WINDOW = 500 + +// Where the last window ended. A keyset beats OFFSET: `OFFSET n` still has to +// produce and throw away n rows, and n grew by a page on every poll that +// dispatched, so the scan got more expensive the longer the rollout ran. +export type IdleRehomeHostCursor = { userId: string; relayHostId: string } | null + +export type IdleRegionalRehomeCandidate = IdleRegionalRehomeRequest & { sourceCellUrl: string } + +export type IdleRegionalRehomeSelection = { + candidates: IdleRegionalRehomeCandidate[] + cursor: IdleRehomeHostCursor +} + +type SourceCell = { + cellId: string + region: string + cellIncarnation: string + startedAt: number + cellUrl: string +} + +type TargetCell = { cellId: string; capacityRequests: number; reservedRequests: number } + +type SelectionInput = { database: RelayDatabase now: number heartbeatTtlMs: number cohortPercent: number - offset: number - connectionHeadroom: Map + preferenceMaxAgeMs: number + hostCooldownMs: number + cursor: IdleRehomeHostCursor + connectionHeadroom: ReadonlyMap cellIsClean: (safety: SqlRow | undefined, runtime: SqlRow, now: number) => boolean -}): Promise> { - const [runtimes, safetyRows] = await Promise.all([ - input.database.query('SELECT * FROM relay_cell_runtime'), - input.database.query('SELECT * FROM relay_cell_rehome_safety') - ]) - const cleanCells = runtimes - .filter((runtime) => - input.cellIsClean( - safetyRows.find((safety) => safety.cell_id === runtime.cell_id), - runtime, - input.now - ) - ) - .map((runtime) => String(runtime.cell_id)) - const targetCells = cleanCells.filter((id) => input.connectionHeadroom.get(id) !== false) - if (!cleanCells.length || !targetCells.length) return [] +} + +const CELL_INVENTORY_QUERY = `SELECT cell.cell_id, cell.cell_url, cell.enabled, + cell.capacity_requests, cell.reserved_requests, region.region, + admission.admission_state, capability.cell_incarnation AS capability_incarnation, + capability.regional_rehome_protocol + FROM relay_cells cell + LEFT JOIN relay_cell_regions region ON region.cell_id = cell.cell_id + LEFT JOIN relay_cell_admission admission ON admission.cell_id = cell.cell_id + LEFT JOIN relay_cell_capabilities capability ON capability.cell_id = cell.cell_id` + +export async function selectIdleRegionalRehomes( + input: SelectionInput +): Promise { + const cells = await readCellInventory(input) + if (!cells.sources.size || !cells.targetsByRegion.size) return { candidates: [], cursor: null } + const sourceRegions = [...new Set([...cells.sources.values()].map((cell) => cell.region))] + const targetRegions = [...cells.targetsByRegion.keys()] + const decisionFilter = `outcome = 'conclusive' AND policy_version = 1 + AND preferred_region IN (${placeholders(targetRegions.length)}) + AND incumbent_region IN (${placeholders(sourceRegions.length)}) + AND preferred_region <> incumbent_region + AND expires_at > ? AND observed_at >= ? AND cohort_bucket < ?` + const decisionParams = [ + ...targetRegions, + ...sourceRegions, + input.now, + input.now - input.preferenceMaxAgeMs, + input.cohortPercent + ] + const after = input.cursor ? [input.cursor.userId, input.cursor.relayHostId] : [] + const afterFilter = input.cursor ? 'AND (user_id, relay_host_id) > (?, ?)' : '' + + // The window is taken first and on its own so the poll knows where it stopped + // reading, not just where it stopped emitting. Every gate below this point can + // reject a host, and a cursor that only advanced past emitted rows would park + // on a rejected host forever. + const window = await input.database.query( + `SELECT user_id, relay_host_id FROM relay_region_decisions + WHERE ${decisionFilter} ${afterFilter} + ORDER BY user_id, relay_host_id LIMIT ?`, + [...decisionParams, ...after, IDLE_REHOME_DECISION_WINDOW] + ) + if (!window.length) return { candidates: [], cursor: null } + const windowEnd = window[window.length - 1]! + const windowWasFull = window.length === IDLE_REHOME_DECISION_WINDOW + + const sourceList = [...cells.sources.values()] const rows = await input.database.query( - `SELECT a.user_id, a.relay_host_id, a.cell_id AS source_cell_id, - a.assignment_epoch, host.generation, r.cell_incarnation, - s.cell_url, target.cell_id AS target_cell_id - FROM relay_region_rehome_control policy - JOIN relay_region_decisions d ON d.outcome = 'conclusive' + // The verification names the window's keys rather than repeating its LIMIT: + // the two reads take separate snapshots, and a decision that turned eligible + // between them would otherwise shift the second LIMIT and push the last host + // out of it while the cursor still advanced past it. + `SELECT d.user_id, d.relay_host_id, d.preferred_region, a.cell_id AS source_cell_id, + a.assignment_epoch, host.generation + FROM (SELECT user_id, relay_host_id, preferred_region, incumbent_region, assignment_epoch + FROM relay_region_decisions + WHERE ${decisionFilter} + AND (user_id, relay_host_id) IN (${Array.from({ length: window.length }, () => '(?,?)').join(',')}) + -- The LIMIT cannot truncate a key set this size; it is here because without + -- it Postgres flattens the subquery, estimates one row out of the join, and + -- drives the whole plan from a sequential scan of the capability table. + ORDER BY user_id, relay_host_id LIMIT ?) d JOIN relay_assignments a ON a.user_id = d.user_id AND a.relay_host_id = d.relay_host_id - JOIN relay_cells s ON s.cell_id = a.cell_id AND s.enabled = 1 - JOIN relay_cell_regions sr ON sr.cell_id = a.cell_id - JOIN relay_cell_admission sa ON sa.cell_id = a.cell_id AND sa.admission_state = 'general' - JOIN relay_cell_runtime r ON r.cell_id = a.cell_id AND r.ready = 1 - JOIN relay_cell_capabilities c ON c.cell_id = r.cell_id AND c.cell_incarnation = r.cell_incarnation - JOIN relay_control_capabilities host ON host.user_id = a.user_id AND host.relay_host_id = a.relay_host_id - AND host.cell_id = a.cell_id AND host.assignment_epoch = a.assignment_epoch - AND host.cell_incarnation = r.cell_incarnation AND host.idle_regional_rehome = 1 - JOIN relay_assignment_activity_leases lease ON lease.user_id = host.user_id - AND lease.relay_host_id = host.relay_host_id AND lease.activity_id = host.activity_id + AND a.assignment_epoch = d.assignment_epoch + JOIN (${inlineRows(SOURCE_CELL_COLUMNS, sourceList.length)}) source + ON source.cell_id = a.cell_id AND source.region = d.incumbent_region + JOIN relay_control_capabilities host ON host.user_id = d.user_id + AND host.relay_host_id = d.relay_host_id AND host.cell_id = a.cell_id + AND host.assignment_epoch = a.assignment_epoch + AND host.cell_incarnation = source.cell_incarnation AND host.idle_regional_rehome = 1 + JOIN relay_assignment_activity_leases lease ON lease.user_id = d.user_id + AND lease.relay_host_id = d.relay_host_id AND lease.activity_id = host.activity_id AND lease.cell_id = a.cell_id AND lease.activity_kind = 'control' - JOIN relay_cell_regions tr ON tr.region = d.preferred_region - JOIN relay_cells target ON target.cell_id = tr.cell_id AND target.enabled = 1 - JOIN relay_cell_admission ta ON ta.cell_id = target.cell_id AND ta.admission_state = 'general' - JOIN relay_cell_runtime rt ON rt.cell_id = target.cell_id AND rt.ready = 1 - JOIN relay_cell_capabilities ct ON ct.cell_id = rt.cell_id AND ct.cell_incarnation = rt.cell_incarnation - WHERE policy.control_id = 'global' AND policy.enabled = 1 AND policy.not_before <= ? - AND d.preferred_region <> sr.region AND d.incumbent_region = sr.region - AND d.assignment_epoch = a.assignment_epoch AND d.policy_version = 1 - AND d.expires_at > ? AND d.observed_at >= ? - policy.preference_max_age_ms - AND d.cohort_bucket < ? AND lease.expires_at > ? AND lease.updated_at >= r.started_at - AND r.last_heartbeat_at > ? AND rt.last_heartbeat_at > ? - AND s.cell_id IN (${cleanCells.map(() => '?').join(',')}) - AND target.cell_id IN (${targetCells.map(() => '?').join(',')}) - -- Reserve the moving host's source activity plus its assignment on the target. - AND target.reserved_requests + 1 + ( - SELECT COALESCE(SUM(activity.request_units), 0) - FROM relay_assignment_activity_leases activity - WHERE activity.user_id = a.user_id AND activity.relay_host_id = a.relay_host_id - AND activity.cell_id = a.cell_id - ) <= target.capacity_requests - AND c.regional_rehome_protocol >= 3 AND ct.regional_rehome_protocol >= 3 - AND NOT EXISTS (SELECT 1 FROM relay_assignment_migrations migration - WHERE migration.user_id = a.user_id AND migration.relay_host_id = a.relay_host_id + AND lease.expires_at > ? AND lease.updated_at >= source.started_at + WHERE NOT EXISTS (SELECT 1 FROM relay_assignment_migrations migration + WHERE migration.user_id = d.user_id AND migration.relay_host_id = d.relay_host_id AND migration.completed_at IS NULL AND migration.aborted_at IS NULL) AND NOT EXISTS (SELECT 1 FROM relay_region_rehome_attempts attempt - WHERE attempt.user_id = a.user_id AND attempt.relay_host_id = a.relay_host_id - AND attempt.created_at > ? - policy.host_cooldown_ms) - ORDER BY a.user_id, a.relay_host_id, host.generation DESC, - (target.reserved_requests + rt.observed_requests) * 1.0 / target.capacity_requests, - target.cell_id - LIMIT ? OFFSET ?`, + WHERE attempt.user_id = d.user_id AND attempt.relay_host_id = d.relay_host_id + AND attempt.created_at > ?) + ORDER BY d.user_id, d.relay_host_id, host.generation DESC + -- Counted in hosts, because a host with one eligible target has to be able + -- to fill a page on its own. A host with many leaves part of this page + -- unread, and the cursor stops where the page stopped, so it is re-read + -- next poll rather than skipped. + LIMIT ?`, [ + ...decisionParams, + ...window.flatMap((row) => [row.user_id, row.relay_host_id]), + IDLE_REHOME_DECISION_WINDOW, + ...sourceList.flatMap((cell) => [cell.cellId, cell.region, cell.cellIncarnation, cell.startedAt]), input.now, - input.now, - input.now, - input.cohortPercent, - input.now, - input.now - input.heartbeatTtlMs, - input.now - input.heartbeatTtlMs, - ...cleanCells, - ...targetCells, - input.now, - IDLE_REHOME_PAGE_SIZE, - input.offset + input.now - input.hostCooldownMs, + IDLE_REHOME_PAGE_SIZE ] ) - return rows.map((row) => { - const request = { - v: 1 as const, - userId: String(row.user_id), - relayHostId: String(row.relay_host_id), - sourceCellId: String(row.source_cell_id), - sourceCellIncarnation: String(row.cell_incarnation), - sourceAssignmentEpoch: Number(row.assignment_epoch), - sourceGeneration: Number(row.generation), - targetCellId: String(row.target_cell_id) + const units = rows.length ? await sourceRequestUnits(input.database, rows) : new Map() + + const candidates: IdleRegionalRehomeCandidate[] = [] + let stoppedAt: IdleRehomeHostCursor = null + for (const row of rows) { + // Whole hosts only: the lower-priority targets are a host's fallbacks when + // the first one defers, and splitting them across pages loses them. + if (candidates.length >= IDLE_REHOME_PAGE_SIZE) { + return { candidates, cursor: stoppedAt } } - // UUIDv5 keeps retries on every director bound to the same source authority and target. - const digest = createHash('sha1') - .update(Buffer.from('0a1c5a9b197b4ea8b6f1f3bcaa3d712c', 'hex')) - .update(JSON.stringify(request)) - .digest() - digest[6] = (digest[6]! & 0x0f) | 0x50 - digest[8] = (digest[8]! & 0x3f) | 0x80 - const hex = digest.subarray(0, 16).toString('hex') - const attemptId = `${hex.slice(0, 8)}-${hex.slice(8, 12)}-${hex.slice(12, 16)}-${hex.slice(16, 20)}-${hex.slice(20)}` - return { ...request, attemptId, sourceCellUrl: String(row.cell_url) } - }) + const source = cells.sources.get(String(row.source_cell_id))! + const sourceUnits = units.get(hostKey(row)) ?? 0 + for (const target of cells.targetsByRegion.get(String(row.preferred_region)) ?? []) { + if (target.reservedRequests + 1 + sourceUnits > target.capacityRequests) continue + candidates.push(idleRehomeCandidate(row, source, target.cellId)) + } + stoppedAt = { userId: String(row.user_id), relayHostId: String(row.relay_host_id) } + } + // A full verification page may have been cut short of the window's end, so only + // a page that ran the window out may wrap to the head of the keyspace. + if (rows.length === IDLE_REHOME_PAGE_SIZE) return { candidates, cursor: stoppedAt } + return { + candidates, + cursor: windowWasFull + ? { userId: String(windowEnd.user_id), relayHostId: String(windowEnd.relay_host_id) } + : null + } +} + +// Every cell predicate the candidate join used to re-evaluate per (host, cell) +// pair. There are tens of cells and tens of thousands of hosts, so this is +// resolved once per poll against the four small inventory tables. +async function readCellInventory( + input: SelectionInput +): Promise<{ sources: Map; targetsByRegion: Map }> { + const { database, now } = input + const [runtimeRows, safetyRows, inventory] = await Promise.all([ + database.query('SELECT * FROM relay_cell_runtime'), + database.query('SELECT * FROM relay_cell_rehome_safety'), + database.query(CELL_INVENTORY_QUERY) + ]) + const runtimes = new Map(runtimeRows.map((row) => [String(row.cell_id), row])) + const safety = new Map(safetyRows.map((row) => [String(row.cell_id), row])) + const sources = new Map() + const targetsByRegion = new Map() + const load = new Map() + for (const cell of inventory) { + const cellId = String(cell.cell_id) + const runtime = runtimes.get(cellId) + if (!runtime || !input.cellIsClean(safety.get(cellId), runtime, now)) continue + if ( + Number(cell.enabled) !== 1 || + cell.admission_state !== 'general' || + cell.region == null || + Number(runtime.ready) !== 1 || + Number(runtime.last_heartbeat_at) <= now - input.heartbeatTtlMs || + cell.capability_incarnation == null || + String(cell.capability_incarnation) !== String(runtime.cell_incarnation) || + Number(cell.regional_rehome_protocol) < 3 + ) { + continue + } + const region = String(cell.region) + sources.set(cellId, { + cellId, + region, + cellIncarnation: String(runtime.cell_incarnation), + startedAt: Number(runtime.started_at), + cellUrl: String(cell.cell_url) + }) + if (input.connectionHeadroom.get(cellId) === false) continue + const capacityRequests = Number(cell.capacity_requests) + const reservedRequests = Number(cell.reserved_requests) + const targets = targetsByRegion.get(region) ?? [] + targets.push({ cellId, capacityRequests, reservedRequests }) + targetsByRegion.set(region, targets) + load.set(cellId, (reservedRequests + Number(runtime.observed_requests)) / capacityRequests) + } + for (const targets of targetsByRegion.values()) { + targets.sort( + (left, right) => + load.get(left.cellId)! - load.get(right.cellId)! || (left.cellId < right.cellId ? -1 : 1) + ) + } + return { sources, targetsByRegion } +} + +// One grouped read for the page instead of a correlated aggregate per (host, cell) pair. +async function sourceRequestUnits( + database: RelayDatabase, + rows: SqlRow[] +): Promise> { + const seen = new Set() + const params: unknown[] = [] + for (const row of rows) { + if (seen.has(hostKey(row))) continue + seen.add(hostKey(row)) + params.push(row.user_id, row.relay_host_id, row.source_cell_id) + } + const sums = await database.query( + `SELECT user_id, relay_host_id, COALESCE(SUM(request_units), 0) AS request_units + FROM relay_assignment_activity_leases + WHERE (user_id, relay_host_id, cell_id) IN (${Array.from({ length: seen.size }, () => '(?,?,?)').join(',')}) + GROUP BY user_id, relay_host_id`, + params + ) + return new Map(sums.map((row) => [hostKey(row), Number(row.request_units)])) +} + +const SOURCE_CELL_COLUMNS = [ + ['cell_id', 'TEXT'], + ['region', 'TEXT'], + ['cell_incarnation', 'TEXT'], + ['started_at', 'BIGINT'] +] as const + +function placeholders(count: number): string { + return Array.from({ length: count }, () => '?').join(',') +} + +// A derived table the planner can hash, in the one syntax both Postgres and the +// SQLite test engine accept (`VALUES ... AS t(col)` and LATERAL are not common to +// both). Only the first branch is cast; both engines take the union's types from it. +function inlineRows(columns: ReadonlyArray, rows: number): string { + const first = columns.map(([name, type]) => `CAST(? AS ${type}) AS ${name}`) + const rest = Array.from({ length: rows - 1 }, () => `UNION ALL SELECT ${placeholders(columns.length)}`) + return `SELECT ${first.join(', ')} ${rest.join(' ')}` +} + +function hostKey(row: SqlRow): string { + return `${String(row.user_id)}${String(row.relay_host_id)}` +} + +function idleRehomeCandidate( + row: SqlRow, + source: SourceCell, + targetCellId: string +): IdleRegionalRehomeCandidate { + const request = { + v: 1 as const, + userId: String(row.user_id), + relayHostId: String(row.relay_host_id), + sourceCellId: source.cellId, + sourceCellIncarnation: source.cellIncarnation, + sourceAssignmentEpoch: Number(row.assignment_epoch), + sourceGeneration: Number(row.generation), + targetCellId + } + // UUIDv5 keeps retries on every director bound to the same source authority and target. + const digest = createHash('sha1') + .update(Buffer.from('0a1c5a9b197b4ea8b6f1f3bcaa3d712c', 'hex')) + .update(JSON.stringify(request)) + .digest() + digest[6] = (digest[6]! & 0x0f) | 0x50 + digest[8] = (digest[8]! & 0x3f) | 0x80 + const hex = digest.subarray(0, 16).toString('hex') + const attemptId = `${hex.slice(0, 8)}-${hex.slice(8, 12)}-${hex.slice(12, 16)}-${hex.slice(16, 20)}-${hex.slice(20)}` + return { ...request, attemptId, sourceCellUrl: source.cellUrl } } diff --git a/cloud/apps/relay/src/idle-regional-rehome-window.test.ts b/cloud/apps/relay/src/idle-regional-rehome-window.test.ts new file mode 100644 index 00000000000..e213b2331f0 --- /dev/null +++ b/cloud/apps/relay/src/idle-regional-rehome-window.test.ts @@ -0,0 +1,310 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.js' +import type { RelayDatabase } from './database.js' +import { IDLE_REHOME_DECISION_WINDOW } from './idle-regional-rehome-selection.js' +import { openIdleRehomeTestDatabase } from './idle-regional-rehome-test-database.js' + +// One source cell and three targets, so a poll has to rank targets per host +// rather than take the single one the two-cell fixture leaves it. +const cells = [ + { id: 'us', url: 'https://us.example.test', region: 'us-central1' as const, capacityRequests: 100 }, + { id: 'asia-busy', url: 'https://asia-busy.example.test', region: 'asia-east2' as const, capacityRequests: 100 }, + { id: 'asia-idle', url: 'https://asia-idle.example.test', region: 'asia-east2' as const, capacityRequests: 100 }, + { id: 'asia-mid', url: 'https://asia-mid.example.test', region: 'asia-east2' as const, capacityRequests: 100 } +] +const incarnations = cells.map((_, index) => `${index + 1}${'1'.repeat(7)}-1111-4111-8111-111111111111`) +const observed = [0, 60, 10, 30] + +const databases: RelayDatabase[] = [] +afterEach(async () => { + vi.restoreAllMocks() + for (const database of databases.splice(0)) await database.close() +}) + +async function setup() { + const database = await openIdleRehomeTestDatabase() + databases.push(database) + let now = 100_000_000 + const store = new RelayAssignmentStore(database, () => now, { regionalRehomeCohortPercent: 100 }) + await store.inspectRegionalRehomeControl() + now += 86_400_000 + await store.applyRegionalRehomeControl({ + expectedGeneration: 0, + enabled: true, + notBefore: now, + ratePerMinute: 10, + preferenceMaxAgeMs: 86_400_000, + hostCooldownMs: 604_800_000, + drainGraceMs: 60_000 + }) + await store.reconcileCells(cells) + const safety = { + observedAt: now, + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + } + for (const [index, cell] of cells.entries()) { + await store.recordCellHeartbeat({ + cellId: cell.id, + cellUrl: cell.url, + region: cell.region, + cellIncarnation: incarnations[index]!, + startedAt: now - 1_000, + ready: true, + observedRequests: observed[index]! + }) + await store.recordCellRegionalRehomeStatus({ + cellId: cell.id, + cellIncarnation: incarnations[index]!, + regionalRehomeProtocol: 3, + safety + }) + } + return { store, database, safety, now } +} + +async function seedHost( + store: RelayAssignmentStore, + identity: { userId: string; relayHostId: string } +): Promise { + const assignment = await store.assign(identity, undefined, 'us-central1') + await store.activateControl(identity, { + cellId: 'us', + assignmentEpoch: assignment.assignmentEpoch, + generation: 7, + cellIncarnation: incarnations[0], + idleRegionalRehome: true + }) + const issued = await store.exchangeRegionCorrection(identity, { v: 1, action: 'issue-window' }, assignment.assignmentEpoch) + await store.exchangeRegionCorrection( + identity, + { + v: 1, + action: 'report', + generation: issued.window!.generation, + assignmentEpoch: assignment.assignmentEpoch, + policyVersion: 1, + outcome: 'conclusive', + measurements: { 'us-central1': 180, 'asia-east2': 40 } + }, + assignment.assignmentEpoch + ) +} + +// Clone one seeded host's rows under new identities, which is far cheaper than +// driving the full activation path thousands of times. +async function cloneHosts( + database: RelayDatabase, + template: { userId: string; relayHostId: string }, + count: number +): Promise { + for (const table of [ + 'relay_assignments', + 'relay_assignment_activity_leases', + 'relay_control_capabilities', + 'relay_region_decisions' + ]) { + const row = ( + await database.query(`SELECT * FROM ${table} WHERE user_id = ? AND relay_host_id = ?`, [ + template.userId, + template.relayHostId + ]) + )[0]! + const columns = Object.keys(row) + const projection = columns.map((column) => + column === 'user_id' || column === 'relay_host_id' ? '?' : column + ) + for (let index = 0; index < count; index++) { + await database.query( + `INSERT INTO ${table} (${columns.join(', ')}) SELECT ${projection.join(', ')} FROM ${table} + WHERE user_id = ? AND relay_host_id = ?`, + [ + `clone-${String(index).padStart(5, '0')}`, + `clonehost${String(index).padStart(7, '0')}`, + template.userId, + template.relayHostId + ] + ) + } + } +} + +describe('idle regional rehome candidate window', () => { + const identity = { userId: 'window-test', relayHostId: 'abcdefghijklmnop' } + + it('offers every eligible target for a host, least loaded first', async () => { + const { store, safety } = await setup() + await seedHost(store, identity) + const candidates = await store.selectIdleRegionalRehomeCandidates(safety) + expect(candidates.map((candidate) => candidate.targetCellId)).toEqual([ + 'asia-idle', + 'asia-mid', + 'asia-busy' + ]) + expect(new Set(candidates.map((candidate) => candidate.sourceCellUrl))).toEqual( + new Set(['https://us.example.test']) + ) + // Every candidate is the same move to a different target, so the attempt ids differ. + expect(new Set(candidates.map((candidate) => candidate.attemptId)).size).toBe(3) + }) + + it('drops only the targets without room for the host plus its source activity', async () => { + const { store, database, safety } = await setup() + await seedHost(store, identity) + await database.query( + 'UPDATE relay_assignment_activity_leases SET request_units = 4 WHERE user_id = ?', + [identity.userId] + ) + // Five units needed: four source units plus the assignment the move reserves. + await database.query("UPDATE relay_cells SET capacity_requests = 4 WHERE cell_id = 'asia-idle'") + const short = await store.selectIdleRegionalRehomeCandidates(safety) + expect(short.map((candidate) => candidate.targetCellId)).toEqual(['asia-mid', 'asia-busy']) + // Exactly enough room is enough; it ranks last because the ratio is per capacity. + await database.query("UPDATE relay_cells SET capacity_requests = 5 WHERE cell_id = 'asia-idle'") + const exact = await store.selectIdleRegionalRehomeCandidates(safety) + expect(exact.map((candidate) => candidate.targetCellId)).toEqual([ + 'asia-mid', + 'asia-busy', + 'asia-idle' + ]) + }) + + it('reads a bounded window of decisions however many hosts are eligible', async () => { + const { store, database, safety } = await setup() + await seedHost(store, identity) + await cloneHosts(database, identity, IDLE_REHOME_DECISION_WINDOW + 200) + const query = vi.spyOn(database, 'query') + await store.selectIdleRegionalRehomeCandidates(safety) + const calls = query.mock.calls.map((call) => call[0]) + const window = calls.findIndex((sql) => /FROM relay_region_decisions\s*$/m.test(sql)) + expect(window).toBeGreaterThanOrEqual(0) + expect(query.mock.calls[window]![1]!.at(-1)).toBe(IDLE_REHOME_DECISION_WINDOW) + // No statement pages by OFFSET any more: that was the cost that grew with the rollout. + expect(calls.some((sql) => /OFFSET/i.test(sql))).toBe(false) + }) + + it('keeps the window\'s last host when a decision turns eligible between the two reads', async () => { + const { store, database, safety, now } = await setup() + await seedHost(store, identity) + // Exactly one full window, whose last key in sort order is the seeded host. + await cloneHosts(database, identity, IDLE_REHOME_DECISION_WINDOW - 1) + await database.query( + "UPDATE relay_assignment_activity_leases SET expires_at = ? WHERE user_id LIKE 'clone-%'", + [now - 1] + ) + const template = ( + await database.query('SELECT * FROM relay_region_decisions WHERE user_id = ?', [ + identity.userId + ]) + )[0]! + const columns = Object.keys(template) + const query = database.query.bind(database) + let inserted = false + vi.spyOn(database, 'query').mockImplementation(async (sql, params) => { + const rows = await query(sql, params) + // A decision that becomes eligible after the window is read and sorts + // inside it: a second LIMIT would push the window's last host out. + if (!inserted && /^SELECT user_id, relay_host_id FROM relay_region_decisions/.test(sql)) { + inserted = true + await query( + `INSERT INTO relay_region_decisions (${columns.join(', ')}) + VALUES (${columns.map(() => '?').join(', ')})`, + columns.map((column) => + column === 'user_id' + ? 'clone-99999' + : column === 'relay_host_id' + ? 'latehost99999999' + : template[column] + ) + ) + } + return rows + }) + const candidates = await store.selectIdleRegionalRehomeCandidates(safety) + expect(candidates.map((candidate) => candidate.userId)).toEqual([ + identity.userId, + identity.userId, + identity.userId + ]) + }) + + it('walks the whole population in bounded pages and wraps only at the end', async () => { + const { store, database, safety } = await setup() + await seedHost(store, identity) + await cloneHosts(database, identity, 120) + const seen = new Set() + let pages = 0 + let wrapped = false + // 121 hosts x 3 targets is 363 candidates, so the page cap has to be hit + // several times before the window runs out and the cursor wraps. + for (let poll = 0; poll < 20 && !wrapped; poll++) { + const page = await store.selectIdleRegionalRehomeCandidates(safety) + pages += 1 + const before = seen.size + for (const candidate of page) seen.add(`${candidate.userId}/${candidate.targetCellId}`) + if (seen.size === before && page.length > 0) wrapped = true + if (page.length < 3) wrapped = true + } + expect(pages).toBeGreaterThan(1) + expect(seen.size).toBe(121 * 3) + }) + + it('does not stall on a host the window found but the join rejected', async () => { + const { store, database, safety, now } = await setup() + await seedHost(store, identity) + await cloneHosts(database, identity, 2) + // The first host in key order loses its control lease, so it can never be a + // candidate; an emitted-rows cursor would sit on it forever. + await database.query('UPDATE relay_assignment_activity_leases SET expires_at = ? WHERE user_id = ?', [ + now - 1, + 'clone-00000' + ]) + const first = await store.selectIdleRegionalRehomeCandidates(safety) + expect(first.map((candidate) => candidate.userId)).not.toContain('clone-00000') + expect(new Set(first.map((candidate) => candidate.userId))).toEqual( + new Set(['clone-00001', identity.userId]) + ) + }) + + it('excludes a host inside its rehome cooldown and takes it back after', async () => { + const { store, database, safety, now } = await setup() + await seedHost(store, identity) + await database.query( + `INSERT INTO relay_region_rehome_attempts + (attempt_id, user_id, relay_host_id, preferred_region, source_cell_id, source_cell_incarnation, + target_cell_id, target_cell_incarnation, previous_epoch, assignment_epoch, drain_grace_ms, + send_attempts, created_at, updated_at) + VALUES ('cooled', ?, ?, 'asia-east2', 'us', ?, 'asia-idle', ?, 0, 9, 0, 0, ?, ?)`, + [identity.userId, identity.relayHostId, incarnations[0], incarnations[2], now - 1_000, now] + ) + expect(await store.selectIdleRegionalRehomeCandidates(safety)).toEqual([]) + await database.query('UPDATE relay_region_rehome_attempts SET created_at = ?', [ + now - 604_800_000 - 1 + ]) + expect(await store.selectIdleRegionalRehomeCandidates(safety)).toHaveLength(3) + }) + + it('excludes a host outside the cohort', async () => { + const { store, database } = await setup() + await seedHost(store, identity) + const now = 100_000_000 + 86_400_000 + const safety = { + observedAt: now, + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + } + await database.query('UPDATE relay_region_decisions SET cohort_bucket = 40') + const narrow = new RelayAssignmentStore(database, () => now, { regionalRehomeCohortPercent: 40 }) + expect(await narrow.selectIdleRegionalRehomeCandidates(safety)).toEqual([]) + const wide = new RelayAssignmentStore(database, () => now, { regionalRehomeCohortPercent: 41 }) + expect(await wide.selectIdleRegionalRehomeCandidates(safety)).toHaveLength(3) + }) +}) diff --git a/cloud/apps/relay/src/regional-rehome-poll-telemetry.test.ts b/cloud/apps/relay/src/regional-rehome-poll-telemetry.test.ts new file mode 100644 index 00000000000..5e0c1104faa --- /dev/null +++ b/cloud/apps/relay/src/regional-rehome-poll-telemetry.test.ts @@ -0,0 +1,58 @@ +import { describe, expect, it } from 'vitest' +import { + REGIONAL_REHOME_POLL_SUMMARY_INTERVAL_MS, + RegionalRehomePollTelemetry +} from './regional-rehome-poll-telemetry.js' + +describe('regional rehome poll telemetry', () => { + it('names the gate that stopped the poll, not just the empty result', () => { + const lines: string[] = [] + const telemetry = new RegionalRehomePollTelemetry((line) => lines.push(line)) + let now = 1_000 + for (let poll = 0; poll < 3; poll++) { + telemetry.record({ now: (now += 6_000), gate: 'budget-closed', candidates: 0 }) + } + telemetry.record({ now: (now += 6_000), gate: 'open', candidates: 0, selectionMs: 12 }) + expect(lines).toEqual([]) + telemetry.record({ + now: now + REGIONAL_REHOME_POLL_SUMMARY_INTERVAL_MS, + gate: 'open', + candidates: 7, + selectionMs: 30 + }) + expect(lines).toHaveLength(1) + expect(JSON.parse(lines[0]!)).toMatchObject({ + event: 'orca_relay_regional_rehome_poll_summary', + polls: 5, + 'budget-closed': 3, + open: 2, + candidates: 7, + selectionMsMax: 30 + }) + }) + + it('starts a fresh window after each summary', () => { + const lines: string[] = [] + const telemetry = new RegionalRehomePollTelemetry((line) => lines.push(line)) + telemetry.record({ now: 0, gate: 'control-closed', candidates: 0 }) + telemetry.record({ + now: REGIONAL_REHOME_POLL_SUMMARY_INTERVAL_MS, + gate: 'control-closed', + candidates: 0 + }) + telemetry.record({ + now: REGIONAL_REHOME_POLL_SUMMARY_INTERVAL_MS * 2, + gate: 'fleet-safety', + candidates: 0 + }) + expect(lines).toHaveLength(2) + expect(JSON.parse(lines[1]!)).toMatchObject({ + polls: 1, + 'control-closed': 0, + 'fleet-safety': 1, + candidates: 0, + selectionMsMax: 0, + selectionMsP95: 0 + }) + }) +}) diff --git a/cloud/apps/relay/src/regional-rehome-poll-telemetry.ts b/cloud/apps/relay/src/regional-rehome-poll-telemetry.ts new file mode 100644 index 00000000000..ca75e1b1505 --- /dev/null +++ b/cloud/apps/relay/src/regional-rehome-poll-telemetry.ts @@ -0,0 +1,69 @@ +// A gated poll and a poll that simply found nobody to move both produce zero +// candidates and no attempt row, so an operator watching a stalled rollout +// cannot tell them apart. One aggregated line a minute per director names the +// gate and prices the selection, at a rate a 50-polls-a-minute worker can afford. +export type RegionalRehomePollGate = + | 'open' + | 'cohort-zero' + | 'process-safety-unavailable' + | 'control-closed' + | 'budget-closed' + | 'fleet-safety' + +export const REGIONAL_REHOME_POLL_SUMMARY_INTERVAL_MS = 60_000 + +const EMPTY_GATES: Record = { + open: 0, + 'cohort-zero': 0, + 'process-safety-unavailable': 0, + 'control-closed': 0, + 'budget-closed': 0, + 'fleet-safety': 0 +} + +export class RegionalRehomePollTelemetry { + private windowStartedAt: number | null = null + private gates = { ...EMPTY_GATES } + private candidates = 0 + private selectionSamplesMs: number[] = [] + + constructor(private readonly write: (line: string) => void = (line) => console.warn(line)) {} + + record(input: { + now: number + gate: RegionalRehomePollGate + candidates: number + selectionMs?: number + }): void { + if (this.windowStartedAt === null) this.windowStartedAt = input.now + this.gates[input.gate] += 1 + this.candidates += input.candidates + if (input.selectionMs !== undefined) this.selectionSamplesMs.push(input.selectionMs) + if (input.now - this.windowStartedAt < REGIONAL_REHOME_POLL_SUMMARY_INTERVAL_MS) return + this.write( + JSON.stringify({ + event: 'orca_relay_regional_rehome_poll_summary', + windowMs: input.now - this.windowStartedAt, + polls: Object.values(this.gates).reduce((total, count) => total + count, 0), + ...this.gates, + candidates: this.candidates, + selectionMsMax: round(Math.max(0, ...this.selectionSamplesMs)), + selectionMsP95: round(percentile(this.selectionSamplesMs, 0.95)) + }) + ) + this.windowStartedAt = input.now + this.gates = { ...EMPTY_GATES } + this.candidates = 0 + this.selectionSamplesMs = [] + } +} + +function percentile(samples: number[], fraction: number): number { + if (!samples.length) return 0 + const sorted = [...samples].sort((a, b) => a - b) + return sorted[Math.min(sorted.length - 1, Math.floor(sorted.length * fraction))]! +} + +function round(value: number): number { + return Math.round(value * 100) / 100 +} diff --git a/cloud/apps/relay/src/regional-rehome-store.test.ts b/cloud/apps/relay/src/regional-rehome-store.test.ts index 9fae78c8afd..2d34661816b 100644 --- a/cloud/apps/relay/src/regional-rehome-store.test.ts +++ b/cloud/apps/relay/src/regional-rehome-store.test.ts @@ -1760,7 +1760,7 @@ function hookAfterCandidateScan( const decorate = (delegate: RelayDatabase): RelayDatabase => ({ query: async (sql, params) => { const rows = await delegate.query(sql, params) - if (!fired && sql.includes('SELECT a.user_id, a.relay_host_id')) { + if (!fired && sql.includes('SELECT d.user_id, d.relay_host_id')) { fired = true await hook(delegate) }