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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

---------

Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Jinwoo Hong
2026-09-18 21:16:19 -04:00
committed by GitHub
co-authored by Claude
parent 742fa638f3
commit 7080eb0604
6 changed files with 754 additions and 109 deletions
+39 -13
View File
@@ -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<Array<IdleRegionalRehomeRequest & { sourceCellUrl: string }>> {
): Promise<IdleRegionalRehomeCandidate[]> {
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(
@@ -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<string, boolean>
preferenceMaxAgeMs: number
hostCooldownMs: number
cursor: IdleRehomeHostCursor
connectionHeadroom: ReadonlyMap<string, boolean>
cellIsClean: (safety: SqlRow | undefined, runtime: SqlRow, now: number) => boolean
}): Promise<Array<IdleRegionalRehomeRequest & { sourceCellUrl: string }>> {
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<IdleRegionalRehomeSelection> {
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<string, number>()
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<string, SourceCell>; targetsByRegion: Map<string, TargetCell[]> }> {
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<string, SourceCell>()
const targetsByRegion = new Map<string, TargetCell[]>()
const load = new Map<string, number>()
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<Map<string, number>> {
const seen = new Set<string>()
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<readonly [string, string]>, 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 }
}
@@ -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<void> {
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<void> {
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<string>()
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)
})
})
@@ -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
})
})
})
@@ -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<RegionalRehomePollGate, number> = {
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
}
@@ -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)
}