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) }