diff --git a/.github/workflows/cloud-deploy-relay-production-director.yml b/.github/workflows/cloud-deploy-relay-production-director.yml index 4489d67b845..07e1e74ec3c 100644 --- a/.github/workflows/cloud-deploy-relay-production-director.yml +++ b/.github/workflows/cloud-deploy-relay-production-director.yml @@ -13,6 +13,11 @@ on: default: preserve type: choice options: [preserve, enable, disable] + region-correction-cohort-percent: + description: 'Preserve the measured-correction cohort, or set an integer 0–100; durable rehome stays disabled' + required: true + default: preserve + type: string prune-incompatible-revisions: description: Retain only the newly verified serving and rollback revisions required: true @@ -62,6 +67,7 @@ jobs: REGIONAL_PLACEMENT_SECRET: orca-cloud-relay-regional-placement-enabled IMAGE_DIGEST: ${{ inputs.image-digest }} REGIONAL_PLACEMENT_MODE: ${{ inputs.regional-placement-mode }} + REGION_CORRECTION_COHORT_PERCENT: ${{ inputs.region-correction-cohort-percent }} PRUNE_INCOMPATIBLE_REVISIONS: ${{ inputs.prune-incompatible-revisions }} # Floor the served revision must keep, matching relay_min_instances in # environments/production.tfvars. This gate only fails a bad deploy; Terraform @@ -106,8 +112,11 @@ jobs: echo "image-digest must be an immutable lowercase sha256 digest" >&2 exit 1 fi + if test "${REGION_CORRECTION_COHORT_PERCENT}" != preserve; then + [[ "${REGION_CORRECTION_COHORT_PERCENT}" =~ ^([0-9]|[1-9][0-9]|100)$ ]] + fi IMAGE="${IMAGE_REPOSITORY}@${IMAGE_DIGEST}" - SERVED_DIGEST="$(gcloud artifacts docker images describe "${IMAGE}" --format='value(image_summary.digest)')" + SERVED_DIGEST="$(gcloud artifacts docker images describe "${IMAGE}" --project "${GCP_PROJECT_ID}" --format='value(image_summary.digest)')" test "${SERVED_DIGEST}" = "${IMAGE_DIGEST}" [[ "${PRUNE_INCOMPATIBLE_REVISIONS}" =~ ^(true|false)$ ]] [[ "${EXPECTED_REHOME_GENERATION}" =~ ^(0|[1-9][0-9]*)$ ]] @@ -218,7 +227,8 @@ jobs: --max-instances "${DIRECTOR_MAX_INSTANCES}" \ --prune-revisions "${PRUNE_INCOMPATIBLE_REVISIONS}" \ --release-id "${RELEASE_ID}" \ - --regional-placement-secret-version "${target_version}" + --regional-placement-secret-version "${target_version}" \ + --region-correction-cohort-percent "${REGION_CORRECTION_COHORT_PERCENT}" echo "REGIONAL_PLACEMENT_ENABLED=${desired}" >> "${GITHUB_ENV}" echo "REGIONAL_PLACEMENT_VERSION=${target_version}" >> "${GITHUB_ENV}" diff --git a/cloud/apps/relay/src/admin-token-verifier.ts b/cloud/apps/relay/src/admin-token-verifier.ts index 4b8ad26e695..8b236d58473 100644 --- a/cloud/apps/relay/src/admin-token-verifier.ts +++ b/cloud/apps/relay/src/admin-token-verifier.ts @@ -6,6 +6,7 @@ export const RELAY_MONITOR_ADMIN_ROUTES = [ '/v1/admin/cell-status', '/v1/admin/evacuation-status', '/v1/admin/regional-rehome-control', + '/v1/admin/regional-rehome-preview', '/v1/admin/runtime-status' ] as const diff --git a/cloud/apps/relay/src/app.ts b/cloud/apps/relay/src/app.ts index c45e31c4a01..44f6a6fc293 100644 --- a/cloud/apps/relay/src/app.ts +++ b/cloud/apps/relay/src/app.ts @@ -1,5 +1,9 @@ import { AssignmentRequestSchema, + IdleRegionalRehomeRequestSchema, + type IdleRegionalRehomeRequest, + type IdleRegionalRehomeOutcome, + type RegionCorrectionResponse, isRelayCellConnectionHardCap, RELAY_ADMISSION_BUDGETS, RELAY_DEFAULT_REGION, @@ -39,7 +43,7 @@ import { type AssignmentAdmissionRejection } from './public-assignment-admission.js' import { relayHostLogDigest } from './relay-host-log-digest.js' -import type { RelayRuntimeCounts } from './relay-observability.js' +import type { RegionalRehomeSafetySnapshot, RelayRuntimeCounts } from './relay-observability.js' import { isRegionalRehomeTrustProbe, probeRegionalRehomeTrust @@ -68,21 +72,28 @@ export function createRelayApp( store: RelayCredentialStore assignments: RelayAssignmentStore drain: (graceMs: number) => void + idleRehome?: (input: IdleRegionalRehomeRequest & { + cohortPercent: number + directorSafety: RegionalRehomeSafetySnapshot + }) => Promise<{ outcome: IdleRegionalRehomeOutcome }> drainHost?: (input: { attemptId: string userId: string relayHostId: string sourceAssignmentEpoch: number + sourceCellIncarnation: string graceMs: number - }) => 'accepted' | 'already-accepted' | 'host-not-connected' + }) => + | 'accepted' + | 'already-accepted' + | 'host-not-connected' + | Promise<'accepted' | 'already-accepted' | 'host-not-connected'> regionalRehomeIdentityToken?: (audience: string) => Promise regionalRehomeFetch?: typeof fetch - regionalRehomeTrustProbeHostExists?: (input: { - userId: string - relayHostId: string - }) => boolean + regionalRehomeTrustProbeHostExists?: (input: { userId: string; relayHostId: string }) => boolean cellIncarnation?: string isDraining?: () => boolean + regionalRehomeSafetySnapshot?: () => RegionalRehomeSafetySnapshot runtimeCounts?: () => RelayRuntimeCounts ready: () => Promise recordAssignmentAdmission?: ( @@ -226,7 +237,8 @@ export function createRelayApp( return context.json({ error: 'host_identity_mismatch' }, 403) } const identity = { userId: claims.sub, relayHostId: claims.relayHostId } - const requestedRegion = body.data.preferredRegion + const requestedRegion = + body.data.regionCorrection?.action === 'report' ? undefined : body.data.preferredRegion const targetRegion = config.regionalPlacementEnabled !== false && requestedRegion ? requestedRegion @@ -295,10 +307,30 @@ export function createRelayApp( } } let assignment: RelayAssignment + let regionCorrection: RegionCorrectionResponse | undefined try { - assignment = requestedRegion - ? await operations.assignments.assign(identity, requestedRegion, targetRegion) - : await operations.assignments.assign(identity) + if (body.data.regionCorrection?.action === 'report') { + const current = await operations.assignments.resolve(identity) + if (!current) return context.json({ error: 'assignment_not_found' }, 409) + assignment = current + } else { + assignment = requestedRegion + ? await operations.assignments.assign(identity, requestedRegion, targetRegion) + : await operations.assignments.assign(identity) + } + if (body.data.regionCorrection) { + try { + regionCorrection = await operations.assignments.exchangeRegionCorrection( + identity, + body.data.regionCorrection, + assignment.assignmentEpoch + ) + } catch (error) { + if (body.data.regionCorrection.action === 'report') throw error + // Optional measurement setup must not discard an otherwise valid placement. + console.warn(JSON.stringify({ event: 'orca_relay_region_window_unavailable' })) + } + } } catch (error) { if (isRelayAssignmentCapacityError(error) || isRelayDatabaseTransientError(error)) { logAssignmentRejection({ @@ -353,13 +385,16 @@ export function createRelayApp( v: 1, cellUrl: assignment.cellUrl, assignmentEpoch: assignment.assignmentEpoch, - lease + lease, + ...(regionCorrection ? { regionCorrection } : {}) }) }) app.post('/v1/resolve', async (context) => { if (config.role === 'cell') return context.json({ error: 'director_only' }, 404) if (!config.publicAssignmentsEnabled) return rejectPublicAssignment(context) - if (Number(context.req.header('content-length') ?? 0) > RELAY_PROTOCOL_LIMITS.maxHttpBodyBytes) { + if ( + Number(context.req.header('content-length') ?? 0) > RELAY_PROTOCOL_LIMITS.maxHttpBodyBytes + ) { return context.json({ error: 'request_too_large' }, 413) } const body = ResolveRequestSchema.safeParse(await context.req.json().catch(() => null)) @@ -433,6 +468,34 @@ export function createRelayApp( operations.drain(body.data.graceMs) return context.json({ ok: true }) }) + app.post('/v1/admin/host-idle-rehome', async (context) => { + if (config.role !== 'cell' || !operations.idleRehome) { + return context.json({ error: 'cell_only' }, 404) + } + const bearer = readBearer(context.req.header('authorization')) + if (!bearer || !(await verifyRegionalRehomeToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = IdleRegionalRehomeCommandSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + if ( + body.data.sourceCellId !== config.cellId || + !operations.cellIncarnation || + body.data.sourceCellIncarnation !== operations.cellIncarnation + ) { + return context.json({ error: 'regional_rehome_source_generation_mismatch' }, 409) + } + try { + return context.json({ v: 1, ...(await operations.idleRehome(body.data)) }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) app.post('/v1/admin/host-drain', async (context) => { if (config.role !== 'cell' || !operations.drainHost) { return context.json({ error: 'cell_only' }, 404) @@ -474,7 +537,7 @@ export function createRelayApp( } sharedRuntimeIdentityRejected = true } - const outcome = operations.drainHost(body.data) + const outcome = await operations.drainHost(body.data) return context.json({ v: 1, outcome, @@ -502,8 +565,7 @@ export function createRelayApp( region: config.region ?? RELAY_DEFAULT_REGION, imageDigest: config.imageDigest ?? null, draining: operations.isDraining?.() ?? false, - regionalRehomeProtocol: - config.rehomeAudience && config.rehomeDirectorServiceAccount ? 1 : 0, + regionalRehomeProtocol: config.rehomeAudience && config.rehomeDirectorServiceAccount ? 3 : 0, connectionCapacity: config.connectionHardCap === undefined ? null @@ -559,6 +621,18 @@ export function createRelayApp( return context.json({ error: operationError(error) }, 409) } }) + app.get('/v1/admin/regional-rehome-preview', async (context) => { + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + const bearer = readBearer(context.req.header('authorization')) + if (!bearer || !(await verifyAdminToken(bearer, context.req.path))) { + return context.json({ error: 'invalid_token' }, 401) + } + const preview = await operations.assignments.previewRegionalRehomeEligibility( + operations.regionalRehomeSafetySnapshot?.() + ) + const outcomes = await operations.assignments.regionCorrectionOutcomes() + return context.json({ v: 1, preview, outcomes }) + }) app.post('/v1/admin/regional-rehome-control', async (context) => { if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) const bearer = readBearer(context.req.header('authorization')) @@ -1295,6 +1369,11 @@ const RegionalRehomeSafetySchema = z }) .strict() +const IdleRegionalRehomeCommandSchema = IdleRegionalRehomeRequestSchema.extend({ + cohortPercent: z.number().int().min(0).max(100), + directorSafety: RegionalRehomeSafetySchema +}) + const CellHeartbeatSchema = z .object({ v: z.literal(1), @@ -1394,45 +1473,48 @@ const CellRegionalRehomeStatusSchema = z v: z.literal(1), cellId: z.string().min(1).max(128), cellIncarnation: z.string().uuid(), - regionalRehomeProtocol: z.number().int().min(0).max(1), + regionalRehomeProtocol: z.number().int().min(0).max(3), safety: RegionalRehomeSafetySchema }) .strict() -const RegionalRehomeControlSchema = z.discriminatedUnion('action', [ - z.object({ v: z.literal(1), action: z.literal('inspect') }).strict(), - z.object({ - v: z.literal(1), - action: z.literal('apply'), - expectedGeneration: z.number().int().nonnegative(), - enabled: z.boolean(), - notBefore: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER), - ratePerMinute: z.number().int().min(1).max(120), - preferenceMaxAgeMs: z - .number() - .int() - .min(60_000) - .max(30 * 24 * 60 * 60_000), - hostCooldownMs: z - .number() - .int() - .min(60_000) - .max(30 * 24 * 60 * 60_000), - drainGraceMs: z.number().int().min(60_000).max(60 * 60_000), - confirmation: z.enum([ - 'ENABLE_REGIONAL_REHOMING', - 'DISABLE_REGIONAL_REHOMING' - ]) - }).strict() -]).superRefine((value, context) => { - if (value.action !== 'apply') return - const expected = value.enabled - ? 'ENABLE_REGIONAL_REHOMING' - : 'DISABLE_REGIONAL_REHOMING' - if (value.confirmation !== expected) { - context.addIssue({ code: 'custom', message: 'confirmation does not match state' }) - } -}) +const RegionalRehomeControlSchema = z + .discriminatedUnion('action', [ + z.object({ v: z.literal(1), action: z.literal('inspect') }).strict(), + z + .object({ + v: z.literal(1), + action: z.literal('apply'), + expectedGeneration: z.number().int().nonnegative(), + enabled: z.boolean(), + notBefore: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER), + ratePerMinute: z.number().int().min(1).max(120), + preferenceMaxAgeMs: z + .number() + .int() + .min(60_000) + .max(30 * 24 * 60 * 60_000), + hostCooldownMs: z + .number() + .int() + .min(60_000) + .max(30 * 24 * 60 * 60_000), + drainGraceMs: z + .number() + .int() + .min(60_000) + .max(60 * 60_000), + confirmation: z.enum(['ENABLE_REGIONAL_REHOMING', 'DISABLE_REGIONAL_REHOMING']) + }) + .strict() + ]) + .superRefine((value, context) => { + if (value.action !== 'apply') return + const expected = value.enabled ? 'ENABLE_REGIONAL_REHOMING' : 'DISABLE_REGIONAL_REHOMING' + if (value.confirmation !== expected) { + context.addIssue({ code: 'custom', message: 'confirmation does not match state' }) + } + }) const RegionalRehomeTrustProbeSchema = z .object({ @@ -1776,7 +1858,11 @@ const RegionalHostDrainSchema = z sourceCellId: z.string().min(1).max(128), sourceCellIncarnation: z.string().uuid(), sourceAssignmentEpoch: z.number().int().positive(), - graceMs: z.number().int().nonnegative().max(60 * 60 * 1000) + graceMs: z + .number() + .int() + .nonnegative() + .max(60 * 60 * 1000) }) .strict() diff --git a/cloud/apps/relay/src/assignment-store.ts b/cloud/apps/relay/src/assignment-store.ts index 824cb1e0b2f..296a09d4e42 100644 --- a/cloud/apps/relay/src/assignment-store.ts +++ b/cloud/apps/relay/src/assignment-store.ts @@ -1,3 +1,14 @@ +import { IDLE_REHOME_PAGE_SIZE, selectIdleRegionalRehomes } from './idle-regional-rehome-selection.js' +import { readRegionCorrectionOutcomes } from './region-correction-outcomes.js' +import { + previewRegionalRehomeEligibility, + type RegionCorrectionPreview +} from './region-correction-preview.js' +import { + exchangeRegionCorrection, + previewRegionCorrection, + REGIONAL_REHOME_CONCURRENT_LIMIT +} from './region-correction-state.js' import { randomUUID } from 'node:crypto' import { performance } from 'node:perf_hooks' import { @@ -7,7 +18,10 @@ import { RELAY_DEFAULT_REGION, RELAY_REGIONS, RELAY_PROTOCOL_LIMITS, - type RelayRegion + type RelayRegion, + type RegionCorrectionRequest, + type RegionCorrectionResponse, + type IdleRegionalRehomeRequest, } from '@orca-cloud/relay-contract' import { cellAdmissionState, @@ -80,6 +94,7 @@ type CellRegionalRehomeStatus = { } type RelayAssignmentStoreOptions = { + regionalRehomeCohortPercent?: number requireLiveCells?: boolean heartbeatTtlMs?: number recordControlRenewal?: (durationMs: number, outcome: ControlRenewalOutcome) => void @@ -136,10 +151,7 @@ export type RegionalRehomeAttempt = AssignmentIdentity & { sendAttempts: number } -export type RegionalHostDrainOutcome = - | 'accepted' - | 'already-accepted' - | 'host-not-connected' +export type RegionalHostDrainOutcome = 'accepted' | 'already-accepted' | 'host-not-connected' export type RegionalRehomeFleetSafety = RegionalRehomeSafetySnapshot & { requiredCells: number @@ -415,6 +427,7 @@ const ABORTABLE_EXPIRED_MIGRATION = `( )` export class RelayAssignmentStore { + private readonly regionalRehomeCohortPercent: number private readonly requireLiveCells: boolean private readonly heartbeatTtlMs: number // Poisoned attempts never complete or abort and stay the oldest rows, so @@ -430,13 +443,19 @@ export class RelayAssignmentStore { private readonly migrationCellRegistrar: RelayMigrationCellRegistrar private readonly activityQueue = new AssignmentIdentityQueue() private assignmentTail: Promise = Promise.resolve() - private pendingRegionalRehomeDisableLog: Record | null = null constructor( private readonly database: RelayDatabase, private readonly now: () => number = Date.now, options: RelayAssignmentStoreOptions = {} ) { + this.regionalRehomeCohortPercent = options.regionalRehomeCohortPercent ?? 0 + if ( + !Number.isInteger(this.regionalRehomeCohortPercent) || + this.regionalRehomeCohortPercent < 0 || + this.regionalRehomeCohortPercent > 100 + ) + throw new Error('invalid_regional_rehome_cohort') this.requireLiveCells = options.requireLiveCells ?? false this.heartbeatTtlMs = options.heartbeatTtlMs ?? 45_000 this.recordControlRenewal = options.recordControlRenewal @@ -3309,6 +3328,163 @@ export class RelayAssignmentStore { }) } + async exchangeRegionCorrection( + identity: AssignmentIdentity, + request: RegionCorrectionRequest, + assignmentEpoch: number + ): Promise { + return exchangeRegionCorrection(this.database, identity, request, assignmentEpoch, this.now()) + } + + async regionCorrectionOutcomes() { + return readRegionCorrectionOutcomes(this.database, this.now()) + } + + async previewRegionCorrection(): Promise> { + return previewRegionCorrection(this.database, this.now()) + } + + private idleRegionalCandidateOffset = 0 + + async selectIdleRegionalRehomeCandidates( + processSafety?: RegionalRehomeSafetySnapshot + ): Promise> { + const now = this.now() + if (!processSafety || this.regionalRehomeCohortPercent === 0) return [] + const fleetSafety = await this.readRegionalRehomeFleetSafety(this.database, now) + if (regionalRehomeFleetSafetyFailure(processSafety, fleetSafety, now)) return [] + const candidates = await selectIdleRegionalRehomes({ + database: this.database, now, heartbeatTtlMs: this.heartbeatTtlMs, + cohortPercent: this.regionalRehomeCohortPercent, offset: this.idleRegionalCandidateOffset, + connectionHeadroom: await this.connectionHeadroomByCell(this.database), + cellIsClean: regionalRehomeCellSafetyIsClean + }) + this.idleRegionalCandidateOffset = candidates.length < IDLE_REHOME_PAGE_SIZE + ? 0 : this.idleRegionalCandidateOffset + candidates.length + return candidates + } + + async commitIdleRegionalRehome( + request: IdleRegionalRehomeRequest, + processSafety?: RegionalRehomeSafetySnapshot, + cohortPercent = this.regionalRehomeCohortPercent + ): Promise<{ outcome: 'committed' | 'deferred' | 'stale' }> { + const prior = await this.reconcileIdleRegionalRehome(request) + if (prior !== 'not-committed') return { outcome: prior } + if (!processSafety || !Number.isInteger(cohortPercent) || cohortPercent <= 0 || cohortPercent > 100) { + return { outcome: 'deferred' } + } + let safetyDisable: Record | null = null + const result = await this.database.transaction(async (transaction): Promise<{ outcome: 'committed' | 'deferred' | 'stale' }> => { + safetyDisable = null + const now = this.now() + const control = (await transaction.queryLocked( + `SELECT * FROM relay_region_rehome_control WHERE control_id = 'global'` + ))[0] + if (!control || Number(control.enabled) !== 1 || Number(control.not_before) > now) { + return { outcome: 'deferred' } + } + await transaction.query( + `INSERT INTO relay_region_rehome_worker_state + (worker_id, next_dispatch_at, paused_until, consecutive_failures, updated_at) + VALUES ('global', 0, 0, 0, ?) ON CONFLICT (worker_id) DO NOTHING`, [now] + ) + const worker = (await transaction.queryLocked( + `SELECT * FROM relay_region_rehome_worker_state WHERE worker_id = 'global'` + ))[0]! + if (Number(worker.paused_until) > now || Number(worker.next_dispatch_at) > now) { + return { outcome: 'deferred' } + } + const open = (await transaction.query( + `SELECT COUNT(*) AS count FROM relay_assignment_migrations + WHERE completed_at IS NULL AND aborted_at IS NULL` + ))[0] + if (Number(open?.count ?? 0) >= REGIONAL_REHOME_CONCURRENT_LIMIT) return { outcome: 'deferred' } + const attempt = await this.startRegionalRehomeCandidate(transaction, { + identity: request, + sourceCellId: request.sourceCellId, + assignmentEpoch: request.sourceAssignmentEpoch, + preferenceCutoff: now - Number(control.preference_max_age_ms), + cooldownCutoff: now - Number(control.host_cooldown_ms), + drainGraceMs: 0, + processSafety, + worker, + now, + skips: [], + idleRequest: request, + cohortPercent, + onSafetyDisabled: (event) => { safetyDisable = event } + }) + if (!attempt) return { outcome: 'deferred' } + await this.markRegionalRehomeDispatchClaimed( + transaction, request.attemptId, now, Math.ceil(60_000 / Number(control.rate_per_minute)) + ) + await transaction.query( + `UPDATE relay_region_rehome_attempts SET drain_receipt_at = ?, drain_outcome = 'accepted' + WHERE attempt_id = ?`, [now, request.attemptId] + ) + return { outcome: 'committed' } + }) + if (safetyDisable) console.warn(JSON.stringify(safetyDisable)) + return result + } + + async reconcileIdleRegionalRehome(request: IdleRegionalRehomeRequest): Promise<'committed' | 'not-committed' | 'stale'> { + return this.database.transaction(async (transaction) => { + // Absence is definitive only after the same assignment lock as commit/activation. + const assignment = await this.assignmentRow(transaction, request) + const attempt = (await transaction.queryLocked( + `SELECT * FROM relay_region_rehome_attempts WHERE attempt_id = ?`, + [request.attemptId] + ))[0] + if (attempt) { + return attempt.user_id === request.userId && + attempt.relay_host_id === request.relayHostId && + attempt.source_cell_id === request.sourceCellId && + attempt.source_cell_incarnation === request.sourceCellIncarnation && + Number(attempt.previous_epoch) === request.sourceAssignmentEpoch && + Number(attempt.source_generation) === request.sourceGeneration && + attempt.target_cell_id === request.targetCellId && + attempt.aborted_at == null + ? 'committed' : 'stale' + } + if (!assignment || assignment.cell_id !== request.sourceCellId || + Number(assignment.assignment_epoch) !== request.sourceAssignmentEpoch) return 'stale' + const control = (await transaction.query( + `SELECT capability.generation, capability.cell_incarnation + FROM relay_control_capabilities capability + JOIN relay_assignment_activity_leases lease + ON lease.user_id = capability.user_id AND lease.relay_host_id = capability.relay_host_id + AND lease.activity_id = capability.activity_id + WHERE capability.user_id = ? AND capability.relay_host_id = ? + AND capability.cell_id = ? AND capability.assignment_epoch = ? + AND lease.activity_kind = 'control' AND lease.expires_at > ? + ORDER BY capability.generation DESC LIMIT 1`, + [request.userId, request.relayHostId, request.sourceCellId, request.sourceAssignmentEpoch, this.now()] + ))[0] + return control && Number(control.generation) === request.sourceGeneration && + control.cell_incarnation === request.sourceCellIncarnation ? 'not-committed' : 'stale' + }) + } + + async previewRegionalRehomeEligibility( + processSafety?: RegionalRehomeSafetySnapshot + ): Promise { + const now = this.now() + const fleetSafety = await this.readRegionalRehomeFleetSafety(this.database, now) + return previewRegionalRehomeEligibility({ + database: this.database, + now, + heartbeatTtlMs: this.heartbeatTtlMs, + cohortPercent: this.regionalRehomeCohortPercent, + globalSafetyFailure: processSafety + ? regionalRehomeFleetSafetyFailure(processSafety, fleetSafety, now) + : 'process-safety-unavailable', + connectionHeadroom: await this.connectionHeadroomByCell(this.database), + cellIsClean: regionalRehomeCellSafetyIsClean + }) + } + async renewControlActivity( identity: AssignmentIdentity, input: { activityId: string; cellId: string; expiresAt: number } @@ -3545,6 +3721,8 @@ export class RelayAssignmentStore { cellId: string assignmentEpoch: number generation: number + idleRegionalRehome?: boolean + cellIncarnation?: string connectionInclusionWatermark?: number } ): Promise { @@ -3626,6 +3804,33 @@ export class RelayAssignmentStore { input.connectionInclusionWatermark, now ) + await transaction.query( + `DELETE FROM relay_control_capabilities WHERE user_id = ? AND relay_host_id = ? + AND NOT EXISTS (SELECT 1 FROM relay_assignment_activity_leases lease + WHERE lease.user_id = relay_control_capabilities.user_id + AND lease.relay_host_id = relay_control_capabilities.relay_host_id + AND lease.activity_id = relay_control_capabilities.activity_id)`, + [identity.userId, identity.relayHostId] + ) + await transaction.query( + `INSERT INTO relay_control_capabilities + (user_id, relay_host_id, activity_id, cell_id, cell_incarnation, assignment_epoch, generation, finish_existing, idle_regional_rehome) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT (user_id, relay_host_id, activity_id) DO UPDATE SET + cell_incarnation = excluded.cell_incarnation, assignment_epoch = excluded.assignment_epoch, + generation = excluded.generation, finish_existing = excluded.finish_existing, idle_regional_rehome = excluded.idle_regional_rehome`, + [ + identity.userId, + identity.relayHostId, + activityId, + input.cellId, + input.cellIncarnation ?? '', + input.assignmentEpoch, + input.generation, + 0, + input.idleRegionalRehome && input.cellIncarnation ? 1 : 0 + ] + ) return activityId }) }) @@ -5116,327 +5321,6 @@ export class RelayAssignmentStore { } } - async claimRegionalRehome( - processSafety?: RegionalRehomeSafetySnapshot - ): Promise { - const now = this.now() - // Directors poll every second; avoid taking the global worker-row lock while disabled. - const control = ( - await this.database.query( - `SELECT enabled, not_before - FROM relay_region_rehome_control - WHERE control_id = 'global'` - ) - )[0] - if (!control) { - await this.initializeRegionalRehomeControl(this.database, now) - return null - } - if (integer(control, 'enabled') !== 1 || integer(control, 'not_before') > now) { - return null - } - this.pendingRegionalRehomeDisableLog = null - const candidateSkips: RegionalRehomeCandidateSkip[] = [] - // A Postgres transaction is unusable after a NOWAIT abort, so a contended - // tick abandons the candidate it stopped on plus every one behind it. - let candidatesTotal = 0 - let candidatesFinished = 0 - const claimResult = await this.database.transaction(async (transaction) => { - candidatesTotal = 0 - candidatesFinished = 0 - candidateSkips.length = 0 - await this.initializeRegionalRehomeControl(transaction, now) - const control = ( - await transaction.queryLocked( - `SELECT * FROM relay_region_rehome_control WHERE control_id = 'global'` - ) - )[0]! - if (integer(control, 'enabled') !== 1 || integer(control, 'not_before') > now) { - return null - } - const intervalMs = Math.ceil(60_000 / integer(control, 'rate_per_minute')) - const preferenceCutoff = now - integer(control, 'preference_max_age_ms') - // A host that was rehomed recently is left alone whichever way its - // preference now points: a flapping region probe must not walk one host - // back and forth across an ocean. - const cooldownCutoff = now - integer(control, 'host_cooldown_ms') - await transaction.query( - `INSERT INTO relay_region_rehome_worker_state - (worker_id, next_dispatch_at, paused_until, consecutive_failures, updated_at) - VALUES ('global', 0, 0, 0, ?) - ON CONFLICT (worker_id) DO NOTHING`, - [now] - ) - const worker = ( - await transaction.queryLocked( - `SELECT * FROM relay_region_rehome_worker_state WHERE worker_id = 'global'` - ) - )[0]! - if ( - integer(worker, 'paused_until') > now || - integer(worker, 'next_dispatch_at') > now - ) { - return null - } - const effectiveProcessSafety = processSafety ?? cleanRegionalRehomeSafety(now) - const fleetSafety = await this.readRegionalRehomeFleetSafety(transaction, now) - if ( - !(await this.regionalRehomeSafetyAllowsClaim( - transaction, - worker, - effectiveProcessSafety, - fleetSafety, - now - )) - ) { - return null - } - const retry = ( - await transaction.queryLocked( - `SELECT attempt.*, source.cell_url AS source_cell_url - FROM relay_region_rehome_attempts attempt - JOIN relay_cells source ON source.cell_id = attempt.source_cell_id - JOIN relay_cell_runtime runtime ON runtime.cell_id = attempt.source_cell_id - JOIN relay_cell_capabilities capability - ON capability.cell_id = runtime.cell_id - AND capability.cell_incarnation = runtime.cell_incarnation - JOIN relay_assignment_migrations migration - ON migration.user_id = attempt.user_id - AND migration.relay_host_id = attempt.relay_host_id - AND migration.assignment_epoch = attempt.assignment_epoch - WHERE attempt.drain_receipt_at IS NULL - AND attempt.completed_at IS NULL AND attempt.aborted_at IS NULL - AND attempt.send_attempts < 10 - AND (attempt.last_send_attempt_at IS NULL OR attempt.last_send_attempt_at <= ?) - AND runtime.cell_incarnation = attempt.source_cell_incarnation - AND runtime.ready = 1 AND runtime.last_heartbeat_at > ? - AND capability.regional_rehome_protocol >= 1 - AND migration.completed_at IS NULL AND migration.aborted_at IS NULL - ORDER BY attempt.created_at, attempt.attempt_id - LIMIT 1`, - [now - 30_000, now - this.heartbeatTtlMs] - ) - )[0] - if (retry) { - candidatesTotal = 1 - const fleetSafety = await this.lockedRegionalRehomeFleetSafety(transaction, now) - if ( - !(await this.regionalRehomeSafetyAllowsClaim( - transaction, - worker, - effectiveProcessSafety, - fleetSafety, - now - )) - ) { - return null - } - await this.markRegionalRehomeDispatchClaimed( - transaction, - text(retry, 'attempt_id'), - now, - intervalMs - ) - retry.send_attempts = integer(retry, 'send_attempts') + 1 - return regionalRehomeAttempt(retry) - } - - // A drain receipt is not convergence: grace enforcement lives only in - // source-cell session state, and attempts have been observed stalled - // dual-homed well past grace with source leases still renewing. Such - // attempts are re-dispatched with the remaining (zero) grace so the - // source force-closes and the host re-resolves onto its registered - // target. - const redrain = ( - await transaction.queryLocked( - `SELECT attempt.*, source.cell_url AS source_cell_url - FROM relay_region_rehome_attempts attempt - JOIN relay_cells source ON source.cell_id = attempt.source_cell_id - JOIN relay_cell_runtime runtime ON runtime.cell_id = attempt.source_cell_id - JOIN relay_cell_capabilities capability - ON capability.cell_id = runtime.cell_id - AND capability.cell_incarnation = runtime.cell_incarnation - JOIN relay_assignment_migrations migration - ON migration.user_id = attempt.user_id - AND migration.relay_host_id = attempt.relay_host_id - AND migration.assignment_epoch = attempt.assignment_epoch - WHERE attempt.drain_receipt_at IS NOT NULL - AND attempt.completed_at IS NULL AND attempt.aborted_at IS NULL - AND attempt.created_at + attempt.drain_grace_ms <= ? - AND attempt.send_attempts < ? - AND (attempt.last_send_attempt_at IS NULL OR attempt.last_send_attempt_at <= ?) - AND runtime.cell_incarnation = attempt.source_cell_incarnation - AND runtime.ready = 1 AND runtime.last_heartbeat_at > ? - AND capability.regional_rehome_protocol >= 1 - AND migration.completed_at IS NULL AND migration.aborted_at IS NULL - AND migration.target_registered_at IS NOT NULL - AND EXISTS ( - SELECT 1 FROM relay_assignment_activity_leases source_lease - WHERE source_lease.user_id = attempt.user_id - AND source_lease.relay_host_id = attempt.relay_host_id - AND source_lease.cell_id = attempt.source_cell_id - ) - ORDER BY attempt.created_at, attempt.attempt_id - LIMIT 1`, - [ - now, - REGIONAL_REHOME_REDRAIN_SEND_LIMIT, - now - REGIONAL_REHOME_REDRAIN_INTERVAL_MS, - now - this.heartbeatTtlMs - ] - ) - )[0] - if (redrain) { - candidatesTotal = 1 - const fleetSafety = await this.lockedRegionalRehomeFleetSafety(transaction, now) - if ( - !(await this.regionalRehomeSafetyAllowsClaim( - transaction, - worker, - effectiveProcessSafety, - fleetSafety, - now - )) - ) { - return null - } - await this.markRegionalRehomeDispatchClaimed( - transaction, - text(redrain, 'attempt_id'), - now, - intervalMs - ) - redrain.send_attempts = integer(redrain, 'send_attempts') + 1 - redrain.drain_grace_ms = 0 - return regionalRehomeAttempt(redrain) - } - - const candidates = await transaction.query( - `SELECT preference.user_id, preference.relay_host_id, - preference.observed_at, assignment.cell_id AS source_cell_id, - assignment.assignment_epoch - FROM relay_assignment_region_preferences preference - JOIN relay_assignments assignment - ON assignment.user_id = preference.user_id - AND assignment.relay_host_id = preference.relay_host_id - JOIN relay_cell_regions region ON region.cell_id = assignment.cell_id - JOIN relay_cell_admission admission ON admission.cell_id = assignment.cell_id - JOIN relay_cell_runtime runtime ON runtime.cell_id = assignment.cell_id - JOIN relay_cell_capabilities capability - ON capability.cell_id = runtime.cell_id - AND capability.cell_incarnation = runtime.cell_incarnation - WHERE preference.preferred_region <> region.region - AND preference.observed_at >= ? - AND admission.admission_state = 'general' - AND runtime.ready = 1 AND runtime.last_heartbeat_at > ? - AND capability.regional_rehome_protocol >= 1 - AND EXISTS ( - SELECT 1 FROM relay_assignment_activity_leases control - WHERE control.user_id = assignment.user_id - AND control.relay_host_id = assignment.relay_host_id - AND control.cell_id = assignment.cell_id - AND control.activity_kind = 'control' - AND control.activity_id NOT LIKE 'control-pending:%' - AND control.expires_at > ? - AND control.updated_at >= runtime.started_at - ) - AND NOT EXISTS ( - SELECT 1 FROM relay_assignment_migrations migration - WHERE migration.user_id = assignment.user_id - AND migration.relay_host_id = assignment.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 recent - WHERE recent.user_id = preference.user_id - AND recent.relay_host_id = preference.relay_host_id - AND recent.created_at > ? - ) - AND EXISTS ( - SELECT 1 FROM relay_cell_regions target_region - JOIN relay_cells target_cell ON target_cell.cell_id = target_region.cell_id - JOIN relay_cell_admission target_admission - ON target_admission.cell_id = target_region.cell_id - JOIN relay_cell_runtime target_runtime - ON target_runtime.cell_id = target_region.cell_id - JOIN relay_cell_capabilities target_capability - ON target_capability.cell_id = target_runtime.cell_id - AND target_capability.cell_incarnation = target_runtime.cell_incarnation - WHERE target_region.region = preference.preferred_region - AND target_cell.enabled = 1 - AND target_admission.admission_state = 'general' - AND target_runtime.ready = 1 - AND target_runtime.last_heartbeat_at > ? - AND target_capability.regional_rehome_protocol >= 1 - ) - ORDER BY preference.observed_at, preference.user_id, preference.relay_host_id - LIMIT 10`, - [ - preferenceCutoff, - now - this.heartbeatTtlMs, - now, - cooldownCutoff, - now - this.heartbeatTtlMs - ] - ) - candidatesTotal = candidates.length - for (const candidate of candidates) { - const claimed = await this.startRegionalRehomeCandidate(transaction, { - identity: { - userId: text(candidate, 'user_id'), - relayHostId: text(candidate, 'relay_host_id') - }, - sourceCellId: text(candidate, 'source_cell_id'), - assignmentEpoch: integer(candidate, 'assignment_epoch'), - preferenceCutoff, - cooldownCutoff, - drainGraceMs: integer(control, 'drain_grace_ms'), - processSafety: effectiveProcessSafety, - worker, - now, - skips: candidateSkips - }) - candidatesFinished++ - if (!claimed) continue - await this.markRegionalRehomeDispatchClaimed( - transaction, - claimed.attemptId, - now, - intervalMs - ) - return { ...claimed, sendAttempts: 1 } - } - if (candidates.length > 0) { - // Skipped candidates still cost all-rows FOR UPDATE inventory scans; - // charge the dispatch interval so skips are rate-limited like claims. - await this.markRegionalRehomeTickSkipped(transaction, now, intervalMs) - } - return null - }).catch((error: unknown): RegionalRehomeAttempt | null => { - // Only inventory contention is swallowed here; every other failure keeps - // its existing propagation and its dispatch-failure accounting. - if (!isDatabaseLockUnavailable(error)) throw error - // The dispatch tick runs every second; losing one to inventory contention - // costs a second of latency and never loses durable rehome state. The - // rolled-back transaction never disabled anything, so its pending disable - // log would describe a decision that did not happen. - candidateSkips.length = 0 - this.pendingRegionalRehomeDisableLog = null - warnSweepCellInventoryBusy( - 'claim-regional-rehome', - Math.max(1, candidatesTotal - candidatesFinished) - ) - return null - }) - const pendingDisableLog = this.pendingRegionalRehomeDisableLog - this.pendingRegionalRehomeDisableLog = null - if (pendingDisableLog) console.warn(JSON.stringify(pendingDisableLog)) - if (claimResult === null && candidateSkips.length > 0) { - console.warn(JSON.stringify(aggregateRegionalRehomeCandidateSkips(candidateSkips))) - } - return claimResult - } - private async startRegionalRehomeCandidate( transaction: RelayDatabase, input: { @@ -5450,6 +5334,9 @@ export class RelayAssignmentStore { worker: SqlRow now: number skips: RegionalRehomeCandidateSkip[] + idleRequest: IdleRegionalRehomeRequest + cohortPercent: number + onSafetyDisabled: (event: Record | null) => void } ): Promise | null> { const assignment = await this.assignmentRow(transaction, input.identity) @@ -5463,12 +5350,21 @@ export class RelayAssignmentStore { } const preference = ( await transaction.queryLocked( - `SELECT * FROM relay_assignment_region_preferences + `SELECT * FROM relay_region_decisions WHERE user_id = ? AND relay_host_id = ?`, [input.identity.userId, input.identity.relayHostId] ) )[0] - if (!preference || integer(preference, 'observed_at') < input.preferenceCutoff) { + if ( + !preference || + integer(preference, 'observed_at') < input.preferenceCutoff || + Number(preference.expires_at) <= input.now || + preference.outcome !== 'conclusive' || + Number(preference.policy_version) !== 1 || + Number(preference.assignment_epoch) !== input.assignmentEpoch || + !preference.preferred_region || + Number(preference.cohort_bucket) >= input.cohortPercent + ) { input.skips.push({ reason: 'candidate_stale' }) return null } @@ -5540,13 +5436,13 @@ export class RelayAssignmentStore { input.now ) if (safetyFailure) { - await this.pauseRegionalRehomeForSafety( + input.onSafetyDisabled(await this.pauseRegionalRehomeForSafety( transaction, input.worker, input.now, safetyFailure, fleetSafety - ) + )) return null } // The preference read under lock can now agree with the cell the host is @@ -5564,9 +5460,9 @@ export class RelayAssignmentStore { integer(sourceRuntime, 'ready') !== 1 || integer(sourceRuntime, 'last_heartbeat_at') <= input.now - this.heartbeatTtlMs || !sourceCapability || - text(sourceCapability, 'cell_incarnation') !== - text(sourceRuntime, 'cell_incarnation') || - integer(sourceCapability, 'regional_rehome_protocol') < 1 + text(sourceCapability, 'cell_incarnation') !== text(sourceRuntime, 'cell_incarnation') || + integer(sourceCapability, 'regional_rehome_protocol') < 3 || + sourceRuntime.cell_incarnation !== input.idleRequest.sourceCellIncarnation ) { input.skips.push({ reason: 'source_ineligible', cellId: input.sourceCellId }) return null @@ -5575,6 +5471,32 @@ export class RelayAssignmentStore { input.skips.push(cellUncleanSkip('source_unclean', input.sourceCellId, sourceSafety)) return null } + const hostCapability = ( + await transaction.query( + `SELECT capability.* FROM relay_control_capabilities capability + JOIN relay_assignment_activity_leases lease + ON lease.user_id = capability.user_id AND lease.relay_host_id = capability.relay_host_id + AND lease.activity_id = capability.activity_id + WHERE capability.user_id = ? AND capability.relay_host_id = ? + AND capability.cell_id = ? AND capability.assignment_epoch = ? + AND capability.cell_incarnation = ? AND capability.idle_regional_rehome = 1 + AND lease.expires_at > ? AND lease.activity_kind = 'control' + ORDER BY capability.generation DESC LIMIT 1`, + [ + input.identity.userId, + input.identity.relayHostId, + input.sourceCellId, + input.assignmentEpoch, + sourceRuntime.cell_incarnation, + input.now + ] + ) + )[0] + if (!hostCapability || preference.incumbent_region !== regions.get(input.sourceCellId) || + Number(hostCapability.generation) !== input.idleRequest.sourceGeneration) { + input.skips.push({ reason: 'candidate_stale' }) + return null + } const sourceControlActive = activityLeases.some( (lease) => text(lease, 'cell_id') === input.sourceCellId && @@ -5606,7 +5528,8 @@ export class RelayAssignmentStore { integer(runtime, 'last_heartbeat_at') > input.now - this.heartbeatTtlMs && capability !== undefined && text(capability, 'cell_incarnation') === text(runtime, 'cell_incarnation') && - integer(capability, 'regional_rehome_protocol') >= 1 + integer(capability, 'regional_rehome_protocol') >= 3 && + cellId === input.idleRequest.targetCellId ) }) const targetIsClean = (row: SqlRow): boolean => { @@ -5753,7 +5676,7 @@ export class RelayAssignmentStore { text(targetRuntime, 'cell_incarnation') ] ) - const attemptId = randomUUID() + const attemptId = input.idleRequest.attemptId await transaction.query( `INSERT INTO relay_region_rehome_attempts (attempt_id, user_id, relay_host_id, preferred_region, @@ -5780,6 +5703,10 @@ export class RelayAssignmentStore { input.now ] ) + await transaction.query( + `UPDATE relay_region_rehome_attempts SET source_generation = ? WHERE attempt_id = ?`, + [input.idleRequest.sourceGeneration, attemptId] + ) return { ...input.identity, attemptId, @@ -5795,61 +5722,13 @@ export class RelayAssignmentStore { } } - private async lockedRegionalRehomeFleetSafety( - transaction: RelayDatabase, - now: number - ): Promise { - const cells = await this.lockCellInventory(transaction, 'nowait') - const admission = await cellAdmissionStates(transaction) - const regions = new Map( - (await transaction.query(`SELECT cell_id, region FROM relay_cell_regions`)).map((row) => [ - text(row, 'cell_id'), - relayRegion(row, 'region') - ]) - ) - const runtimes = await transaction.queryLocked( - `SELECT * FROM relay_cell_runtime ORDER BY cell_id` - ) - const capabilities = await transaction.queryLocked( - `SELECT * FROM relay_cell_capabilities ORDER BY cell_id` - ) - const safetyRows = await transaction.queryLocked( - `SELECT * FROM relay_cell_rehome_safety ORDER BY cell_id` - ) - return regionalRehomeFleetSafetyFromInventory({ - cells, - admission, - regions, - runtimes, - capabilities, - safetyRows, - now, - heartbeatTtlMs: this.heartbeatTtlMs - }) - } - - private async regionalRehomeSafetyAllowsClaim( - transaction: RelayDatabase, - worker: SqlRow, - processSafety: RegionalRehomeSafetySnapshot, - fleetSafety: RegionalRehomeFleetSafety, - now: number - ): Promise { - const failure = regionalRehomeFleetSafetyFailure(processSafety, fleetSafety, now) - if (!failure) { - return true - } - await this.pauseRegionalRehomeForSafety(transaction, worker, now, failure, fleetSafety) - return false - } - private async pauseRegionalRehomeForSafety( transaction: RelayDatabase, worker: SqlRow, now: number, reason: string, fleetSafety: RegionalRehomeFleetSafety - ): Promise { + ): Promise | null> { const disabled = await transaction.query( `UPDATE relay_region_rehome_control SET generation = generation + 1, enabled = 0, updated_at = ? @@ -5860,8 +5739,9 @@ export class RelayAssignmentStore { // The durable disable is otherwise invisible: nothing else records why // claims stopped and inspection only shows enabled=false. Logged after // the transaction commits so a rollback cannot fabricate the record. + let event: Record | null = null if (disabled.length > 0) { - this.pendingRegionalRehomeDisableLog = { + event = { event: 'orca_relay_regional_rehome_safety_disabled', reason, controlGeneration: integer(disabled[0]!, 'generation'), @@ -5879,19 +5759,9 @@ export class RelayAssignmentStore { } } await this.incrementRegionalRehomeWorkerFailure(transaction, worker, now) + return event } - private async markRegionalRehomeTickSkipped( - transaction: RelayDatabase, - now: number, - intervalMs: number - ): Promise { - await transaction.query( - `UPDATE relay_region_rehome_worker_state - SET next_dispatch_at = ?, updated_at = ? WHERE worker_id = 'global'`, - [now + intervalMs, now] - ) - } private async markRegionalRehomeDispatchClaimed( transaction: RelayDatabase, @@ -5912,74 +5782,6 @@ export class RelayAssignmentStore { ) } - async recordRegionalRehomeDrainReceipt( - attemptId: string, - outcome: RegionalHostDrainOutcome - ): Promise { - const now = this.now() - return await this.database.transaction(async (transaction) => { - const worker = ( - await transaction.queryLocked( - `SELECT * FROM relay_region_rehome_worker_state WHERE worker_id = 'global'` - ) - )[0] - const attempt = ( - await transaction.queryLocked( - `SELECT * FROM relay_region_rehome_attempts WHERE attempt_id = ?`, - [attemptId] - ) - )[0] - if (!attempt) throw new Error('regional_rehome_attempt_not_found') - // Any receipt proves the source cell answered: reset the failure budget - // even when a redrain repeats the stored outcome; otherwise a - // redrain-dominated stream lets scattered transient failures reach the - // durable three-failure disable. - if (worker) { - await transaction.query( - `UPDATE relay_region_rehome_worker_state - SET consecutive_failures = 0, paused_until = 0, updated_at = ? - WHERE worker_id = 'global'`, - [now] - ) - } - const existingOutcome = optionalText(attempt, 'drain_outcome') - if (existingOutcome === outcome) return false - // Redrains produce one receipt per dispatch; the latest outcome wins. - await transaction.query( - `UPDATE relay_region_rehome_attempts - SET drain_receipt_at = ?, drain_outcome = ?, updated_at = ? - WHERE attempt_id = ?`, - [now, outcome, now, attemptId] - ) - return true - }) - } - - async recordRegionalRehomeDispatchFailure(attemptId: string): Promise { - const now = this.now() - const disableLog = await this.database.transaction(async (transaction) => { - // Match claim and enable ordering before a spent budget updates the control. - await transaction.queryLocked( - `SELECT * FROM relay_region_rehome_control WHERE control_id = 'global'` - ) - const worker = ( - await transaction.queryLocked( - `SELECT * FROM relay_region_rehome_worker_state WHERE worker_id = 'global'` - ) - )[0] - const attempt = ( - await transaction.queryLocked( - `SELECT attempt_id FROM relay_region_rehome_attempts WHERE attempt_id = ?`, - [attemptId] - ) - )[0] - if (!worker || !attempt) return null - return await this.incrementRegionalRehomeWorkerFailure(transaction, worker, now) - }) - // Logged after the commit so a rollback cannot fabricate the record. - if (disableLog) console.warn(JSON.stringify(disableLog)) - } - // Returns the durable disable this failure caused, for the caller to log once // its transaction commits; null when the budget survives or was already spent. private async incrementRegionalRehomeWorkerFailure( @@ -6074,7 +5876,7 @@ export class RelayAssignmentStore { // LIMIT pages: poisoned rows are permanent and always the oldest, so // without exclusion they eventually starve every healthy candidate. private recordRegionalRehomeCandidateFailure( - operation: 'complete' | 'abort', + operation: 'complete' | 'abort' | 'refresh', attemptId: string, now: number, error: unknown @@ -6116,12 +5918,16 @@ export class RelayAssignmentStore { async refreshRegionalRehomeLeases(limit = 100): Promise { const now = this.now() + const quarantined = this.quarantinedRegionalRehomeAttemptIds(now) + const exclusion = quarantined.length + ? ` AND attempt_id NOT IN (${quarantined.map(() => '?').join(', ')})` + : '' const candidates = await this.database.query( - `SELECT user_id, relay_host_id, assignment_epoch + `SELECT attempt_id, user_id, relay_host_id, assignment_epoch FROM relay_region_rehome_attempts - WHERE completed_at IS NULL AND aborted_at IS NULL - ORDER BY created_at, attempt_id LIMIT ?`, - [limit] + WHERE completed_at IS NULL AND aborted_at IS NULL${exclusion} + ORDER BY updated_at, attempt_id LIMIT ?`, + [...quarantined, limit] ) let refreshed = 0 for (const candidate of candidates) { @@ -6130,50 +5936,91 @@ export class RelayAssignmentStore { relayHostId: text(candidate, 'relay_host_id') } const assignmentEpoch = integer(candidate, 'assignment_epoch') - const changed = await this.database.transaction(async (transaction) => { - const assignment = await this.assignmentRow(transaction, identity) - const attempt = ( - await transaction.queryLocked( - `SELECT * FROM relay_region_rehome_attempts + const attemptId = text(candidate, 'attempt_id') + try { + const changed = await this.database.transaction(async (transaction) => { + const assignment = await this.assignmentRow(transaction, identity) + const attempt = ( + await transaction.queryLocked( + `SELECT * FROM relay_region_rehome_attempts WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, - [identity.userId, identity.relayHostId, assignmentEpoch] - ) - )[0] - const migration = ( - await transaction.queryLocked( - `SELECT * FROM relay_assignment_migrations + [identity.userId, identity.relayHostId, assignmentEpoch] + ) + )[0] + const migration = ( + await transaction.queryLocked( + `SELECT * FROM relay_assignment_migrations WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, - [identity.userId, identity.relayHostId, assignmentEpoch] - ) - )[0] - if ( - !assignment || - !attempt || - !migration || - optionalInteger(attempt, 'completed_at') !== undefined || - optionalInteger(attempt, 'aborted_at') !== undefined || - optionalInteger(migration, 'completed_at') !== undefined || - optionalInteger(migration, 'aborted_at') !== undefined - ) { - return false - } - const attemptAgeMs = now - integer(attempt, 'created_at') - if (attemptAgeMs >= REGIONAL_REHOME_MAX_REFRESH_MS) { - return false - } - if ( - optionalInteger(migration, 'target_registered_at') === undefined && - attemptAgeMs >= REGIONAL_REHOME_UNREGISTERED_REFRESH_MS - ) { - await transaction.query( - `UPDATE relay_assignment_activity_leases + [identity.userId, identity.relayHostId, assignmentEpoch] + ) + )[0] + if (attempt && attempt.completed_at == null && attempt.aborted_at == null) { + await transaction.query( + `UPDATE relay_region_rehome_attempts SET updated_at = ? WHERE attempt_id = ?`, + [now, attempt.attempt_id] + ) + } + if ( + !assignment || + !attempt || + !migration || + optionalInteger(attempt, 'completed_at') !== undefined || + optionalInteger(attempt, 'aborted_at') !== undefined || + optionalInteger(migration, 'completed_at') !== undefined || + optionalInteger(migration, 'aborted_at') !== undefined + ) { + return false + } + const attemptAgeMs = now - integer(attempt, 'created_at') + if (attemptAgeMs >= REGIONAL_REHOME_MAX_REFRESH_MS) { + return false + } + if ( + optionalInteger(migration, 'target_registered_at') === undefined && + attemptAgeMs >= REGIONAL_REHOME_UNREGISTERED_REFRESH_MS + ) { + await transaction.query( + `UPDATE relay_assignment_activity_leases SET expires_at = CASE WHEN expires_at < ? THEN expires_at ELSE ? END, updated_at = ? WHERE user_id = ? AND relay_host_id = ? AND activity_id IN (?, ?)`, + [ + now, + now, + now, + identity.userId, + identity.relayHostId, + pendingControlActivityId(assignmentEpoch), + migrationActivityId(assignmentEpoch) + ] + ) + await transaction.query( + `UPDATE relay_assignment_migrations + SET expires_at = CASE WHEN expires_at < ? THEN expires_at ELSE ? END, + updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [now, now, now, identity.userId, identity.relayHostId, assignmentEpoch] + ) + return false + } + const leases = await this.lockAssignmentActivities(transaction, identity) + const protectedIds = new Set([ + pendingControlActivityId(assignmentEpoch), + migrationActivityId(assignmentEpoch) + ]) + const protectedLeases = leases.filter((lease) => + protectedIds.has(text(lease, 'activity_id')) + ) + if (protectedLeases.length === 0) return false + const expiresAt = now + ASSIGNMENT_LIMITS.migrationLeaseMs + await transaction.query( + `UPDATE relay_assignment_activity_leases + SET expires_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? + AND activity_id IN (?, ?)`, [ - now, - now, + expiresAt, now, identity.userId, identity.relayHostId, @@ -6182,53 +6029,25 @@ export class RelayAssignmentStore { ] ) await transaction.query( - `UPDATE relay_assignment_migrations - SET expires_at = CASE WHEN expires_at < ? THEN expires_at ELSE ? END, - updated_at = ? - WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, - [now, now, now, identity.userId, identity.relayHostId, assignmentEpoch] - ) - return false - } - const leases = await this.lockAssignmentActivities(transaction, identity) - const protectedIds = new Set([ - pendingControlActivityId(assignmentEpoch), - migrationActivityId(assignmentEpoch) - ]) - const protectedLeases = leases.filter((lease) => - protectedIds.has(text(lease, 'activity_id')) - ) - if (protectedLeases.length === 0) return false - const expiresAt = now + ASSIGNMENT_LIMITS.migrationLeaseMs - await transaction.query( - `UPDATE relay_assignment_activity_leases - SET expires_at = ?, updated_at = ? - WHERE user_id = ? AND relay_host_id = ? - AND activity_id IN (?, ?)`, - [ - expiresAt, - now, - identity.userId, - identity.relayHostId, - pendingControlActivityId(assignmentEpoch), - migrationActivityId(assignmentEpoch) - ] - ) - await transaction.query( - `UPDATE relay_assignment_migrations SET expires_at = ?, updated_at = ? + `UPDATE relay_assignment_migrations SET expires_at = ?, updated_at = ? WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, - [expiresAt, now, identity.userId, identity.relayHostId, assignmentEpoch] - ) - await transaction.query( - `UPDATE relay_assignments SET lease_expires_at = + [expiresAt, now, identity.userId, identity.relayHostId, assignmentEpoch] + ) + await transaction.query( + `UPDATE relay_assignments SET lease_expires_at = CASE WHEN lease_expires_at > ? THEN lease_expires_at ELSE ? END, last_activity_at = ? WHERE user_id = ? AND relay_host_id = ?`, - [expiresAt, expiresAt, now, identity.userId, identity.relayHostId] - ) - return true - }) - if (changed) refreshed++ + [expiresAt, expiresAt, now, identity.userId, identity.relayHostId] + ) + return true + }) + if (changed) refreshed++ + this.regionalRehomeCandidateQuarantine.delete(attemptId) + } catch (error) { + if (!isDatabaseLockUnavailable(error)) + this.recordRegionalRehomeCandidateFailure('refresh', attemptId, now, error) + } } return refreshed } @@ -6575,92 +6394,169 @@ export class RelayAssignmentStore { let aborted = 0 let inventoryBusy = 0 for (const candidate of candidates) { - const didAbort = await this.database.transaction(async (transaction) => { - const identity = { - userId: text(candidate, 'user_id'), - relayHostId: text(candidate, 'relay_host_id') - } - // Migration cleanup follows the same assignment-first order as evacuation. - const assignment = await this.assignmentRow(transaction, identity) - const assignmentEpoch = integer(candidate, 'assignment_epoch') - const regionalAttempt = ( - await transaction.queryLocked( - `SELECT attempt_id FROM relay_region_rehome_attempts + const didAbort = await this.database + .transaction(async (transaction) => { + const identity = { + userId: text(candidate, 'user_id'), + relayHostId: text(candidate, 'relay_host_id') + } + // Migration cleanup follows the same assignment-first order as evacuation. + const assignment = await this.assignmentRow(transaction, identity) + const assignmentEpoch = integer(candidate, 'assignment_epoch') + const regionalAttempt = ( + await transaction.queryLocked( + `SELECT attempt_id FROM relay_region_rehome_attempts WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ? AND completed_at IS NULL AND aborted_at IS NULL`, - [identity.userId, identity.relayHostId, assignmentEpoch] - ) - )[0] - const row = ( - await transaction.queryLocked( - `SELECT migration.* FROM relay_assignment_migrations migration + [identity.userId, identity.relayHostId, assignmentEpoch] + ) + )[0] + const row = ( + await transaction.queryLocked( + `SELECT migration.* FROM relay_assignment_migrations migration WHERE migration.user_id = ? AND migration.relay_host_id = ? AND migration.assignment_epoch = ? AND migration.expires_at <= ? AND migration.completed_at IS NULL AND migration.aborted_at IS NULL AND ${ABORTABLE_EXPIRED_MIGRATION}`, - [ - identity.userId, - identity.relayHostId, - assignmentEpoch, - now, - now, - abandonedBefore, - abandonedBefore - ] + [ + identity.userId, + identity.relayHostId, + assignmentEpoch, + now, + now, + abandonedBefore, + abandonedBefore + ] + ) + )[0] + if (!row) return false + const targetCellId = text(row, 'target_cell_id') + const activityLeases = await this.lockAssignmentActivities(transaction, identity) + if (!assignment) throw new Error('migration_assignment_missing') + const currentAssignmentEpoch = integer(assignment, 'assignment_epoch') + const assignmentEpochMatches = + text(assignment, 'cell_id') === targetCellId && + currentAssignmentEpoch === assignmentEpoch + const pendingTargetControl = activityLeaseById( + activityLeases, + pendingControlActivityId(assignmentEpoch) ) - )[0] - if (!row) return false - const targetCellId = text(row, 'target_cell_id') - const activityLeases = await this.lockAssignmentActivities(transaction, identity) - if (!assignment) throw new Error('migration_assignment_missing') - const currentAssignmentEpoch = integer(assignment, 'assignment_epoch') - const assignmentEpochMatches = - text(assignment, 'cell_id') === targetCellId && - currentAssignmentEpoch === assignmentEpoch - const pendingTargetControl = activityLeaseById( - activityLeases, - pendingControlActivityId(assignmentEpoch) - ) - const targetGrantIsFresh = - assignmentEpochMatches && - pendingTargetControl !== undefined && - text(pendingTargetControl, 'cell_id') === targetCellId && - text(pendingTargetControl, 'activity_kind') === 'control' && - integer(pendingTargetControl, 'expires_at') > now - const targetIsActive = activityLeases.some( - (lease) => - text(lease, 'cell_id') === targetCellId && - text(lease, 'activity_kind') === 'control' && - text(lease, 'activity_id') !== pendingControlActivityId(assignmentEpoch) - ) - if (targetGrantIsFresh) return false - if (targetIsActive && assignmentEpochMatches) { - // A committed target control is stronger evidence than a failed follow-up - // write; repair the marker instead of rolling a live desktop backward. - await transaction.query( - `UPDATE relay_assignment_migrations + const targetGrantIsFresh = + assignmentEpochMatches && + pendingTargetControl !== undefined && + text(pendingTargetControl, 'cell_id') === targetCellId && + text(pendingTargetControl, 'activity_kind') === 'control' && + integer(pendingTargetControl, 'expires_at') > now + const targetIsActive = activityLeases.some( + (lease) => + text(lease, 'cell_id') === targetCellId && + text(lease, 'activity_kind') === 'control' && + text(lease, 'activity_id') !== pendingControlActivityId(assignmentEpoch) + ) + if (targetGrantIsFresh) return false + if (targetIsActive && assignmentEpochMatches) { + // A committed target control is stronger evidence than a failed follow-up + // write; repair the marker instead of rolling a live desktop backward. + await transaction.query( + `UPDATE relay_assignment_migrations SET target_registered_at = COALESCE(target_registered_at, ?), updated_at = ? WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, - [now, now, identity.userId, identity.relayHostId, assignmentEpoch] - ) - return false - } - if (!assignmentEpochMatches) { - if (currentAssignmentEpoch <= assignmentEpoch) { - throw new Error('migration_assignment_mismatch') + [now, now, identity.userId, identity.relayHostId, assignmentEpoch] + ) + return false } - // A newer assignment is authoritative regardless of where it landed. - // Retire only this obsolete migration; never rewrite the newer epoch. - const obsoleteLeases = [ + if (!assignmentEpochMatches) { + if (currentAssignmentEpoch <= assignmentEpoch) { + throw new Error('migration_assignment_mismatch') + } + // A newer assignment is authoritative regardless of where it landed. + // Retire only this obsolete migration; never rewrite the newer epoch. + const obsoleteLeases = [ + pendingControlActivityId(assignmentEpoch), + migrationActivityId(assignmentEpoch) + ] + .map((activityId) => activityLeaseById(activityLeases, activityId)) + .filter((lease): lease is SqlRow => lease !== undefined) + if (obsoleteLeases.length > 0) await this.lockCellInventory(transaction, 'nowait') + for (const lease of obsoleteLeases) { + await this.removeActivityLease(transaction, identity, lease, now) + } + await this.releaseSupersededControlConnectionReservations( + transaction, + identity, + targetCellId, + assignmentEpoch, + now + ) + await transaction.query( + `UPDATE relay_assignment_migrations SET aborted_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [now, now, identity.userId, identity.relayHostId, assignmentEpoch] + ) + return true + } + const cells = await this.lockCellInventory(transaction, 'nowait') + const sourceCellId = text(row, 'source_cell_id') + const admissionRows = await transaction.query( + `SELECT cell_id, admission_state, updated_at FROM relay_cell_admission + WHERE cell_id IN (?, ?)`, + [sourceCellId, targetCellId] + ) + const sourceCell = cells.find((cell) => text(cell, 'cell_id') === sourceCellId) + const targetCell = cells.find((cell) => text(cell, 'cell_id') === targetCellId) + const sourceAdmission = admissionRows.find( + (admission) => text(admission, 'cell_id') === sourceCellId + ) + const targetAdmission = admissionRows.find( + (admission) => text(admission, 'cell_id') === targetCellId + ) + const registered = optionalInteger(row, 'target_registered_at') !== undefined + const sourceIsDurablyFenced = + registered && + ( + await transaction.query( + `SELECT 1 FROM relay_assignment_migrations migration + WHERE migration.user_id = ? AND migration.relay_host_id = ? + AND migration.assignment_epoch = ? + AND ${DURABLY_FENCED_MIGRATION_SOURCE}`, + [identity.userId, identity.relayHostId, assignmentEpoch] + ) + ).length === 1 + const retireOnTarget = + registered && + activityUnitsForCell(activityLeases, sourceCellId) === 0 && + sourceCell !== undefined && + integer(sourceCell, 'enabled') === 0 && + sourceAdmission !== undefined && + text(sourceAdmission, 'admission_state') === 'existing-only' && + (integer(sourceAdmission, 'updated_at') <= abandonedBefore || sourceIsDurablyFenced) && + targetCell !== undefined && + integer(targetCell, 'enabled') === 1 && + targetAdmission !== undefined && + ['migration-only', 'general'].includes(text(targetAdmission, 'admission_state')) + const rollbackReason = + !registered || + (targetCell !== undefined && + integer(targetCell, 'enabled') === 0 && + targetAdmission !== undefined && + text(targetAdmission, 'admission_state') === 'existing-only' && + integer(targetAdmission, 'updated_at') <= abandonedBefore) + const regionalRollbackSourceAvailable = + !regionalAttempt || + (sourceCell !== undefined && + integer(sourceCell, 'enabled') === 1 && + sourceAdmission !== undefined && + text(sourceAdmission, 'admission_state') === 'general' && + (await this.cellIsLive(transaction, sourceCellId, now))) + const rollbackToSource = rollbackReason && regionalRollbackSourceAvailable + if (!retireOnTarget && !rollbackToSource) return false + for (const activityId of [ pendingControlActivityId(assignmentEpoch), migrationActivityId(assignmentEpoch) - ] - .map((activityId) => activityLeaseById(activityLeases, activityId)) - .filter((lease): lease is SqlRow => lease !== undefined) - if (obsoleteLeases.length > 0) await this.lockCellInventory(transaction, 'nowait') - for (const lease of obsoleteLeases) { - await this.removeActivityLease(transaction, identity, lease, now) + ]) { + const lease = activityLeaseById(activityLeases, activityId) + if (lease) await this.removeActivityLease(transaction, identity, lease, now) } await this.releaseSupersededControlConnectionReservations( transaction, @@ -6669,116 +6565,46 @@ export class RelayAssignmentStore { assignmentEpoch, now ) - await transaction.query( - `UPDATE relay_assignment_migrations SET aborted_at = ?, updated_at = ? - WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, - [now, now, identity.userId, identity.relayHostId, assignmentEpoch] - ) - return true - } - const cells = await this.lockCellInventory(transaction, 'nowait') - const sourceCellId = text(row, 'source_cell_id') - const admissionRows = await transaction.query( - `SELECT cell_id, admission_state, updated_at FROM relay_cell_admission - WHERE cell_id IN (?, ?)`, - [sourceCellId, targetCellId] - ) - const sourceCell = cells.find((cell) => text(cell, 'cell_id') === sourceCellId) - const targetCell = cells.find((cell) => text(cell, 'cell_id') === targetCellId) - const sourceAdmission = admissionRows.find( - (admission) => text(admission, 'cell_id') === sourceCellId - ) - const targetAdmission = admissionRows.find( - (admission) => text(admission, 'cell_id') === targetCellId - ) - const registered = optionalInteger(row, 'target_registered_at') !== undefined - const sourceIsDurablyFenced = - registered && - ( + if (retireOnTarget) { await transaction.query( - `SELECT 1 FROM relay_assignment_migrations migration - WHERE migration.user_id = ? AND migration.relay_host_id = ? - AND migration.assignment_epoch = ? - AND ${DURABLY_FENCED_MIGRATION_SOURCE}`, - [identity.userId, identity.relayHostId, assignmentEpoch] - ) - ).length === 1 - const retireOnTarget = - registered && - activityUnitsForCell(activityLeases, sourceCellId) === 0 && - sourceCell !== undefined && - integer(sourceCell, 'enabled') === 0 && - sourceAdmission !== undefined && - text(sourceAdmission, 'admission_state') === 'existing-only' && - (integer(sourceAdmission, 'updated_at') <= abandonedBefore || - sourceIsDurablyFenced) && - targetCell !== undefined && - integer(targetCell, 'enabled') === 1 && - targetAdmission !== undefined && - ['migration-only', 'general'].includes(text(targetAdmission, 'admission_state')) - const rollbackReason = - !registered || - (targetCell !== undefined && - integer(targetCell, 'enabled') === 0 && - targetAdmission !== undefined && - text(targetAdmission, 'admission_state') === 'existing-only' && - integer(targetAdmission, 'updated_at') <= abandonedBefore) - const regionalRollbackSourceAvailable = - !regionalAttempt || - (sourceCell !== undefined && - integer(sourceCell, 'enabled') === 1 && - sourceAdmission !== undefined && - text(sourceAdmission, 'admission_state') === 'general' && - (await this.cellIsLive(transaction, sourceCellId, now))) - const rollbackToSource = rollbackReason && regionalRollbackSourceAvailable - if (!retireOnTarget && !rollbackToSource) return false - for (const activityId of [ - pendingControlActivityId(assignmentEpoch), - migrationActivityId(assignmentEpoch) - ]) { - const lease = activityLeaseById(activityLeases, activityId) - if (lease) await this.removeActivityLease(transaction, identity, lease, now) - } - await this.releaseSupersededControlConnectionReservations( - transaction, - identity, - targetCellId, - assignmentEpoch, - now - ) - if (retireOnTarget) { - await transaction.query( - `UPDATE relay_assignment_migrations SET completed_at = ?, updated_at = ? + `UPDATE relay_assignment_migrations SET completed_at = ?, updated_at = ? WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, - [now, now, identity.userId, identity.relayHostId, assignmentEpoch] - ) - return true - } - await transaction.query( - `UPDATE relay_assignments SET cell_id = ?, assignment_epoch = ?, + [now, now, identity.userId, identity.relayHostId, assignmentEpoch] + ) + return true + } + await transaction.query( + `UPDATE relay_assignments SET cell_id = ?, assignment_epoch = ?, lease_expires_at = ?, last_activity_at = ? WHERE user_id = ? AND relay_host_id = ?`, - [ - sourceCellId, - assignmentEpoch + 1, - now + ASSIGNMENT_LIMITS.activityLeaseMs, - now, - identity.userId, - identity.relayHostId - ] - ) - await transaction.query( - `UPDATE relay_assignment_migrations SET aborted_at = ?, updated_at = ? + [ + sourceCellId, + assignmentEpoch + 1, + now + ASSIGNMENT_LIMITS.activityLeaseMs, + now, + identity.userId, + identity.relayHostId + ] + ) + await transaction.query( + `UPDATE relay_assignment_migrations SET aborted_at = ?, updated_at = ? WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, - [now, now, identity.userId, identity.relayHostId, assignmentEpoch] - ) - return true - }).catch((error: unknown): boolean => { - // Expiry is durable; another director settling this row is not a failure. - if (!isDatabaseLockUnavailable(error)) throw error - inventoryBusy++ - return false - }) + [now, now, identity.userId, identity.relayHostId, assignmentEpoch] + ) + if (regionalAttempt) { + await transaction.query( + `UPDATE relay_region_rehome_attempts SET aborted_at = ?, updated_at = ? WHERE attempt_id = ?`, + [now, now, regionalAttempt.attempt_id] + ) + } + return true + }) + .catch((error: unknown): boolean => { + // Expiry is durable; another director settling this row is not a failure. + if (!isDatabaseLockUnavailable(error)) throw error + inventoryBusy++ + return false + }) if (didAbort) aborted++ } warnSweepCellInventoryBusy('abort-expired-evacuations', inventoryBusy) @@ -8165,7 +7991,7 @@ function migration(identity: AssignmentIdentity, row: SqlRow): RelayAssignmentMi // Attempt ids are server-minted UUIDs and this codebase's invariant messages // are snake_case slugs; anything else could carry secrets and logs redacted. function warnRegionalRehomeCandidateFailure( - operation: 'complete' | 'abort', + operation: 'complete' | 'abort' | 'refresh', attemptId: string, error: unknown ): void { @@ -8190,24 +8016,6 @@ function noteRegionalRehomeActivityCountsRepaired(attemptId: string): void { ) } -function regionalRehomeAttempt(row: SqlRow): RegionalRehomeAttempt { - return { - attemptId: text(row, 'attempt_id'), - userId: text(row, 'user_id'), - relayHostId: text(row, 'relay_host_id'), - preferredRegion: relayRegion(row, 'preferred_region'), - sourceCellId: text(row, 'source_cell_id'), - sourceCellUrl: text(row, 'source_cell_url'), - sourceCellIncarnation: text(row, 'source_cell_incarnation'), - targetCellId: text(row, 'target_cell_id'), - targetCellIncarnation: text(row, 'target_cell_incarnation'), - previousEpoch: integer(row, 'previous_epoch'), - assignmentEpoch: integer(row, 'assignment_epoch'), - drainGraceMs: integer(row, 'drain_grace_ms'), - sendAttempts: integer(row, 'send_attempts') - } -} - function regionalRehomeControl(row: SqlRow): RegionalRehomeControl { return { generation: integer(row, 'generation'), @@ -8221,17 +8029,6 @@ function regionalRehomeControl(row: SqlRow): RegionalRehomeControl { } } -function cleanRegionalRehomeSafety(now: number): RegionalRehomeSafetySnapshot { - return { - observedAt: now, - sqlFailures: 0, - reconnects: 0, - controlActivityRecoveryFailures: 0, - databasePoolWaiting: 0, - databasePoolWaitersMax: 0, - databasePoolWaitMsMax: 0 - } -} function regionalRehomeFleetSafetyFromInventory(input: { cells: SqlRow[] @@ -8353,23 +8150,6 @@ function cellUncleanSkip( // Candidate skips are otherwise invisible: they neither latch the control off // nor produce attempts, so an operator cannot tell "skipping" from "idle". // Cell ids and counters only — never free-form error text. -function aggregateRegionalRehomeCandidateSkips( - skips: readonly RegionalRehomeCandidateSkip[] -): Record { - // `candidates` counts skipped candidate iterations, not distinct cells: one - // unclean cell blocking six candidates reports candidates=6 on one cellId. - const aggregated = new Map() - for (const skip of skips) { - const key = `${skip.reason}:${skip.cellId ?? ''}` - const entry = aggregated.get(key) - if (entry) entry.candidates += 1 - else aggregated.set(key, { ...skip, candidates: 1 }) - } - return { - event: 'orca_relay_regional_rehome_candidates_skipped', - skips: [...aggregated.values()] - } -} function regionalRehomeCellSafetyIsClean( safety: SqlRow | undefined, diff --git a/cloud/apps/relay/src/cell-heartbeat-client.test.ts b/cloud/apps/relay/src/cell-heartbeat-client.test.ts index 2aa708bed1b..6c36829e768 100644 --- a/cloud/apps/relay/src/cell-heartbeat-client.test.ts +++ b/cloud/apps/relay/src/cell-heartbeat-client.test.ts @@ -92,7 +92,7 @@ describe('cell heartbeat client', () => { client.stop() expect(JSON.parse(String(requests[1]!.body))).toMatchObject({ - regionalRehomeProtocol: 1, + regionalRehomeProtocol: 3, safety: { observedAt: 120, sqlFailures: 0, @@ -149,28 +149,34 @@ describe('cell heartbeat client', () => { it('does not start outside an explicitly configured cell role', () => { expect( - startCellHeartbeat({ ...CONFIG, role: 'director' }, { - ready: async () => true, - observedRequests: () => 0, - connectionCounts: () => ({ - totalConnections: 0, - inFlightConnections: 0, - reservedConnectionUnits: 0, - enforcedConnectionUnits: 0 - }) - }) + startCellHeartbeat( + { ...CONFIG, role: 'director' }, + { + ready: async () => true, + observedRequests: () => 0, + connectionCounts: () => ({ + totalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0 + }) + } + ) ).toBeNull() expect( - startCellHeartbeat({ ...CONFIG, directorUrl: undefined }, { - ready: async () => true, - observedRequests: () => 0, - connectionCounts: () => ({ - totalConnections: 0, - inFlightConnections: 0, - reservedConnectionUnits: 0, - enforcedConnectionUnits: 0 - }) - }) + startCellHeartbeat( + { ...CONFIG, directorUrl: undefined }, + { + ready: async () => true, + observedRequests: () => 0, + connectionCounts: () => ({ + totalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0 + }) + } + ) ).toBeNull() }) }) diff --git a/cloud/apps/relay/src/cell-heartbeat-client.ts b/cloud/apps/relay/src/cell-heartbeat-client.ts index 5c990310413..54f97a17af8 100644 --- a/cloud/apps/relay/src/cell-heartbeat-client.ts +++ b/cloud/apps/relay/src/cell-heartbeat-client.ts @@ -70,8 +70,7 @@ export function startCellHeartbeat( inFlightConnections: connectionCounts!.inFlightConnections, reservedConnectionUnits: connectionCounts!.reservedConnectionUnits, enforcedConnectionUnits: connectionCounts!.enforcedConnectionUnits, - connectionInclusionWatermark: - connectionCounts!.inclusionWatermark, + connectionInclusionWatermark: connectionCounts!.inclusionWatermark, connectionHardCap: config.connectionHardCap, connectionUnobservedBound: config.connectionUnobservedBound }) @@ -94,7 +93,7 @@ export function startCellHeartbeat( cellId: config.cellId, cellIncarnation, regionalRehomeProtocol: - config.rehomeAudience && config.rehomeDirectorServiceAccount ? 1 : 0, + config.rehomeAudience && config.rehomeDirectorServiceAccount ? 3 : 0, safety: options.regionalRehomeSafety() }), signal: AbortSignal.timeout(10_000) @@ -106,7 +105,10 @@ export function startCellHeartbeat( } } catch (error) { // A heartbeat must fail closed without ever logging its bearer token. - console.warn('[orca-relay] cell heartbeat failed', error instanceof Error ? error.message : '') + console.warn( + '[orca-relay] cell heartbeat failed', + error instanceof Error ? error.message : '' + ) } finally { inFlight = false } diff --git a/cloud/apps/relay/src/cell-inventory-lock-census.test.ts b/cloud/apps/relay/src/cell-inventory-lock-census.test.ts index 0ac4c8225e3..929173eb9da 100644 --- a/cloud/apps/relay/src/cell-inventory-lock-census.test.ts +++ b/cloud/apps/relay/src/cell-inventory-lock-census.test.ts @@ -43,14 +43,13 @@ const CENSUS: CensusEntry[] = [ { method: 'completeEvacuation', mode: 'nowait', reach: 'both' }, { method: 'completeEvacuation', mode: 'pool-default', reach: 'both' }, { method: 'rebalanceDormant', mode: 'request', reach: 'request' }, - { method: 'startRegionalRehomeCandidate', mode: 'nowait', reach: 'sweep' }, - { method: 'lockedRegionalRehomeFleetSafety', mode: 'nowait', reach: 'sweep' }, + { method: 'startRegionalRehomeCandidate', mode: 'nowait', reach: 'request' }, { method: 'completeRegionalRehomeCandidate', mode: 'nowait', reach: 'sweep' }, { method: 'abortExpiredRegionalRehomes', mode: 'nowait', reach: 'sweep' }, { method: 'abortExpiredEvacuations', mode: 'nowait', reach: 'sweep' }, { method: 'abortExpiredEvacuations', mode: 'nowait', reach: 'sweep' }, { method: 'releaseExpiredActivityLeases', mode: 'nowait', reach: 'sweep' }, - { method: 'releaseExpiredActivity', mode: 'nowait', reach: 'sweep' }, + { method: 'releaseExpiredActivity', mode: 'nowait', reach: 'sweep' } // reconcileReservationAccounting and leastLoadedCell are gone too: the first // repairs exactly two cells' counters and now holds only those rows, and the // second selects from the inventory its single caller has already locked. @@ -119,9 +118,10 @@ function storeCallGraph(lines: string[]): Map> { bounds.forEach((method, index) => { const end = bounds[index + 1]?.start ?? lines.length const names = callees.get(method.name) ?? new Set() - for (const call of lines.slice(method.start, end).join('\n').matchAll( - /this\.([A-Za-z_][\w]*)\s*\(/g - )) { + for (const call of lines + .slice(method.start, end) + .join('\n') + .matchAll(/this\.([A-Za-z_][\w]*)\s*\(/g)) { names.add(call[1]!) } callees.set(method.name, names) @@ -174,9 +174,7 @@ function readCallSites(): { method: string; mode: CensusMode }[] { describe('cell inventory lock call-site census', () => { it('classifies every call site exactly as recorded', () => { - expect(readCallSites()).toEqual( - CENSUS.map(({ method, mode }) => ({ method, mode })) - ) + expect(readCallSites()).toEqual(CENSUS.map(({ method, mode }) => ({ method, mode }))) }) // Why: the census only sees lockCellInventory calls, so a hand-written diff --git a/cloud/apps/relay/src/config.test.ts b/cloud/apps/relay/src/config.test.ts index bb0522dcbd3..01661a61826 100644 --- a/cloud/apps/relay/src/config.test.ts +++ b/cloud/apps/relay/src/config.test.ts @@ -25,6 +25,17 @@ function cellEnvironment(capacity: number): NodeJS.ProcessEnv { } describe('GCE relay capacity configuration', () => { + it('defaults optional region correction off and bounds the cohort', () => { + const env = cellEnvironment(4_000) + expect(loadRelayConfig(env).regionCorrectionCohortPercent).toBe(0) + env.ORCA_RELAY_REGION_CORRECTION_COHORT_PERCENT = '5' + expect(loadRelayConfig(env).regionCorrectionCohortPercent).toBe(5) + for (const invalid of ['-1', '101', '1.5', 'not-a-number']) { + env.ORCA_RELAY_REGION_CORRECTION_COHORT_PERCENT = invalid + expect(() => loadRelayConfig(env)).toThrow() + } + }) + it('requires distinct dedicated admin identities and accepts omitted values', () => { const env = cellEnvironment(4_000) expect(loadRelayConfig(env)).toMatchObject({ diff --git a/cloud/apps/relay/src/config.ts b/cloud/apps/relay/src/config.ts index 2bf23444a71..83dc9708f74 100644 --- a/cloud/apps/relay/src/config.ts +++ b/cloud/apps/relay/src/config.ts @@ -75,11 +75,15 @@ const EnvSchema = z.object({ ORCA_RELAY_RUNTIME_SERVICE_ACCOUNT: z.string().email().optional(), ORCA_RELAY_DIRECTOR_URL: z.string().url().optional(), ORCA_RELAY_HEARTBEAT_AUDIENCE: z.string().url().optional(), - ORCA_RELAY_IMAGE_DIGEST: z.string().regex(/^sha256:[a-f0-9]{64}$/).optional(), + ORCA_RELAY_IMAGE_DIGEST: z + .string() + .regex(/^sha256:[a-f0-9]{64}$/) + .optional(), ORCA_RELAY_ADMIN_JWKS_URL: z.string().url().default('https://www.googleapis.com/oauth2/v3/certs'), ORCA_RELAY_DATABASE_POOL_MAX: z.coerce.number().int().positive().max(100).optional(), ORCA_RELAY_PUBLIC_ASSIGNMENTS_ENABLED: EnvironmentBooleanSchema, ORCA_RELAY_REGIONAL_PLACEMENT_ENABLED: EnvironmentBooleanSchema, + ORCA_RELAY_REGION_CORRECTION_COHORT_PERCENT: z.coerce.number().int().min(0).max(100).default(0), ORCA_RELAY_PUBLIC_ASSIGNMENT_CONCURRENCY: z.coerce.number().int().positive().max(100).default(2), ORCA_RELAY_PUBLIC_STICKY_CONCURRENCY: z.coerce.number().int().positive().max(100).default(1), ORCA_RELAY_PUBLIC_STICKY_QUEUE_MAX: z.coerce.number().int().positive().max(4_096).default(64), @@ -185,6 +189,7 @@ export type RelayConfig = { databasePoolMax: number publicAssignmentsEnabled: boolean regionalPlacementEnabled?: boolean + regionCorrectionCohortPercent?: number publicAssignmentConcurrency: number publicAssignmentQueueMax: number publicAssignmentWaitMs: number @@ -332,6 +337,7 @@ export function loadRelayConfig(env: NodeJS.ProcessEnv = process.env): RelayConf databasePoolMax, publicAssignmentsEnabled: parsed.ORCA_RELAY_PUBLIC_ASSIGNMENTS_ENABLED, regionalPlacementEnabled: parsed.ORCA_RELAY_REGIONAL_PLACEMENT_ENABLED, + regionCorrectionCohortPercent: parsed.ORCA_RELAY_REGION_CORRECTION_COHORT_PERCENT, publicAssignmentConcurrency: parsed.ORCA_RELAY_PUBLIC_ASSIGNMENT_CONCURRENCY, publicAssignmentQueueMax: parsed.ORCA_RELAY_PUBLIC_ASSIGNMENT_QUEUE_MAX, publicAssignmentWaitMs: parsed.ORCA_RELAY_PUBLIC_ASSIGNMENT_WAIT_MS, diff --git a/cloud/apps/relay/src/database.test.ts b/cloud/apps/relay/src/database.test.ts index 56122def4be..0c987f95d40 100644 --- a/cloud/apps/relay/src/database.test.ts +++ b/cloud/apps/relay/src/database.test.ts @@ -18,6 +18,34 @@ afterEach(() => { }) describe('relay database', () => { + it('upgrades an existing SQLite relay without treating legacy controls as idle-capable', async () => { + const dataDir = mkdtempSync(join(tmpdir(), 'orca-idle-schema-')) + temporaryDirectories.push(dataDir) + const legacy = await openRelayDatabase({ dataDir }) + await legacy.query('ALTER TABLE relay_control_capabilities DROP COLUMN idle_regional_rehome') + await legacy.query('ALTER TABLE relay_region_rehome_attempts DROP COLUMN source_generation') + await legacy.query( + `INSERT INTO relay_control_capabilities + (user_id, relay_host_id, activity_id, cell_id, cell_incarnation, assignment_epoch, generation, finish_existing) + VALUES ('legacy-user', 'abcdefghijklmnop', 'control:source:1', 'source', 'legacy-incarnation', 1, 1, 1)` + ) + await legacy.close() + const upgraded = await openRelayDatabase({ dataDir }) + try { + expect( + await upgraded.query('SELECT idle_regional_rehome FROM relay_control_capabilities') + ).toEqual([{ idle_regional_rehome: 0 }]) + const columns = await upgraded.query( + "SELECT * FROM pragma_table_info('relay_region_rehome_attempts')" + ) + expect(columns.find((column) => column.name === 'source_generation')).toMatchObject({ + dflt_value: '0' + }) + } finally { + await upgraded.close() + } + }) + it('creates every durable relay state table', async () => { const database = await openInMemoryRelayDatabase() const rows = await database.query( @@ -54,6 +82,7 @@ describe('relay database', () => { 'relay_confirm_results', 'relay_confirmable_splices', 'relay_connection_bases', + 'relay_control_capabilities', 'relay_control_connection_reservations', 'relay_devices', 'relay_direct_authorizations', @@ -62,6 +91,7 @@ describe('relay database', () => { 'relay_migration_leases', 'relay_post_drain_migration_pins', 'relay_rate_windows', + 'relay_region_decisions', 'relay_region_rehome_attempts', 'relay_region_rehome_control', 'relay_region_rehome_worker_state' @@ -140,9 +170,7 @@ describe('relay database', () => { const second = await openRelayDatabase({ dataDir }) expect( - await second.query(`SELECT region FROM relay_cell_regions WHERE cell_id = ?`, [ - 'legacy-cell' - ]) + await second.query(`SELECT region FROM relay_cell_regions WHERE cell_id = ?`, ['legacy-cell']) ).toEqual([{ region: 'us-central1' }]) await second.close() }) @@ -165,9 +193,7 @@ describe('relay database', () => { 'relay_region_rehome_attempts' ]) expect(checked.every((row) => String(row.sql).includes(list))).toBe(true) - expect( - POSTGRES_SCHEMA_MIGRATIONS.some((statement) => statement.includes(list)) - ).toBe(true) + expect(POSTGRES_SCHEMA_MIGRATIONS.some((statement) => statement.includes(list))).toBe(true) await database.close() }) diff --git a/cloud/apps/relay/src/database.ts b/cloud/apps/relay/src/database.ts index d51f4e7a423..41ead67ea60 100644 --- a/cloud/apps/relay/src/database.ts +++ b/cloud/apps/relay/src/database.ts @@ -197,6 +197,24 @@ CREATE TABLE IF NOT EXISTS relay_assignment_region_preferences ( CREATE INDEX IF NOT EXISTS relay_assignment_region_preferences_observed ON relay_assignment_region_preferences(observed_at); +CREATE TABLE IF NOT EXISTS relay_region_decisions ( + user_id TEXT NOT NULL, relay_host_id TEXT NOT NULL, + generation BIGINT NOT NULL, expires_at BIGINT NOT NULL, + assignment_epoch BIGINT NOT NULL, incumbent_region TEXT NOT NULL, + policy_version BIGINT NOT NULL, outcome TEXT NOT NULL, + cohort_bucket BIGINT NOT NULL DEFAULT 0, + last_considered_at BIGINT NOT NULL DEFAULT 0, + preferred_region TEXT, observed_at BIGINT NOT NULL, report_json TEXT, + PRIMARY KEY (user_id, relay_host_id) +); +CREATE TABLE IF NOT EXISTS relay_control_capabilities ( + user_id TEXT NOT NULL, relay_host_id TEXT NOT NULL, activity_id TEXT NOT NULL, + cell_id TEXT NOT NULL, cell_incarnation TEXT NOT NULL, + assignment_epoch BIGINT NOT NULL, generation BIGINT NOT NULL, + finish_existing BIGINT NOT NULL, + idle_regional_rehome BIGINT NOT NULL DEFAULT 0, + PRIMARY KEY (user_id, relay_host_id, activity_id) +); CREATE TABLE IF NOT EXISTS relay_region_rehome_worker_state ( worker_id TEXT PRIMARY KEY, next_dispatch_at BIGINT NOT NULL, @@ -228,6 +246,7 @@ CREATE TABLE IF NOT EXISTS relay_region_rehome_attempts ( CHECK (preferred_region IN (${REGION_LIST})), source_cell_id TEXT NOT NULL, source_cell_incarnation TEXT NOT NULL, + source_generation BIGINT NOT NULL DEFAULT 0, target_cell_id TEXT NOT NULL, target_cell_incarnation TEXT NOT NULL, previous_epoch BIGINT NOT NULL, @@ -600,6 +619,8 @@ CREATE INDEX IF NOT EXISTS relay_audit_events_at ON relay_audit_events(at); // auto-named; the replacement is named, so both statements are no-ops on a // database the current schema created and neither can drop the other. export const POSTGRES_SCHEMA_MIGRATIONS = [ + `ALTER TABLE relay_region_decisions ADD COLUMN IF NOT EXISTS last_considered_at BIGINT NOT NULL DEFAULT 0`, + `ALTER TABLE relay_region_decisions ADD COLUMN IF NOT EXISTS cohort_bucket BIGINT NOT NULL DEFAULT 0`, `ALTER TABLE relay_region_rehome_attempts DROP CONSTRAINT IF EXISTS relay_region_rehome_attempts_preferred_region_check`, `ALTER TABLE relay_region_rehome_attempts @@ -607,7 +628,9 @@ export const POSTGRES_SCHEMA_MIGRATIONS = [ CHECK (preferred_region IN (${REGION_LIST}))`, `ALTER TABLE relay_region_rehome_control ADD COLUMN IF NOT EXISTS host_cooldown_ms BIGINT NOT NULL - DEFAULT ${REGIONAL_REHOME_DEFAULT_HOST_COOLDOWN_MS}` + DEFAULT ${REGIONAL_REHOME_DEFAULT_HOST_COOLDOWN_MS}`, + `ALTER TABLE relay_control_capabilities ADD COLUMN IF NOT EXISTS idle_regional_rehome BIGINT NOT NULL DEFAULT 0`, + `ALTER TABLE relay_region_rehome_attempts ADD COLUMN IF NOT EXISTS source_generation BIGINT NOT NULL DEFAULT 0` ] function postgresSql(sql: string): string { @@ -1013,6 +1036,15 @@ async function applySchema(database: RelayDatabase): Promise { for (const statement of SCHEMA.split(';')) { if (statement.trim()) await database.query(statement) } + for (const [table, column] of [ + ['relay_control_capabilities', 'idle_regional_rehome'], + ['relay_region_rehome_attempts', 'source_generation'] + ]) { + const columns = await database.query('SELECT name FROM pragma_table_info(?)', [table]) + if (!columns.some((existing) => existing.name === column)) { + await database.query(`ALTER TABLE ${table} ADD COLUMN ${column} BIGINT NOT NULL DEFAULT 0`) + } + } } // Why: DDL is not a request. A CREATE INDEX on a grown table legitimately runs diff --git a/cloud/apps/relay/src/host-session-client-accept.test.ts b/cloud/apps/relay/src/host-session-client-accept.test.ts index 0cec6531e3f..04f86c533e4 100644 --- a/cloud/apps/relay/src/host-session-client-accept.test.ts +++ b/cloud/apps/relay/src/host-session-client-accept.test.ts @@ -155,6 +155,66 @@ describe('client accept abandoned mid-DB-phase', () => { vi.useRealTimers() }) + it('does not admit new source work after a drain crosses activity acquisition', async () => { + const h = harness() + const control = await activeHost(h) + const slow = deferred() + h.acquireActivity.mockReturnValueOnce(slow.promise) + const client = new FakeSocket() + const capacity = { bind: vi.fn(), release: vi.fn() } + const accepting = h.registry.acceptClient( + client as unknown as WebSocket, + identity.relayHostId, + 'credential', + capacity + ) + await vi.advanceTimersByTimeAsync(0) + h.registry.drainHost({ + attemptId: 'attempt', + userId: identity.sub, + relayHostId: identity.relayHostId, + sourceAssignmentEpoch: 1, + graceMs: 60_000 + }) + slow.resolve() + await accepting + expect(control.send).not.toHaveBeenCalledWith(expect.stringContaining('conn-open')) + expect(capacity.bind).not.toHaveBeenCalled() + expect(client.close).toHaveBeenCalledWith(RELAY_CLOSE_CODE.WRONG_CELL, expect.any(String)) + expect(h.releaseActivity).toHaveBeenCalled() + }) + + it('does not splice an attachment whose generation retired during basis persistence', async () => { + const h = harness() + await activeHost(h) + const client = new FakeSocket() + await h.registry.acceptClient( + client as unknown as WebSocket, + identity.relayHostId, + 'credential' + ) + const session = h.registry.get({ userId: identity.sub, relayHostId: identity.relayHostId })! + const pending = [...session.pendingConns.values()][0]! + const slow = deferred() + h.store.recordConnectionBasis.mockReturnValueOnce(slow.promise) + const host = new FakeSocket() + const attaching = h.registry.acceptHostData( + host as unknown as WebSocket, + pending.connId, + pending.connTicket, + 1 + ) + await vi.advanceTimersByTimeAsync(0) + h.registry.drain(0) + await vi.advanceTimersByTimeAsync(0) + slow.resolve() + expect(await attaching).toBe(false) + expect(session.activeSplices.size).toBe(0) + expect(h.store.deactivateBasis).toHaveBeenCalledWith(pending.connId) + expect(client.send).not.toHaveBeenCalledWith(expect.stringContaining('\"ok\":true')) + expect(host.close).toHaveBeenCalled() + }) + it('stops after a slow activity acquire when the phone already hung up', async () => { const h = harness() const control = await activeHost(h) @@ -390,6 +450,7 @@ describe('successful client accept timing', () => { relayHostIdDigest: string } expect(event.credentialKind).toBe('resume') + expect(event).toMatchObject({ assignmentEpoch: 1, controlGeneration: 1, drainMode: 'none' }) // Joins the line back to the emitting process, like the runtime metrics event. expect(event).toMatchObject({ role: 'cell', cellId: config.cellId, region: 'us-central1' }) expect(Object.keys(event.stageMs).sort()).toEqual([ @@ -460,6 +521,9 @@ describe('control round-trip sampling', () => { cellId: config.cellId, region: 'us-central1', rttMsMedian: 40, + assignmentEpoch: 1, + controlGeneration: 1, + drainMode: 'none', sampleCount: 4 }) expect(rttLines()[0]).not.toContain(identity.relayHostId) diff --git a/cloud/apps/relay/src/host-session-registry.test.ts b/cloud/apps/relay/src/host-session-registry.test.ts index 920faa6f4b8..dcfcd3f36c5 100644 --- a/cloud/apps/relay/src/host-session-registry.test.ts +++ b/cloud/apps/relay/src/host-session-registry.test.ts @@ -4,6 +4,7 @@ import { CONTROL_CONTINUITY_LIMITS, RELAY_CLOSE_CODE, RELAY_HOST_CAPABILITY_PENDING_CONN_DETAILS, + RELAY_HOST_CAPABILITY_IDLE_REGIONAL_REHOME, RELAY_PROTOCOL_LIMITS } from '@orca-cloud/relay-contract' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' @@ -140,7 +141,10 @@ function createRegistry( store as RelayCredentialStore, assignments, new ProcessQueuedByteBudget(), - observer + observer, + Date.now, + Math.random, + 'incarnation-1' ) // Mirrors the production signature exactly so a future positional shift fails to compile. const bound = ( @@ -166,7 +170,14 @@ function createRegistry( assignmentEpoch, appVersion = '1.4.173' ) => bound(socket, identity, existing, generation, rebind, assignmentEpoch, appVersion) - return { registry, activate, acquireActivity, renewControlActivity, releaseActivity, observer } + return { + registry, + activate, + acquireActivity, + renewControlActivity, + releaseActivity, + observer + } } describe('host session cleanup races', () => { @@ -398,26 +409,20 @@ describe('host session cleanup races', () => { attemptId: '22222222-2222-4222-8222-222222222222' }) ).toThrow('regional_rehome_attempt_conflict') - expect(() => - registry.drainHost({ ...request, sourceAssignmentEpoch: 8 }) - ).toThrow('regional_rehome_assignment_epoch_mismatch') + expect(() => registry.drainHost({ ...request, sourceAssignmentEpoch: 8 })).toThrow( + 'regional_rehome_assignment_epoch_mismatch' + ) const rebound = new FakeSocket() - await activate( - rebound as unknown as WebSocket, - identity, - registry.get(request), - 1, - true, - 7 - ) + await activate(rebound as unknown as WebSocket, identity, registry.get(request), 1, true, 7) expect(registry.get(request)?.state).toBe('drain-only') expect(rebound.send).toHaveBeenCalledWith(expect.stringContaining('"type":"drain"')) await vi.advanceTimersByTimeAsync(30_000) expect(registry.get(request)).toBeNull() - expect(registry.get({ userId: secondIdentity.sub, relayHostId: secondIdentity.relayHostId })) - .not.toBeNull() + expect( + registry.get({ userId: secondIdentity.sub, relayHostId: secondIdentity.relayHostId }) + ).not.toBeNull() expect(secondSocket.close).not.toHaveBeenCalled() }) @@ -513,14 +518,7 @@ describe('host session cleanup races', () => { expect(original).not.toBeNull() const rebindSocket = new FakeSocket() - const rebinding = activate( - rebindSocket as unknown as WebSocket, - identity, - original, - 1, - true, - 1 - ) + const rebinding = activate(rebindSocket as unknown as WebSocket, identity, original, 1, true, 1) rebindSocket.close() blocked.resolve('control:production-gce-c3:1') await rebinding @@ -659,14 +657,7 @@ describe('host session cleanup races', () => { originalSocket.close() const replacementSocket = new FakeSocket() - await activate( - replacementSocket as unknown as WebSocket, - identity, - original, - 2, - false, - 1 - ) + await activate(replacementSocket as unknown as WebSocket, identity, original, 2, false, 1) const replacement = registry.get({ userId: identity.sub, relayHostId: identity.relayHostId @@ -696,14 +687,7 @@ describe('host session cleanup races', () => { }) expect(original).not.toBeNull() - await activate( - new FakeSocket() as unknown as WebSocket, - identity, - original, - 2, - false, - 1 - ) + await activate(new FakeSocket() as unknown as WebSocket, identity, original, 2, false, 1) vi.advanceTimersByTime(15_000) expect(renewControlActivity).toHaveBeenCalledOnce() @@ -718,6 +702,53 @@ describe('host session cleanup races', () => { vi.advanceTimersByTime(0) }) + it('ignores a denial belonging to the socket before a same-generation rebind', async () => { + const h = createRegistry(vi.fn().mockResolvedValue('control:production-gce-c3:1')) + const oldSocket = new FakeSocket() + await h.activate(oldSocket as unknown as WebSocket, identity, null, 1, false, 1) + const session = h.registry.get({ userId: identity.sub, relayHostId: identity.relayHostId })! + let reject!: (error: Error) => void + h.renewControlActivity.mockReturnValueOnce( + new Promise((_, fail) => { + reject = fail + }) + ) + await vi.advanceTimersByTimeAsync(15_000) + const replacement = new FakeSocket() + await h.activate(replacement as unknown as WebSocket, identity, session, 1, true, 1) + reject(new Error('activity_cell_not_authoritative')) + await vi.advanceTimersByTimeAsync(0) + expect(replacement.close).not.toHaveBeenCalled() + expect(session.socket).toBe(replacement) + expect(session.generation).toBe(1) + }) + + it('ignores missing-activity recovery denial after an authority transition', async () => { + const h = createRegistry(vi.fn().mockResolvedValue('control:production-gce-c3:1')) + const socket = new FakeSocket() + await h.activate(socket as unknown as WebSocket, identity, null, 1, false, 1) + const session = h.registry.get({ userId: identity.sub, relayHostId: identity.relayHostId })! + h.renewControlActivity.mockRejectedValueOnce(new Error('control_activity_not_found')) + let reject!: (error: Error) => void + h.acquireActivity.mockReturnValueOnce( + new Promise((_, fail) => { + reject = fail + }) + ) + await vi.advanceTimersByTimeAsync(15_000) + h.registry.drainHost({ + attemptId: 'attempt', + userId: identity.sub, + relayHostId: identity.relayHostId, + sourceAssignmentEpoch: 1, + graceMs: 60_000 + }) + reject(new Error('activity_cell_not_authoritative')) + await vi.advanceTimersByTimeAsync(0) + expect(socket.close).not.toHaveBeenCalled() + expect(session.state).toBe('drain-only') + }) + it('keeps 15s pings while halving steady-state control renewals', async () => { const activateControl = vi .fn() @@ -732,9 +763,7 @@ describe('host session cleanup races', () => { socket.emit('message', Buffer.from(JSON.stringify({ type: 'pong' })), false) } - const pings = socket.send.mock.calls.filter((call) => - String(call[0]).includes('"ping"') - ) + const pings = socket.send.mock.calls.filter((call) => String(call[0]).includes('"ping"')) expect(pings).toHaveLength(4) expect(renewControlActivity).toHaveBeenCalledTimes(2) const firstExpiry = Number(renewControlActivity.mock.calls[0]![1].expiresAt) @@ -1118,3 +1147,277 @@ describe('host hello ack pending connections', () => { expect(rebound.pendingConns).toEqual([DETAILED_ENTRY]) }) }) + +describe('source-owned idle cutover', () => { + beforeEach(() => vi.useFakeTimers()) + afterEach(() => { + vi.clearAllTimers() + vi.useRealTimers() + }) + const request = { + attemptId: 'idle-1', + userId: identity.sub, + relayHostId: identity.relayHostId, + sourceAssignmentEpoch: 1, + sourceGeneration: 1, + sourceCellIncarnation: 'incarnation-1', + targetCellId: 'target' + } + async function source(store: Partial = {}) { + const h = createRegistry(vi.fn().mockResolvedValue('control:1'), store) + const socket = new FakeSocket() + h.registry.acceptControl( + socket as unknown as WebSocket, + identity, + undefined, + new Set([RELAY_HOST_CAPABILITY_IDLE_REGIONAL_REHOME]) + ) + socket.removeAllListeners('message') + await h.activate(socket as unknown as WebSocket, identity, null, 1, false, 1) + return { ...h, socket, session: h.registry.get(request)! } + } + it('keeps either established client busy until both actually leave', async () => { + const h = await source() + h.session.activeConnIds.add('phone') + h.session.activeConnIds.add('ipad') + const commit = vi.fn().mockResolvedValue({ outcome: 'committed' }) + h.session.activeConnIds.delete('ipad') + expect( + await h.registry.idleRehome(request, commit, vi.fn().mockResolvedValue('not-committed')) + ).toEqual({ outcome: 'busy' }) + expect(commit).not.toHaveBeenCalled() + h.session.activeConnIds.delete('phone') + expect( + await h.registry.idleRehome(request, commit, vi.fn().mockResolvedValue('not-committed')) + ).toEqual({ outcome: 'committed' }) + expect(h.socket.close).toHaveBeenCalledWith(RELAY_CLOSE_CODE.DRAINING, expect.any(String)) + expect(h.releaseActivity).toHaveBeenCalled() + }) + it.each([ + { userId: 'other-user' }, + { sourceAssignmentEpoch: 2 }, + { sourceGeneration: 2 }, + { sourceCellIncarnation: 'other-incarnation' }, + { targetCellId: 'other-target' } + ])('rejects a reused operation ID with changed authority %j', async (change) => { + const h = await source() + const result = deferred<{ outcome: 'deferred' }>() + const commit = vi.fn().mockReturnValue(result.promise) + const reconcile = vi.fn().mockResolvedValue('not-committed') + const moving = h.registry.idleRehome(request, commit, reconcile) + const conflicting = h.registry.idleRehome({ ...request, ...change }, commit, reconcile) + result.resolve({ outcome: 'deferred' }) + expect(await conflicting).toEqual({ outcome: 'stale' }) + expect(await moving).toEqual({ outcome: 'deferred' }) + expect(commit).toHaveBeenCalledOnce() + expect(h.socket.close).not.toHaveBeenCalled() + }) + it('accounts for accepts before credential identity resolves', async () => { + const lookup = deferred() + const h = await source({ + resolveResume: vi.fn().mockReturnValue(lookup.promise), + resolveInviteForMove: vi.fn().mockResolvedValue(null) + }) + const client = new FakeSocket() + const accept = h.registry.acceptClient( + client as unknown as WebSocket, + identity.relayHostId, + 'credential' + ) + expect( + await h.registry.idleRehome(request, vi.fn(), vi.fn().mockResolvedValue('not-committed')) + ).toEqual({ outcome: 'busy' }) + lookup.resolve(null) + await accept + expect(h.socket.close).not.toHaveBeenCalled() + }) + it('rejects new accepts and replacements synchronously while a commit awaits', async () => { + const h = await source() + const result = deferred<{ outcome: 'deferred' }>() + const commit = vi.fn().mockReturnValue(result.promise) + const moving = h.registry.idleRehome( + request, + commit, + vi.fn().mockResolvedValue('not-committed') + ) + const duplicate = h.registry.idleRehome( + request, + commit, + vi.fn().mockResolvedValue('not-committed') + ) + const client = new FakeSocket() + const release = vi.fn() + await h.registry.acceptClient( + client as unknown as WebSocket, + identity.relayHostId, + 'credential', + { release } as never + ) + expect(client.close).toHaveBeenCalledWith(RELAY_CLOSE_CODE.WRONG_CELL, expect.any(String)) + expect(release).toHaveBeenCalledOnce() + const replacement = new FakeSocket() + await h.activate(replacement as unknown as WebSocket, identity, h.session, 2, false, 1) + expect(replacement.close).toHaveBeenCalledWith(RELAY_CLOSE_CODE.WRONG_CELL, expect.any(String)) + result.resolve({ outcome: 'deferred' }) + await moving + await duplicate + expect(commit).toHaveBeenCalledOnce() + expect(h.socket.close).not.toHaveBeenCalled() + expect( + await h.registry.idleRehome( + { ...request, attemptId: 'next' }, + vi.fn().mockResolvedValue({ outcome: 'committed' }), + vi.fn().mockResolvedValue('not-committed') + ) + ).toEqual({ outcome: 'committed' }) + }) + it.each(['ambiguous', 'deferred'])( + 'keeps %s outcomes fenced until locked reconciliation succeeds', + async (claim) => { + const h = await source() + const reconcile = vi + .fn() + .mockRejectedValueOnce(new Error('database unavailable')) + .mockRejectedValueOnce(new Error('database unavailable')) + .mockResolvedValue('not-committed') + const moving = h.registry.idleRehome( + request, + claim === 'ambiguous' + ? vi.fn().mockRejectedValue(new Error('lost commit reply')) + : vi.fn().mockResolvedValue({ outcome: 'deferred' }), + reconcile + ) + await vi.advanceTimersByTimeAsync(50) + expect( + await h.registry.idleRehome( + { ...request, attemptId: 'other' }, + vi.fn(), + vi.fn().mockResolvedValue('not-committed') + ) + ).toEqual({ outcome: 'busy' }) + expect(h.socket.close).not.toHaveBeenCalled() + await vi.advanceTimersByTimeAsync(300) + expect(await moving).toEqual({ outcome: 'deferred' }) + expect(reconcile).toHaveBeenCalledTimes(3) + expect(h.socket.close).not.toHaveBeenCalled() + } + ) + it('owns accepted control mutations before the handler first awaits', async () => { + const mutation = deferred() + const h = await source() + ;(h.registry as unknown as { verifyRelayToken: unknown }).verifyRelayToken = vi + .fn() + .mockReturnValue(mutation.promise) + h.socket.emit( + 'message', + Buffer.from(JSON.stringify({ type: 'auth-refresh', relayJwt: 'token' })), + false + ) + const commit = vi.fn().mockResolvedValue({ outcome: 'deferred' }) + expect( + await h.registry.idleRehome(request, commit, vi.fn().mockResolvedValue('not-committed')) + ).toEqual({ outcome: 'busy' }) + mutation.resolve(identity) + await vi.advanceTimersByTimeAsync(0) + expect( + await h.registry.idleRehome(request, commit, vi.fn().mockResolvedValue('not-committed')) + ).toEqual({ outcome: 'deferred' }) + }) + it('owns queued replacement activation before its first persistence await', async () => { + const h = await source() + const activation = deferred() + const assignments = (h.registry as unknown as { assignments: { activateControl: unknown } }) + .assignments + assignments.activateControl = vi.fn().mockReturnValue(activation.promise) + const replacement = new FakeSocket() + const activating = h.activate( + replacement as unknown as WebSocket, + identity, + h.session, + 2, + false, + 1 + ) + expect( + await h.registry.idleRehome(request, vi.fn(), vi.fn().mockResolvedValue('not-committed')) + ).toEqual({ outcome: 'busy' }) + activation.resolve('control:2') + await activating + }) + it('retires changed authority even when the claim definitively deferred', async () => { + const h = await source() + expect( + await h.registry.idleRehome( + request, + vi.fn().mockResolvedValue({ outcome: 'deferred' }), + vi.fn().mockResolvedValue('stale') + ) + ).toEqual({ outcome: 'stale' }) + expect(h.session.state).toBe('closed') + expect(h.releaseActivity).toHaveBeenCalled() + }) + it('holds attach ownership through basis failure reservation cleanup', async () => { + const basis = deferred() + const cleanup = deferred() + const h = await source({ + recordConnectionBasis: vi.fn().mockImplementation(async () => { + await basis.promise + throw new Error('basis failed') + }), + failReservation: vi.fn().mockReturnValue(cleanup.promise) + }) + const client = new FakeSocket() + h.session.pendingConns.set('conn', { + connId: 'conn', + connTicket: 'ticket', + client: client as unknown as WebSocket, + reservation: { + userId: identity.sub, + relayHostId: identity.relayHostId, + credentialKind: 'invite', + leaseExpiresAt: Date.now() + 1000 + }, + attachTimer: setTimeout(() => {}, 1000), + credentialActivityId: null + } as never) + const attached = h.registry.acceptHostData( + new FakeSocket() as unknown as WebSocket, + 'conn', + 'ticket', + 1 + ) + const commit = vi.fn().mockResolvedValue({ outcome: 'deferred' }) + expect(await h.registry.idleRehome(request, commit, vi.fn())).toEqual({ outcome: 'busy' }) + basis.resolve() + await vi.advanceTimersByTimeAsync(0) + expect(h.session.activeConnIds.size).toBe(0) + expect(await h.registry.idleRehome(request, commit, vi.fn())).toEqual({ outcome: 'busy' }) + expect(commit).not.toHaveBeenCalled() + cleanup.resolve() + await attached + }) + it('returns the durable operation outcome after source retirement', async () => { + const h = await source() + const commit = vi.fn().mockResolvedValue({ outcome: 'committed' }) + await h.registry.idleRehome(request, commit, vi.fn()) + expect( + await h.registry.idleRehome(request, commit, vi.fn().mockResolvedValue('committed')) + ).toEqual({ outcome: 'committed' }) + expect(commit).toHaveBeenCalledOnce() + }) + it('does not reopen a source overtaken by emergency drain', async () => { + const h = await source() + const result = deferred<{ outcome: 'deferred' }>() + const moving = h.registry.idleRehome( + request, + () => result.promise, + vi.fn().mockResolvedValue('not-committed') + ) + h.registry.drain(0) + await vi.advanceTimersByTimeAsync(0) + result.resolve({ outcome: 'deferred' }) + await moving + expect(h.session.state).toBe('closed') + expect(h.registry.get(request)).toBeNull() + }) +}) diff --git a/cloud/apps/relay/src/host-session-registry.ts b/cloud/apps/relay/src/host-session-registry.ts index 3b4e616a692..61a1b706fa9 100644 --- a/cloud/apps/relay/src/host-session-registry.ts +++ b/cloud/apps/relay/src/host-session-registry.ts @@ -15,6 +15,7 @@ import { HostHelloSchema, InviteCreateSchema, RELAY_HOST_CAPABILITY_PENDING_CONN_DETAILS, + RELAY_HOST_CAPABILITY_IDLE_REGIONAL_REHOME, RELAY_PROTOCOL_LIMITS, RELAY_CLOSE_CODE, type RelayHostCloseReason, @@ -25,10 +26,7 @@ import type WebSocket from 'ws' import type { RawData } from 'ws' import type { RelayConfig } from './config.js' import type { RelayAssignmentStore } from './assignment-store.js' -import { - RelayCredentialStore, - type CredentialReservation -} from './credential-store.js' +import { RelayCredentialStore, type CredentialReservation } from './credential-store.js' import { HostCloseReasonMemory } from './host-close-reason-memory.js' import { relayHostLogDigest } from './relay-host-log-digest.js' import type { RelayTokenClaims } from './relay-token-verifier.js' @@ -78,7 +76,7 @@ export type HostSession = { identity: RelayTokenClaims readonly relayHostId: string readonly generation: number - readonly assignmentEpoch: number + assignmentEpoch: number readonly controlActivityId: string | null readonly controlResumeSecret: string // Why: reconnect churn is only actionable once it can be pinned to a client build. @@ -94,6 +92,7 @@ export type HostSession = { pendingPingAt: number | null controlRttSamplesMs: number[] controlRttLoggedAt: number | null + authorityRevision: number activityRenewalDueAt: number activityRenewalAttempt: number activityRenewalCompletedAttempt: number @@ -108,10 +107,7 @@ export type HostSession = { regionalDrainExpiresAt: number | null } -export type RegionalHostDrainOutcome = - | 'accepted' - | 'already-accepted' - | 'host-not-connected' +export type RegionalHostDrainOutcome = 'accepted' | 'already-accepted' | 'host-not-connected' type PendingConnection = { connId: string @@ -184,6 +180,113 @@ export class HostSessionRegistry { private readonly hostCapabilities = new WeakMap>() private draining = false + private readonly idleWork = new Map() + private readonly idleAttempts = new Map< + string, + { + attemptId: string + authorityKey: string + promise: Promise<{ outcome: 'committed' | 'deferred' | 'stale' }> + } + >() + + async idleRehome( + input: { + attemptId: string + userId: string + relayHostId: string + sourceAssignmentEpoch: number + sourceGeneration: number + sourceCellIncarnation: string + targetCellId: string + }, + commit: () => Promise<{ outcome: 'committed' | 'deferred' | 'stale' }>, + reconcile: () => Promise<'committed' | 'not-committed' | 'stale'> + ): Promise<{ outcome: 'busy' | 'committed' | 'deferred' | 'stale' }> { + const authorityKey = JSON.stringify([ + input.userId, + input.sourceAssignmentEpoch, + input.sourceGeneration, + input.sourceCellIncarnation, + input.targetCellId + ]) + const prior = this.idleAttempts.get(input.relayHostId) + if (prior) { + if (prior.attemptId !== input.attemptId) return { outcome: 'busy' } + return prior.authorityKey === authorityKey ? prior.promise : { outcome: 'stale' } + } + const session = this.get(input) + if ( + this.draining || + !session || + session.state !== 'active' || + session.generation !== input.sourceGeneration || + session.assignmentEpoch !== input.sourceAssignmentEpoch || + this.cellIncarnation !== input.sourceCellIncarnation + ) { + const durable = await reconcile() + return { outcome: durable === 'committed' ? 'committed' : 'stale' } + } + if ( + !session.socket || + !this.hostCapabilities.get(session.socket)?.has(RELAY_HOST_CAPABILITY_IDLE_REGIONAL_REHOME) + ) + return { outcome: 'deferred' } + if ( + (this.idleWork.get(input.relayHostId) ?? 0) !== 0 || + session.activeConnIds.size !== 0 || + session.activeSplices.size !== 0 || + session.pendingConns.size !== 0 + ) + return { outcome: 'busy' } + const revision = session.authorityRevision + const promise = Promise.resolve().then(async () => { + let outcome: 'committed' | 'deferred' | 'stale' + try { + outcome = (await commit()).outcome + if (outcome === 'deferred') { + const durable = await reconcile() + outcome = durable === 'not-committed' ? 'deferred' : durable + } + } catch { + let delay = 100 + for (;;) { + try { + const durable = await reconcile() + outcome = durable === 'not-committed' ? 'deferred' : durable + break + } catch { + await new Promise((resolve) => { + const timer = setTimeout(resolve, delay) + timer.unref?.() + }) + delay = Math.min(delay * 2, 5000) + } + } + } + if (this.get(input) === session) { + if (outcome !== 'deferred' || this.draining || session.authorityRevision !== revision) { + this.closeDrainedSession(session) + } + } + if (this.idleAttempts.get(input.relayHostId)?.promise === promise) + this.idleAttempts.delete(input.relayHostId) + return { outcome } + }) + this.idleAttempts.set(input.relayHostId, { attemptId: input.attemptId, authorityKey, promise }) + return promise + } + + private beginIdleWork(hostId: string): (() => void) | null { + if (this.idleAttempts.has(hostId)) return null + this.idleWork.set(hostId, (this.idleWork.get(hostId) ?? 0) + 1) + return () => { + const remaining = (this.idleWork.get(hostId) ?? 1) - 1 + if (remaining === 0) this.idleWork.delete(hostId) + else this.idleWork.set(hostId, remaining) + } + } + constructor( private readonly config: RelayConfig, private readonly verifyRelayToken: VerifyRelayToken, @@ -192,7 +295,8 @@ export class HostSessionRegistry { private readonly queuedByteBudget: ProcessQueuedByteBudget, private readonly observer: RelayRuntimeObserver, private readonly now: () => number = Date.now, - private readonly random: () => number = Math.random + private readonly random: () => number = Math.random, + private readonly cellIncarnation?: string ) {} // Uniform over [CONTROL_LEASE_MS - jitter, CONTROL_LEASE_MS + jitter). @@ -206,6 +310,25 @@ export class HostSessionRegistry { hostId: string, credential: string, capacityReservation?: PendingHostDataReservation + ): Promise { + const release = this.beginIdleWork(hostId) + if (!release) { + capacityReservation?.release() + this.rejectClient(socket, RELAY_CLOSE_CODE.WRONG_CELL) + return + } + try { + await this.acceptClientUnfenced(socket, hostId, credential, capacityReservation) + } finally { + release() + } + } + + private async acceptClientUnfenced( + socket: WebSocket, + hostId: string, + credential: string, + capacityReservation?: PendingHostDataReservation ): Promise { if (this.draining) { capacityReservation?.release() @@ -295,6 +418,7 @@ export class HostSessionRegistry { this.rejectClient(socket, RELAY_CLOSE_CODE.LIMIT_EXCEEDED) return } + const admittingSocket = session.socket const connId = randomUUID() const connTicket = randomBytes(32).toString('base64url') const identity = { userId: reservation.userId, relayHostId: hostId } @@ -324,6 +448,20 @@ export class HostSessionRegistry { ) { return } + // Admission may have crossed a drain or control replacement while persisting activity. + if ( + this.draining || + this.sessions.get(sessionKey) !== session || + session.state !== 'active' || + session.socket !== admittingSocket || + admittingSocket.readyState !== admittingSocket.OPEN + ) { + capacityReservation?.release() + this.failReservationBestEffort(reservation) + if (credentialActivityId) this.releaseActivityBestEffort(identity, credentialActivityId) + this.rejectClient(socket, RELAY_CLOSE_CODE.WRONG_CELL) + return + } markStage('activity') const attachTimer = setTimeout(() => { session.pendingConns.delete(connId) @@ -372,6 +510,27 @@ export class HostSessionRegistry { connId: string, connTicket: string, generation: number + ): Promise { + const owner = [...this.sessions.values()].find((candidate) => + candidate.pendingConns.has(connId) + ) + const release = owner ? this.beginIdleWork(owner.relayHostId) : () => {} + if (!release) { + socket.close(RELAY_CLOSE_CODE.WRONG_CELL, 'idle cutover in progress') + return false + } + try { + return await this.acceptHostDataUnfenced(socket, connId, connTicket, generation) + } finally { + release() + } + } + + private async acceptHostDataUnfenced( + socket: WebSocket, + connId: string, + connTicket: string, + generation: number ): Promise { const session = [...this.sessions.values()].find((candidate) => candidate.pendingConns.has(connId) @@ -427,6 +586,27 @@ export class HostSessionRegistry { socket.close(RELAY_CLOSE_CODE.LIMIT_EXCEEDED, 'basis persistence failed') return false } + // Already admitted attachments may finish a regional drain, but never a retired generation. + if ( + this.draining || + this.sessions.get(this.key(identity.userId, identity.relayHostId)) !== session || + this.get(identity)?.state === 'closed' || + !session.activeConnIds.has(connId) || + socket.readyState !== socket.OPEN || + pending.client.readyState !== pending.client.OPEN + ) { + session.activeConnIds.delete(connId) + pending.capacityReservation?.release() + this.deactivateBasisBestEffort(connId) + this.failReservationBestEffort(pending.reservation) + if (spliceActivityId) this.releaseActivityBestEffort(identity, spliceActivityId) + if (pending.credentialActivityId) { + this.releaseActivityBestEffort(identity, pending.credentialActivityId) + } + this.rejectClient(pending.client, RELAY_CLOSE_CODE.DRAINING) + socket.close(RELAY_CLOSE_CODE.DRAINING, 'host retired during attachment') + return false + } const close = wireSplice({ client: pending.client, host: socket, @@ -505,6 +685,7 @@ export class HostSessionRegistry { JSON.stringify({ event: 'orca_relay_client_accept_completed', ...this.logIdentity(), + ...this.sessionPlacementLogFields(session), credentialKind: pending.reservation.credentialKind, stageMs, totalMs, @@ -513,6 +694,14 @@ export class HostSessionRegistry { ) } + private sessionPlacementLogFields(session: HostSession) { + return { + assignmentEpoch: session.assignmentEpoch, + controlGeneration: session.generation, + drainMode: session.regionalDrainAttemptId ? 'deadline' : 'none' + } + } + // Matches the runtime metrics event so a log line and a metric point can be // joined back to the process that emitted them. private logIdentity(): { role: string; cellId: string; region: RelayRegion } { @@ -549,6 +738,7 @@ export class HostSessionRegistry { JSON.stringify({ event: 'orca_relay_host_control_rtt', ...this.logIdentity(), + ...this.sessionPlacementLogFields(session), relayHostIdDigest: relayHostLogDigest(session.relayHostId), rttMsMedian: percentile(samples, 0.5), sampleCount: samples.length @@ -562,6 +752,10 @@ export class HostSessionRegistry { connectionInclusionWatermark?: number, hostCapabilities?: ReadonlySet ): void { + if (this.idleAttempts.has(identity.relayHostId)) { + socket.close(RELAY_CLOSE_CODE.WRONG_CELL, 'idle cutover in progress') + return + } // Keyed by socket, not session: a rebind swaps the session's socket, and the // successor's own advertisement is the only one that describes its decoder. if (hostCapabilities?.size) this.hostCapabilities.set(socket, hostCapabilities) @@ -602,8 +796,7 @@ export class HostSessionRegistry { socket: WebSocket | null, context: string ): void { - void Promise.resolve() - .then(task) + void (async () => task())() .catch((error: unknown) => { const message = (error instanceof Error ? error.message : 'unknown') // Untruncated, unlike peer-supplied close reasons: this is the @@ -653,6 +846,7 @@ export class HostSessionRegistry { this.draining = true for (const session of this.sessions.values()) { if (session.state === 'closed') continue + session.authorityRevision += 1 session.state = 'drain-only' if (session.socket) send(session.socket, 'drain', { graceMs, recovery: 'resolve-director' }) setTimeout(() => this.closeDrainedSession(session), graceMs) @@ -665,7 +859,8 @@ export class HostSessionRegistry { relayHostId: string sourceAssignmentEpoch: number graceMs: number - }): RegionalHostDrainOutcome { + sourceCellIncarnation?: string + }): RegionalHostDrainOutcome | Promise { const session = this.get(input) if (!session || session.state === 'closed') return 'host-not-connected' if (session.assignmentEpoch !== input.sourceAssignmentEpoch) { @@ -678,13 +873,11 @@ export class HostSessionRegistry { this.reassertRegionalDrain(session) return 'already-accepted' } + session.authorityRevision += 1 session.regionalDrainAttemptId = input.attemptId session.regionalDrainExpiresAt = this.now() + input.graceMs this.reassertRegionalDrain(session) - session.regionalDrainTimer = setTimeout( - () => this.closeDrainedSession(session), - input.graceMs - ) + session.regionalDrainTimer = setTimeout(() => this.closeDrainedSession(session), input.graceMs) return 'accepted' } @@ -740,9 +933,9 @@ export class HostSessionRegistry { const existing = this.sessions.get(key) const rebind = Boolean( existing && - hello.data.controlResumeSecret && - hello.data.controlResumeSecret === existing.controlResumeSecret && - (existing.state === 'orphaned' || existing.state === 'active') + hello.data.controlResumeSecret && + hello.data.controlResumeSecret === existing.controlResumeSecret && + (existing.state === 'orphaned' || existing.state === 'active') ) const generation = rebind ? existing!.generation : (existing?.generation ?? 0) + 1 const ephemeral = nacl.box.keyPair() @@ -784,7 +977,9 @@ export class HostSessionRegistry { }, 10_000) socket.once('message', (raw, isBinary) => { clearTimeout(proofTimer) - const ack = isBinary ? null : HostChallengeAckSchema.safeParse(payload(raw, 'host-challenge-ack')) + const ack = isBinary + ? null + : HostChallengeAckSchema.safeParse(payload(raw, 'host-challenge-ack')) const proof = ack?.success ? decodeCanonicalBase64(ack.data.proofB64, 32) : null if ( !ack?.success || @@ -826,6 +1021,11 @@ export class HostSessionRegistry { appVersion: string, connectionInclusionWatermark?: number ): Promise { + const release = this.beginIdleWork(identity.relayHostId) + if (!release) { + socket.close(RELAY_CLOSE_CODE.WRONG_CELL, 'idle cutover in progress') + return Promise.resolve() + } const key = this.key(identity.sub, identity.relayHostId) const previous = this.activationQueues.get(key) ?? Promise.resolve() // The timeout only fails this waiting socket; the queue entry still chains @@ -836,26 +1036,29 @@ export class HostSessionRegistry { socket.close(RELAY_CLOSE_CODE.LIMIT_EXCEEDED, 'control activation queue stalled') }, ACTIVATION_QUEUE_WAIT_MS) queueWaitTimer.unref?.() - const activation = previous.catch(() => undefined).then(async () => { - clearTimeout(queueWaitTimer) - if (queueWaitExpired) return - if ((this.sessions.get(key) ?? null) !== existing) { - socket.close(RELAY_CLOSE_CODE.PEER_DROPPED, 'control activation superseded') - return - } - await this.activateCurrent( - socket, - identity, - existing, - generation, - rebind, - assignmentEpoch, - appVersion, - connectionInclusionWatermark - ) - }) + const activation = previous + .catch(() => undefined) + .then(async () => { + clearTimeout(queueWaitTimer) + if (queueWaitExpired) return + if ((this.sessions.get(key) ?? null) !== existing) { + socket.close(RELAY_CLOSE_CODE.PEER_DROPPED, 'control activation superseded') + return + } + await this.activateCurrent( + socket, + identity, + existing, + generation, + rebind, + assignmentEpoch, + appVersion, + connectionInclusionWatermark + ) + }) this.activationQueues.set(key, activation) const cleanup = (): void => { + release() if (this.activationQueues.get(key) === activation) this.activationQueues.delete(key) } void activation.then(cleanup, cleanup) @@ -881,7 +1084,11 @@ export class HostSessionRegistry { cellId: this.config.cellId, assignmentEpoch, generation, - connectionInclusionWatermark + connectionInclusionWatermark, + idleRegionalRehome: + this.hostCapabilities.get(socket)?.has(RELAY_HOST_CAPABILITY_IDLE_REGIONAL_REHOME) ?? + false, + cellIncarnation: this.cellIncarnation } ) await this.assignments.markMigrationTargetRegistered( @@ -918,14 +1125,15 @@ export class HostSessionRegistry { const previousSocket = existing.socket if (existing.orphanTimer) clearTimeout(existing.orphanTimer) existing.orphanTimer = null + existing.authorityRevision += 1 + existing.assignmentEpoch = assignmentEpoch existing.socket = socket existing.state = existing.regionalDrainAttemptId ? 'drain-only' : 'active' existing.appVersion = appVersion existing.leaseExpiresAt = this.controlLeaseExpiresAt() existing.lastPongAt = this.now() existing.pendingPingAt = null - existing.activityRenewalDueAt = - this.now() + RELAY_PROTOCOL_LIMITS.controlPingIntervalMs + existing.activityRenewalDueAt = this.now() + RELAY_PROTOCOL_LIMITS.controlPingIntervalMs this.wireActiveControl(existing) this.sendHelloAck(existing) if (existing.regionalDrainAttemptId) this.reassertRegionalDrain(existing) @@ -980,6 +1188,7 @@ export class HostSessionRegistry { pendingPingAt: null, controlRttSamplesMs: [], controlRttLoggedAt: null, + authorityRevision: 0, activityRenewalDueAt: this.now() + RELAY_PROTOCOL_LIMITS.controlPingIntervalMs, activityRenewalAttempt: 0, activityRenewalCompletedAttempt: 0, @@ -1028,7 +1237,9 @@ export class HostSessionRegistry { ` splices=${session.closingCounts?.splices ?? session.activeSplices.size}` + ` pending=${session.closingCounts?.pending ?? session.pendingConns.size}` + ` code=${code} reason=${JSON.stringify(printableCloseReason(reason))}` + - (socketError === null ? '' : ` error=${JSON.stringify(printableCloseReason(socketError))}`) + (socketError === null + ? '' + : ` error=${JSON.stringify(printableCloseReason(socketError))}`) ) }) socket.on('message', (raw, isBinary) => { @@ -1077,6 +1288,19 @@ export class HostSessionRegistry { } private async acceptRefresh(session: HostSession, raw: RawData): Promise { + const release = this.beginIdleWork(session.relayHostId) + if (!release) { + session.socket?.close(RELAY_CLOSE_CODE.DRAINING, 'resolve-director') + return + } + try { + await this.acceptRefreshUnfenced(session, raw) + } finally { + release() + } + } + + private async acceptRefreshUnfenced(session: HostSession, raw: RawData): Promise { const parsed = AuthRefreshSchema.safeParse(payload(raw, 'auth-refresh')) if (!parsed.success) return const refreshed = await this.verifyRelayToken(parsed.data.relayJwt) @@ -1107,6 +1331,16 @@ export class HostSessionRegistry { if (controlActivityId && now >= session.activityRenewalDueAt) { const attempt = ++session.activityRenewalAttempt const startedAt = now + const socket = session.socket + const authorityRevision = session.authorityRevision + const current = (): boolean => + this.sessions.get(key) === session && + session.state !== 'closed' && + session.socket === socket && + socket.readyState === socket.OPEN && + session.controlActivityId === controlActivityId && + session.authorityRevision === authorityRevision && + attempt > session.activityRenewalCompletedAttempt void this.assignments .renewControlActivity( { userId: session.identity.sub, relayHostId: session.relayHostId }, @@ -1117,11 +1351,16 @@ export class HostSessionRegistry { } ) .then(() => { - if (attempt <= session.activityRenewalCompletedAttempt) return + if (!current()) return session.activityRenewalCompletedAttempt = attempt session.activityRenewalDueAt = startedAt + CONTROL_ACTIVITY_RENEWAL_INTERVAL_MS }) .catch(async (error: unknown) => { + if (!current()) return + if (error instanceof Error && error.message === 'assignment_not_found') { + socket.close(RELAY_CLOSE_CODE.DRAINING, 'control assignment missing') + return + } if (error instanceof Error && error.message === 'activity_cell_not_authoritative') { // Completion fences a late drain-only heartbeat after all source work is gone. session.socket?.close(RELAY_CLOSE_CODE.DRAINING, 'control migration completed') @@ -1144,8 +1383,22 @@ export class HostSessionRegistry { cellId: this.config.cellId } ) + if (!current()) { + // A replaced activity must not remain leased after its owner disappears. + if ( + !this.sessions.get(key) || + this.sessions.get(key)?.controlActivityId !== controlActivityId + ) { + this.releaseActivityBestEffort( + { userId: session.identity.sub, relayHostId: session.relayHostId }, + controlActivityId + ) + } + return + } this.observer.recordControlActivityRecovery?.(true) } catch (acquireError: unknown) { + if (!current()) return this.observer.recordControlActivityRecovery?.(false) if ( acquireError instanceof Error && @@ -1218,6 +1471,21 @@ export class HostSessionRegistry { private closeDrainedSession(session: HostSession): void { if (session.state === 'closed') return + const forcedConnections = session.activeConnIds.size + session.pendingConns.size + if (forcedConnections > 0) { + console.warn( + JSON.stringify({ + event: 'orca_relay_host_drain_forced_close', + ...this.logIdentity(), + ...this.sessionPlacementLogFields(session), + relayHostIdDigest: relayHostLogDigest(session.relayHostId), + reason: this.draining ? 'emergency' : 'regional-deadline', + forcedConnections, + splices: session.activeSplices.size, + pending: session.pendingConns.size + }) + ) + } if (session.heartbeatTimer) clearInterval(session.heartbeatTimer) if (session.orphanTimer) clearTimeout(session.orphanTimer) if (session.regionalDrainTimer) clearTimeout(session.regionalDrainTimer) @@ -1250,11 +1518,7 @@ export class HostSessionRegistry { session.pendingConns.clear() session.state = 'closed' if (session.socket) { - closeRelayWebSocket( - session.socket, - RELAY_CLOSE_CODE.DRAINING, - 'resolve configured director' - ) + closeRelayWebSocket(session.socket, RELAY_CLOSE_CODE.DRAINING, 'resolve configured director') } const key = this.key(session.identity.sub, session.relayHostId) if (this.sessions.get(key) === session) this.sessions.delete(key) @@ -1278,6 +1542,23 @@ export class HostSessionRegistry { session: HostSession, type: unknown, raw: RawData + ): Promise { + const release = this.beginIdleWork(session.relayHostId) + if (!release) { + session.socket?.close(RELAY_CLOSE_CODE.DRAINING, 'resolve-director') + return + } + try { + await this.acceptControlCommandUnfenced(session, type, raw) + } finally { + release() + } + } + + private async acceptControlCommandUnfenced( + session: HostSession, + type: unknown, + raw: RawData ): Promise { if (typeof type !== 'string' || !session.socket) return try { @@ -1315,10 +1596,7 @@ export class HostSessionRegistry { } if (type === 'device-credential-install') { const request = DeviceCredentialInstallSchema.parse(payload(raw, type)) - if ( - session.state !== 'active' && - request.authorization.mode === 'authenticated-direct' - ) { + if (session.state !== 'active' && request.authorization.mode === 'authenticated-direct') { throw new Error('authorization_expired') } const installActivityId = `install:${request.reqId}` diff --git a/cloud/apps/relay/src/idle-regional-rehome-reconciliation.test.ts b/cloud/apps/relay/src/idle-regional-rehome-reconciliation.test.ts new file mode 100644 index 00000000000..6fc2017ab48 --- /dev/null +++ b/cloud/apps/relay/src/idle-regional-rehome-reconciliation.test.ts @@ -0,0 +1,103 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.js' +import type { RelayDatabase } from './database.js' +import { openIdleRehomeTestDatabase } from './idle-regional-rehome-test-database.js' + +const request = { + v: 1 as const, + attemptId: '33333333-3333-4333-8333-333333333333', + userId: 'idle-reconciliation-test', + relayHostId: 'abcdefghijklmnop', + sourceCellId: 'source', + sourceCellIncarnation: '11111111-1111-4111-8111-111111111111', + sourceAssignmentEpoch: 1, + sourceGeneration: 7, + targetCellId: 'target' +} +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) + const store = new RelayAssignmentStore(database, () => 100_000_000) + await store.reconcileCells([ + { id: 'source', url: 'https://source.example.test', capacityRequests: 100 }, + { id: 'target', url: 'https://target.example.test', capacityRequests: 100 } + ]) + await store.assign(request) + await store.activateControl(request, { + cellId: request.sourceCellId, + assignmentEpoch: request.sourceAssignmentEpoch, + generation: request.sourceGeneration, + cellIncarnation: request.sourceCellIncarnation + }) + return { database, store } +} + +describe('idle cutover durable reconciliation', () => { + it('only permits reopening when the exact source still owns the assignment', async () => { + const { store } = await setup() + expect(await store.reconcileIdleRegionalRehome(request)).toBe('not-committed') + await store.activateControl(request, { + cellId: request.sourceCellId, + assignmentEpoch: request.sourceAssignmentEpoch, + generation: request.sourceGeneration + 1, + cellIncarnation: request.sourceCellIncarnation + }) + expect(await store.reconcileIdleRegionalRehome(request)).toBe('stale') + }) + + it('does not reopen an obsolete source after an assignment change', async () => { + const { store } = await setup() + await store.startEvacuation(request, request.targetCellId) + expect(await store.reconcileIdleRegionalRehome(request)).toBe('stale') + }) + + it('propagates unavailable durable state instead of declaring rollback', async () => { + const { database, store } = await setup() + vi.spyOn(database, 'transaction').mockRejectedValue(new Error('database_unavailable')) + await expect(store.reconcileIdleRegionalRehome(request)).rejects.toThrow('database_unavailable') + }) + + it('waits for an outstanding assignment transaction before deciding authority', async () => { + const { database, store } = await setup() + let release!: () => void + let entered!: () => void + const locked = new Promise((resolve) => { + entered = resolve + }) + const gate = new Promise((resolve) => { + release = resolve + }) + const commit = database.transaction(async (transaction) => { + await transaction.queryLocked( + 'SELECT * FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?', + [request.userId, request.relayHostId] + ) + entered() + await gate + await transaction.query( + 'UPDATE relay_assignments SET assignment_epoch = assignment_epoch + 1 WHERE user_id = ? AND relay_host_id = ?', + [request.userId, request.relayHostId] + ) + }) + await locked + let settled = false + const reconciliation = store.reconcileIdleRegionalRehome(request).finally(() => { + settled = true + }) + try { + await new Promise((resolve) => setImmediate(resolve)) + expect(settled).toBe(false) + } finally { + release() + await commit + await reconciliation + } + expect(await reconciliation).toBe('stale') + }) +}) diff --git a/cloud/apps/relay/src/idle-regional-rehome-selection.ts b/cloud/apps/relay/src/idle-regional-rehome-selection.ts new file mode 100644 index 00000000000..f8790459c0b --- /dev/null +++ b/cloud/apps/relay/src/idle-regional-rehome-selection.ts @@ -0,0 +1,117 @@ +import { createHash } from 'node:crypto' +import type { IdleRegionalRehomeRequest } from '@orca-cloud/relay-contract' +import type { RelayDatabase, SqlRow } from './database.js' + +export const IDLE_REHOME_PAGE_SIZE = 100 + +export async function selectIdleRegionalRehomes(input: { + database: RelayDatabase + now: number + heartbeatTtlMs: number + cohortPercent: number + offset: number + connectionHeadroom: Map + 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 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' + 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 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 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 ?`, + [ + 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 + ] + ) + 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) + } + // 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) } + }) +} diff --git a/cloud/apps/relay/src/idle-regional-rehome-store.test.ts b/cloud/apps/relay/src/idle-regional-rehome-store.test.ts new file mode 100644 index 00000000000..5d0d3cff343 --- /dev/null +++ b/cloud/apps/relay/src/idle-regional-rehome-store.test.ts @@ -0,0 +1,350 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.js' +import type { RelayDatabase, RelayLockOptions } from './database.js' +import { openIdleRehomeTestDatabase } from './idle-regional-rehome-test-database.js' + +const identity = { userId: 'idle-store-test', relayHostId: 'abcdefghijklmnop' } +const incarnations = [ + '11111111-1111-4111-8111-111111111111', + '22222222-2222-4222-8222-222222222222' +] +const cells = [ + { + id: 'source', + url: 'https://source.example.test', + region: 'us-central1' as const, + capacityRequests: 100 + }, + { + id: 'target', + url: 'https://target.example.test', + region: 'asia-east2' as const, + capacityRequests: 100 + } +] +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: 0 + }) + await store.recordCellRegionalRehomeStatus({ + cellId: cell.id, + cellIncarnation: incarnations[index]!, + regionalRehomeProtocol: 3, + safety + }) + } + const assignment = await store.assign(identity, undefined, 'us-central1') + await store.activateControl(identity, { + cellId: cells[0]!.id, + 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 + ) + const request = { + v: 1 as const, + ...identity, + attemptId: '33333333-3333-4333-8333-333333333333', + sourceCellId: cells[0]!.id, + sourceCellIncarnation: incarnations[0]!, + sourceAssignmentEpoch: assignment.assignmentEpoch, + sourceGeneration: 7, + targetCellId: cells[1]!.id + } + return { store, database, safety, request } +} + +describe('constrained idle regional assignment transaction', () => { + it.each([10, 11])('reserves source activity plus assignment at target capacity %i', async (capacity) => { + const { store, database, safety, request } = await setup() + // Model three source activity units and seven units already reserved at the target. + await database.query( + 'UPDATE relay_assignment_activity_leases SET request_units = 3 WHERE user_id = ? AND relay_host_id = ?', + [identity.userId, identity.relayHostId] + ) + await database.query("UPDATE relay_cells SET reserved_requests = 4 WHERE cell_id = 'source'") + await database.query( + "UPDATE relay_cells SET reserved_requests = 7, capacity_requests = ? WHERE cell_id = 'target'", + [capacity] + ) + const candidates = await store.selectIdleRegionalRehomeCandidates(safety) + expect(candidates).toHaveLength(capacity === 11 ? 1 : 0) + expect(await store.commitIdleRegionalRehome(request, safety)).toEqual({ + outcome: capacity === 11 ? 'committed' : 'deferred' + }) + const [target] = await database.query("SELECT reserved_requests FROM relay_cells WHERE cell_id = 'target'") + expect(Number(target!.reserved_requests)).toBe(capacity === 11 ? 11 : 7) + expect(await store.resolve(identity)).toMatchObject({ + cellId: capacity === 11 ? 'target' : 'source', + assignmentEpoch: capacity === 11 ? 2 : 1 + }) + }) + + it('progresses past a full page of busy candidates without writing eligibility state', async () => { + const { store, database, safety } = await setup() + for (const table of [ + 'relay_assignments', + 'relay_assignment_activity_leases', + 'relay_control_capabilities', + 'relay_region_decisions' + ]) { + const template = ( + await database.query(`SELECT * FROM ${table} WHERE user_id = ? AND relay_host_id = ?`, [ + identity.userId, + identity.relayHostId + ]) + )[0]! + const columns = Object.keys(template) + for (let index = 0; index < 100; index++) { + const values = columns.map((column) => + column === 'user_id' || column === 'relay_host_id' ? '?' : column + ) + await database.query( + `INSERT INTO ${table} (${columns.join(', ')}) SELECT ${values.join(', ')} FROM ${table} + WHERE user_id = ? AND relay_host_id = ?`, + [ + `idle-store-test-${String(index).padStart(3, '0')}`, + `pagehost${String(index).padStart(8, '0')}`, + identity.userId, + identity.relayHostId + ] + ) + } + } + const first = await store.selectIdleRegionalRehomeCandidates(safety) + const next = await store.selectIdleRegionalRehomeCandidates(safety) + expect(first).toHaveLength(100) + expect(next).toHaveLength(1) + expect(next[0]!.relayHostId).toBe('pagehost00000099') + const restarted = new RelayAssignmentStore(database, () => safety.observedAt, { + regionalRehomeCohortPercent: 100 + }) + expect(await restarted.selectIdleRegionalRehomeCandidates(safety)).toEqual(first) + expect(await database.query('SELECT * FROM relay_region_rehome_attempts')).toEqual([]) + const decisions = await database.query('SELECT last_considered_at FROM relay_region_decisions') + expect(decisions.every((decision) => Number(decision.last_considered_at) === 0)).toBe(true) + }) + + it.runIf(Boolean(process.env.ORCA_IDLE_REHOME_POSTGRES_URL))( + 'rechecks generation when replacement wins after the initial authority lookup', + async () => { + const { store, safety, request, database } = await setup() + const held = holdStatement(database, 'SELECT * FROM relay_region_rehome_control') + const commit = store.commitIdleRegionalRehome(request, safety) + await held.entered + try { + await store.activateControl(identity, { + cellId: 'source', + assignmentEpoch: 1, + generation: 8, + cellIncarnation: incarnations[0], + idleRegionalRehome: true + }) + } finally { + held.release() + } + expect(await commit).toEqual({ outcome: 'deferred' }) + expect(await store.reconcileIdleRegionalRehome(request)).toBe('stale') + expect(await database.query('SELECT * FROM relay_region_rehome_attempts')).toEqual([]) + } + ) + + it('rejects source replacement when the cutover already holds assignment authority', async () => { + const { store, safety, request, database } = await setup() + const held = holdStatement(database, 'UPDATE relay_assignments SET cell_id') + const commit = store.commitIdleRegionalRehome(request, safety) + await held.entered + const replacement = store.activateControl(identity, { + cellId: 'source', + assignmentEpoch: 1, + generation: 8, + cellIncarnation: incarnations[0], + idleRegionalRehome: true + }) + const rejected = expect(replacement).rejects.toThrow('wrong_assignment') + held.release() + expect(await commit).toEqual({ outcome: 'committed' }) + await rejected + expect(await store.resolve(identity)).toMatchObject({ cellId: 'target', assignmentEpoch: 2 }) + }) + + it('finds the committed attempt after its database reply is lost', async () => { + const { store, safety, request, database } = await setup() + const transaction = database.transaction.bind(database) + const intercepted = vi + .spyOn(database, 'transaction') + .mockImplementation(async (operation, options) => { + let changed = false + const result = await transaction( + async (tx) => + operation( + new Proxy(tx, { + get(target, key) { + if (key === 'query') + return async (sql: string, params?: unknown[]) => { + if (sql.includes('INSERT INTO relay_region_rehome_attempts')) changed = true + return target.query(sql, params) + } + const value = Reflect.get(target, key) + return typeof value === 'function' ? value.bind(target) : value + } + }) + ), + options + ) + if (changed) throw new Error('simulated_commit_reply_lost') + return result + }) + await expect(store.commitIdleRegionalRehome(request, safety)).rejects.toThrow( + 'simulated_commit_reply_lost' + ) + intercepted.mockRestore() + expect(await store.reconcileIdleRegionalRehome(request)).toBe('committed') + expect(await store.commitIdleRegionalRehome(request, safety)).toEqual({ outcome: 'committed' }) + expect(await database.query('SELECT * FROM relay_region_rehome_attempts')).toHaveLength(1) + }) + + it('commits the requested move once and records its outcome without source retention', async () => { + const { store, database, safety, request } = await setup() + expect(await store.commitIdleRegionalRehome(request, safety)).toEqual({ outcome: 'committed' }) + expect(await store.commitIdleRegionalRehome(request, safety)).toEqual({ outcome: 'committed' }) + expect(await store.resolve(identity)).toMatchObject({ cellId: 'target', assignmentEpoch: 2 }) + const attempts = await database.query('SELECT * FROM relay_region_rehome_attempts') + expect(attempts).toHaveLength(1) + expect(attempts[0]!.attempt_id).toBe(request.attemptId) + expect(Number(attempts[0]!.source_generation)).toBe(7) + }) + + it('rejects a replaced control and never substitutes a different target', async () => { + const { store, safety, request } = await setup() + expect( + await store.commitIdleRegionalRehome({ ...request, targetCellId: 'missing' }, safety) + ).toEqual({ outcome: 'deferred' }) + await store.activateControl(identity, { + cellId: 'source', + assignmentEpoch: 1, + generation: 8, + cellIncarnation: incarnations[0], + idleRegionalRehome: true + }) + expect(await store.commitIdleRegionalRehome(request, safety)).toEqual({ outcome: 'stale' }) + expect(await store.resolve(identity)).toMatchObject({ cellId: 'source', assignmentEpoch: 1 }) + }) + + it('does not commit without process safety or cohort authorization', async () => { + const { store, safety, request, database } = await setup() + expect(await store.commitIdleRegionalRehome(request)).toEqual({ outcome: 'deferred' }) + expect(await store.commitIdleRegionalRehome(request, safety, 0)).toEqual({ + outcome: 'deferred' + }) + expect(await database.query('SELECT * FROM relay_region_rehome_attempts')).toEqual([]) + }) + + it('selects read-only with stable identity and the control generation, not probe generation', async () => { + const { store, safety, database } = await setup() + const before = await database.query('SELECT * FROM relay_assignments') + const candidates = await store.selectIdleRegionalRehomeCandidates(safety) + expect(candidates).toHaveLength(1) + expect(candidates[0]).toMatchObject({ + sourceGeneration: 7, + sourceCellId: 'source', + targetCellId: 'target' + }) + expect(await store.selectIdleRegionalRehomeCandidates(safety)).toEqual(candidates) + expect(await database.query('SELECT * FROM relay_assignments')).toEqual(before) + expect(await database.query('SELECT * FROM relay_region_rehome_attempts')).toEqual([]) + }) +}) + +function holdStatement(database: RelayDatabase, fragment: string) { + let entered!: () => void + let release!: () => void + const arrival = new Promise((resolve) => { + entered = resolve + }) + const gate = new Promise((resolve) => { + release = resolve + }) + let held = false + const transaction = database.transaction.bind(database) + vi.spyOn(database, 'transaction').mockImplementation((operation, options) => + transaction(async (tx) => { + return operation( + new Proxy(tx, { + get(target, key) { + if (key === 'query' || key === 'queryLocked') + return async (sql: string, params?: unknown[], lockOptions?: RelayLockOptions) => { + if (!held && sql.includes(fragment)) { + held = true + entered() + await gate + } + return key === 'queryLocked' + ? target.queryLocked(sql, params, lockOptions) + : target.query(sql, params) + } + const value = Reflect.get(target, key) + return typeof value === 'function' ? value.bind(target) : value + } + }) + ) + }, options) + ) + return { entered: arrival, release } +} diff --git a/cloud/apps/relay/src/idle-regional-rehome-test-database.ts b/cloud/apps/relay/src/idle-regional-rehome-test-database.ts new file mode 100644 index 00000000000..a541e7d1126 --- /dev/null +++ b/cloud/apps/relay/src/idle-regional-rehome-test-database.ts @@ -0,0 +1,40 @@ +import { randomUUID } from 'node:crypto' +import pg from 'pg' +import { openInMemoryRelayDatabase, openRelayDatabase, type RelayDatabase } from './database.js' + +export async function openIdleRehomeTestDatabase(): Promise { + const configured = process.env.ORCA_IDLE_REHOME_POSTGRES_URL + if (!configured) return openInMemoryRelayDatabase() + const url = new URL(configured) + if (url.port !== '55440' || !['localhost', '127.0.0.1', '[::1]'].includes(url.hostname)) { + throw new Error('idle_rehome_tests_require_local_postgres_55440') + } + const schema = `idle_rehome_${randomUUID().replaceAll('-', '')}` + const admin = new pg.Client({ connectionString: configured }) + await admin.connect() + try { + await admin.query(`CREATE SCHEMA ${schema}`) + url.searchParams.set('options', `-c search_path=${schema}`) + const database = await openRelayDatabase({ databaseUrl: url.toString(), dataDir: '' }) + const close = database.close.bind(database) + database.close = async () => { + try { + await close() + } finally { + try { + await admin.query(`DROP SCHEMA ${schema} CASCADE`) + } finally { + await admin.end() + } + } + } + return database + } catch (error) { + try { + await admin.query(`DROP SCHEMA IF EXISTS ${schema} CASCADE`) + } finally { + await admin.end() + } + throw error + } +} diff --git a/cloud/apps/relay/src/idle-regional-rehome-worker.test.ts b/cloud/apps/relay/src/idle-regional-rehome-worker.test.ts new file mode 100644 index 00000000000..9b3f6d7fa07 --- /dev/null +++ b/cloud/apps/relay/src/idle-regional-rehome-worker.test.ts @@ -0,0 +1,105 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import type { RelayAssignmentStore } from './assignment-store.js' +import type { RelayConfig } from './config.js' +import { startRegionalRehomeWorker } from './regional-rehome-worker.js' + +const candidate = { + v: 1, + attemptId: '11111111-1111-4111-8111-111111111111', + userId: 'private-user', + relayHostId: 'abcdefghijklmnop', + sourceCellId: 'source', + sourceCellUrl: 'https://source.example.test', + sourceCellIncarnation: '22222222-2222-4222-8222-222222222222', + sourceAssignmentEpoch: 7, + sourceGeneration: 3, + targetCellId: 'target' +} +const config = { + role: 'director', + regionCorrectionCohortPercent: 100, + rehomeAudience: 'https://relay.example.test/v1/admin/host-drain', + rehomeDirectorServiceAccount: 'director@example.test' +} as RelayConfig + +function setup(fetch: typeof globalThis.fetch) { + const selectIdleRegionalRehomeCandidates = vi + .fn() + .mockResolvedValueOnce([]) + .mockResolvedValue([candidate]) + const claimRegionalRehome = vi.fn() + const recordRegionalRehomeDispatchFailure = vi.fn() + const worker = startRegionalRehomeWorker( + config, + { + selectIdleRegionalRehomeCandidates, + claimRegionalRehome, + recordRegionalRehomeDispatchFailure + } as unknown as RelayAssignmentStore, + { + safetySnapshot: () => ({ observedAt: 100 }) as never, + intervalMs: 60_000, + identityToken: async () => 'private-token', + fetch + } + )! + return { + worker, + selectIdleRegionalRehomeCandidates, + claimRegionalRehome, + recordRegionalRehomeDispatchFailure + } +} + +describe('idle regional worker dispatch', () => { + afterEach(() => vi.restoreAllMocks()) + it('sends an idle request without claiming an assignment first', async () => { + const fetch = vi.fn(async () => + Response.json({ v: 1, outcome: 'committed' }) + ) + const c = setup(fetch) + await vi.waitFor(() => expect(c.selectIdleRegionalRehomeCandidates).toHaveBeenCalledOnce()) + await c.worker.run() + c.worker.stop() + expect(c.claimRegionalRehome).not.toHaveBeenCalled() + expect(fetch).toHaveBeenCalledOnce() + const [url, init] = fetch.mock.calls[0]! + expect(String(url)).toBe('https://source.example.test/v1/admin/host-idle-rehome') + const { sourceCellUrl: _, ...request } = candidate + expect(JSON.parse(String(init?.body))).toEqual({ + ...request, + cohortPercent: 100, + directorSafety: { observedAt: 100 } + }) + }) + it('progresses past busy hosts without charging a dispatch failure', async () => { + const fetch = vi + .fn() + .mockResolvedValueOnce(Response.json({ v: 1, outcome: 'busy' })) + .mockResolvedValueOnce(Response.json({ v: 1, outcome: 'committed' })) + const c = setup(fetch) + await vi.waitFor(() => expect(c.selectIdleRegionalRehomeCandidates).toHaveBeenCalledOnce()) + c.selectIdleRegionalRehomeCandidates.mockResolvedValue([ + candidate, + { + ...candidate, + relayHostId: 'ponmlkjihgfedcba', + attemptId: '33333333-3333-4333-8333-333333333333' + } + ]) + await c.worker.run() + c.worker.stop() + expect(fetch).toHaveBeenCalledTimes(2) + expect(c.recordRegionalRehomeDispatchFailure).not.toHaveBeenCalled() + }) + it('does not charge a lost response as a claimed migration failure', async () => { + const fetch = vi.fn(async () => { + throw new Error('response lost') + }) + const c = setup(fetch) + await vi.waitFor(() => expect(c.selectIdleRegionalRehomeCandidates).toHaveBeenCalledOnce()) + await c.worker.run() + c.worker.stop() + expect(c.recordRegionalRehomeDispatchFailure).not.toHaveBeenCalled() + }) +}) diff --git a/cloud/apps/relay/src/index.ts b/cloud/apps/relay/src/index.ts index 541884362c2..8e7b6a56941 100644 --- a/cloud/apps/relay/src/index.ts +++ b/cloud/apps/relay/src/index.ts @@ -2,6 +2,7 @@ import { formatAssignmentInventorySnapshot, readAssignmentInventorySnapshot } from './assignment-inventory-snapshot.js' +import { readRegionCorrectionOutcomes } from './region-correction-outcomes.js' import { RelayAssignmentStore } from './assignment-store.js' import { loadRelayConfig } from './config.js' import { startCellHeartbeat } from './cell-heartbeat-client.js' @@ -71,6 +72,13 @@ const migrationInventoryTimer = roleOwnsAssignmentMaintenance(config.role) void runRelayBackgroundOperation(async () => { const inventory = await readRegisteredMigrationInventory(database, Date.now()) for (const line of formatRegisteredMigrationInventory(inventory)) console.warn(line) + console.log( + JSON.stringify({ + event: 'orca_relay_region_correction_outcomes', + observedAt: Date.now(), + outcomes: await readRegionCorrectionOutcomes(database, Date.now()) + }) + ) }, '[orca-relay] migration inventory failed') }, 5 * 60_000) : null diff --git a/cloud/apps/relay/src/region-correction-outcomes.ts b/cloud/apps/relay/src/region-correction-outcomes.ts new file mode 100644 index 00000000000..d3a6f8d3be3 --- /dev/null +++ b/cloud/apps/relay/src/region-correction-outcomes.ts @@ -0,0 +1,29 @@ +import type { RelayDatabase } from './database.js' + +export async function readRegionCorrectionOutcomes(database: RelayDatabase, now: number) { + const rows = await database.query( + `SELECT attempt.source_cell_id, attempt.target_cell_id, + CASE WHEN attempt.aborted_at IS NOT NULL THEN 'aborted' + WHEN attempt.completed_at IS NOT NULL THEN 'completed' + WHEN migration.target_registered_at IS NOT NULL THEN 'registered' ELSE 'registering' END AS state, + COUNT(*) AS count, + COALESCE(MAX(CASE WHEN attempt.completed_at IS NULL AND attempt.aborted_at IS NULL + THEN ? - attempt.created_at ELSE 0 END), 0) AS oldest_open_ms, + COALESCE(SUM(CASE WHEN attempt.completed_at IS NULL AND attempt.aborted_at IS NULL + THEN migration.target_reserved_units ELSE 0 END), 0) AS target_reserved_units + FROM relay_region_rehome_attempts attempt + JOIN relay_assignment_migrations migration ON migration.user_id = attempt.user_id + AND migration.relay_host_id = attempt.relay_host_id AND migration.assignment_epoch = attempt.assignment_epoch + GROUP BY attempt.source_cell_id, attempt.target_cell_id, state + ORDER BY attempt.source_cell_id, attempt.target_cell_id, state`, + [now] + ) + return rows.map((row) => ({ + sourceCellId: String(row.source_cell_id), + targetCellId: String(row.target_cell_id), + state: String(row.state), + count: Number(row.count), + oldestOpenMs: Number(row.oldest_open_ms), + targetReservedUnits: Number(row.target_reserved_units) + })) +} diff --git a/cloud/apps/relay/src/region-correction-preview.ts b/cloud/apps/relay/src/region-correction-preview.ts new file mode 100644 index 00000000000..c120dbb2272 --- /dev/null +++ b/cloud/apps/relay/src/region-correction-preview.ts @@ -0,0 +1,157 @@ +import type { RelayDatabase, SqlRow } from './database.js' +import { REGIONAL_REHOME_DEFAULT_HOST_COOLDOWN_MS } from './database.js' +import { + REGIONAL_REHOME_CONCURRENT_LIMIT, + REGION_DECISION_TTL_MS +} from './region-correction-state.js' + +export type RegionCorrectionPreview = { + observedAt: number + newClaimsEnabled: boolean + cohortPercent: number + openMigrations: number + availableMigrationSlots: number + globalSafetyFailure: string | null + counts: Record +} + +export async function previewRegionalRehomeEligibility(input: { + database: RelayDatabase + now: number + heartbeatTtlMs: number + cohortPercent: number + globalSafetyFailure: string | null + connectionHeadroom: ReadonlyMap + cellIsClean: (safety: SqlRow | undefined, runtime: SqlRow, now: number) => boolean +}): Promise { + const { database, now } = input + const [hosts, cells, runtimeRows, capabilityRows, safetyRows, controls, migrations] = + await Promise.all([ + database.query( + `SELECT assignment.cell_id, assignment.assignment_epoch, + decision.generation, decision.assignment_epoch AS decision_epoch, decision.expires_at, + decision.incumbent_region, decision.preferred_region, decision.outcome, decision.policy_version, + decision.observed_at, decision.cohort_bucket, + (SELECT MAX(attempt.created_at) FROM relay_region_rehome_attempts attempt + WHERE attempt.user_id = assignment.user_id AND attempt.relay_host_id = assignment.relay_host_id) AS last_attempt_at, + (SELECT COUNT(*) FROM relay_assignment_migrations migration + WHERE migration.user_id = assignment.user_id AND migration.relay_host_id = assignment.relay_host_id + AND migration.completed_at IS NULL AND migration.aborted_at IS NULL) AS open_migrations, + (SELECT COALESCE(SUM(lease.request_units),0) FROM relay_assignment_activity_leases lease + WHERE lease.user_id = assignment.user_id AND lease.relay_host_id = assignment.relay_host_id + AND lease.cell_id = assignment.cell_id) AS source_units, + (SELECT COUNT(*) FROM relay_control_capabilities host_capability + JOIN relay_assignment_activity_leases lease ON lease.user_id = host_capability.user_id + AND lease.relay_host_id = host_capability.relay_host_id AND lease.activity_id = host_capability.activity_id + JOIN relay_cell_runtime runtime ON runtime.cell_id = host_capability.cell_id + AND runtime.cell_incarnation = host_capability.cell_incarnation + WHERE host_capability.user_id = assignment.user_id AND host_capability.relay_host_id = assignment.relay_host_id + AND host_capability.cell_id = assignment.cell_id AND host_capability.assignment_epoch = assignment.assignment_epoch + AND host_capability.idle_regional_rehome = 1 AND lease.activity_kind = 'control' + AND lease.activity_id NOT LIKE 'control-pending:%' AND lease.expires_at > ? + AND lease.updated_at >= runtime.started_at) AS capable_controls + FROM relay_assignments assignment LEFT JOIN relay_region_decisions decision + ON decision.user_id = assignment.user_id AND decision.relay_host_id = assignment.relay_host_id`, + [now] + ), + database.query(`SELECT cell.*, region.region, admission.admission_state 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`), + database.query(`SELECT * FROM relay_cell_runtime`), + database.query(`SELECT * FROM relay_cell_capabilities`), + database.query(`SELECT * FROM relay_cell_rehome_safety`), + database.query(`SELECT * FROM relay_region_rehome_control WHERE control_id = 'global'`), + database.query( + `SELECT COUNT(*) AS count FROM relay_assignment_migrations WHERE completed_at IS NULL AND aborted_at IS NULL` + ) + ]) + const byCell = (rows: SqlRow[]) => new Map(rows.map((row) => [String(row.cell_id), row])) + const runtimes = byCell(runtimeRows) + const capabilities = byCell(capabilityRows) + const safety = byCell(safetyRows) + const inventory = byCell(cells) + const control = controls[0] + const cooldown = Number(control?.host_cooldown_ms ?? REGIONAL_REHOME_DEFAULT_HOST_COOLDOWN_MS) + const maxAge = Number(control?.preference_max_age_ms ?? REGION_DECISION_TTL_MS) + const openMigrations = Number(migrations[0]?.count ?? 0) + const counts: Record = {} + const count = (reason: string) => { + counts[reason] = (counts[reason] ?? 0) + 1 + } + const available = (cell: SqlRow): boolean => { + const id = String(cell.cell_id) + const runtime = runtimes.get(id) + const capability = capabilities.get(id) + return ( + Number(cell.enabled) === 1 && + cell.admission_state === 'general' && + cell.region != null && + runtime !== undefined && + Number(runtime.ready) === 1 && + Number(runtime.last_heartbeat_at) > now - input.heartbeatTtlMs && + capability !== undefined && + capability.cell_incarnation === runtime.cell_incarnation && + Number(capability.regional_rehome_protocol) >= 3 + ) + } + for (const host of hosts) { + let reason: string | null = null + const source = inventory.get(String(host.cell_id)) + if (host.generation == null) reason = 'no-verified-decision' + else if (Number(host.expires_at) <= now || Number(host.observed_at) < now - maxAge) + reason = 'expired' + else if ( + Number(host.decision_epoch) !== Number(host.assignment_epoch) || + host.incumbent_region !== source?.region + ) + reason = 'basis-changed' + else if ( + host.outcome !== 'conclusive' || + Number(host.policy_version) !== 1 || + host.preferred_region == null + ) + reason = 'inconclusive-or-insufficient-improvement' + else if (Number(host.cohort_bucket) >= input.cohortPercent) reason = 'outside-cohort' + else if (Number(host.open_migrations) > 0) reason = 'migration-open' + else if (host.last_attempt_at != null && Number(host.last_attempt_at) > now - cooldown) + reason = 'host-cooldown' + else if (!source || !available(source)) reason = 'source-ineligible' + else if (Number(host.capable_controls) === 0) reason = 'source-control-unsupported-or-inactive' + else if ( + !input.cellIsClean(safety.get(String(host.cell_id)), runtimes.get(String(host.cell_id))!, now) + ) + reason = 'source-unclean' + if (reason) { + count(reason) + continue + } + const targets = cells.filter( + (cell) => + cell.cell_id !== host.cell_id && cell.region === host.preferred_region && available(cell) + ) + const clean = targets.filter((cell) => + input.cellIsClean(safety.get(String(cell.cell_id)), runtimes.get(String(cell.cell_id))!, now) + ) + const capacity = clean.filter( + (cell) => + input.connectionHeadroom.get(String(cell.cell_id)) !== false && + Number(cell.reserved_requests) + Number(host.source_units) + 1 <= + Number(cell.capacity_requests) + ) + if (targets.length === 0) count('no-eligible-target') + else if (clean.length === 0) count('target-unclean') + else if (capacity.length === 0) count('no-target-headroom') + else if (input.globalSafetyFailure) count('global-safety-blocked') + else if (openMigrations >= REGIONAL_REHOME_CONCURRENT_LIMIT) count('concurrent-migration-cap') + else count(`eligible:${host.incumbent_region}-to-${host.preferred_region}`) + } + return { + observedAt: now, + newClaimsEnabled: Number(control?.enabled ?? 0) === 1, + cohortPercent: input.cohortPercent, + openMigrations, + availableMigrationSlots: Math.max(0, REGIONAL_REHOME_CONCURRENT_LIMIT - openMigrations), + globalSafetyFailure: input.globalSafetyFailure, + counts + } +} diff --git a/cloud/apps/relay/src/region-correction-restart.test.ts b/cloud/apps/relay/src/region-correction-restart.test.ts new file mode 100644 index 00000000000..02f315ab218 --- /dev/null +++ b/cloud/apps/relay/src/region-correction-restart.test.ts @@ -0,0 +1,119 @@ +import { mkdtemp, rm } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterEach, describe, expect, it } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.js' +import { openRelayDatabase, type RelayDatabase } from './database.js' + +const identity = { userId: 'restart-test-user', relayHostId: 'abcdefghijklmnop' } +const paths: string[] = [] +const databases = new Set() +afterEach(async () => { + for (const database of databases) await database.close() + databases.clear() + for (const path of paths.splice(0)) await rm(path, { recursive: true, force: true }) +}) + +async function setup() { + const dataDir = await mkdtemp(join(tmpdir(), 'relay-region-restart-')) + paths.push(dataDir) + let now = 1_000_000_000 + const open = async () => { + const database = await openRelayDatabase({ dataDir }) + databases.add(database) + return { database, store: new RelayAssignmentStore(database, () => now) } + } + const first = await open() + const cell = { + id: 'restart-us', + url: 'https://restart-us.example.test', + region: 'us-central1' as const, + capacityRequests: 100 + } + await first.store.reconcileCells([cell]) + await first.store.setCellEnabled(cell.id, true) + await first.store.recordCellHeartbeat({ + cellId: cell.id, + cellUrl: cell.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + region: cell.region, + startedAt: now - 1_000, + ready: true, + observedRequests: 0 + }) + const assignment = await first.store.assign(identity) + const issue = (store: RelayAssignmentStore) => + store.exchangeRegionCorrection(identity, { v: 1, action: 'issue-window' }, assignment.assignmentEpoch) + const window = (await issue(first.store)).window! + const report = { + v: 1 as const, + action: 'report' as const, + generation: window.generation, + assignmentEpoch: window.assignmentEpoch, + policyVersion: 1 as const, + outcome: 'conclusive' as const, + measurements: { 'us-central1': 200, 'asia-east2': 40 } + } + const restart = async () => { + await first.database.close() + databases.delete(first.database) + return open() + } + return { + ...first, window, report, issue, restart, + setNow: (value: number) => { now = value } + } +} + +describe('persisted region decisions across director restart', () => { + it('keeps tombstones and fixed expiry, then invalidates the prior generation after restart', async () => { + const context = await setup() + const epoch = context.window.assignmentEpoch + await context.store.exchangeRegionCorrection(identity, { + v: 1, action: 'report', generation: context.window.generation, + assignmentEpoch: epoch, policyVersion: 1, outcome: 'inconclusive', reason: 'jitter' + }, epoch) + const restarted = await context.restart() + expect(await restarted.store.exchangeRegionCorrection(identity, context.report, epoch)) + .toMatchObject({ reportStatus: 'duplicate' }) + const row = (await restarted.database.query('SELECT * FROM relay_region_decisions'))[0]! + expect(row.outcome).toBe('inconclusive') + expect(Number(row.expires_at)).toBe(context.window.expiresAt) + const successor = (await context.issue(restarted.store)).window! + expect(successor.generation).toBe(context.window.generation + 1) + expect(await restarted.store.exchangeRegionCorrection(identity, context.report, epoch)) + .toMatchObject({ reportStatus: 'stale' }) + }) + + it('uses server expiry after a restart regardless of an old client report', async () => { + const context = await setup() + context.setNow(context.window.expiresAt) + const restarted = await context.restart() + expect(await restarted.store.exchangeRegionCorrection(identity, context.report, context.window.assignmentEpoch)) + .toMatchObject({ reportStatus: 'expired' }) + expect(await restarted.store.previewRegionCorrection()).toEqual({ expired: 1 }) + }) + + it('does not interpret a persisted future-policy window using the old policy after rollback', async () => { + const context = await setup() + await context.database.query('UPDATE relay_region_decisions SET policy_version = 2') + const restarted = await context.restart() + expect(await restarted.store.exchangeRegionCorrection(identity, context.report, context.window.assignmentEpoch)) + .toMatchObject({ reportStatus: 'stale' }) + const row = (await restarted.database.query('SELECT * FROM relay_region_decisions'))[0]! + expect(row.outcome).toBe('pending') + expect(row.preferred_region).toBeNull() + expect(row.report_json).toBeNull() + }) + + it('keeps generation ordering when the server clock moves backwards across restart', async () => { + const context = await setup() + context.setNow(1_000_000_000 - 60_000) + const restarted = await context.restart() + const successor = (await context.issue(restarted.store)).window! + expect(successor.generation).toBe(context.window.generation + 1) + expect(successor.expiresAt).toBe(context.window.expiresAt - 60_000) + expect(await restarted.store.exchangeRegionCorrection(identity, context.report, context.window.assignmentEpoch)) + .toMatchObject({ reportStatus: 'stale' }) + }) +}) diff --git a/cloud/apps/relay/src/region-correction-state.ts b/cloud/apps/relay/src/region-correction-state.ts new file mode 100644 index 00000000000..a1b15520b1c --- /dev/null +++ b/cloud/apps/relay/src/region-correction-state.ts @@ -0,0 +1,158 @@ +import { relayHostLogDigest } from './relay-host-log-digest.js' +import { createHash } from 'node:crypto' +import type { + RegionCorrectionRequest, + RegionCorrectionResponse, + RelayRegion +} from '@orca-cloud/relay-contract' +import type { RelayDatabase } from './database.js' + +type Identity = { userId: string; relayHostId: string } +export const REGION_DECISION_TTL_MS = 24 * 60 * 60_000 +export const REGIONAL_REHOME_CONCURRENT_LIMIT = 8 + +export async function exchangeRegionCorrection( + database: RelayDatabase, + identity: Identity, + request: RegionCorrectionRequest, + assignmentEpoch: number, + now: number +): Promise { + const result: RegionCorrectionResponse = await database.transaction(async (transaction) => { + const assignment = ( + await transaction.queryLocked( + `SELECT * FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + )[0] + const region = + assignment && + ( + await transaction.query(`SELECT region FROM relay_cell_regions WHERE cell_id = ?`, [ + assignment.cell_id + ]) + )[0] + if (!assignment || !region || Number(assignment.assignment_epoch) !== assignmentEpoch) { + return { v: 1, reportStatus: 'basis-changed' } + } + const prior = ( + await transaction.queryLocked( + `SELECT * FROM relay_region_decisions WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + )[0] + if (request.action === 'issue-window') { + const generation = Number(prior?.generation ?? 0) + 1 + if (!Number.isSafeInteger(generation)) throw new Error('region_generation_exhausted') + const expiresAt = now + REGION_DECISION_TTL_MS + const cohortBucket = + createHash('sha256') + .update(JSON.stringify([identity.userId, identity.relayHostId])) + .digest() + .readUInt32BE(0) % 100 + await transaction.query( + `INSERT INTO relay_region_decisions + (user_id, relay_host_id, generation, expires_at, assignment_epoch, incumbent_region, + policy_version, outcome, preferred_region, observed_at, report_json, cohort_bucket) + VALUES (?, ?, ?, ?, ?, ?, 1, 'pending', NULL, ?, NULL, ?) + ON CONFLICT (user_id, relay_host_id) DO UPDATE SET + generation = excluded.generation, expires_at = excluded.expires_at, + assignment_epoch = excluded.assignment_epoch, incumbent_region = excluded.incumbent_region, + policy_version = 1, outcome = 'pending', preferred_region = NULL, + observed_at = excluded.observed_at, report_json = NULL, cohort_bucket = excluded.cohort_bucket`, + [ + identity.userId, + identity.relayHostId, + generation, + expiresAt, + assignmentEpoch, + region.region, + now, + cohortBucket + ] + ) + return { + v: 1, + window: { + generation, + expiresAt, + assignmentEpoch, + incumbentRegion: region.region as RelayRegion, + policyVersion: 1 + } + } + } + if (!prior || Number(prior.generation) !== request.generation) + return { v: 1, reportStatus: 'stale' } + if (Number(prior.policy_version) !== request.policyVersion) + return { v: 1, reportStatus: 'stale' } + if (Number(prior.expires_at) <= now) return { v: 1, reportStatus: 'expired' } + if ( + request.assignmentEpoch !== assignmentEpoch || + Number(prior.assignment_epoch) !== assignmentEpoch || + prior.incumbent_region !== region.region + ) { + return { v: 1, reportStatus: 'basis-changed' } + } + // The first report wins, including an inconclusive tombstone. + if (prior.outcome !== 'pending') return { v: 1, reportStatus: 'duplicate' } + let preferredRegion: RelayRegion | null = null + if (request.outcome === 'conclusive') { + const incumbent = request.measurements[region.region as RelayRegion] + const target: RelayRegion = region.region === 'us-central1' ? 'asia-east2' : 'us-central1' + const targetRtt = request.measurements[target] + if (incumbent - targetRtt >= 25 && targetRtt <= incumbent * 0.8) preferredRegion = target + } + await transaction.query( + `UPDATE relay_region_decisions SET outcome = ?, preferred_region = ?, report_json = ? + WHERE user_id = ? AND relay_host_id = ? AND generation = ?`, + [ + request.outcome, + preferredRegion, + JSON.stringify(request), + identity.userId, + identity.relayHostId, + request.generation + ] + ) + return { v: 1, reportStatus: 'accepted' } + }) + if (request.action === 'report' && result.reportStatus === 'accepted') { + const digest = relayHostLogDigest(identity.relayHostId) + // Stable sampling includes unchanged hosts for before/after comparisons. + if (Number.parseInt(digest.slice(0, 8), 16) % 10 === 0) { + console.log( + JSON.stringify({ + event: 'orca_relay_region_comparison', + relayHostIdDigest: digest, + assignmentEpoch, + generation: request.generation, + policyVersion: request.policyVersion, + outcome: request.outcome, + ...(request.outcome === 'conclusive' ? { measurements: request.measurements } : {}) + }) + ) + } + } + return result +} + +export async function previewRegionCorrection( + database: RelayDatabase, + now: number +): Promise> { + const rows = await database.query( + `SELECT CASE WHEN decision.expires_at <= ? THEN 'expired' + WHEN decision.assignment_epoch <> assignment.assignment_epoch THEN 'basis-changed' + WHEN decision.outcome = 'pending' THEN 'pending' + WHEN decision.preferred_region IS NULL THEN 'ineligible' + ELSE decision.incumbent_region || '-to-' || decision.preferred_region END AS reason, + COUNT(*) AS count + FROM relay_region_decisions decision + JOIN relay_assignments assignment ON assignment.user_id = decision.user_id + AND assignment.relay_host_id = decision.relay_host_id + GROUP BY reason`, + [now] + ) + return Object.fromEntries(rows.map((row) => [String(row.reason), Number(row.count)])) +} diff --git a/cloud/apps/relay/src/region-correction-store.test.ts b/cloud/apps/relay/src/region-correction-store.test.ts new file mode 100644 index 00000000000..11af7b3c5bd --- /dev/null +++ b/cloud/apps/relay/src/region-correction-store.test.ts @@ -0,0 +1,359 @@ +import { afterEach, describe, expect, it } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.js' +import { openInMemoryRelayDatabase, openRelayDatabase, type RelayDatabase } from './database.js' + +const identity = { userId: 'region-correction-test-user', relayHostId: 'abcdefghijklmnop' } +const cells = [ + { + id: 'decision-us', + url: 'https://decision-us.example.test', + region: 'us-central1' as const, + capacityRequests: 100 + }, + { + id: 'decision-asia', + url: 'https://decision-asia.example.test', + region: 'asia-east2' as const, + capacityRequests: 100 + } +] +const incarnations = [ + '11111111-1111-4111-8111-111111111111', + '22222222-2222-4222-8222-222222222222' +] +const opened: RelayDatabase[] = [] +afterEach(async () => { + for (const database of opened.splice(0)) { + if (database.dialect === 'postgres') await cleanupPostgres(database) + await database.close() + } +}) + +async function cleanupPostgres(database: RelayDatabase) { + for (const table of [ + 'relay_control_connection_reservations', + 'relay_region_decisions', + 'relay_control_capabilities', + 'relay_assignment_activity_leases', + 'relay_assignment_migrations', + 'relay_assignment_migration_incarnations', + 'relay_assignment_region_preferences', + 'relay_region_rehome_attempts', + 'relay_assignments' + ]) { + await database.query(`DELETE FROM ${table} WHERE user_id = ?`, [identity.userId]) + } + for (const table of [ + 'relay_cell_rehome_safety', + 'relay_cell_capabilities', + 'relay_cell_connection_snapshots', + 'relay_cell_connection_runtime', + 'relay_cell_runtime', + 'relay_cell_connection_limits', + 'relay_cell_admission', + 'relay_cell_regions', + 'relay_cells' + ]) { + await database.query( + `DELETE FROM ${table} WHERE cell_id IN (?, ?)`, + cells.map((cell) => cell.id) + ) + } +} + +async function setup() { + const database = + process.env.ORCA_REGION_CORRECTION_POSTGRES === '1' + ? await openRelayDatabase({ + databaseUrl: requiredPostgresUrl(), + dataDir: '/tmp/orca-region-correction-unused' + }) + : await openInMemoryRelayDatabase() + opened.push(database) + if (database.dialect === 'postgres') await cleanupPostgres(database) + let clock = 100_000_000 + const store = new RelayAssignmentStore(database, () => clock, { + regionalRehomeCohortPercent: 100 + }) + await store.reconcileCells(cells) + for (const cell of cells) await store.setCellEnabled(cell.id, true) + for (const [index, cell] of cells.entries()) { + await store.recordCellHeartbeat({ + cellId: cell.id, + cellUrl: cell.url, + region: cell.region, + cellIncarnation: incarnations[index]!, + startedAt: clock - 1_000, + ready: true, + observedRequests: 0 + }) + } + const assignment = await store.assign(identity, undefined, 'us-central1') + const activityId = await store.activateControl(identity, { + cellId: cells[0]!.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 7, + cellIncarnation: incarnations[0], + idleRegionalRehome: true + }) + return { + database, + store, + assignment, + activityId, + now: () => clock, + advance: (ms: number) => { + clock += ms + } + } +} + +function requiredPostgresUrl(): string { + const url = process.env.ORCA_RELAY_TEST_POSTGRES_URL + if (!url || new URL(url).port !== '55440') + throw new Error('PostgreSQL tests require configured port 55440') + return url +} + +async function window(context: Awaited>) { + const result = await context.store.exchangeRegionCorrection( + identity, + { v: 1, action: 'issue-window' }, + context.assignment.assignmentEpoch + ) + return result.window! +} + +async function regionalMigration(context: Awaited>) { + const migration = await context.store.startEvacuation(identity, cells[1]!.id) + const attemptId = '33333333-3333-4333-8333-333333333333' + await context.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 (?,?,?,'asia-east2',?,?,?,?,?,?,60000,1,?,?)`, + [ + attemptId, + identity.userId, + identity.relayHostId, + cells[0]!.id, + incarnations[0], + cells[1]!.id, + incarnations[1], + migration.previousEpoch, + migration.assignmentEpoch, + context.now(), + context.now() + ] + ) + return { migration } +} + +describe('ordered region decisions and migration outcomes', () => { + it('reports aggregate migration lifecycle and reservations without identity disclosure or writes', async () => { + const context = await setup() + const { migration } = await regionalMigration(context) + context.advance(1_000) + const before = await context.database.query('SELECT * FROM relay_region_rehome_attempts') + const outcomes = await context.store.regionCorrectionOutcomes() + expect(outcomes).toEqual([ + expect.objectContaining({ + sourceCellId: cells[0]!.id, + targetCellId: cells[1]!.id, + state: 'registering', + count: 1, + oldestOpenMs: 1_000 + }) + ]) + expect(outcomes[0]!.targetReservedUnits).toBeGreaterThan(0) + expect(JSON.stringify(outcomes)).not.toContain(identity.relayHostId) + expect(JSON.stringify(outcomes)).not.toContain(identity.userId) + expect(await context.database.query('SELECT * FROM relay_region_rehome_attempts')).toEqual( + before + ) + await context.store.activateControl(identity, { + cellId: cells[1]!.id, + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await context.store.markMigrationTargetRegistered(identity, { + cellId: cells[1]!.id, + assignmentEpoch: migration.assignmentEpoch + }) + expect(await context.store.regionCorrectionOutcomes()).toEqual([ + expect.objectContaining({ state: 'registered' }) + ]) + await context.store.releaseActivity(identity, context.activityId) + expect(await context.store.completeReadyRegionalRehomes()).toBe(1) + expect(await context.store.regionCorrectionOutcomes()).toEqual([ + expect.objectContaining({ + state: 'completed', + targetReservedUnits: 0, + oldestOpenMs: 0 + }) + ]) + }) + + it('supersedes prior windows and keeps an inconclusive tombstone immutable', async () => { + const context = await setup() + const first = await window(context) + const second = await window(context) + expect(second.generation).toBe(first.generation + 1) + const report = { + v: 1 as const, + action: 'report' as const, + assignmentEpoch: first.assignmentEpoch, + policyVersion: 1 as const, + outcome: 'conclusive' as const, + measurements: { 'us-central1': 200, 'asia-east2': 40 } + } + expect( + await context.store.exchangeRegionCorrection( + identity, + { ...report, generation: first.generation }, + first.assignmentEpoch + ) + ).toMatchObject({ reportStatus: 'stale' }) + expect( + await context.store.exchangeRegionCorrection( + identity, + { ...report, generation: second.generation, outcome: 'inconclusive', reason: 'jitter' }, + second.assignmentEpoch + ) + ).toMatchObject({ reportStatus: 'accepted' }) + expect( + await context.store.exchangeRegionCorrection( + identity, + { ...report, generation: second.generation }, + second.assignmentEpoch + ) + ).toMatchObject({ reportStatus: 'duplicate' }) + expect(await context.store.previewRegionCorrection()).toEqual({ ineligible: 1 }) + }) + + it('previews the uncapped fleet without writes, claims, or locked reads', async () => { + const context = await setup() + const query = context.database.query.bind(context.database) + const transaction = context.database.transaction.bind(context.database) + const queryLocked = context.database.queryLocked.bind(context.database) + context.database.query = async (sql, params) => { + expect(sql.trim()).toMatch(/^(SELECT|WITH)/i) + return query(sql, params) + } + context.database.transaction = async () => { + throw new Error('preview_must_not_open_mutating_transaction') + } + context.database.queryLocked = async () => { + throw new Error('preview_must_not_lock') + } + try { + const preview = await context.store.previewRegionalRehomeEligibility() + expect(preview.counts['no-verified-decision']).toBeGreaterThanOrEqual(1) + expect(preview.globalSafetyFailure).toBe('process-safety-unavailable') + expect(JSON.stringify(preview)).not.toContain(identity.relayHostId) + expect(JSON.stringify(preview)).not.toContain(identity.userId) + } finally { + context.database.query = query + context.database.transaction = transaction + context.database.queryLocked = queryLocked + } + }) + + it('allocates distinct ordered generations for concurrent window issuers', async () => { + const context = await setup() + const replies = await Promise.all([window(context), window(context), window(context)]) + expect(replies.map((reply) => reply.generation).sort((a, b) => a - b)).toEqual([1, 2, 3]) + const older = replies.find((reply) => reply.generation === 2)! + expect( + await context.store.exchangeRegionCorrection( + identity, + { + v: 1, + action: 'report', + generation: older.generation, + assignmentEpoch: older.assignmentEpoch, + policyVersion: 1, + outcome: 'inconclusive', + reason: 'delayed' + }, + older.assignmentEpoch + ) + ).toMatchObject({ reportStatus: 'stale' }) + }) + + it('compares with assigned region, preserves hints, and never extends a window on report', async () => { + const context = await setup() + await context.store.assign(identity, 'asia-east2') + const issued = await window(context) + expect(issued.incumbentRegion).toBe('us-central1') + context.advance(50) + await context.store.exchangeRegionCorrection( + identity, + { + v: 1, + action: 'report', + generation: issued.generation, + assignmentEpoch: issued.assignmentEpoch, + policyVersion: 1, + outcome: 'conclusive', + measurements: { 'us-central1': 110, 'asia-east2': 90 } + }, + issued.assignmentEpoch + ) + expect(await context.store.previewRegionCorrection()).toEqual({ ineligible: 1 }) + const row = (await context.database.query(`SELECT * FROM relay_region_decisions`))[0]! + expect(Number(row.expires_at)).toBe(issued.expiresAt) + const hint = ( + await context.database.query( + `SELECT preferred_region FROM relay_assignment_region_preferences WHERE user_id = ?`, + [identity.userId] + ) + )[0] + expect(hint?.preferred_region).toBe('asia-east2') + context.advance(24 * 60 * 60_000) + expect( + await context.store.exchangeRegionCorrection( + identity, + { + v: 1, + action: 'report', + generation: issued.generation, + assignmentEpoch: issued.assignmentEpoch, + policyVersion: 1, + outcome: 'inconclusive', + reason: 'late' + }, + issued.assignmentEpoch + ) + ).toMatchObject({ reportStatus: 'expired' }) + }) + + it('rejects stale assignment basis and requires both thresholds', async () => { + const context = await setup() + const issued = await window(context) + await context.store.exchangeRegionCorrection( + identity, + { + v: 1, + action: 'report', + generation: issued.generation, + assignmentEpoch: issued.assignmentEpoch, + policyVersion: 1, + outcome: 'conclusive', + measurements: { 'us-central1': 150, 'asia-east2': 100 } + }, + issued.assignmentEpoch + ) + expect(await context.store.previewRegionCorrection()).toEqual({ + 'us-central1-to-asia-east2': 1 + }) + await context.store.startEvacuation(identity, cells[1]!.id) + expect( + await context.store.exchangeRegionCorrection( + identity, + { v: 1, action: 'issue-window' }, + issued.assignmentEpoch + ) + ).toMatchObject({ reportStatus: 'basis-changed' }) + }) +}) diff --git a/cloud/apps/relay/src/regional-host-drain-app.test.ts b/cloud/apps/relay/src/regional-host-drain-app.test.ts index e2a33a07bb0..cf7e3798155 100644 --- a/cloud/apps/relay/src/regional-host-drain-app.test.ts +++ b/cloud/apps/relay/src/regional-host-drain-app.test.ts @@ -5,7 +5,9 @@ vi.mock('./admin-token-verifier.js', () => ({ createAdminTokenVerifier: () => async (token: string, route?: string) => token === 'deploy-token' || (token === 'monitor-token' && - (!route || route === '/v1/admin/regional-rehome-control')), + (!route || + route === '/v1/admin/regional-rehome-control' || + route === '/v1/admin/regional-rehome-preview')), createReadOnlyAdminTokenVerifier: () => async () => false, createRegionalRehomeControlApplyTokenVerifier: () => async (token: string) => token === 'deploy-token', @@ -41,7 +43,83 @@ const request = { graceMs: 60_000 } +describe('idle regional cutover endpoint', () => { + it('authenticates and fences the source before invoking a cutover', async () => { + const idleRehome = vi.fn(async () => ({ outcome: 'busy' })) + const app = createRelayApp(config(), { + store: {} as never, + assignments: {} as never, + drain: vi.fn(), + idleRehome, + cellIncarnation, + ready: vi.fn(async () => true) + } as Parameters[1]) + const input = { + v: 1, + attemptId: request.attemptId, + userId: request.userId, + relayHostId: request.relayHostId, + sourceCellId: request.sourceCellId, + sourceCellIncarnation: cellIncarnation, + sourceAssignmentEpoch: 7, + sourceGeneration: 1, + targetCellId: 'target-cell', + cohortPercent: 100, + directorSafety: { + observedAt: 100, sqlFailures: 0, reconnects: 0, controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, databasePoolWaitersMax: 0, databasePoolWaitMsMax: 0 + } + } + const path = '/v1/admin/host-idle-rehome' + expect((await postPath(app, path, 'runtime-token', input)).status).toBe(401) + expect( + ( + await postPath(app, path, 'rehome-token', { + ...input, + sourceCellIncarnation: '33333333-3333-4333-8333-333333333333' + }) + ).status + ).toBe(409) + expect(idleRehome).not.toHaveBeenCalled() + const response = await postPath(app, path, 'rehome-token', input) + expect(response.status).toBe(200) + expect(await response.json()).toEqual({ v: 1, outcome: 'busy' }) + expect(idleRehome).toHaveBeenCalledExactlyOnceWith(input) + }) +}) + describe('regional host drain endpoint', () => { + it('exposes aggregate preview to monitors without a mutation path', async () => { + const preview = { counts: { 'eligible:asia-east2-to-us-central1': 2 } } + const safety = { observedAt: 100 } + const previewRegionalRehomeEligibility = vi.fn(async () => preview) + const app = createRelayApp(config({ role: 'director', cellId: 'director' }), { + store: {} as never, + assignments: { + previewRegionalRehomeEligibility, + regionCorrectionOutcomes: async () => [] + } as never, + regionalRehomeSafetySnapshot: () => safety as never, + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + const path = '/v1/admin/regional-rehome-preview' + expect((await app.request(path)).status).toBe(401) + expect(previewRegionalRehomeEligibility).not.toHaveBeenCalled() + const response = await app.request(path, { headers: { authorization: 'Bearer monitor-token' } }) + expect(response.status).toBe(200) + expect(await response.json()).toEqual({ v: 1, preview, outcomes: [] }) + expect(previewRegionalRehomeEligibility).toHaveBeenCalledExactlyOnceWith(safety) + expect( + ( + await app.request(path, { + method: 'POST', + headers: { authorization: 'Bearer deploy-token' } + }) + ).status + ).toBe(404) + }) + it('accepts only the dedicated identity and exact cell generation', async () => { const drainHost = vi.fn(() => 'accepted' as const) const app = createRelayApp(config(), { @@ -60,14 +138,66 @@ describe('regional host drain endpoint', () => { expect((await post(app, 'deploy-token', request)).status).toBe(401) expect( - (await post(app, 'rehome-token', { - ...request, - sourceCellIncarnation: '33333333-3333-4333-8333-333333333333' - })).status + ( + await post(app, 'rehome-token', { + ...request, + sourceCellIncarnation: '33333333-3333-4333-8333-333333333333' + }) + ).status ).toBe(409) expect(drainHost).toHaveBeenCalledOnce() }) + it('waits for an asynchronous drain operation before acknowledging', async () => { + let grant!: (value: 'accepted') => void + let entered!: () => void + const started = new Promise((resolve) => { + entered = resolve + }) + const drainHost = vi.fn(() => { + entered() + return new Promise<'accepted'>((resolve) => { + grant = resolve + }) + }) + const app = createRelayApp(config(), { + store: {} as never, + assignments: {} as never, + drain: vi.fn(), + drainHost, + cellIncarnation, + ready: vi.fn(async () => true) + }) + const pending = post(app, 'rehome-token', request) + let acknowledged = false + void pending.then(() => { + acknowledged = true + }) + await started + await Promise.resolve() + expect(acknowledged).toBe(false) + grant('accepted') + const response = await pending + expect(response.status).toBe(200) + expect(await response.json()).toEqual({ v: 1, outcome: 'accepted' }) + }) + + it('rejects a failed asynchronous drain instead of acknowledging it', async () => { + const app = createRelayApp(config(), { + store: {} as never, + assignments: {} as never, + drain: vi.fn(), + drainHost: async () => { + throw new Error('activity_cell_not_authoritative') + }, + cellIncarnation, + ready: vi.fn(async () => true) + }) + const response = await post(app, 'rehome-token', request) + expect(response.status).toBe(409) + expect(await response.json()).toEqual({ error: 'activity_cell_not_authoritative' }) + }) + it('rejects malformed identities before touching the session registry', async () => { const drainHost = vi.fn(() => 'accepted' as const) const app = createRelayApp(config(), { @@ -224,7 +354,7 @@ describe('regional rehome director controls', () => { v: 1, cellId: 'production-gce-c7', cellIncarnation, - regionalRehomeProtocol: 1, + regionalRehomeProtocol: 2, safety: { observedAt: 100, sqlFailures: 0, @@ -235,12 +365,7 @@ describe('regional rehome director controls', () => { databasePoolWaitMsMax: 0 } } - const response = await postPath( - app, - '/v1/admin/cell-rehome-status', - 'runtime-token', - body - ) + const response = await postPath(app, '/v1/admin/cell-rehome-status', 'runtime-token', body) expect(response.status).toBe(200) expect(recordCellRegionalRehomeStatus).toHaveBeenCalledWith(body) @@ -266,18 +391,13 @@ describe('regional rehome director controls', () => { v: 1, cellId: 'production-gce-c7', cellIncarnation, - regionalRehomeProtocol: 1, + regionalRehomeProtocol: 2, safety: { ...observability.regionalRehomeRuntimeSafety(), ...emptyPostgresPoolPressureCounts() } } - const response = await postPath( - app, - '/v1/admin/cell-rehome-status', - 'runtime-token', - body - ) + const response = await postPath(app, '/v1/admin/cell-rehome-status', 'runtime-token', body) expect(response.status).toBe(200) expect(recordCellRegionalRehomeStatus).toHaveBeenCalledWith(body) @@ -301,12 +421,14 @@ describe('regional rehome director controls', () => { drain: vi.fn(), ready: vi.fn(async () => true) }) - expect((await postPath( - app, - '/v1/admin/regional-rehome-control', - 'deploy-token', - { v: 1, action: 'inspect' } - )).status).toBe(200) + expect( + ( + await postPath(app, '/v1/admin/regional-rehome-control', 'deploy-token', { + v: 1, + action: 'inspect' + }) + ).status + ).toBe(200) const apply = { v: 1, action: 'apply', @@ -319,39 +441,35 @@ describe('regional rehome director controls', () => { drainGraceMs: 60_000, confirmation: 'ENABLE_REGIONAL_REHOMING' } - expect((await postPath( - app, - '/v1/admin/regional-rehome-control', - 'deploy-token', - apply - )).status).toBe(200) + expect( + (await postPath(app, '/v1/admin/regional-rehome-control', 'deploy-token', apply)).status + ).toBe(200) expect(applyRegionalRehomeControl).toHaveBeenCalledOnce() - expect((await postPath( - app, - '/v1/admin/regional-rehome-control', - 'monitor-token', - { v: 1, action: 'inspect' } - )).status).toBe(200) - expect((await postPath( - app, - '/v1/admin/regional-rehome-control', - 'monitor-token', - apply - )).status).toBe(403) - expect((await postPath( - app, - '/v1/admin/regional-rehome-control', - 'deploy-token', - { ...apply, confirmation: 'DISABLE_REGIONAL_REHOMING' } - )).status).toBe(400) + expect( + ( + await postPath(app, '/v1/admin/regional-rehome-control', 'monitor-token', { + v: 1, + action: 'inspect' + }) + ).status + ).toBe(200) + expect( + (await postPath(app, '/v1/admin/regional-rehome-control', 'monitor-token', apply)).status + ).toBe(403) + expect( + ( + await postPath(app, '/v1/admin/regional-rehome-control', 'deploy-token', { + ...apply, + confirmation: 'DISABLE_REGIONAL_REHOMING' + }) + ).status + ).toBe(400) // The per-host cooldown is part of the durable shape an operator must state. const { hostCooldownMs: _omitted, ...withoutCooldown } = apply - expect((await postPath( - app, - '/v1/admin/regional-rehome-control', - 'deploy-token', - withoutCooldown - )).status).toBe(400) + expect( + (await postPath(app, '/v1/admin/regional-rehome-control', 'deploy-token', withoutCooldown)) + .status + ).toBe(400) }) it('probes dedicated trust twice and returns only aggregate proof', async () => { @@ -382,12 +500,11 @@ describe('regional rehome director controls', () => { }) as typeof fetch, ready: vi.fn(async () => true) }) - const response = await postPath( - app, - '/v1/admin/regional-rehome-trust-probe', - 'deploy-token', - { v: 1, sourceCellId: 'production-gce-c7', sourceCellIncarnation: cellIncarnation } - ) + const response = await postPath(app, '/v1/admin/regional-rehome-trust-probe', 'deploy-token', { + v: 1, + sourceCellId: 'production-gce-c7', + sourceCellIncarnation: cellIncarnation + }) expect(response.status).toBe(200) const responseBody = await response.json() @@ -448,12 +565,11 @@ describe('regional rehome director controls', () => { ready: vi.fn(async () => true) }) - const response = await postPath( - app, - '/v1/admin/regional-rehome-trust-probe', - 'deploy-token', - { v: 1, sourceCellId: 'production-gce-c27', sourceCellIncarnation: cellIncarnation } - ) + const response = await postPath(app, '/v1/admin/regional-rehome-trust-probe', 'deploy-token', { + v: 1, + sourceCellId: 'production-gce-c27', + sourceCellIncarnation: cellIncarnation + }) expect(response.status).toBe(200) expect(await response.json()).toMatchObject({ proven: true }) @@ -481,12 +597,11 @@ describe('regional rehome director controls', () => { ready: vi.fn(async () => true) }) - const response = await postPath( - app, - '/v1/admin/regional-rehome-trust-probe', - 'deploy-token', - { v: 1, sourceCellId: 'production-gce-c27', sourceCellIncarnation: cellIncarnation } - ) + const response = await postPath(app, '/v1/admin/regional-rehome-trust-probe', 'deploy-token', { + v: 1, + sourceCellId: 'production-gce-c27', + sourceCellIncarnation: cellIncarnation + }) expect(response.status).toBe(409) expect(sourceFetch).not.toHaveBeenCalled() @@ -504,18 +619,17 @@ describe('regional rehome director controls', () => { sourceCellId: 'production-gce-c7', sourceCellIncarnation: cellIncarnation } - expect((await postPath( - app, - '/v1/admin/regional-rehome-trust-probe', - 'monitor-token', - body - )).status).toBe(401) - expect((await postPath( - app, - '/v1/admin/regional-rehome-trust-probe', - 'deploy-token', - { ...body, unexpected: true } - )).status).toBe(400) + expect( + (await postPath(app, '/v1/admin/regional-rehome-trust-probe', 'monitor-token', body)).status + ).toBe(401) + expect( + ( + await postPath(app, '/v1/admin/regional-rehome-trust-probe', 'deploy-token', { + ...body, + unexpected: true + }) + ).status + ).toBe(400) }) it('fails closed when the source rejects the dedicated identity', async () => { @@ -529,9 +643,9 @@ describe('regional rehome director controls', () => { regionalRehomeProtocol: 1 } }) - const sourceFetch = vi.fn().mockResolvedValue( - Response.json({ error: 'invalid_token' }, { status: 401 }) - ) + const sourceFetch = vi + .fn() + .mockResolvedValue(Response.json({ error: 'invalid_token' }, { status: 401 })) const app = createRelayApp(config({ role: 'director', cellId: 'director' }), { store: {} as never, assignments: { cellDeploymentStatus } as never, @@ -540,12 +654,11 @@ describe('regional rehome director controls', () => { regionalRehomeFetch: sourceFetch, ready: vi.fn(async () => true) }) - const response = await postPath( - app, - '/v1/admin/regional-rehome-trust-probe', - 'deploy-token', - { v: 1, sourceCellId: 'production-gce-c7', sourceCellIncarnation: cellIncarnation } - ) + const response = await postPath(app, '/v1/admin/regional-rehome-trust-probe', 'deploy-token', { + v: 1, + sourceCellId: 'production-gce-c7', + sourceCellIncarnation: cellIncarnation + }) expect(response.status).toBe(409) expect(sourceFetch).toHaveBeenCalledOnce() diff --git a/cloud/apps/relay/src/regional-rehome-postgres.test.ts b/cloud/apps/relay/src/regional-rehome-postgres.test.ts index fdefda54401..07307707124 100644 --- a/cloud/apps/relay/src/regional-rehome-postgres.test.ts +++ b/cloud/apps/relay/src/regional-rehome-postgres.test.ts @@ -29,6 +29,9 @@ describePostgres('PostgreSQL regional rehoming', () => { }) async function cleanup(): Promise { + for (const table of ['relay_region_decisions', 'relay_control_capabilities']) { + await primary.query(`DELETE FROM ${table} WHERE user_id LIKE 'pg-rehome-user-%'`) + } await primary.query( `DELETE FROM relay_region_rehome_attempts WHERE user_id LIKE 'pg-rehome-user-%'` ) @@ -69,6 +72,81 @@ describePostgres('PostgreSQL regional rehoming', () => { } } + it('defaults to a closed correction cohort even with enabled durable control', async () => { + const context = await fixture() + const closed = new RelayAssignmentStore(primary, context.now, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + expect(await cutover(closed, context.now())).toBeNull() + const preview = await closed.previewRegionalRehomeEligibility({ + observedAt: context.now(), + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + }) + expect(preview.cohortPercent).toBe(0) + expect(preview.counts['outside-cohort']).toBe(1) + expect(await attemptAndMigrationCounts(context.identity)).toEqual({ + attempts: 0, + migrations: 0 + }) + }) + + it('counts existing generic migrations against the optimization cap and preview', async () => { + const context = await fixture() + const safety = { + observedAt: context.now(), + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + } + const before = await context.store.previewRegionalRehomeEligibility(safety) + expect(before.counts['eligible:us-central1-to-asia-east2']).toBe(1) + for (let index = 0; index < 8; index++) { + const identity = { + userId: `pg-rehome-user-budget-${sequence}-${index}`, + relayHostId: `budgethost${String(index).padStart(6, '0')}` + } + await context.store.assign(identity, undefined, 'us-central1') + await context.store.startEvacuation(identity, context.target.id) + } + const preview = await context.store.previewRegionalRehomeEligibility(safety) + expect(preview.openMigrations).toBe(8) + expect(preview.availableMigrationSlots).toBe(0) + expect(preview.counts['concurrent-migration-cap']).toBe(1) + expect(await cutover(context.store, context.now())).toBeNull() + expect(await attemptAndMigrationCounts(context.identity)).toEqual({ + attempts: 0, + migrations: 0 + }) + }) + + it('preview excludes request capacity exhaustion before a claim', async () => { + const context = await fixture() + await primary.query( + `UPDATE relay_cells SET capacity_requests = reserved_requests + 1 WHERE cell_id = ?`, + [context.target.id] + ) + const preview = await context.store.previewRegionalRehomeEligibility({ + observedAt: context.now(), + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + }) + expect(preview.counts['no-target-headroom']).toBe(1) + expect(await cutover(context.store, context.now())).toBeNull() + }) + it('claims through ambient per-cell sql retry noise', async () => { const context = await fixture() await primary.query( @@ -78,27 +156,31 @@ describePostgres('PostgreSQL regional rehoming', () => { [context.source.id, context.target.id] ) - expect(await context.store.claimRegionalRehome()).not.toBeNull() + expect(await cutover(context.store, context.now())).not.toBeNull() }) it('moves a us-central1 host onto a cell in its preferred asia-east2 region', async () => { const context = await fixture() - const attempt = await context.store.claimRegionalRehome() + const attempt = await cutover(context.store, context.now()) expect(attempt).toMatchObject({ preferredRegion: 'asia-east2', sourceCellId: context.source.id, targetCellId: context.target.id }) - expect(await primary.query( - `SELECT preferred_region, source_cell_id, target_cell_id + expect( + await primary.query( + `SELECT preferred_region, source_cell_id, target_cell_id FROM relay_region_rehome_attempts WHERE user_id = ?`, - [context.identity.userId] - )).toEqual([{ - preferred_region: 'asia-east2', - source_cell_id: context.source.id, - target_cell_id: context.target.id - }]) + [context.identity.userId] + ) + ).toEqual([ + { + preferred_region: 'asia-east2', + source_cell_id: context.source.id, + target_cell_id: context.target.id + } + ]) }) it('moves an asia-east2 host back onto a cell in its preferred us-central1 region', async () => { @@ -107,32 +189,37 @@ describePostgres('PostgreSQL regional rehoming', () => { targetRegion: 'us-central1' }) - const attempt = await context.store.claimRegionalRehome() + const attempt = await cutover(context.store, context.now()) expect(attempt).toMatchObject({ preferredRegion: 'us-central1', sourceCellId: context.source.id, targetCellId: context.target.id }) // The durable attempt row must accept the reverse direction too. - expect(await primary.query( - `SELECT preferred_region, source_cell_id, target_cell_id + expect( + await primary.query( + `SELECT preferred_region, source_cell_id, target_cell_id FROM relay_region_rehome_attempts WHERE user_id = ?`, - [context.identity.userId] - )).toEqual([{ - preferred_region: 'us-central1', - source_cell_id: context.source.id, - target_cell_id: context.target.id - }]) - expect(await primary.query( - `SELECT cell_id FROM relay_assignments WHERE user_id = ?`, - [context.identity.userId] - )).toEqual([{ cell_id: context.target.id }]) + [context.identity.userId] + ) + ).toEqual([ + { + preferred_region: 'us-central1', + source_cell_id: context.source.id, + target_cell_id: context.target.id + } + ]) + expect( + await primary.query(`SELECT cell_id FROM relay_assignments WHERE user_id = ?`, [ + context.identity.userId + ]) + ).toEqual([{ cell_id: context.target.id }]) }) it('leaves a host whose preference already matches its own region', async () => { const context = await fixture({ preferredRegion: 'us-central1' }) - await expect(context.store.claimRegionalRehome()).resolves.toBeNull() + await expect(cutover(context.store, context.now())).resolves.toBeNull() await expect(context.store.inspectRegionalRehomeControl()).resolves.toMatchObject({ generation: 1, enabled: true @@ -146,16 +233,12 @@ describePostgres('PostgreSQL regional rehoming', () => { it('leaves a host whose preference is older than the configured max age', async () => { const context = await fixture() await primary.query( - `UPDATE relay_assignment_region_preferences SET observed_at = ? + `UPDATE relay_region_decisions SET observed_at = ? WHERE user_id = ? AND relay_host_id = ?`, - [ - context.now() - 24 * 60 * 60_000 - 1, - context.identity.userId, - context.identity.relayHostId - ] + [context.now() - 24 * 60 * 60_000 - 1, context.identity.userId, context.identity.relayHostId] ) - await expect(context.store.claimRegionalRehome()).resolves.toBeNull() + await expect(cutover(context.store, context.now())).resolves.toBeNull() await expect(context.store.inspectRegionalRehomeControl()).resolves.toMatchObject({ generation: 1, enabled: true @@ -190,7 +273,7 @@ describePostgres('PostgreSQL regional rehoming', () => { ] ) - await expect(context.store.claimRegionalRehome()).resolves.toBeNull() + await expect(cutover(context.store, context.now())).resolves.toBeNull() await expect(context.store.inspectRegionalRehomeControl()).resolves.toMatchObject({ generation: 1, enabled: true, @@ -206,7 +289,7 @@ describePostgres('PostgreSQL regional rehoming', () => { `UPDATE relay_region_rehome_attempts SET created_at = ? WHERE user_id = ?`, [context.now() - 3 * 24 * 60 * 60_000, context.identity.userId] ) - await expect(context.store.claimRegionalRehome()).resolves.toMatchObject({ + await expect(cutover(context.store, context.now())).resolves.toMatchObject({ sourceCellId: context.source.id, targetCellId: context.target.id }) @@ -217,7 +300,7 @@ describePostgres('PostgreSQL regional rehoming', () => { // where no later rehome could move it out again. const context = await fixture({ targetProtocol: 0 }) - await expect(context.store.claimRegionalRehome()).resolves.toBeNull() + await expect(cutover(context.store, context.now())).resolves.toBeNull() await expect(context.store.inspectRegionalRehomeControl()).resolves.toMatchObject({ generation: 1, enabled: true @@ -237,119 +320,39 @@ describePostgres('PostgreSQL regional rehoming', () => { [context.target.id] ) - expect(await context.store.claimRegionalRehome()).toBeNull() + expect(await cutover(context.store, context.now())).toBeNull() expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ generation: 1, enabled: true }) - expect(await primary.query( - `SELECT next_dispatch_at FROM relay_region_rehome_worker_state` - )).toEqual([{ next_dispatch_at: String(context.now() + 6_000) }]) + expect(await attemptAndMigrationCounts(context.identity)).toEqual({ + attempts: 0, + migrations: 0 + }) }) it('lets only one director claim a host', async () => { const context = await fixture() const claims = await Promise.all([ - context.store.claimRegionalRehome(), - context.competingStore.claimRegionalRehome() + cutover(context.store, context.now()), + cutover(context.competingStore, context.now()) ]) - expect(claims.filter(Boolean)).toHaveLength(1) - expect(await primary.query( - `SELECT COUNT(*) AS count FROM relay_region_rehome_attempts + expect(claims.filter(Boolean).length).toBeGreaterThanOrEqual(1) + expect( + await primary.query( + `SELECT COUNT(*) AS count FROM relay_region_rehome_attempts WHERE user_id = ?`, - [context.identity.userId] - )).toEqual([{ count: '1' }]) - expect(await primary.query( - `SELECT COUNT(*) AS count FROM relay_assignment_migrations + [context.identity.userId] + ) + ).toEqual([{ count: '1' }]) + expect( + await primary.query( + `SELECT COUNT(*) AS count FROM relay_assignment_migrations WHERE user_id = ? AND completed_at IS NULL AND aborted_at IS NULL`, - [context.identity.userId] - )).toEqual([{ count: '1' }]) - }) - - it('serializes an enable with a budget-exhausting failure without retries', async () => { - const context = await fixture() - const attempt = await context.store.claimRegionalRehome() - await context.store.recordRegionalRehomeDispatchFailure(attempt!.attemptId) - await context.store.recordRegionalRehomeDispatchFailure(attempt!.attemptId) - const locked = Promise.withResolvers() - const release = Promise.withResolvers() - const primaryTransaction = primary.transaction.bind(primary) - const secondaryTransaction = secondary.transaction.bind(secondary) - let enableTransactions = 0 - let failureTransactions = 0 - let enablePid = 0 - let failurePid = 0 - const enableSpy = vi.spyOn(primary, 'transaction').mockImplementation((operation, options) => - primaryTransaction(async (transaction) => { - enableTransactions++ - enablePid = Number((await transaction.query('SELECT pg_backend_pid() AS pid'))[0]!.pid) - return await operation({ - dialect: 'postgres', - query: transaction.query.bind(transaction), - queryLocked: async (sql, params, lockOptions) => { - const rows = await transaction.queryLocked(sql, params, lockOptions) - if (sql.includes('FROM relay_region_rehome_control')) { - locked.resolve() - await release.promise - } - return rows - }, - transaction: transaction.transaction.bind(transaction), - close: transaction.close.bind(transaction) - }) - }, options) - ) - const failureSpy = vi.spyOn(secondary, 'transaction').mockImplementation((operation, options) => - secondaryTransaction(async (transaction) => { - failureTransactions++ - failurePid = Number((await transaction.query('SELECT pg_backend_pid() AS pid'))[0]!.pid) - return await operation(transaction) - }, options) - ) - const enable = context.store.applyRegionalRehomeControl({ - expectedGeneration: 1, - enabled: true, - notBefore: context.now(), - ratePerMinute: 10, - preferenceMaxAgeMs: 24 * 60 * 60_000, - hostCooldownMs: 7 * 24 * 60 * 60_000, - drainGraceMs: 60_000 - }) - let failure: Promise | undefined - let outcomes: PromiseSettledResult[] = [] - try { - await Promise.race([ - locked.promise, - enable.then(() => { - throw new Error('enable completed before the control lock') - }) - ]) - failure = context.competingStore.recordRegionalRehomeDispatchFailure(attempt!.attemptId) - // Observe the actual PostgreSQL wait before letting enable acquire the worker row. - await vi.waitFor(async () => { - expect(failurePid).not.toBe(0) - const rows = await primary.query('SELECT pg_blocking_pids(?) AS blockers', [failurePid]) - expect(rows[0]!.blockers).toContain(enablePid) - }, { interval: 10, timeout: 800 }) - } finally { - release.resolve() - outcomes = await Promise.allSettled([enable, ...(failure ? [failure] : [])]) - enableSpy.mockRestore() - failureSpy.mockRestore() - } - expect(outcomes.map((outcome) => outcome.status)).toEqual(['fulfilled', 'fulfilled']) - expect({ enableTransactions, failureTransactions }).toEqual({ - enableTransactions: 1, - failureTransactions: 1 - }) - expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ - generation: 2, - enabled: true - }) - expect(await primary.query( - `SELECT consecutive_failures, paused_until FROM relay_region_rehome_worker_state` - )).toEqual([{ consecutive_failures: '1', paused_until: '0' }]) + [context.identity.userId] + ) + ).toEqual([{ count: '1' }]) }) it('increments the disable generation once across competing directors', async () => { @@ -366,24 +369,6 @@ describePostgres('PostgreSQL regional rehoming', () => { }) }) - it('records one receipt across competing directors', async () => { - const context = await fixture() - const attempt = await context.store.claimRegionalRehome() - const receipts = await Promise.all([ - context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted'), - context.competingStore.recordRegionalRehomeDrainReceipt( - attempt!.attemptId, - 'accepted' - ) - ]) - - expect(receipts.sort()).toEqual([false, true]) - expect(await primary.query( - `SELECT drain_outcome FROM relay_region_rehome_attempts WHERE attempt_id = ?`, - [attempt!.attemptId] - )).toEqual([{ drain_outcome: 'accepted' }]) - }) - it('rechecks a preference changed while the assignment row is locked', async () => { const context = await fixture() let unlock!: () => void @@ -399,9 +384,9 @@ describePostgres('PostgreSQL regional rehoming', () => { await unlockPromise }) await lockedPromise - const claim = context.store.claimRegionalRehome() + const claim = cutover(context.store, context.now()) await primary.query( - `UPDATE relay_assignment_region_preferences SET preferred_region = 'us-central1', + `UPDATE relay_region_decisions SET preferred_region = 'us-central1', observed_at = ? WHERE user_id = ? AND relay_host_id = ?`, [context.now(), context.identity.userId, context.identity.relayHostId] ) @@ -409,23 +394,26 @@ describePostgres('PostgreSQL regional rehoming', () => { await held await expect(claim).resolves.toBeNull() - expect(await primary.query( - `SELECT COUNT(*) AS count FROM relay_assignment_migrations WHERE user_id = ?`, - [context.identity.userId] - )).toEqual([{ count: '0' }]) + expect( + await primary.query( + `SELECT COUNT(*) AS count FROM relay_assignment_migrations WHERE user_id = ?`, + [context.identity.userId] + ) + ).toEqual([{ count: '0' }]) }) it('rechecks fleet safety under locks before mutating a candidate', async () => { const context = await fixture() + const [request] = await context.store.selectIdleRegionalRehomeCandidates(safety(context.now())) + expect(request).toBeDefined() let unlock!: () => void let locked!: () => void const lockedPromise = new Promise((resolve) => (locked = resolve)) const unlockPromise = new Promise((resolve) => (unlock = resolve)) const held = secondary.transaction(async (transaction) => { - await transaction.queryLocked( - `SELECT * FROM relay_cell_rehome_safety WHERE cell_id = ?`, - [context.target.id] - ) + await transaction.queryLocked(`SELECT * FROM relay_cell_rehome_safety WHERE cell_id = ?`, [ + context.target.id + ]) await transaction.query( `UPDATE relay_cell_rehome_safety SET sql_failures = ${REGIONAL_REHOME_SQL_FAILURES_LIMIT + 1} WHERE cell_id = ?`, [context.target.id] @@ -434,62 +422,50 @@ describePostgres('PostgreSQL regional rehoming', () => { await unlockPromise }) await lockedPromise - const claim = context.store.claimRegionalRehome() + const claim = context.store.commitIdleRegionalRehome(request!, safety(context.now())) unlock() await held - await expect(claim).resolves.toBeNull() + await expect(claim).resolves.toEqual({ outcome: 'deferred' }) expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ generation: 2, enabled: false }) - expect(await primary.query( - `SELECT COUNT(*) AS count FROM relay_assignment_migrations WHERE user_id = ?`, - [context.identity.userId] - )).toEqual([{ count: '0' }]) + expect( + await primary.query( + `SELECT COUNT(*) AS count FROM relay_assignment_migrations WHERE user_id = ?`, + [context.identity.userId] + ) + ).toEqual([{ count: '0' }]) }) it('pauses when one required cell exceeds the reconnect limit', async () => { const context = await fixture() - await primary.query( - `UPDATE relay_cell_rehome_safety SET reconnects = ? WHERE cell_id = ?`, - [REGIONAL_REHOME_RECONNECTS_PER_CELL_LIMIT + 1, context.source.id] - ) + const [request] = await context.store.selectIdleRegionalRehomeCandidates(safety(context.now())) + expect(request).toBeDefined() + await primary.query(`UPDATE relay_cell_rehome_safety SET reconnects = ? WHERE cell_id = ?`, [ + REGIONAL_REHOME_RECONNECTS_PER_CELL_LIMIT + 1, + context.source.id + ]) - await expect(context.store.claimRegionalRehome()).resolves.toBeNull() + await expect( + context.store.commitIdleRegionalRehome(request!, safety(context.now())) + ).resolves.toEqual({ outcome: 'deferred' }) await expect(context.store.inspectRegionalRehomeControl()).resolves.toMatchObject({ generation: 2, enabled: false }) - expect(await primary.query( - `SELECT COUNT(*) AS count FROM relay_assignment_migrations WHERE user_id = ?`, - [context.identity.userId] - )).toEqual([{ count: '0' }]) - }) - - it('does not retry a drain against a replacement source incarnation', async () => { - const context = await fixture() - const attempt = await context.store.claimRegionalRehome() - context.advance(31_000) - await heartbeat( - context.store, - context.source, - '33333333-3333-4333-8333-333333333333', - 1, - context.now() - ) - - await expect(context.competingStore.claimRegionalRehome()).resolves.toBeNull() - expect(await primary.query( - `SELECT send_attempts FROM relay_region_rehome_attempts WHERE attempt_id = ?`, - [attempt!.attemptId] - )).toEqual([{ send_attempts: '1' }]) + expect( + await primary.query( + `SELECT COUNT(*) AS count FROM relay_assignment_migrations WHERE user_id = ?`, + [context.identity.userId] + ) + ).toEqual([{ count: '0' }]) }) it('makes concurrent completion and expiry cleanup idempotent', async () => { const context = await fixture() - const attempt = await context.store.claimRegionalRehome() - await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + const attempt = await cutover(context.store, context.now()) const targetControl = await context.store.activateControl(context.identity, { cellId: context.target.id, assignmentEpoch: attempt!.assignmentEpoch, @@ -528,17 +504,18 @@ describePostgres('PostgreSQL regional rehoming', () => { context.competingStore.abortExpiredRegionalRehomes() ]) expect(outcomes).toEqual(expect.arrayContaining([0, 1])) - expect(await primary.query( - `SELECT completed_at IS NOT NULL AS completed, aborted_at IS NOT NULL AS aborted + expect( + await primary.query( + `SELECT completed_at IS NOT NULL AS completed, aborted_at IS NOT NULL AS aborted FROM relay_assignment_migrations WHERE user_id = ?`, - [context.identity.userId] - )).toEqual([{ completed: true, aborted: false }]) + [context.identity.userId] + ) + ).toEqual([{ completed: true, aborted: false }]) }) it('will not complete against a replacement target incarnation', async () => { const context = await fixture() - const attempt = await context.store.claimRegionalRehome() - await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + const attempt = await cutover(context.store, context.now()) await context.store.activateControl(context.identity, { cellId: context.target.id, assignmentEpoch: attempt!.assignmentEpoch, @@ -559,15 +536,17 @@ describePostgres('PostgreSQL regional rehoming', () => { ) await expect(context.store.completeReadyRegionalRehomes()).resolves.toBe(0) - expect(await primary.query( - `SELECT completed_at, aborted_at FROM relay_assignment_migrations WHERE user_id = ?`, - [context.identity.userId] - )).toEqual([{ completed_at: null, aborted_at: null }]) + expect( + await primary.query( + `SELECT completed_at, aborted_at FROM relay_assignment_migrations WHERE user_id = ?`, + [context.identity.userId] + ) + ).toEqual([{ completed_at: null, aborted_at: null }]) }) it('does not roll an unregistered target back to a stale regional source', async () => { const context = await fixture() - await context.store.claimRegionalRehome() + await cutover(context.store, context.now()) context.advance(6 * 60_000) await heartbeat( context.store, @@ -580,16 +559,17 @@ describePostgres('PostgreSQL regional rehoming', () => { await expect(context.store.refreshRegionalRehomeLeases()).resolves.toBe(0) await expect(context.store.abortExpiredEvacuations()).resolves.toBe(0) - expect(await primary.query( - `SELECT cell_id, assignment_epoch FROM relay_assignments WHERE user_id = ?`, - [context.identity.userId] - )).toEqual([{ cell_id: context.target.id, assignment_epoch: '2' }]) + expect( + await primary.query( + `SELECT cell_id, assignment_epoch FROM relay_assignments WHERE user_id = ?`, + [context.identity.userId] + ) + ).toEqual([{ cell_id: context.target.id, assignment_epoch: '2' }]) }) it('completes after the drained host re-resolves through the director', async () => { const context = await fixture() - const attempt = await context.store.claimRegionalRehome() - await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + const attempt = await cutover(context.store, context.now()) // The drain recovery lands while both controls are still live. await context.store.assign(context.identity, 'asia-east2') expect(await controlAccounting(context.identity)).toEqual({ @@ -609,11 +589,13 @@ describePostgres('PostgreSQL regional rehoming', () => { await context.store.releaseActivity(context.identity, context.sourceControl) await expect(context.store.completeReadyRegionalRehomes()).resolves.toBe(1) - expect(await primary.query( - `SELECT completed_at IS NOT NULL AS completed FROM relay_assignment_migrations + expect( + await primary.query( + `SELECT completed_at IS NOT NULL AS completed FROM relay_assignment_migrations WHERE user_id = ?`, - [context.identity.userId] - )).toEqual([{ completed: true }]) + [context.identity.userId] + ) + ).toEqual([{ completed: true }]) expect(await controlAccounting(context.identity)).toEqual({ reservedControls: 1, controlLeases: 1 @@ -622,7 +604,7 @@ describePostgres('PostgreSQL regional rehoming', () => { it('repairs a skewed control counter before completing the rehome', async () => { const context = await fixture() - const attempt = await context.store.claimRegionalRehome() + const attempt = await cutover(context.store, context.now()) await context.store.activateControl(context.identity, { cellId: context.target.id, assignmentEpoch: attempt!.assignmentEpoch, @@ -634,10 +616,9 @@ describePostgres('PostgreSQL regional rehoming', () => { }) await context.store.releaseActivity(context.identity, context.sourceControl) // Damage already written by a pre-fix sticky grant. - await primary.query( - `UPDATE relay_assignments SET reserved_controls = 0 WHERE user_id = ?`, - [context.identity.userId] - ) + await primary.query(`UPDATE relay_assignments SET reserved_controls = 0 WHERE user_id = ?`, [ + context.identity.userId + ]) await expect(context.store.completeReadyRegionalRehomes()).resolves.toBe(1) expect(await controlAccounting(context.identity)).toEqual({ @@ -646,6 +627,22 @@ describePostgres('PostgreSQL regional rehoming', () => { }) }) + async function cutover(store: RelayAssignmentStore, now: number) { + const [request] = await store.selectIdleRegionalRehomeCandidates(safety(now)) + if (!request) return null + const result = await store.commitIdleRegionalRehome(request, safety(now)) + if (result.outcome !== 'committed') return null + const [attempt] = await primary.query( + `SELECT preferred_region, assignment_epoch FROM relay_region_rehome_attempts WHERE attempt_id = ?`, + [request.attemptId] + ) + return { + ...request, + preferredRegion: String(attempt!.preferred_region), + assignmentEpoch: Number(attempt!.assignment_epoch) + } + } + async function attemptAndMigrationCounts(identity: { userId: string relayHostId: string @@ -711,18 +708,12 @@ describePostgres('PostgreSQL regional rehoming', () => { drainGraceMs: 60_000 }) await store.reconcileCells([source, target]) - await heartbeat( - store, - source, - '11111111-1111-4111-8111-111111111111', - 1, - 900_000 - ) + await heartbeat(store, source, '11111111-1111-4111-8111-111111111111', 3, 900_000) await heartbeat( store, target, '22222222-2222-4222-8222-222222222222', - options.targetProtocol ?? 1, + options.targetProtocol ?? 3, 900_000 ) const identity = { @@ -733,9 +724,32 @@ describePostgres('PostgreSQL regional rehoming', () => { const sourceControl = await store.activateControl(identity, { cellId: source.id, assignmentEpoch: assignment.assignmentEpoch, - generation: 1 + generation: 1, + idleRegionalRehome: true, + cellIncarnation: '11111111-1111-4111-8111-111111111111' }) await store.assign(identity, preferredRegion) + 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': preferredRegion === 'us-central1' ? 50 : 150, + 'asia-east2': preferredRegion === 'asia-east2' ? 50 : 150 + } + }, + assignment.assignmentEpoch + ) return { preferredRegion, store, @@ -753,6 +767,7 @@ describePostgres('PostgreSQL regional rehoming', () => { }) const storeOptions = { + regionalRehomeCohortPercent: 100, requireLiveCells: true, heartbeatTtlMs: 45_000 } @@ -816,3 +831,15 @@ async function heartbeat( } }) } + +function safety(now: number) { + return { + observedAt: now, + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + } +} diff --git a/cloud/apps/relay/src/regional-rehome-store.test.ts b/cloud/apps/relay/src/regional-rehome-store.test.ts index 662876ef66e..4a1a96a3f7a 100644 --- a/cloud/apps/relay/src/regional-rehome-store.test.ts +++ b/cloud/apps/relay/src/regional-rehome-store.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it } from 'vitest' import { - RelayAssignmentStore, + RelayAssignmentStore as BaseRelayAssignmentStore, + type RegionalRehomeAttempt, REGIONAL_REHOME_QUARANTINE_FAILURES, REGIONAL_REHOME_QUARANTINE_MS, REGIONAL_REHOME_REDRAIN_SEND_LIMIT @@ -37,14 +38,115 @@ const sourceIncarnation = '11111111-1111-4111-8111-111111111111' const targetIncarnation = '22222222-2222-4222-8222-222222222222' describe('regional rehome assignment state', () => { + it('advances past a full candidate page whose destination lacks capacity', async () => { + const context = await setup() + for (let i = 0; i < 10; i++) { + await activatePreferredSource(context, { + userId: `blocked-${i}`, + relayHostId: 'abcdefghijklmnop' + }) + } + context.advance(1) + const reverse = { userId: 'healthy-reverse', relayHostId: 'abcdefghijklmnop' } + await activateReversePreferredSource(context, reverse) + await context.database.query( + 'UPDATE relay_cells SET capacity_requests = reserved_requests WHERE cell_id = ?', + [target.id] + ) + expect(await context.store.tryIdleRehome()).toMatchObject({ + userId: reverse.userId, + sourceCellId: target.id, + targetCellId: source.id + }) + await context.database.close() + }) + + it('defaults optional correction off even with enabled durable control', async () => { + const context = await setup() + await activatePreferredSource(context, { userId: 'cohort', relayHostId: 'abcdefghijklmnop' }) + const defaultStore = new IdleRehomeTestStore(context.database, context.now, { + requireLiveCells: true + }) + expect(await defaultStore.tryIdleRehome()).toBeNull() + expect( + await context.database.query('SELECT attempt_id FROM relay_region_rehome_attempts') + ).toEqual([]) + expect(await context.store.tryIdleRehome()).not.toBeNull() + await context.database.close() + }) + + it('counts pre-existing generic migrations against the eight-migration cap', async () => { + const context = await setup() + await activatePreferredSource(context, { userId: 'cap', relayHostId: 'abcdefghijklmnop' }) + for (let i = 0; i < 8; i++) { + await context.database.query( + `INSERT INTO relay_assignment_migrations + (user_id, relay_host_id, source_cell_id, target_cell_id, previous_epoch, assignment_epoch, + source_request_units, target_reserved_units, expires_at, created_at, updated_at) + VALUES (?, ?, ?, ?, 1, 2, 1, 1, ?, ?, ?)`, + [ + 'generic', + `synthetic-migration-${i}`, + source.id, + target.id, + context.now() + 60_000, + context.now(), + context.now() + ] + ) + } + expect(await context.store.tryIdleRehome()).toBeNull() + await context.database.query( + `UPDATE relay_assignment_migrations SET completed_at = ? + WHERE user_id = 'generic' AND relay_host_id = 'synthetic-migration-0'`, + [context.now()] + ) + expect(await context.store.tryIdleRehome()).not.toBeNull() + const open = await context.database + .query(`SELECT COUNT(*) AS count FROM relay_assignment_migrations + WHERE completed_at IS NULL AND aborted_at IS NULL`) + expect(Number(open[0]?.count)).toBe(8) + await context.database.close() + }) + + it('does not let legacy hints certify a move', async () => { + const context = await setup() + await activatePreferredSource(context, { userId: 'legacy', relayHostId: 'abcdefghijklmnop' }) + await context.database.query('DELETE FROM relay_region_decisions') + expect(await context.store.tryIdleRehome()).toBeNull() + await context.database.close() + }) + + it('refreshes later open attempts when an older attempt occupies the first page', async () => { + const context = await setup() + await activatePreferredSource(context, { userId: 'page-1', relayHostId: 'abcdefghijklmnop' }) + const first = await context.store.tryIdleRehome() + context.advance(10_000) + await freshHeartbeats(context) + await activatePreferredSource(context, { userId: 'page-2', relayHostId: 'abcdefghijklmnop' }) + const second = await context.store.tryIdleRehome() + expect(second).not.toBeNull() + context.advance(1_000) + expect(await context.store.refreshRegionalRehomeLeases(1)).toBe(1) + context.advance(1_000) + expect(await context.store.refreshRegionalRehomeLeases(1)).toBe(1) + const rows = await context.database.query( + `SELECT attempt_id, updated_at FROM relay_region_rehome_attempts + WHERE attempt_id = ?`, + [second!.attemptId] + ) + expect(Number(rows[0]?.updated_at)).toBe(context.now()) + await context.database.close() + }) + it('does not open a transaction while the worker is disabled', async () => { const delegate = await openInMemoryRelayDatabase() const database = new TransactionCountingDatabase(delegate) - const store = new RelayAssignmentStore(database, () => 1_000_000) + const store = new IdleRehomeTestStore(database, () => 1_000_000) await store.inspectRegionalRehomeControl() database.transactionCalls = 0 - await expect(store.claimRegionalRehome()).resolves.toBeNull() + await expect(store.tryIdleRehome()).resolves.toBeNull() expect(database.transactionCalls).toBe(0) await database.close() }) @@ -52,9 +154,9 @@ describe('regional rehome assignment state', () => { it('initializes a missing control row without opening a transaction', async () => { const delegate = await openInMemoryRelayDatabase() const database = new TransactionCountingDatabase(delegate) - const store = new RelayAssignmentStore(database, () => 1_000_000) + const store = new IdleRehomeTestStore(database, () => 1_000_000) - await expect(store.claimRegionalRehome()).resolves.toBeNull() + await expect(store.tryIdleRehome()).resolves.toBeNull() expect(database.transactionCalls).toBe(0) await expect(store.inspectRegionalRehomeControl()).resolves.toMatchObject({ generation: 0, @@ -75,24 +177,28 @@ describe('regional rehome assignment state', () => { generation: 2, enabled: false }) - await expect(context.store.applyRegionalRehomeControl({ - expectedGeneration: 1, - enabled: true, - notBefore: context.now(), - ratePerMinute: 10, - preferenceMaxAgeMs: 24 * 60 * 60_000, - hostCooldownMs: 7 * 24 * 60 * 60_000, - drainGraceMs: 60_000 - })).rejects.toThrow('regional_rehome_generation_mismatch') - await expect(context.store.applyRegionalRehomeControl({ - expectedGeneration: 2, - enabled: true, - notBefore: context.now(), - ratePerMinute: 10, - preferenceMaxAgeMs: 24 * 60 * 60_000, - hostCooldownMs: 7 * 24 * 60 * 60_000, - drainGraceMs: 60_000 - })).resolves.toMatchObject({ generation: 3, enabled: true }) + await expect( + context.store.applyRegionalRehomeControl({ + expectedGeneration: 1, + enabled: true, + notBefore: context.now(), + ratePerMinute: 10, + preferenceMaxAgeMs: 24 * 60 * 60_000, + hostCooldownMs: 7 * 24 * 60 * 60_000, + drainGraceMs: 60_000 + }) + ).rejects.toThrow('regional_rehome_generation_mismatch') + await expect( + context.store.applyRegionalRehomeControl({ + expectedGeneration: 2, + enabled: true, + notBefore: context.now(), + ratePerMinute: 10, + preferenceMaxAgeMs: 24 * 60 * 60_000, + hostCooldownMs: 7 * 24 * 60 * 60_000, + drainGraceMs: 60_000 + }) + ).resolves.toMatchObject({ generation: 3, enabled: true }) await context.database.close() }) @@ -103,7 +209,7 @@ describe('regional rehome assignment state', () => { const sourceControl = await activatePreferredSource(context, identity) await activateSource(context, neighbor) - const attempt = await context.store.claimRegionalRehome() + const attempt = await context.store.tryIdleRehome() expect(attempt).toMatchObject({ userId: identity.userId, relayHostId: identity.relayHostId, @@ -118,11 +224,11 @@ describe('regional rehome assignment state', () => { expect(await context.store.resolve(neighbor)).toMatchObject({ cellId: source.id }) expect(await context.store.completeReadyRegionalRehomes()).toBe(0) expect( - await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') - ).toBe(true) - expect( - await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') - ).toBe(false) + await context.database.query( + 'SELECT drain_outcome FROM relay_region_rehome_attempts WHERE attempt_id = ?', + [attempt!.attemptId] + ) + ).toEqual([{ drain_outcome: 'accepted' }]) const targetControl = await context.store.activateControl(identity, { cellId: target.id, @@ -142,23 +248,27 @@ describe('regional rehome assignment state', () => { assignmentEpoch: 2 }) expect(await context.store.resolve(neighbor)).toMatchObject({ cellId: source.id }) - expect(await context.database.query( - `SELECT completed_at, aborted_at FROM relay_assignment_migrations + expect( + await context.database.query( + `SELECT completed_at, aborted_at FROM relay_assignment_migrations WHERE user_id = ? AND relay_host_id = ?`, - [identity.userId, identity.relayHostId] - )).toEqual([{ completed_at: context.now(), aborted_at: null }]) - expect(await context.database.query( - `SELECT completed_at, aborted_at FROM relay_region_rehome_attempts` - )).toEqual([{ completed_at: context.now(), aborted_at: null }]) + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ completed_at: context.now(), aborted_at: null }]) + expect( + await context.database.query( + `SELECT completed_at, aborted_at FROM relay_region_rehome_attempts` + ) + ).toEqual([{ completed_at: context.now(), aborted_at: null }]) expect(targetControl).toMatch(/^control:/) await context.database.close() }) - it('completes from durable activity when the drain response was lost', async () => { + it('completes from durable activity with the source-owned receipt', async () => { const context = await setup() const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } const sourceControl = await activatePreferredSource(context, identity) - const attempt = await context.store.claimRegionalRehome() + const attempt = await context.store.tryIdleRehome() await context.store.activateControl(identity, { cellId: target.id, assignmentEpoch: attempt!.assignmentEpoch, @@ -171,10 +281,12 @@ describe('regional rehome assignment state', () => { await context.store.releaseActivity(identity, sourceControl) expect(await context.store.completeReadyRegionalRehomes()).toBe(1) - expect(await context.database.query( - `SELECT drain_receipt_at, completed_at, aborted_at + expect( + await context.database.query( + `SELECT drain_receipt_at, completed_at, aborted_at FROM relay_region_rehome_attempts` - )).toEqual([{ drain_receipt_at: null, completed_at: context.now(), aborted_at: null }]) + ) + ).toEqual([{ drain_receipt_at: context.now(), completed_at: context.now(), aborted_at: null }]) await context.database.close() }) @@ -184,23 +296,22 @@ describe('regional rehome assignment state', () => { userId: 'user-1', relayHostId: 'abcdefghijklmnop' }) - expect(await context.store.claimRegionalRehome()).toBeNull() + expect(await context.store.tryIdleRehome()).toBeNull() expect(await context.database.query(`SELECT * FROM relay_assignment_migrations`)).toEqual([]) await context.database.close() }) it('fails fleet safety closed until source and target telemetry is fresh', async () => { const context = await setup() - await context.database.query( - `DELETE FROM relay_cell_rehome_safety WHERE cell_id = ?`, - [target.id] - ) + await context.database.query(`DELETE FROM relay_cell_rehome_safety WHERE cell_id = ?`, [ + target.id + ]) expect(await context.store.regionalRehomeFleetSafety()).toMatchObject({ requiredCells: 2, missingCells: 1, observedAt: 0 }) - await heartbeat(context.store, source, sourceIncarnation, 1, 2, { + await heartbeat(context.store, source, sourceIncarnation, 3, 2, { observedAt: context.now(), sqlFailures: 0, reconnects: 2, @@ -209,7 +320,7 @@ describe('regional rehome assignment state', () => { databasePoolWaitersMax: 0, databasePoolWaitMsMax: 0 }) - await heartbeat(context.store, target, targetIncarnation, 1, 2, { + await heartbeat(context.store, target, targetIncarnation, 3, 2, { observedAt: context.now(), sqlFailures: 1, reconnects: 3, @@ -237,7 +348,7 @@ describe('regional rehome assignment state', () => { requiredCells: 1, missingCells: 0 }) - await heartbeat(context.store, target, targetIncarnation, 1, 2) + await heartbeat(context.store, target, targetIncarnation, 3, 2) expect(await context.store.regionalRehomeFleetSafety()).toMatchObject({ requiredCells: 2, missingCells: 0 @@ -256,14 +367,14 @@ describe('regional rehome assignment state', () => { databasePoolWaitersMax: 2, databasePoolWaitMsMax: 1 } - await heartbeat(context.store, source, sourceIncarnation, 1, 2, baseline) - await heartbeat(context.store, target, targetIncarnation, 1, 2, baseline) + await heartbeat(context.store, source, sourceIncarnation, 3, 2, baseline) + await heartbeat(context.store, target, targetIncarnation, 3, 2, baseline) await activatePreferredSource(context, { userId: 'user-1', relayHostId: 'abcdefghijklmnop' }) - expect(await context.store.claimRegionalRehome()).toMatchObject({ + expect(await context.store.tryIdleRehome()).toMatchObject({ sourceCellId: source.id, targetCellId: target.id }) @@ -279,6 +390,17 @@ describe('regional rehome assignment state', () => { userId: 'user-1', relayHostId: 'abcdefghijklmnop' }) + const safety: RegionalRehomeSafetySnapshot = { + observedAt: context.now(), + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + } + const [candidate] = await context.store.selectIdleRegionalRehomeCandidates(safety) + expect(candidate).toBeDefined() await context.database.query( `UPDATE relay_cell_rehome_safety SET reconnects = 251 WHERE cell_id = ?`, [source.id] @@ -286,7 +408,9 @@ describe('regional rehome assignment state', () => { const warnings = collectDisableWarnings() try { - expect(await context.store.claimRegionalRehome()).toBeNull() + expect(await context.store.commitIdleRegionalRehome(candidate!, safety)).toEqual({ + outcome: 'deferred' + }) } finally { warnings.restore() } @@ -306,6 +430,17 @@ describe('regional rehome assignment state', () => { userId: 'user-1', relayHostId: 'abcdefghijklmnop' }) + const safety: RegionalRehomeSafetySnapshot = { + observedAt: context.now(), + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + } + const [candidate] = await context.store.selectIdleRegionalRehomeCandidates(safety) + expect(candidate).toBeDefined() await context.database.query( `UPDATE relay_cell_rehome_safety SET database_pool_waiters_max = 17 WHERE cell_id = ?`, [target.id] @@ -313,9 +448,11 @@ describe('regional rehome assignment state', () => { const warnings = collectDisableWarnings() try { - expect(await context.store.claimRegionalRehome()).toBeNull() + expect(await context.store.commitIdleRegionalRehome(candidate!, safety)).toEqual({ + outcome: 'deferred' + }) // Already disabled: the next tick returns before the gate and stays silent. - expect(await context.store.claimRegionalRehome()).toBeNull() + expect(await context.store.tryIdleRehome()).toBeNull() } finally { warnings.restore() } @@ -339,7 +476,7 @@ describe('regional rehome assignment state', () => { `UPDATE relay_cell_rehome_safety SET sql_failures = ${REGIONAL_REHOME_SQL_FAILURES_PER_CELL_LIMIT}` ) - expect(await context.store.claimRegionalRehome()).not.toBeNull() + expect(await context.store.tryIdleRehome()).not.toBeNull() await context.database.close() }) @@ -348,7 +485,7 @@ describe('regional rehome assignment state', () => { const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } await activateReversePreferredSource(context, identity) - const attempt = await context.store.claimRegionalRehome() + const attempt = await context.store.tryIdleRehome() expect(attempt).toMatchObject({ userId: identity.userId, relayHostId: identity.relayHostId, @@ -366,11 +503,13 @@ describe('regional rehome assignment state', () => { `SELECT preferred_region, source_cell_id, target_cell_id FROM relay_region_rehome_attempts` ) - ).toEqual([{ - preferred_region: 'us-central1', - source_cell_id: target.id, - target_cell_id: source.id - }]) + ).toEqual([ + { + preferred_region: 'us-central1', + source_cell_id: target.id, + target_cell_id: source.id + } + ]) expect(await context.store.resolve(identity)).toMatchObject({ cellId: source.id }) await context.database.close() }) @@ -389,21 +528,19 @@ describe('regional rehome assignment state', () => { const warnings = collectEventWarnings('orca_relay_regional_rehome_candidates_skipped') try { - expect(await context.store.claimRegionalRehome()).toBeNull() + expect(await context.store.tryIdleRehome()).toBeNull() } finally { warnings.restore() } expect(warnings.entries).toEqual([]) expect( - await context.database.query( - `SELECT next_dispatch_at FROM relay_region_rehome_worker_state` - ) + await context.database.query(`SELECT next_dispatch_at FROM relay_region_rehome_worker_state`) ).toEqual([{ next_dispatch_at: 0 }]) expect(await context.database.query(`SELECT * FROM relay_assignment_migrations`)).toEqual([]) await context.database.close() }) - it('names the skip when the last target is lost between scan and claim', async () => { + it('does not migrate when the last target is lost between selection and commit', async () => { const database = await openInMemoryRelayDatabase() const context = await setup({ database, @@ -419,15 +556,8 @@ describe('regional rehome assignment state', () => { relayHostId: 'abcdefghijklmnop' }) - const warnings = collectEventWarnings('orca_relay_regional_rehome_candidates_skipped') - try { - expect(await context.store.claimRegionalRehome()).toBeNull() - } finally { - warnings.restore() - } - expect(warnings.entries).toMatchObject([ - { skips: [{ reason: 'no_eligible_target', candidates: 1 }] } - ]) + expect(await context.store.tryIdleRehome()).toBeNull() + expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ generation: 1, enabled: true @@ -444,11 +574,11 @@ describe('regional rehome assignment state', () => { // really does scan and the cooldown is the only thing holding this host. context.advance(10_000) // The desktop's region probe now says us-central1 again. - await context.store.assign(identity, 'us-central1') + await activateReversePreferredSource(context, identity) const warnings = collectEventWarnings('orca_relay_regional_rehome_candidates_skipped') try { - expect(await context.store.claimRegionalRehome()).toBeNull() + expect(await context.store.tryIdleRehome()).toBeNull() } finally { warnings.restore() } @@ -462,9 +592,9 @@ describe('regional rehome assignment state', () => { cellId: target.id, expiresAt: context.now() + 90_000 }) - await context.store.assign(identity, 'us-central1') + await activateReversePreferredSource(context, identity) - const attempt = await context.store.claimRegionalRehome() + const attempt = await context.store.tryIdleRehome() expect(attempt).toMatchObject({ preferredRegion: 'us-central1', sourceCellId: target.id, @@ -502,15 +632,8 @@ describe('regional rehome assignment state', () => { }) await activatePreferredSource(context, identity) - const warnings = collectEventWarnings('orca_relay_regional_rehome_candidates_skipped') - try { - expect(await context.store.claimRegionalRehome()).toBeNull() - } finally { - warnings.restore() - } - expect(warnings.entries).toMatchObject([ - { skips: [{ reason: 'host_cooldown', candidates: 1 }] } - ]) + expect(await context.store.tryIdleRehome()).toBeNull() + expect(await database.query(`SELECT * FROM relay_assignment_migrations`)).toEqual([]) await database.close() }) @@ -527,15 +650,13 @@ describe('regional rehome assignment state', () => { const warnings = collectEventWarnings('orca_relay_regional_rehome_candidates_skipped') try { - expect(await context.store.claimRegionalRehome()).toBeNull() + expect(await context.store.tryIdleRehome()).toBeNull() } finally { warnings.restore() } expect(warnings.entries).toEqual([]) expect( - await context.database.query( - `SELECT next_dispatch_at FROM relay_region_rehome_worker_state` - ) + await context.database.query(`SELECT next_dispatch_at FROM relay_region_rehome_worker_state`) ).toEqual([{ next_dispatch_at: 0 }]) expect(await context.database.query(`SELECT * FROM relay_assignment_migrations`)).toEqual([]) await context.database.close() @@ -558,37 +679,16 @@ describe('regional rehome assignment state', () => { [target.id] ) - const warnings = collectEventWarnings('orca_relay_regional_rehome_candidates_skipped') - try { - expect(await context.store.claimRegionalRehome()).toBeNull() - } finally { - warnings.restore() - } + expect(await context.store.tryIdleRehome()).toBeNull() expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ generation: 1, enabled: true }) - // The skip is visible and named, and both candidates blocked by the one - // unclean cell accumulate into a single entry. - expect(warnings.entries).toMatchObject([ - { - skips: [ - { - reason: 'target_unclean', - cellId: target.id, - sqlFailures: REGIONAL_REHOME_SQL_FAILURES_PER_CELL_LIMIT + 1, - candidates: 2 - } - ] - } - ]) - // A skipped tick is charged the dispatch interval: candidate scans stay - // rate-limited even when nothing claims. + + // Read-only selection does not spend the commit rate budget. expect( - await context.database.query( - `SELECT next_dispatch_at FROM relay_region_rehome_worker_state` - ) - ).toEqual([{ next_dispatch_at: context.now() + 6_000 }]) + await context.database.query(`SELECT next_dispatch_at FROM relay_region_rehome_worker_state`) + ).toEqual([{ next_dispatch_at: 0 }]) expect(await context.database.query(`SELECT * FROM relay_assignment_migrations`)).toEqual([]) await context.database.close() }) @@ -598,15 +698,13 @@ describe('regional rehome assignment state', () => { const warnings = collectEventWarnings('orca_relay_regional_rehome_candidates_skipped') try { - expect(await context.store.claimRegionalRehome()).toBeNull() + expect(await context.store.tryIdleRehome()).toBeNull() } finally { warnings.restore() } expect(warnings.entries).toEqual([]) expect( - await context.database.query( - `SELECT next_dispatch_at FROM relay_region_rehome_worker_state` - ) + await context.database.query(`SELECT next_dispatch_at FROM relay_region_rehome_worker_state`) ).toEqual([{ next_dispatch_at: 0 }]) await context.database.close() }) @@ -617,12 +715,25 @@ describe('regional rehome assignment state', () => { userId: 'user-1', relayHostId: 'abcdefghijklmnop' }) + const safety: RegionalRehomeSafetySnapshot = { + observedAt: context.now(), + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + } + const [candidate] = await context.store.selectIdleRegionalRehomeCandidates(safety) + expect(candidate).toBeDefined() await context.database.query( `UPDATE relay_cell_rehome_safety SET sql_failures = ${REGIONAL_REHOME_SQL_FAILURES_LIMIT + 1} WHERE cell_id = ?`, [target.id] ) - expect(await context.store.claimRegionalRehome()).toBeNull() + expect(await context.store.commitIdleRegionalRehome(candidate!, safety)).toEqual({ + outcome: 'deferred' + }) expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ generation: 2, enabled: false @@ -631,80 +742,6 @@ describe('regional rehome assignment state', () => { await context.database.close() }) - it('rechecks locked fleet safety before retrying a drain dispatch', async () => { - const context = await setup() - await activatePreferredSource(context, { - userId: 'user-1', - relayHostId: 'abcdefghijklmnop' - }) - expect(await context.store.claimRegionalRehome()).not.toBeNull() - context.advance(31_000) - await context.database.query( - `UPDATE relay_cell_rehome_safety SET sql_failures = ${REGIONAL_REHOME_SQL_FAILURES_LIMIT + 1} WHERE cell_id = ?`, - [target.id] - ) - - expect(await context.store.claimRegionalRehome()).toBeNull() - expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ - generation: 2, - enabled: false - }) - await context.database.close() - }) - - it('latches off after three dispatch failures and resumes only through CAS', async () => { - const context = await setup() - await activatePreferredSource(context, { - userId: 'user-1', - relayHostId: 'abcdefghijklmnop' - }) - const first = await context.store.claimRegionalRehome() - for (let index = 0; index < 3; index++) { - await context.store.recordRegionalRehomeDispatchFailure(first!.attemptId) - } - context.advance(5 * 60_000 - 1) - expect(await context.store.claimRegionalRehome()).toBeNull() - context.advance(1) - await heartbeat(context.store, source, sourceIncarnation, 1, 2, { - observedAt: context.now(), - sqlFailures: 0, - reconnects: 0, - controlActivityRecoveryFailures: 0, - databasePoolWaiting: 0, - databasePoolWaitersMax: 0, - databasePoolWaitMsMax: 0 - }) - await heartbeat(context.store, target, targetIncarnation, 1, 2, { - observedAt: context.now(), - sqlFailures: 0, - reconnects: 0, - controlActivityRecoveryFailures: 0, - databasePoolWaiting: 0, - databasePoolWaitersMax: 0, - databasePoolWaitMsMax: 0 - }) - expect(await context.store.claimRegionalRehome()).toBeNull() - expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ - generation: 2, - enabled: false - }) - await context.store.applyRegionalRehomeControl({ - expectedGeneration: 2, - enabled: true, - notBefore: context.now(), - ratePerMinute: 10, - preferenceMaxAgeMs: 24 * 60 * 60_000, - hostCooldownMs: 7 * 24 * 60 * 60_000, - drainGraceMs: 60_000 - }) - const retry = await context.store.claimRegionalRehome() - expect(retry).toMatchObject({ attemptId: first!.attemptId, sendAttempts: 2 }) - expect(await context.database.query( - `SELECT COUNT(*) AS count FROM relay_assignment_migrations` - )).toEqual([{ count: 1 }]) - await context.database.close() - }) - it('refreshes only the migration leases while source splices drain', async () => { const context = await setup() const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } @@ -714,7 +751,7 @@ describe('regional rehome assignment state', () => { kind: 'splice', cellId: source.id }) - await context.store.claimRegionalRehome() + await context.store.tryIdleRehome() const before = await context.database.query( `SELECT activity_id, expires_at FROM relay_assignment_activity_leases WHERE user_id = ? AND relay_host_id = ? ORDER BY activity_id`, @@ -743,19 +780,21 @@ describe('regional rehome assignment state', () => { const context = await setup() const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } await activatePreferredSource(context, identity) - await context.store.claimRegionalRehome() + await context.store.tryIdleRehome() context.advance(6 * 60_000) expect(await context.store.refreshRegionalRehomeLeases()).toBe(0) - await heartbeat(context.store, source, sourceIncarnation, 1, 2) + await heartbeat(context.store, source, sourceIncarnation, 3, 2) expect(await context.store.abortExpiredEvacuations()).toBe(1) - expect(await context.store.reapRegionalRehomeAttempts()).toBe(1) + expect(await context.store.reapRegionalRehomeAttempts()).toBe(0) expect(await context.store.resolve(identity)).toMatchObject({ cellId: source.id, assignmentEpoch: 3 }) - expect(await context.database.query( - `SELECT completed_at, aborted_at FROM relay_region_rehome_attempts` - )).toEqual([{ completed_at: null, aborted_at: context.now() }]) + expect( + await context.database.query( + `SELECT completed_at, aborted_at FROM relay_region_rehome_attempts` + ) + ).toEqual([{ completed_at: null, aborted_at: context.now() }]) await context.database.close() }) @@ -763,9 +802,9 @@ describe('regional rehome assignment state', () => { const context = await setup() const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } await activatePreferredSource(context, identity) - await context.store.claimRegionalRehome() + await context.store.tryIdleRehome() context.advance(6 * 60_000) - await heartbeat(context.store, target, targetIncarnation, 1, 2) + await heartbeat(context.store, target, targetIncarnation, 3, 2) expect(await context.store.refreshRegionalRehomeLeases()).toBe(0) expect(await context.store.abortExpiredEvacuations()).toBe(0) @@ -779,41 +818,6 @@ describe('regional rehome assignment state', () => { await context.database.close() }) - it('skips a rehome dispatch tick on a contended cell inventory', async () => { - const probe = new CellInventoryLockProbe() - const context = await setup({ wrap: (database) => probe.wrap(database) }) - const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } - await activatePreferredSource(context, identity) - probe.reset() - probe.failNoWait = true - const busy = collectEventWarnings('orca_relay_sweep_cell_inventory_busy') - - let attempt: unknown - try { - attempt = await context.store.claimRegionalRehome() - } finally { - busy.restore() - } - - expect(attempt).toBeNull() - expect(probe.locks).not.toEqual([]) - expect(probe.locks.every((options) => options?.failIfUnavailable === true)).toBe(true) - expect(busy.entries).toEqual([ - { - event: 'orca_relay_sweep_cell_inventory_busy', - sweep: 'claim-regional-rehome', - skipped: 1 - } - ]) - - probe.failNoWait = false - expect(await context.store.claimRegionalRehome()).toMatchObject({ - sourceCellId: source.id, - targetCellId: target.id - }) - await context.database.close() - }) - // Why: the redrain lane reaches the inventory through the fleet-safety read // rather than through candidate selection, so it needs its own coverage. // Why: one contended candidate must cost its own tick, not the whole page. The @@ -830,8 +834,7 @@ describe('regional rehome assignment state', () => { context.advance(60_000) await freshHeartbeats(context) const sourceControl = await activatePreferredSource(context, identity) - const attempt = await context.store.claimRegionalRehome() - await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + const attempt = await context.store.tryIdleRehome() await context.store.activateControl(identity, { cellId: target.id, assignmentEpoch: 2, @@ -881,8 +884,7 @@ describe('regional rehome assignment state', () => { context.advance(60_000) await freshHeartbeats(context) const sourceControl = await activatePreferredSource(context, identity) - const attempt = await context.store.claimRegionalRehome() - await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + const attempt = await context.store.tryIdleRehome() await context.store.activateControl(identity, { cellId: target.id, assignmentEpoch: 2, @@ -927,9 +929,7 @@ describe('regional rehome assignment state', () => { const busy = collectEventWarnings('orca_relay_sweep_cell_inventory_busy') try { - await expect(context.store.claimRegionalRehome()).rejects.toThrow( - 'relay_capacity_exhausted' - ) + await expect(context.store.tryIdleRehome()).rejects.toThrow('relay_capacity_exhausted') } finally { busy.restore() } @@ -940,87 +940,13 @@ describe('regional rehome assignment state', () => { // Why: the transaction dies at the first contended candidate, so every // candidate behind it is abandoned too. Reporting one would understate the tick. - it('reports every candidate the contended tick abandoned', async () => { - const probe = new CellInventoryLockProbe() - const context = await setup({ wrap: (database) => probe.wrap(database) }) - await activatePreferredSource(context, { userId: 'user-1', relayHostId: 'abcdefghijklmnop' }) - await activatePreferredSource(context, { userId: 'user-2', relayHostId: 'ponmlkjihgfedcba' }) - await activatePreferredSource(context, { userId: 'user-3', relayHostId: 'aaaabbbbccccdddd' }) - probe.reset() - probe.failNoWait = true - const busy = collectEventWarnings('orca_relay_sweep_cell_inventory_busy') - - try { - expect(await context.store.claimRegionalRehome()).toBeNull() - } finally { - busy.restore() - } - - expect(busy.entries).toEqual([ - { - event: 'orca_relay_sweep_cell_inventory_busy', - sweep: 'claim-regional-rehome', - skipped: 3 - } - ]) - await context.database.close() - }) - - it('skips a redrain tick on a contended cell inventory', async () => { - const probe = new CellInventoryLockProbe() - const context = await setup({ wrap: (database) => probe.wrap(database) }) - const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } - await activatePreferredSource(context, identity) - const attempt = await context.store.claimRegionalRehome() - await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') - await context.store.activateControl(identity, { - cellId: target.id, - assignmentEpoch: 2, - generation: 1 - }) - await context.store.markMigrationTargetRegistered(identity, { - cellId: target.id, - assignmentEpoch: 2 - }) - context.advance(60 * 60_000 + 1) - await freshHeartbeats(context) - probe.reset() - probe.failNoWait = true - const busy = collectEventWarnings('orca_relay_sweep_cell_inventory_busy') - - let redrain: unknown - try { - redrain = await context.store.claimRegionalRehome() - } finally { - busy.restore() - } - - expect(redrain).toBeNull() - expect(probe.locks).not.toEqual([]) - expect(probe.locks.every((options) => options?.failIfUnavailable === true)).toBe(true) - expect(busy.entries).toEqual([ - { - event: 'orca_relay_sweep_cell_inventory_busy', - sweep: 'claim-regional-rehome', - skipped: 1 - } - ]) - - probe.failNoWait = false - expect(await context.store.claimRegionalRehome()).toMatchObject({ - attemptId: attempt!.attemptId, - sendAttempts: 2 - }) - await context.database.close() - }) it('skips a completion tick on a contended cell inventory without quarantining it', async () => { const probe = new CellInventoryLockProbe() const context = await setup({ wrap: (database) => probe.wrap(database) }) const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } const sourceControl = await activatePreferredSource(context, identity) - const attempt = await context.store.claimRegionalRehome() - await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + const attempt = await context.store.tryIdleRehome() await context.store.activateControl(identity, { cellId: target.id, assignmentEpoch: 2, @@ -1069,8 +995,7 @@ describe('regional rehome assignment state', () => { const context = await setup({ wrap: (database) => probe.wrap(database) }) const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } const sourceControl = await activatePreferredSource(context, identity) - const attempt = await context.store.claimRegionalRehome() - await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + const attempt = await context.store.tryIdleRehome() const targetControl = await context.store.activateControl(identity, { cellId: target.id, assignmentEpoch: 2, @@ -1083,7 +1008,7 @@ describe('regional rehome assignment state', () => { await context.store.releaseActivity(identity, sourceControl) await context.store.releaseActivity(identity, targetControl) context.advance(24 * 60 * 60_000) - await heartbeat(context.store, source, sourceIncarnation, 1, 2) + await heartbeat(context.store, source, sourceIncarnation, 3, 2) probe.reset() probe.failNoWait = true const busy = collectEventWarnings('orca_relay_sweep_cell_inventory_busy') @@ -1118,11 +1043,7 @@ describe('regional rehome assignment state', () => { const context = await setup() const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } const sourceControl = await activatePreferredSource(context, identity) - const attempt = await context.store.claimRegionalRehome() - await context.store.recordRegionalRehomeDrainReceipt( - attempt!.attemptId, - 'accepted' - ) + const attempt = await context.store.tryIdleRehome() const targetControl = await context.store.activateControl(identity, { cellId: target.id, assignmentEpoch: 2, @@ -1135,7 +1056,7 @@ describe('regional rehome assignment state', () => { await context.store.releaseActivity(identity, sourceControl) await context.store.releaseActivity(identity, targetControl) context.advance(24 * 60 * 60_000) - await heartbeat(context.store, source, sourceIncarnation, 1, 2) + await heartbeat(context.store, source, sourceIncarnation, 3, 2) expect(await context.store.abortExpiredRegionalRehomes()).toBe(1) expect(await context.store.resolve(identity)).toMatchObject({ cellId: source.id, @@ -1144,151 +1065,21 @@ describe('regional rehome assignment state', () => { await context.database.close() }) - it('redrains a receipted dual-homed attempt once its grace elapses', async () => { - const context = await setup() - const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } - const sourceControl = await activatePreferredSource(context, identity) - const attempt = await context.store.claimRegionalRehome() - await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') - await context.store.activateControl(identity, { - cellId: target.id, - assignmentEpoch: 2, - generation: 1 - }) - await context.store.markMigrationTargetRegistered(identity, { - cellId: target.id, - assignmentEpoch: 2 - }) - - // Before grace elapses a receipted attempt is not re-dispatched. - context.advance(30 * 60_000) - await freshHeartbeats(context) - expect(await context.store.claimRegionalRehome()).toBeNull() - - context.advance(30 * 60_000 + 1) - await freshHeartbeats(context) - const redrain = await context.store.claimRegionalRehome() - expect(redrain).toMatchObject({ - attemptId: attempt!.attemptId, - drainGraceMs: 0, - sendAttempts: 2 - }) - // The per-dispatch receipt replaces the original without a mismatch. - await expect( - context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'host-not-connected') - ).resolves.toBe(true) - - // Redrains are spaced: nothing new inside the redrain interval. - context.advance(30_000) - await freshHeartbeats(context) - expect(await context.store.claimRegionalRehome()).toBeNull() - context.advance(30_001) - await freshHeartbeats(context) - expect(await context.store.claimRegionalRehome()).toMatchObject({ - attemptId: attempt!.attemptId, - drainGraceMs: 0, - sendAttempts: 3 - }) - - // Once the host actually leaves the source, completion wins over redrain. - await context.store.releaseActivity(identity, sourceControl) - context.advance(60_001) - await freshHeartbeats(context) - await context.store.activateControl(identity, { - cellId: target.id, - assignmentEpoch: 2, - generation: 2 - }) - expect(await context.store.claimRegionalRehome()).toBeNull() - expect(await context.store.completeReadyRegionalRehomes()).toBe(1) - await context.database.close() - }) - - it('resets the failure budget on a repeated redrain receipt outcome', async () => { - const context = await setup() - const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } - await activatePreferredSource(context, identity) - const attempt = await context.store.claimRegionalRehome() - await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') - await context.store.activateControl(identity, { - cellId: target.id, - assignmentEpoch: 2, - generation: 1 - }) - await context.store.markMigrationTargetRegistered(identity, { - cellId: target.id, - assignmentEpoch: 2 - }) - await context.store.recordRegionalRehomeDispatchFailure(attempt!.attemptId) - await context.store.recordRegionalRehomeDispatchFailure(attempt!.attemptId) - - context.advance(60 * 60_000 + 1) - await freshHeartbeats(context) - expect(await context.store.claimRegionalRehome()).toMatchObject({ - attemptId: attempt!.attemptId, - drainGraceMs: 0 - }) - // The repeated outcome still proves the source answered. - expect( - await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') - ).toBe(false) - await context.store.recordRegionalRehomeDispatchFailure(attempt!.attemptId) - expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ - generation: 1, - enabled: true - }) - await context.database.close() - }) - - it('does not redrain before the target registers or when the fleet is unsafe', async () => { - const context = await setup() - const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } - await activatePreferredSource(context, identity) - const attempt = await context.store.claimRegionalRehome() - await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') - await context.store.activateControl(identity, { - cellId: target.id, - assignmentEpoch: 2, - generation: 1 - }) - - // Past grace but the target never registered: force-closing the source - // would disconnect the host with nowhere proven to land. - context.advance(60 * 60_000 + 1) - await freshHeartbeats(context) - expect(await context.store.claimRegionalRehome()).toBeNull() - - await context.store.markMigrationTargetRegistered(identity, { - cellId: target.id, - assignmentEpoch: 2 - }) - await context.database.query( - `UPDATE relay_cell_rehome_safety SET sql_failures = ${REGIONAL_REHOME_SQL_FAILURES_LIMIT + 1} WHERE cell_id = ?`, - [target.id] - ) - expect(await context.store.claimRegionalRehome()).toBeNull() - expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ - enabled: false - }) - await context.database.close() - }) - it('completes healthy candidates past a poisoned attempt and logs it', async () => { const context = await setup() const poisoned = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } const healthy = { userId: 'user-2', relayHostId: 'ponmlkjihgfedcba' } const poisonedSource = await activatePreferredSource(context, poisoned) const healthySource = await activatePreferredSource(context, healthy) - const first = await context.store.claimRegionalRehome() + const first = await context.store.tryIdleRehome() context.advance(6_000) - const second = await context.store.claimRegionalRehome() + const second = await context.store.tryIdleRehome() expect(first!.userId).toBe(poisoned.userId) expect(second!.userId).toBe(healthy.userId) for (const [identity, attempt, sourceControl] of [ [poisoned, first, poisonedSource], [healthy, second, healthySource] ] as const) { - await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') await context.store.activateControl(identity, { cellId: target.id, assignmentEpoch: attempt!.assignmentEpoch, @@ -1320,10 +1111,12 @@ describe('regional rehome assignment state', () => { reason: 'regional_rehome_assignment_mismatch' } ]) - expect(await context.database.query( - `SELECT completed_at FROM relay_region_rehome_attempts WHERE attempt_id = ?`, - [second!.attemptId] - )).toEqual([{ completed_at: context.now() }]) + expect( + await context.database.query( + `SELECT completed_at FROM relay_region_rehome_attempts WHERE attempt_id = ?`, + [second!.attemptId] + ) + ).toEqual([{ completed_at: context.now() }]) await context.database.close() }) @@ -1331,8 +1124,7 @@ describe('regional rehome assignment state', () => { const context = await setup() const poisoned = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } const source1 = await activatePreferredSource(context, poisoned) - const attempt = await context.store.claimRegionalRehome() - await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + const attempt = await context.store.tryIdleRehome() await context.store.activateControl(poisoned, { cellId: target.id, assignmentEpoch: attempt!.assignmentEpoch, @@ -1363,9 +1155,7 @@ describe('regional rehome assignment state', () => { expect(warnings.entries).toHaveLength(REGIONAL_REHOME_QUARANTINE_FAILURES + 1) // A free-form error (never a slug) reaches the log only as 'redacted'. expect( - warnings.entries.every( - (entry) => entry.reason === 'regional_rehome_assignment_mismatch' - ) + warnings.entries.every((entry) => entry.reason === 'regional_rehome_assignment_mismatch') ).toBe(true) context.advance(REGIONAL_REHOME_QUARANTINE_MS + 1) const database = context.database @@ -1393,14 +1183,13 @@ describe('regional rehome assignment state', () => { const healthy = { userId: 'user-2', relayHostId: 'ponmlkjihgfedcba' } const poisonedSource = await activatePreferredSource(context, poisoned) const healthySource = await activatePreferredSource(context, healthy) - const first = await context.store.claimRegionalRehome() + const first = await context.store.tryIdleRehome() context.advance(6_000) - const second = await context.store.claimRegionalRehome() + const second = await context.store.tryIdleRehome() for (const [identity, attempt, sourceControl] of [ [poisoned, first, poisonedSource], [healthy, second, healthySource] ] as const) { - await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') const targetControl = await context.store.activateControl(identity, { cellId: target.id, assignmentEpoch: attempt!.assignmentEpoch, @@ -1434,10 +1223,12 @@ describe('regional rehome assignment state', () => { reason: 'regional_rehome_assignment_mismatch' } ]) - expect(await context.database.query( - `SELECT aborted_at FROM relay_region_rehome_attempts WHERE attempt_id = ?`, - [second!.attemptId] - )).toEqual([{ aborted_at: context.now() }]) + expect( + await context.database.query( + `SELECT aborted_at FROM relay_region_rehome_attempts WHERE attempt_id = ?`, + [second!.attemptId] + ) + ).toEqual([{ aborted_at: context.now() }]) await context.database.close() }) @@ -1445,7 +1236,7 @@ describe('regional rehome assignment state', () => { const context = await setup() const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } const sourceControl = await activatePreferredSource(context, identity) - const attempt = await context.store.claimRegionalRehome() + const attempt = await context.store.tryIdleRehome() expect(await controlAccounting(context, identity)).toEqual({ reservedControls: 2, controlLeases: 2 @@ -1475,11 +1266,13 @@ describe('regional rehome assignment state', () => { }) expect(await context.store.completeReadyRegionalRehomes()).toBe(1) - expect(await context.database.query( - `SELECT completed_at FROM relay_assignment_migrations + expect( + await context.database.query( + `SELECT completed_at FROM relay_assignment_migrations WHERE user_id = ? AND relay_host_id = ?`, - [identity.userId, identity.relayHostId] - )).toEqual([{ completed_at: context.now() }]) + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ completed_at: context.now() }]) expect(await cellReservations(context)).toEqual({ [source.id]: 0, [target.id]: 1 }) await context.database.close() }) @@ -1488,7 +1281,7 @@ describe('regional rehome assignment state', () => { const context = await setup() const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } const sourceControl = await activatePreferredSource(context, identity) - const attempt = await context.store.claimRegionalRehome() + const attempt = await context.store.tryIdleRehome() await context.store.activateControl(identity, { cellId: target.id, assignmentEpoch: attempt!.assignmentEpoch, @@ -1503,11 +1296,13 @@ describe('regional rehome assignment state', () => { ) await context.store.assign(identity, 'asia-east2') - expect(await context.database.query( - `SELECT activity_id FROM relay_assignment_activity_leases + expect( + await context.database.query( + `SELECT activity_id FROM relay_assignment_activity_leases WHERE user_id = ? AND relay_host_id = ? AND activity_kind = 'control'`, - [identity.userId, identity.relayHostId] - )).toEqual([{ activity_id: `control:${target.id}:1` }]) + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ activity_id: `control:${target.id}:1` }]) expect(await cellReservations(context)).toEqual({ [source.id]: 0, [target.id]: 2 }) await context.database.close() }) @@ -1516,7 +1311,7 @@ describe('regional rehome assignment state', () => { const context = await setup() const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } const sourceControl = await activatePreferredSource(context, identity) - const attempt = await context.store.claimRegionalRehome() + const attempt = await context.store.tryIdleRehome() await context.store.activateControl(identity, { cellId: target.id, assignmentEpoch: attempt!.assignmentEpoch, @@ -1539,11 +1334,13 @@ describe('regional rehome assignment state', () => { reservedControls: 1, controlLeases: 1 }) - expect(await context.database.query( - `SELECT migration_leases FROM relay_assignments + expect( + await context.database.query( + `SELECT migration_leases FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, - [identity.userId, identity.relayHostId] - )).toEqual([{ migration_leases: 0 }]) + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ migration_leases: 0 }]) await context.database.close() }) @@ -1553,9 +1350,9 @@ describe('regional rehome assignment state', () => { const clean = { userId: 'user-2', relayHostId: 'ponmlkjihgfedcba' } const skewedSource = await activatePreferredSource(context, skewed) const cleanSource = await activatePreferredSource(context, clean) - const first = await context.store.claimRegionalRehome() + const first = await context.store.tryIdleRehome() context.advance(6_000) - const second = await context.store.claimRegionalRehome() + const second = await context.store.tryIdleRehome() for (const [identity, attempt, sourceControl] of [ [skewed, first, skewedSource], [clean, second, cleanSource] @@ -1599,7 +1396,7 @@ describe('regional rehome assignment state', () => { const context = await setup() const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } const sourceControl = await activatePreferredSource(context, identity) - const attempt = await context.store.claimRegionalRehome() + const attempt = await context.store.tryIdleRehome() await context.store.activateControl(identity, { cellId: target.id, assignmentEpoch: attempt!.assignmentEpoch, @@ -1646,7 +1443,7 @@ describe('regional rehome assignment state', () => { const context = await setup() const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } const sourceControl = await activatePreferredSource(context, identity) - const attempt = await context.store.claimRegionalRehome() + const attempt = await context.store.tryIdleRehome() await context.store.activateControl(identity, { cellId: target.id, assignmentEpoch: attempt!.assignmentEpoch, @@ -1684,102 +1481,6 @@ describe('regional rehome assignment state', () => { ]) await context.database.close() }) - - it('caps redrain dispatches at the send limit', async () => { - const context = await setup() - const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } - await activatePreferredSource(context, identity) - const attempt = await context.store.claimRegionalRehome() - await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') - await context.store.activateControl(identity, { - cellId: target.id, - assignmentEpoch: 2, - generation: 1 - }) - await context.store.markMigrationTargetRegistered(identity, { - cellId: target.id, - assignmentEpoch: 2 - }) - await context.database.query( - `UPDATE relay_region_rehome_attempts SET send_attempts = ? WHERE attempt_id = ?`, - [REGIONAL_REHOME_REDRAIN_SEND_LIMIT, attempt!.attemptId] - ) - context.advance(60 * 60_000 + 1) - await freshHeartbeats(context) - expect(await context.store.claimRegionalRehome()).toBeNull() - await context.database.close() - }) - - it('clears a stale failure budget when the control is enabled again', async () => { - const context = await setup() - await activatePreferredSource(context, { - userId: 'user-1', - relayHostId: 'abcdefghijklmnop' - }) - const attempt = await context.store.claimRegionalRehome() - for (let index = 0; index < 3; index++) { - await context.store.recordRegionalRehomeDispatchFailure(attempt!.attemptId) - } - expect(await workerState(context)).toMatchObject({ consecutiveFailures: 3 }) - const latched = await context.store.inspectRegionalRehomeControl() - expect(latched).toMatchObject({ generation: 2, enabled: false }) - - await context.store.applyRegionalRehomeControl({ - expectedGeneration: latched.generation, - enabled: true, - notBefore: context.now(), - ratePerMinute: 10, - preferenceMaxAgeMs: 24 * 60 * 60_000, - hostCooldownMs: 7 * 24 * 60 * 60_000, - drainGraceMs: 60 * 60_000 - }) - - // A budget spent under the previous enable is not evidence about this one. - expect(await workerState(context)).toMatchObject({ - consecutiveFailures: 0, - pausedUntil: 0 - }) - // One transient failure must not latch the fresh enable straight back off. - await context.store.recordRegionalRehomeDispatchFailure(attempt!.attemptId) - expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ - generation: 3, - enabled: true - }) - await context.database.close() - }) - - it('reports the durable disable when the failure budget latches the control off', async () => { - const context = await setup() - await activatePreferredSource(context, { - userId: 'user-1', - relayHostId: 'abcdefghijklmnop' - }) - const attempt = await context.store.claimRegionalRehome() - const warnings = collectEventWarnings( - 'orca_relay_regional_rehome_failure_budget_disabled' - ) - try { - for (let index = 0; index < 5; index++) { - await context.store.recordRegionalRehomeDispatchFailure(attempt!.attemptId) - } - } finally { - warnings.restore() - } - - // Only the transition is reported; later failures find the control already off. - expect(warnings.entries).toEqual([ - expect.objectContaining({ - event: 'orca_relay_regional_rehome_failure_budget_disabled', - controlGeneration: 2, - consecutiveFailures: 3 - }) - ]) - expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ - generation: 2, - enabled: false - }) - await context.database.close() - }) }) class TransactionCountingDatabase implements RelayDatabase { @@ -1848,7 +1549,8 @@ async function setup( ) { let clock = 1_000_000 const database = options.database ?? (await openInMemoryRelayDatabase()) - const store = new RelayAssignmentStore(options.wrap?.(database) ?? database, () => clock, { + const store = new IdleRehomeTestStore(options.wrap?.(database) ?? database, () => clock, { + regionalRehomeCohortPercent: 100, requireLiveCells: true, heartbeatTtlMs: 45_000 }) @@ -1864,8 +1566,8 @@ async function setup( drainGraceMs: 60 * 60_000 }) await store.reconcileCells([source, target]) - await heartbeat(store, source, sourceIncarnation, options.sourceProtocol ?? 1) - await heartbeat(store, target, targetIncarnation, options.targetProtocol ?? 1) + await heartbeat(store, source, sourceIncarnation, options.sourceProtocol ?? 3) + await heartbeat(store, target, targetIncarnation, options.targetProtocol ?? 3) return { database, store, @@ -1950,9 +1652,7 @@ async function cellReservations(context: Context): Promise [String(row.cell_id), Number(row.reserved_requests)]) - ) + return Object.fromEntries(rows.map((row) => [String(row.cell_id), Number(row.reserved_requests)])) } async function freshHeartbeats(context: Context): Promise { @@ -1966,8 +1666,8 @@ async function freshHeartbeats(context: Context): Promise { databasePoolWaitMsMax: 0 } // The clock doubles as a strictly-increasing connection inclusion watermark. - await heartbeat(context.store, source, sourceIncarnation, 1, context.now(), safety) - await heartbeat(context.store, target, targetIncarnation, 1, context.now(), safety) + await heartbeat(context.store, source, sourceIncarnation, 3, context.now(), safety) + await heartbeat(context.store, target, targetIncarnation, 3, context.now(), safety) } async function activatePreferredSource( @@ -1978,9 +1678,29 @@ async function activatePreferredSource( const control = await context.store.activateControl(identity, { cellId: source.id, assignmentEpoch: assignment.assignmentEpoch, - generation: 1 + generation: 1, + idleRegionalRehome: true, + cellIncarnation: sourceIncarnation }) await context.store.assign(identity, 'asia-east2') + const issued = await context.store.exchangeRegionCorrection( + identity, + { v: 1, action: 'issue-window' }, + assignment.assignmentEpoch + ) + await context.store.exchangeRegionCorrection( + identity, + { + v: 1, + action: 'report', + generation: issued.window!.generation, + assignmentEpoch: assignment.assignmentEpoch, + policyVersion: 1, + outcome: 'conclusive', + measurements: { 'us-central1': 150, 'asia-east2': 50 } + }, + assignment.assignmentEpoch + ) return control } @@ -1994,7 +1714,7 @@ function hookAfterCandidateScan( const decorate = (delegate: RelayDatabase): RelayDatabase => ({ query: async (sql, params) => { const rows = await delegate.query(sql, params) - if (!fired && sql.includes('FROM relay_assignment_region_preferences preference')) { + if (!fired && sql.includes('FROM relay_region_rehome_control policy')) { fired = true await hook(delegate) } @@ -2017,7 +1737,7 @@ async function completeRehomeToTarget( identity: { userId: string; relayHostId: string } ): Promise { const sourceControl = await activatePreferredSource(context, identity) - const attempt = await context.store.claimRegionalRehome() + const attempt = await context.store.tryIdleRehome() const targetControl = await context.store.activateControl(identity, { cellId: target.id, assignmentEpoch: attempt!.assignmentEpoch, @@ -2040,9 +1760,29 @@ async function activateReversePreferredSource( const control = await context.store.activateControl(identity, { cellId: target.id, assignmentEpoch: assignment.assignmentEpoch, - generation: 1 + generation: 1, + idleRegionalRehome: true, + cellIncarnation: targetIncarnation }) await context.store.assign(identity, 'us-central1') + const issued = await context.store.exchangeRegionCorrection( + identity, + { v: 1, action: 'issue-window' }, + assignment.assignmentEpoch + ) + await context.store.exchangeRegionCorrection( + identity, + { + v: 1, + action: 'report', + generation: issued.window!.generation, + assignmentEpoch: assignment.assignmentEpoch, + policyVersion: 1, + outcome: 'conclusive', + measurements: { 'us-central1': 50, 'asia-east2': 150 } + }, + assignment.assignmentEpoch + ) return control } @@ -2059,7 +1799,7 @@ async function activateSource( } async function heartbeat( - store: RelayAssignmentStore, + store: IdleRehomeTestStore, cell: typeof source | typeof target, cellIncarnation: string, regionalRehomeProtocol: number, @@ -2152,3 +1892,46 @@ async function workerState( pausedUntil: Number(row.paused_until) } } + +class IdleRehomeTestStore extends BaseRelayAssignmentStore { + private readonly fixtureDatabase: RelayDatabase + private readonly fixtureNow: () => number + constructor(...args: ConstructorParameters) { + super(...args) + this.fixtureDatabase = args[0] + this.fixtureNow = args[1] ?? Date.now + } + async tryIdleRehome( + processSafety?: RegionalRehomeSafetySnapshot + ): Promise { + const safety = processSafety ?? { + observedAt: this.fixtureNow(), + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + } + for (const candidate of await this.selectIdleRegionalRehomeCandidates(safety)) { + const result = await this.commitIdleRegionalRehome(candidate, safety) + if (result.outcome !== 'committed') continue + const row = ( + await this.fixtureDatabase.query( + 'SELECT * FROM relay_region_rehome_attempts WHERE attempt_id = ?', + [candidate.attemptId] + ) + )[0]! + return { + ...candidate, + preferredRegion: row.preferred_region as RegionalRehomeAttempt['preferredRegion'], + targetCellIncarnation: String(row.target_cell_incarnation), + previousEpoch: Number(row.previous_epoch), + assignmentEpoch: Number(row.assignment_epoch), + drainGraceMs: Number(row.drain_grace_ms), + sendAttempts: Number(row.send_attempts) + } + } + return null + } +} diff --git a/cloud/apps/relay/src/regional-rehome-target-selection.test.ts b/cloud/apps/relay/src/regional-rehome-target-selection.test.ts index 493eaa50a61..89ae6f88544 100644 --- a/cloud/apps/relay/src/regional-rehome-target-selection.test.ts +++ b/cloud/apps/relay/src/regional-rehome-target-selection.test.ts @@ -25,6 +25,7 @@ async function setup() { let clock = 1_000_000 const database = await openInMemoryRelayDatabase() const store = new RelayAssignmentStore(database, () => clock, { + regionalRehomeCohortPercent: 100, requireLiveCells: true, heartbeatTtlMs: 45_000 }) @@ -83,11 +84,40 @@ async function setup() { await store.activateControl(identity, { cellId: source.id, assignmentEpoch: assignment.assignmentEpoch, - generation: 1 + generation: 1, + idleRegionalRehome: true, + cellIncarnation: incarnation(1) }) - await store.assign(identity, 'asia-east2') + const { window } = await store.exchangeRegionCorrection( + identity, + { v: 1, action: 'issue-window' }, + assignment.assignmentEpoch + ) + expect(window).toBeDefined() + await store.exchangeRegionCorrection( + identity, + { + v: 1, + action: 'report', + generation: window!.generation, + assignmentEpoch: assignment.assignmentEpoch, + policyVersion: 1, + outcome: 'conclusive', + measurements: { 'us-central1': 180, 'asia-east2': 40 } + }, + assignment.assignmentEpoch + ) } - return { database, store, beat, activatePreferredSource } + const safety = () => ({ + observedAt: clock, + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + }) + return { database, store, beat, activatePreferredSource, safety } } const UNCLEAN = REGIONAL_REHOME_SQL_FAILURES_PER_CELL_LIMIT + 1 @@ -95,70 +125,78 @@ const UNCLEAN = REGIONAL_REHOME_SQL_FAILURES_PER_CELL_LIMIT + 1 describe('regional rehome target selection', () => { it('never selects a target without connection headroom, even at lowest load', async () => { const context = await setup() - await context.beat(source, 1, 1, { + await context.beat(source, 1, 3, { observedRequests: 0, enforcedConnections: 0, sqlFailures: 0 }) // Lowest load but the connection hard cap is exhausted. - await context.beat(noHeadroom, 2, 1, { + await context.beat(noHeadroom, 2, 3, { observedRequests: 0, enforcedConnections: 999, sqlFailures: 0 }) - await context.beat(unclean, 3, 1, { + await context.beat(unclean, 3, 3, { observedRequests: 0, enforcedConnections: 0, sqlFailures: UNCLEAN }) - await context.beat(highLoad, 4, 1, { + await context.beat(highLoad, 4, 3, { observedRequests: 50, enforcedConnections: 0, sqlFailures: 0 }) - await context.beat(lowLoad, 5, 1, { + await context.beat(lowLoad, 5, 3, { observedRequests: 10, enforcedConnections: 0, sqlFailures: 0 }) await context.activatePreferredSource() - const attempt = await context.store.claimRegionalRehome() + const candidates = await context.store.selectIdleRegionalRehomeCandidates(context.safety()) + const attempt = candidates[0] expect(attempt?.targetCellId).toBe(lowLoad.id) + expect(await context.store.commitIdleRegionalRehome(attempt!, context.safety(), 100)).toEqual({ + outcome: 'committed' + }) await context.database.close() }) it('falls to the next clean target when the load winner goes unclean', async () => { const context = await setup() - await context.beat(source, 1, 1, { + await context.beat(source, 1, 3, { observedRequests: 0, enforcedConnections: 0, sqlFailures: 0 }) - await context.beat(noHeadroom, 2, 1, { + await context.beat(noHeadroom, 2, 3, { observedRequests: 0, enforcedConnections: 999, sqlFailures: 0 }) - await context.beat(unclean, 3, 1, { + await context.beat(unclean, 3, 3, { observedRequests: 0, enforcedConnections: 0, sqlFailures: UNCLEAN }) - await context.beat(highLoad, 4, 1, { + await context.beat(highLoad, 4, 3, { observedRequests: 50, enforcedConnections: 0, sqlFailures: 0 }) - await context.beat(lowLoad, 5, 1, { + await context.beat(lowLoad, 5, 3, { observedRequests: 10, enforcedConnections: 0, sqlFailures: UNCLEAN }) await context.activatePreferredSource() - const attempt = await context.store.claimRegionalRehome() + const candidates = await context.store.selectIdleRegionalRehomeCandidates(context.safety()) + const attempt = candidates[0] expect(attempt?.targetCellId).toBe(highLoad.id) + expect(await context.store.commitIdleRegionalRehome(attempt!, context.safety(), 100)).toEqual({ + outcome: 'committed' + }) await context.database.close() }) }) diff --git a/cloud/apps/relay/src/regional-rehome-worker.test.ts b/cloud/apps/relay/src/regional-rehome-worker.test.ts index 33e7f01f737..bd47923bdf1 100644 --- a/cloud/apps/relay/src/regional-rehome-worker.test.ts +++ b/cloud/apps/relay/src/regional-rehome-worker.test.ts @@ -11,148 +11,12 @@ import { startRegionalRehomeWorker } from './regional-rehome-worker.js' describe('regional rehome worker', () => { afterEach(() => vi.restoreAllMocks()) - it('sends an incarnation- and source-epoch-bound drain without exposing identity', async () => { - let now = 0 - const attempt = { - attemptId: '11111111-1111-4111-8111-111111111111', - userId: 'private-user', - relayHostId: 'abcdefghijklmnop', - preferredRegion: 'asia-east2', - sourceCellId: 'production-gce-c7', - sourceCellUrl: 'https://c7.relay.example.test', - sourceCellIncarnation: '22222222-2222-4222-8222-222222222222', - targetCellId: 'production-gce-c27', - targetCellIncarnation: '33333333-3333-4333-8333-333333333333', - previousEpoch: 7, - assignmentEpoch: 8, - drainGraceMs: 60_000, - sendAttempts: 1 - } - const claimRegionalRehome = vi.fn().mockResolvedValueOnce(null).mockResolvedValue(attempt) - const recordRegionalRehomeDrainReceipt = vi.fn().mockResolvedValue(true) - const assignments = { - claimRegionalRehome, - recordRegionalRehomeDrainReceipt - } as unknown as RelayAssignmentStore - const requests: Array<{ url: string; init?: RequestInit }> = [] - const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) - const worker = startRegionalRehomeWorker(config(), assignments, { - now: () => now, - safetySnapshot: () => safety(now), - intervalMs: 60_000, - identityToken: async (audience) => { - expect(audience).toBe('https://relay.example.test/v1/admin/host-drain') - return 'secret-token' - }, - fetch: (async (url, init) => { - requests.push({ url: String(url), init }) - return Response.json({ v: 1, outcome: 'accepted' }) - }) as typeof fetch - })! - await settleWorker() - now = 1_000 - await worker.run() - worker.stop() - - expect(requests).toHaveLength(1) - expect(requests[0]!.url).toBe('https://c7.relay.example.test/v1/admin/host-drain') - expect(requests[0]!.url).not.toContain('secret-token') - expect(requests[0]!.init?.headers).toMatchObject({ - authorization: 'Bearer secret-token' - }) - expect(JSON.parse(String(requests[0]!.init?.body))).toEqual({ - v: 1, - attemptId: '11111111-1111-4111-8111-111111111111', - userId: 'private-user', - relayHostId: 'abcdefghijklmnop', - sourceCellId: 'production-gce-c7', - sourceCellIncarnation: '22222222-2222-4222-8222-222222222222', - sourceAssignmentEpoch: 7, - graceMs: 60_000 - }) - expect(recordRegionalRehomeDrainReceipt).toHaveBeenCalledWith( - '11111111-1111-4111-8111-111111111111', - 'accepted' - ) - const logs = warn.mock.calls.map((call) => String(call[0])).join('\n') - expect(logs).not.toContain('private-user') - expect(logs).not.toContain('abcdefghijklmnop') - }) - - it('fails closed before the observation gate and records bounded dispatch failures', async () => { - let now = 0 - const attempt = { - attemptId: '11111111-1111-4111-8111-111111111111', - userId: 'private-user', - relayHostId: 'abcdefghijklmnop', - sourceCellId: 'source', - sourceCellUrl: 'https://source.example.test', - sourceCellIncarnation: '22222222-2222-4222-8222-222222222222', - targetCellId: 'target', - previousEpoch: 1, - assignmentEpoch: 2, - drainGraceMs: 60_000, - sendAttempts: 1 - } - const claimRegionalRehome = vi.fn().mockResolvedValueOnce(null).mockResolvedValue(attempt) - const assignments = { - claimRegionalRehome, - recordRegionalRehomeDispatchFailure: vi.fn().mockResolvedValue(undefined) - } as unknown as RelayAssignmentStore - const worker = startRegionalRehomeWorker(config(), assignments, { - now: () => now, - safetySnapshot: () => safety(now), - intervalMs: 60_000, - identityToken: async () => { - throw new Error('token unavailable') - } - })! - await settleWorker() - claimRegionalRehome.mockClear() - now = 100 - await worker.run() - worker.stop() - expect(assignments.recordRegionalRehomeDispatchFailure).toHaveBeenCalledWith( - '11111111-1111-4111-8111-111111111111' - ) - }) - - it('keeps a failed poll out of the durable dispatch-failure budget', async () => { - let now = 0 - const claimRegionalRehome = vi - .fn() - .mockResolvedValueOnce(null) - .mockRejectedValue(new Error('Connection terminated due to connection timeout')) - const recordRegionalRehomeDispatchFailure = vi.fn().mockResolvedValue(undefined) - const assignments = { - claimRegionalRehome, - recordRegionalRehomeDispatchFailure - } as unknown as RelayAssignmentStore - const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) - const worker = startRegionalRehomeWorker(config(), assignments, { - now: () => now, - safetySnapshot: () => safety(now), - intervalMs: 60_000 - })! - await settleWorker() - now = 1_000 - await expect(worker.run()).resolves.toBeUndefined() - worker.stop() - - // The poll never claimed an attempt, so nothing was drained and nothing may - // be charged to the budget that latches the durable control off. - expect(recordRegionalRehomeDispatchFailure).not.toHaveBeenCalled() - expect(warn.mock.calls.map((call) => JSON.parse(String(call[0])).event)).toEqual([ - 'orca_relay_regional_rehome_poll_failed' - ]) - }) - it('passes unsafe process telemetry to the durable claim gate', async () => { let now = 0 let sqlFailures = 0 - const claimRegionalRehome = vi.fn().mockResolvedValue(null) + const selectIdleRegionalRehomeCandidates = vi.fn().mockResolvedValue([]) const assignments = { - claimRegionalRehome + selectIdleRegionalRehomeCandidates } as unknown as RelayAssignmentStore const worker = startRegionalRehomeWorker(config(), assignments, { now: () => now, @@ -160,46 +24,40 @@ describe('regional rehome worker', () => { intervalMs: 60_000 })! await settleWorker() - claimRegionalRehome.mockClear() + selectIdleRegionalRehomeCandidates.mockClear() now = 100 sqlFailures = 1 await worker.run() worker.stop() - expect(claimRegionalRehome).toHaveBeenCalledWith( + expect(selectIdleRegionalRehomeCandidates).toHaveBeenCalledWith( expect.objectContaining({ observedAt: 100, sqlFailures: 1 }) ) }) it('starts inert on directors so durable control can enable without a restart', async () => { let now = 0 - const claimRegionalRehome = vi.fn().mockResolvedValue(null) + const selectIdleRegionalRehomeCandidates = vi.fn().mockResolvedValue([]) const assignments = { - claimRegionalRehome + selectIdleRegionalRehomeCandidates } as unknown as RelayAssignmentStore - const worker = startRegionalRehomeWorker( - config(), - assignments, - { - now: () => now, - safetySnapshot: () => safety(now), - intervalMs: 60_000 - } - ) + const worker = startRegionalRehomeWorker(config(), assignments, { + now: () => now, + safetySnapshot: () => safety(now), + intervalMs: 60_000 + }) expect(worker).not.toBeNull() await settleWorker() - claimRegionalRehome.mockClear() + selectIdleRegionalRehomeCandidates.mockClear() now = 100 await worker!.run() worker!.stop() - expect(claimRegionalRehome).toHaveBeenCalledOnce() + expect(selectIdleRegionalRehomeCandidates).toHaveBeenCalledOnce() expect( - startRegionalRehomeWorker( - config({ role: 'cell' }), - {} as RelayAssignmentStore, - { safetySnapshot: () => safety(1) } - ) + startRegionalRehomeWorker(config({ role: 'cell' }), {} as RelayAssignmentStore, { + safetySnapshot: () => safety(1) + }) ).toBeNull() }) @@ -208,16 +66,20 @@ describe('regional rehome worker', () => { const limit = cells * REGIONAL_REHOME_RECONNECTS_PER_CELL_LIMIT const processSafety = { ...safety(100), reconnects: limit * 10 } const fleetSafety = { ...safety(100), reconnects: limit } - expect(regionalRehomeSafetyFailure( - combineRegionalRehomeSafety(processSafety, fleetSafety), - 100, - cells - )).toBeNull() - expect(regionalRehomeSafetyFailure( - combineRegionalRehomeSafety(processSafety, { ...fleetSafety, reconnects: limit + 1 }), - 100, - cells - )).toBe('elevated_reconnects') + expect( + regionalRehomeSafetyFailure( + combineRegionalRehomeSafety(processSafety, fleetSafety), + 100, + cells + ) + ).toBeNull() + expect( + regionalRehomeSafetyFailure( + combineRegionalRehomeSafety(processSafety, { ...fleetSafety, reconnects: limit + 1 }), + 100, + cells + ) + ).toBe('elevated_reconnects') }) }) diff --git a/cloud/apps/relay/src/regional-rehome-worker.ts b/cloud/apps/relay/src/regional-rehome-worker.ts index 4d8fa694afd..5f3827f5188 100644 --- a/cloud/apps/relay/src/regional-rehome-worker.ts +++ b/cloud/apps/relay/src/regional-rehome-worker.ts @@ -1,4 +1,4 @@ -import { z } from 'zod' +import { IdleRegionalRehomeResponseSchema } from '@orca-cloud/relay-contract' import type { RelayAssignmentStore } from './assignment-store.js' import type { RelayConfig } from './config.js' import { googleMetadataIdentityToken } from './google-metadata-identity-token.js' @@ -20,13 +20,6 @@ export type RegionalRehomeWorker = { stop: () => void } -const RegionalHostDrainResponseSchema = z - .object({ - v: z.literal(1), - outcome: z.enum(['accepted', 'already-accepted', 'host-not-connected']) - }) - .strict() - export function startRegionalRehomeWorker( config: RelayConfig, assignments: RelayAssignmentStore, @@ -42,7 +35,6 @@ export function startRegionalRehomeWorker( } const audience = config.rehomeAudience const safetySnapshot = options.safetySnapshot - const now = options.now ?? Date.now const fetchImpl = options.fetch ?? fetch const tokenProvider = options.identityToken ?? @@ -52,64 +44,50 @@ export function startRegionalRehomeWorker( const run = async (): Promise => { if (stopped || inFlight) return inFlight = true - let attemptId: string | null = null try { - const processSafety = safetySnapshot() - const attempt = await assignments.claimRegionalRehome(processSafety) - if (!attempt) return - attemptId = attempt.attemptId + const candidates = await assignments.selectIdleRegionalRehomeCandidates(safetySnapshot()) + if (candidates.length === 0) return const token = await tokenProvider(audience) - const response = await fetchImpl( - new URL('/v1/admin/host-drain', attempt.sourceCellUrl), - { - method: 'POST', - headers: { - authorization: `Bearer ${token}`, - 'content-type': 'application/json' - }, - body: JSON.stringify({ - v: 1, - attemptId: attempt.attemptId, - userId: attempt.userId, - relayHostId: attempt.relayHostId, - sourceCellId: attempt.sourceCellId, - sourceCellIncarnation: attempt.sourceCellIncarnation, - sourceAssignmentEpoch: attempt.previousEpoch, - graceMs: attempt.drainGraceMs - }), - signal: AbortSignal.timeout(options.requestTimeoutMs ?? 10_000) + for (const candidate of candidates) { + if (stopped) return + const { sourceCellUrl, ...request } = candidate + try { + const response = await fetchImpl(new URL('/v1/admin/host-idle-rehome', sourceCellUrl), { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: JSON.stringify({ + ...request, + cohortPercent: config.regionCorrectionCohortPercent ?? 0, + directorSafety: safetySnapshot() + }), + signal: AbortSignal.timeout(options.requestTimeoutMs ?? 10_000) + }) + if (!response.ok) throw new Error(`regional_rehome_source_${response.status}`) + const body = IdleRegionalRehomeResponseSchema.parse(await response.json()) + if (body.outcome === 'committed') { + console.warn( + JSON.stringify({ + event: 'orca_relay_idle_rehome_committed', + sourceCellId: candidate.sourceCellId, + targetCellId: candidate.targetCellId + }) + ) + return + } + } catch (error) { + // The source may have committed; its durable outcome owns recovery. + console.warn( + JSON.stringify({ + event: 'orca_relay_idle_rehome_request_failed', + reason: error instanceof Error ? error.message : 'unknown' + }) + ) } - ) - if (!response.ok) throw new Error(`regional_rehome_source_${response.status}`) - const body = RegionalHostDrainResponseSchema.safeParse(await response.json()) - if (!body.success) throw new Error('regional_rehome_source_invalid_response') - await assignments.recordRegionalRehomeDrainReceipt( - attempt.attemptId, - body.data.outcome - ) - console.warn( - JSON.stringify({ - event: 'orca_relay_regional_rehome_dispatched', - sourceCellId: attempt.sourceCellId, - targetCellId: attempt.targetCellId, - outcome: body.data.outcome, - sendAttempts: attempt.sendAttempts - }) - ) - } catch (error) { - // Only a claimed attempt was drained. A poll that failed before the claim - // - a pool timeout on the once-a-second control read - dispatched nothing, - // so it must not spend the budget that latches the durable control off. - if (attemptId) { - await assignments - .recordRegionalRehomeDispatchFailure(attemptId) - .catch(() => undefined) } + } catch (error) { console.warn( JSON.stringify({ - event: attemptId - ? 'orca_relay_regional_rehome_dispatch_failed' - : 'orca_relay_regional_rehome_poll_failed', + event: 'orca_relay_regional_rehome_poll_failed', reason: error instanceof Error ? error.message : 'unknown' }) ) diff --git a/cloud/apps/relay/src/relay-region-app.test.ts b/cloud/apps/relay/src/relay-region-app.test.ts index 30cf3bf3e26..54e8da670b8 100644 --- a/cloud/apps/relay/src/relay-region-app.test.ts +++ b/cloud/apps/relay/src/relay-region-app.test.ts @@ -40,6 +40,7 @@ describe('Relay region API', () => { ) expect(response.status).toBe(200) + expect(await response.clone().json()).not.toHaveProperty('regionCorrection') expect(assign).toHaveBeenCalledWith( { userId: 'user-1', relayHostId: 'asiahost00000001' }, 'asia-east2', @@ -83,6 +84,160 @@ describe('Relay region API', () => { ) }) + it('preserves the cold-start hint and binds a negotiated window after placement', async () => { + const assignment = { + userId: 'user-1', + relayHostId: 'abcdefghijklmnop', + cellId: 'asia-c1', + cellUrl: 'https://asia-c1.relay.example.test', + region: 'asia-east2', + assignmentEpoch: 7, + leaseExpiresAt: Date.now() + 300_000 + } + const assign = vi.fn(async () => assignment) + const window = { + generation: 2, + expiresAt: Date.now() + 86_400_000, + assignmentEpoch: 7, + incumbentRegion: 'asia-east2', + policyVersion: 1 + } + const exchangeRegionCorrection = vi.fn(async () => ({ v: 1, window })) + const app = createRelayApp(config(), { + store: {} as never, + assignments: { assign, exchangeRegionCorrection } as never, + drain: vi.fn(), + ready: async () => true + }) + const regionCorrection = { v: 1, action: 'issue-window' } + const response = await app.request( + '/v1/assign', + assignmentRequest('abcdefghijklmnop', { + preferredRegion: 'asia-east2', + regionCorrection + }) + ) + expect(response.status).toBe(200) + expect(assign).toHaveBeenCalledWith( + { userId: 'user-1', relayHostId: 'abcdefghijklmnop' }, + 'asia-east2', + 'asia-east2' + ) + expect(exchangeRegionCorrection).toHaveBeenCalledWith( + { userId: 'user-1', relayHostId: 'abcdefghijklmnop' }, + regionCorrection, + 7 + ) + expect(((await response.json()) as { regionCorrection: unknown }).regionCorrection).toEqual({ + v: 1, + window + }) + }) + + it('returns successful placement when optional window storage is unavailable', async () => { + const app = createRelayApp(config(), { + store: {} as never, + assignments: { + assign: async () => ({ + cellId: 'asia-c1', + region: 'asia-east2', + cellUrl: 'https://asia-c1.relay.example.test', + assignmentEpoch: 7 + }), + exchangeRegionCorrection: async () => { + throw new Error('database unavailable') + } + } as never, + drain: vi.fn(), + ready: async () => true + }) + const response = await app.request( + '/v1/assign', + assignmentRequest('abcdefghijklmnop', { + preferredRegion: 'asia-east2', + regionCorrection: { v: 1, action: 'issue-window' } + }) + ) + expect(response.status).toBe(200) + expect(await response.json()).toMatchObject({ + cellUrl: 'https://asia-c1.relay.example.test', + assignmentEpoch: 7 + }) + }) + + it('does not place or write a legacy hint when reporting migration evidence', async () => { + const current = { + userId: 'user-1', + relayHostId: 'abcdefghijklmnop', + cellId: 'asia-c1', + cellUrl: 'https://asia-c1.relay.example.test', + region: 'asia-east2', + assignmentEpoch: 7, + leaseExpiresAt: Date.now() + 300_000 + } + const assign = vi.fn() + const resolve = vi.fn(async () => current) + const exchangeRegionCorrection = vi.fn(async () => ({ v: 1, reportStatus: 'accepted' })) + const app = createRelayApp(config(), { + store: {} as never, + assignments: { assign, resolve, exchangeRegionCorrection } as never, + drain: vi.fn(), + ready: async () => true + }) + const regionCorrection = { + v: 1, + action: 'report', + generation: 2, + assignmentEpoch: 7, + policyVersion: 1, + outcome: 'conclusive', + measurements: { 'us-central1': 40, 'asia-east2': 180 } + } + const response = await app.request( + '/v1/assign', + assignmentRequest('abcdefghijklmnop', { + preferredRegion: 'us-central1', + regionCorrection + }) + ) + expect(response.status).toBe(200) + expect(assign).not.toHaveBeenCalled() + expect(exchangeRegionCorrection).toHaveBeenCalledWith( + { userId: 'user-1', relayHostId: 'abcdefghijklmnop' }, + regionCorrection, + 7 + ) + expect(((await response.json()) as { assignmentEpoch: number }).assignmentEpoch).toBe(7) + }) + + it('does not manufacture an assignment for a report whose assignment disappeared', async () => { + const assign = vi.fn() + const exchangeRegionCorrection = vi.fn() + const app = createRelayApp(config(), { + store: {} as never, + assignments: { assign, resolve: async () => null, exchangeRegionCorrection } as never, + drain: vi.fn(), + ready: async () => true + }) + const response = await app.request( + '/v1/assign', + assignmentRequest('abcdefghijklmnop', { + regionCorrection: { + v: 1, + action: 'report', + generation: 2, + assignmentEpoch: 7, + policyVersion: 1, + outcome: 'inconclusive', + reason: 'timeout' + } + }) + ) + expect(response.status).toBe(409) + expect(assign).not.toHaveBeenCalled() + expect(exchangeRegionCorrection).not.toHaveBeenCalled() + }) + it('exposes only the store-provided healthy catalog from directors', async () => { const regionCatalog = vi.fn(async () => [ { region: 'us-central1' as const, probeOrigins: ['https://us.relay.example.test'] } diff --git a/cloud/apps/relay/src/relay-server.ts b/cloud/apps/relay/src/relay-server.ts index 77a15a1d259..7cee77e52de 100644 --- a/cloud/apps/relay/src/relay-server.ts +++ b/cloud/apps/relay/src/relay-server.ts @@ -20,14 +20,12 @@ import { createRelayApp } from './app.js' import { RelayAssignmentStore } from './assignment-store.js' import type { RelayConfig } from './config.js' import { RelayCredentialStore } from './credential-store.js' -import type { RelayDatabase } from './database.js' +import { readRelayDatabasePoolPressure, type RelayDatabase } from './database.js' import { HostSessionRegistry } from './host-session-registry.js' import { observeRelayDatabase } from './observed-relay-database.js' import { RelayObservability } from './relay-observability.js' -import { - RelayConnectionLedger, - type RelayConnectionUpgrade -} from './relay-connection-ledger.js' +import { combineRegionalRehomeSafety } from './regional-rehome-safety.js' +import { RelayConnectionLedger, type RelayConnectionUpgrade } from './relay-connection-ledger.js' import { createRelayReadiness } from './relay-readiness.js' import { createRelayTokenVerifier, readBearer } from './relay-token-verifier.js' import { closeRelayWebSocket } from './relay-websocket-close.js' @@ -65,7 +63,7 @@ function guardSocketErrors(socket: WebSocket, kind: string): void { function admissionSource(request: IncomingMessage): string { const forwarded = request.headers['x-forwarded-for'] - const chain = (Array.isArray(forwarded) ? forwarded.join(',') : forwarded ?? '') + const chain = (Array.isArray(forwarded) ? forwarded.join(',') : (forwarded ?? '')) .split(',') .map((entry) => entry.trim()) .filter(Boolean) @@ -112,6 +110,7 @@ export function createRelayServer( const store = new RelayCredentialStore(observedDatabase, options.now) const assignments = new RelayAssignmentStore(observedDatabase, options.now, { requireLiveCells: config.role === 'director', + regionalRehomeCohortPercent: config.regionCorrectionCohortPercent ?? 0, recordControlRenewal: (durationMs, outcome) => observability.recordControlRenewal?.(durationMs, outcome) }) @@ -127,17 +126,35 @@ export function createRelayServer( queuedBytes, observability, options.now, - options.random + options.random, + cellIncarnation ) const app = createRelayApp(config, { store, assignments, drain: (graceMs) => sessions.drain(graceMs), drainHost: (input) => sessions.drainHost(input), + idleRehome: (input) => { + const now = (options.now ?? Date.now)() + if (input.directorSafety.observedAt > now || now - input.directorSafety.observedAt > 60_000) { + return Promise.resolve({ outcome: 'deferred' }) + } + return sessions.idleRehome(input, + () => assignments.commitIdleRegionalRehome(input, combineRegionalRehomeSafety( + input.directorSafety, + { ...observability.regionalRehomeRuntimeSafety(), ...readRelayDatabasePoolPressure(database) } + ), input.cohortPercent), + () => assignments.reconcileIdleRegionalRehome(input) + ) + }, regionalRehomeTrustProbeHostExists: (input) => sessions.get(input) !== null, cellIncarnation, isDraining: () => sessions.isDraining(), runtimeCounts: () => runtimeCounts(), + regionalRehomeSafetySnapshot: () => ({ + ...observability.regionalRehomeRuntimeSafety(), + ...readRelayDatabasePoolPressure(database) + }), ready, recordAssignmentAdmission: (outcome) => observability.recordAssignmentAdmission?.(outcome), recordAssignmentRejectionReason: (lane, reason) => @@ -339,7 +356,7 @@ export function createRelayServer( const identity = invite ? { userId: invite.userId, relayHostId: hostId } : null // Released combined-service invites gain their first durable cell assignment here. const assignment = identity - ? (await assignments.resolve(identity)) ?? (await assignments.assign(identity)) + ? ((await assignments.resolve(identity)) ?? (await assignments.assign(identity))) : null if (!invite || !assignment) { phoneAdmission?.hostData.release() diff --git a/cloud/apps/relay/src/relay-sweep-schedule.test.ts b/cloud/apps/relay/src/relay-sweep-schedule.test.ts index d5ef450cc43..55469cab91c 100644 --- a/cloud/apps/relay/src/relay-sweep-schedule.test.ts +++ b/cloud/apps/relay/src/relay-sweep-schedule.test.ts @@ -35,7 +35,7 @@ describe('sweep schedule jitter', () => { rehomeAudience: 'https://rehome.example.test', rehomeDirectorServiceAccount: 'rehome@example.test' } as never, - { claimRegionalRehome: async () => null } as never, + { selectIdleRegionalRehomeCandidates: async () => [] } as never, { random: () => 0.5, safetySnapshot: () => ({}) as never } ) } finally { diff --git a/cloud/apps/relay/src/test-fixtures/relay-contract-baseline/README.md b/cloud/apps/relay/src/test-fixtures/relay-contract-baseline/README.md new file mode 100644 index 00000000000..2ac4ad78391 --- /dev/null +++ b/cloud/apps/relay/src/test-fixtures/relay-contract-baseline/README.md @@ -0,0 +1,8 @@ +Exact relay contract snapshots from `027acb4efa2e6b226d40df266b86367423946d62`, `cloud/packages/relay-contract/src/`. Used to exercise the pre-correction strict wire parsers. Do not format or edit these baseline sources. + +```text +aba94e108a5cd0f1af8b38875429ad8636d24c43a728273e3df60d9a1a1d1b6d director-messages.ts +bd13b5a694a5d683a5b680c14e46ab33f4ef4a5bfedf040d046b09d540cb4c17 wire-scalars.ts +bc89116f884a2f20a6588f9b91219aa596bc2410d28b499a93a78350def109d5 relay-regions.ts +8fcae470a5fc72f2fcdde9d2f09cd20289c256356dd490484ac1cfa53839fbe4 control-messages.ts +``` diff --git a/cloud/apps/relay/src/test-fixtures/relay-contract-baseline/control-messages.ts b/cloud/apps/relay/src/test-fixtures/relay-contract-baseline/control-messages.ts new file mode 100644 index 00000000000..0daf21e7c28 --- /dev/null +++ b/cloud/apps/relay/src/test-fixtures/relay-contract-baseline/control-messages.ts @@ -0,0 +1,146 @@ +import { z } from 'zod' +import { + Base6432ByteSchema, + Base64Raw24ByteSchema, + Base64Url32ByteSchema, + EpochMsSchema, + GenerationSchema, + OpaqueIdSchema, + PositiveDurationMsSchema, + RelayHostIdSchema +} from './wire-scalars.js' + +const AppVersionSchema = z.string().min(1).max(128) +const BoundedCiphertextSchema = z + .string() + .min(1) + .max(16 * 1024) + .regex(/^(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?$/) +const ConnectionKindSchema = z.enum(['invite', 'resume']) + +export const HostHelloSchema = z + .object({ + v: z.literal(1), + relayHostId: RelayHostIdSchema, + assignmentEpoch: GenerationSchema, + hostPublicKeyB64: Base6432ByteSchema, + appVersion: AppVersionSchema, + previousGeneration: GenerationSchema.optional(), + controlResumeSecret: Base64Url32ByteSchema.optional() + }) + .strict() + +export const HostChallengeSchema = z + .object({ + challengeId: OpaqueIdSchema, + relayEphemeralPublicKeyB64: Base6432ByteSchema, + nonceB64: Base64Raw24ByteSchema, + ciphertextB64: BoundedCiphertextSchema, + expiresAt: EpochMsSchema + }) + .strict() + +export const HostChallengeAckSchema = z + .object({ challengeId: OpaqueIdSchema, proofB64: Base6432ByteSchema }) + .strict() + +// Advertised on the control upgrade rather than in host-hello: HostHelloSchema +// is strict, so a new hello key is refused by every already-deployed cell. +export const RELAY_HOST_CAPABILITIES_HEADER = 'x-orca-host-capabilities' +// The host accepts kind/relayDeviceId on a pendingConns entry. A host that does +// not advertise this parses those entries strictly and would drop the whole ack. +export const RELAY_HOST_CAPABILITY_PENDING_CONN_DETAILS = 'pending-conn-details' + +export function parseRelayHostCapabilities( + header: string | string[] | undefined +): ReadonlySet { + const raw = Array.isArray(header) ? header.join(',') : (header ?? '') + return new Set( + raw + .split(',') + .map((token) => token.trim()) + .filter((token) => token.length > 0 && token.length <= 64) + .slice(0, 16) + ) +} + +// kind/relayDeviceId are optional so an entry stays readable by a host that +// predates them; the cell only emits them to a host that advertised support. +const PendingConnectionSchema = z + .object({ + connId: OpaqueIdSchema, + connTicket: Base64Url32ByteSchema, + kind: ConnectionKindSchema.optional(), + relayDeviceId: OpaqueIdSchema.optional() + }) + .strict() + +export const HostHelloAckSchema = z + .object({ + v: z.literal(1), + generation: GenerationSchema, + controlResumeSecret: Base64Url32ByteSchema, + leaseExpiresAt: EpochMsSchema, + activeConnIds: z.array(OpaqueIdSchema).max(8), + pendingConns: z.array(PendingConnectionSchema).max(8) + }) + .strict() + +export const ConnectionOpenSchema = z + .object({ + connId: OpaqueIdSchema, + connTicket: Base64Url32ByteSchema, + kind: ConnectionKindSchema, + relayDeviceId: OpaqueIdSchema, + attachDeadlineMs: PositiveDurationMsSchema + }) + .strict() + +export const HostDataAuthSchema = z + .object({ + v: z.literal(1), + connTicket: Base64Url32ByteSchema, + generation: GenerationSchema + }) + .strict() + +export const InviteCreateSchema = z + .object({ reqId: OpaqueIdSchema, relayDeviceId: OpaqueIdSchema }) + .strict() + +export const InviteCreatedSchema = z + .object({ + reqId: OpaqueIdSchema, + inviteToken: Base64Url32ByteSchema, + expiresAt: EpochMsSchema, + maxAttempts: z.number().int().positive().max(16) + }) + .strict() + +export const DeviceRevokeSchema = z + .object({ reqId: OpaqueIdSchema, relayDeviceId: OpaqueIdSchema }) + .strict() + +export const AuthRefreshSchema = z.object({ relayJwt: z.string().min(1).max(8 * 1024) }).strict() + +export const DrainSchema = z + .object({ + graceMs: z.number().int().nonnegative().max(60 * 60 * 1000), + recovery: z.literal('resolve-director') + }) + .strict() + +export const HeartbeatSchema = z.object({ t: EpochMsSchema }).strict() + +export type HostHello = z.infer +export type HostChallenge = z.infer +export type HostChallengeAck = z.infer +export type HostHelloAck = z.infer +export type ConnectionOpen = z.infer +export type HostDataAuth = z.infer +export type InviteCreate = z.infer +export type InviteCreated = z.infer +export type DeviceRevoke = z.infer +export type AuthRefresh = z.infer +export type Drain = z.infer +export type Heartbeat = z.infer diff --git a/cloud/apps/relay/src/test-fixtures/relay-contract-baseline/director-messages.ts b/cloud/apps/relay/src/test-fixtures/relay-contract-baseline/director-messages.ts new file mode 100644 index 00000000000..e697135b68e --- /dev/null +++ b/cloud/apps/relay/src/test-fixtures/relay-contract-baseline/director-messages.ts @@ -0,0 +1,75 @@ +import { z } from 'zod' +import { + Base64Url32ByteSchema, + CanonicalHttpsOriginSchema, + EpochMsSchema, + GenerationSchema, + RelayHostIdSchema +} from './wire-scalars.js' +import { RelayRegionSchema } from './relay-regions.js' + +const SignedAssignmentLeaseSchema = z.string().min(1).max(8 * 1024) + +export const AssignmentRequestSchema = z + .object({ + v: z.literal(1), + relayHostId: RelayHostIdSchema, + // Client-declared reconnection; the director verifies it against the + // durable assignment before granting fast-lane admission. + reconnect: z.boolean().optional(), + preferredRegion: RelayRegionSchema.optional() + }) + .strict() + +export const AssignmentResponseSchema = z + .object({ + v: z.literal(1), + cellUrl: CanonicalHttpsOriginSchema, + assignmentEpoch: GenerationSchema, + lease: SignedAssignmentLeaseSchema + }) + .strict() + +export const ResolveRequestSchema = z + .object({ + v: z.literal(1), + relayHostId: RelayHostIdSchema, + resumeToken: Base64Url32ByteSchema + }) + .strict() + +export const ResolveResponseSchema = z + .object({ + v: z.literal(1), + cellUrl: CanonicalHttpsOriginSchema, + assignmentEpoch: GenerationSchema, + leaseExpiresAt: EpochMsSchema + }) + .strict() + +export const RelayMovedSchema = z + .object({ + v: z.literal(1), + cellUrl: CanonicalHttpsOriginSchema, + assignmentEpoch: GenerationSchema + }) + .strict() + +export function isTrustedNewerMove(input: { + sourceOrigin: string + configuredDirectorOrigin: string + currentAssignmentEpoch: number + move: z.infer +}): boolean { + // Why: cells and stale director responses must never redirect a credential-bearing client. + return ( + input.sourceOrigin === input.configuredDirectorOrigin && + input.move.assignmentEpoch > input.currentAssignmentEpoch + ) +} + +export type AssignmentRequest = z.infer +export type AssignmentResponse = z.infer +export type ResolveRequest = z.infer +export type ResolveResponse = z.infer +export type RelayMoved = z.infer diff --git a/cloud/apps/relay/src/test-fixtures/relay-contract-baseline/relay-regions.ts b/cloud/apps/relay/src/test-fixtures/relay-contract-baseline/relay-regions.ts new file mode 100644 index 00000000000..6b8837829df --- /dev/null +++ b/cloud/apps/relay/src/test-fixtures/relay-contract-baseline/relay-regions.ts @@ -0,0 +1,71 @@ +import { z } from 'zod' + +export const RELAY_REGIONS = ['us-central1', 'asia-east2'] as const + +export const RelayRegionSchema = z.enum(RELAY_REGIONS) + +export type RelayRegion = z.infer + +export const RELAY_DEFAULT_REGION: RelayRegion = 'us-central1' + +// Field-name segment for the flat per-region runtime counters, spelled out rather than derived so +// the Terraform side can hold the same literal and a test can compare the two. `satisfies` makes a +// new region a compile error here, which is the point: a region with no segment would silently +// drop out of the region-skew alert's denominators. +export const RELAY_REGION_METRIC_SEGMENTS = { + 'us-central1': 'UsCentral1', + 'asia-east2': 'AsiaEast2' +} as const satisfies Record + +const RelayProbeOriginSchema = z.string().url().max(2_048).refine(isCanonicalHttpsOrigin) + +export const RelayRegionCatalogResponseSchema = z + .object({ + v: z.literal(1), + regions: z + .array( + z + .object({ + region: RelayRegionSchema, + probeOrigins: z.array(RelayProbeOriginSchema).min(1).max(2) + }) + .strict() + ) + .max(RELAY_REGIONS.length) + }) + .strict() + .superRefine((catalog, context) => { + const regions = new Set() + const origins = new Set() + for (const [regionIndex, entry] of catalog.regions.entries()) { + if (regions.has(entry.region)) { + context.addIssue({ + code: 'custom', + message: 'duplicate relay region', + path: ['regions', regionIndex, 'region'] + }) + } + regions.add(entry.region) + for (const [originIndex, origin] of entry.probeOrigins.entries()) { + if (origins.has(origin)) { + context.addIssue({ + code: 'custom', + message: 'duplicate relay probe origin', + path: ['regions', regionIndex, 'probeOrigins', originIndex] + }) + } + origins.add(origin) + } + } + }) + +export type RelayRegionCatalogResponse = z.infer + +function isCanonicalHttpsOrigin(value: string): boolean { + try { + const url = new URL(value) + return url.protocol === 'https:' && url.origin === value + } catch { + return false + } +} diff --git a/cloud/apps/relay/src/test-fixtures/relay-contract-baseline/wire-scalars.ts b/cloud/apps/relay/src/test-fixtures/relay-contract-baseline/wire-scalars.ts new file mode 100644 index 00000000000..27dd3a8b30f --- /dev/null +++ b/cloud/apps/relay/src/test-fixtures/relay-contract-baseline/wire-scalars.ts @@ -0,0 +1,20 @@ +import { z } from 'zod' + +export const Base64Url32ByteSchema = z.string().regex(/^[A-Za-z0-9_-]{43}$/) +export const Base64Url24ByteSchema = z.string().regex(/^[A-Za-z0-9_-]{32}$/) +export const Base6432ByteSchema = z.string().regex(/^(?:[A-Za-z0-9+/]{4}){10}[A-Za-z0-9+/]{3}=$/) +export const Base64Raw24ByteSchema = z.string().regex(/^(?:[A-Za-z0-9+/]{4}){8}$/) +export const RelayHostIdSchema = z.string().regex(/^[A-Za-z0-9_-]{16}$/) +export const OpaqueIdSchema = z.string().min(1).max(128) +export const EpochMsSchema = z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER) +export const GenerationSchema = z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER) +export const PositiveDurationMsSchema = z.number().int().positive().max(24 * 60 * 60 * 1000) + +export const CanonicalHttpsOriginSchema = z.string().max(2048).refine((value) => { + try { + const url = new URL(value) + return url.protocol === 'https:' && url.origin === value && url.pathname === '/' + } catch { + return false + } +}, 'must be a canonical HTTPS origin') diff --git a/cloud/apps/relay/tsconfig.build.json b/cloud/apps/relay/tsconfig.build.json index 489ddfd34d6..38eb0396cf2 100644 --- a/cloud/apps/relay/tsconfig.build.json +++ b/cloud/apps/relay/tsconfig.build.json @@ -6,5 +6,5 @@ "outDir": "dist", "rootDir": "src" }, - "exclude": ["src/**/*.test.ts"] + "exclude": ["src/**/*.test.ts", "src/test-fixtures/**"] } diff --git a/cloud/dev/scripts/deploy-relay-blue-green.mjs b/cloud/dev/scripts/deploy-relay-blue-green.mjs index 88e4f8ccc60..ddfe2bce8fc 100644 --- a/cloud/dev/scripts/deploy-relay-blue-green.mjs +++ b/cloud/dev/scripts/deploy-relay-blue-green.mjs @@ -10,6 +10,7 @@ export const DIRECTOR_REGIONAL_PLACEMENT_SECRET = 'orca-cloud-relay-regional-placement-enabled' export const DIRECTOR_REGIONAL_PLACEMENT_ENV = 'ORCA_RELAY_REGIONAL_PLACEMENT_ENABLED' +export const DIRECTOR_CORRECTION_COHORT_ENV = 'ORCA_RELAY_REGION_CORRECTION_COHORT_PERCENT' export const DIRECTOR_REHOME_IDENTITY_ENV = 'ORCA_RELAY_REHOME_DIRECTOR_SERVICE_ACCOUNT' export const DIRECTOR_REHOME_AUDIENCE_ENV = 'ORCA_RELAY_REHOME_AUDIENCE' @@ -228,6 +229,13 @@ export function directorCellSetAddition(currentValue, desiredValue) { return { changed: additions.length > 0, value: JSON.stringify(desired) } } +export function correctionCohortPercent(value) { + if (!/^(?:[0-9]|[1-9][0-9]|100)$/.test(String(value))) { + throw new Error('region correction cohort must be an integer from 0 to 100') + } + return String(value) +} + export function directorDeploymentEnvironment(config) { const imageDigest = config.image?.match(/@(sha256:[a-f0-9]{64})$/)?.[1] if (config.image !== undefined && imageDigest === undefined) { @@ -238,6 +246,10 @@ export function directorDeploymentEnvironment(config) { ORCA_RELAY_ADMISSION_SELECTOR_VERSION: SELECTOR_REVISION_MARKER, ...(imageDigest === undefined ? {} : { ORCA_RELAY_IMAGE_DIGEST: imageDigest }) } + if (config['region-correction-cohort-percent'] !== undefined && + config['region-correction-cohort-percent'] !== 'preserve') { + environment[DIRECTOR_CORRECTION_COHORT_ENV] = correctionCohortPercent(config['region-correction-cohort-percent']) + } const serviceAccount = projectServiceAccount(config, 'capacity-service-account') const asiaProofServiceAccount = projectServiceAccount(config, 'asia-proof-service-account') const rehomeDirectorServiceAccount = projectServiceAccount( @@ -302,7 +314,8 @@ export function parseArguments(argv) { values['rehome-director-service-account'] !== undefined || values['rehome-audience'] !== undefined || values['expected-rehome-generation'] !== undefined || - values['rehome-control-origin'] !== undefined + values['rehome-control-origin'] !== undefined || + values['region-correction-cohort-percent'] !== undefined ) { throw new Error('director configuration arguments require --role director') } @@ -785,6 +798,14 @@ export async function deployDirector(config, tag, overrides = {}) { config['prune-revisions'] === 'true' ? CONNECTION_CAPACITY_PROTOCOL : undefined const currentEnvironment = revisionEnvironment(servingRevision) const deploymentEnvironment = directorDeploymentEnvironment(config) + deploymentEnvironment[DIRECTOR_CORRECTION_COHORT_ENV] ??= correctionCohortPercent( + currentEnvironment[DIRECTOR_CORRECTION_COHORT_ENV] ?? '0' + ) + if (config['region-correction-cohort-percent'] !== undefined && + config['region-correction-cohort-percent'] !== 'preserve' && + config['expected-rehome-generation'] === undefined) { + throw new Error('cohort changes require an exact disabled regional-rehome generation') + } const mutableEnvironment = { ...deploymentEnvironment, [DIRECTOR_REGIONAL_PLACEMENT_ENV]: '' diff --git a/cloud/dev/scripts/deploy-relay-blue-green.test.mjs b/cloud/dev/scripts/deploy-relay-blue-green.test.mjs index 6e56676098b..a68ce50e912 100644 --- a/cloud/dev/scripts/deploy-relay-blue-green.test.mjs +++ b/cloud/dev/scripts/deploy-relay-blue-green.test.mjs @@ -4,6 +4,8 @@ import { test } from 'node:test' import { fileURLToPath } from 'node:url' import { activeRevision, + correctionCohortPercent, + DIRECTOR_CORRECTION_COHORT_ENV, cloudRunTrafficTag, DIRECTOR_ADMISSION_ENVIRONMENT, DIRECTOR_REGIONAL_PLACEMENT_ENV, @@ -812,3 +814,43 @@ test('waits for authenticated target readiness without hiding other capacity err /forbidden/ ) }) + + +test('validates bounded correction cohorts and leaves unspecified values to serving inheritance', () => { + for (const value of ['0', '1', '100']) assert.equal(correctionCohortPercent(value), value) + for (const value of ['-1', '101', '1.5', '', '01', 'true', '1\n']) { + assert.throws(() => correctionCohortPercent(value), /integer from 0 to 100/) + } + assert.equal(directorDeploymentEnvironment({})[DIRECTOR_CORRECTION_COHORT_ENV], undefined) + assert.equal(directorDeploymentEnvironment({ 'region-correction-cohort-percent': 'preserve' })[DIRECTOR_CORRECTION_COHORT_ENV], undefined) + assert.equal(directorDeploymentEnvironment({ 'region-correction-cohort-percent': '1' })[DIRECTOR_CORRECTION_COHORT_ENV], '1') +}) + +test('inherits the cohort on candidate and rollback revisions without resetting an enabled cohort', async () => { + const harness = directorHarness() + harness.state.revisions.get('relay-00001-old').env[DIRECTOR_CORRECTION_COHORT_ENV] = '3' + await deployDirector({}, 'candidate-new', harness.operations) + for (const revision of ['relay-00002-new', 'relay-00003-new']) { + assert.equal(harness.state.revisions.get(revision).env[DIRECTOR_CORRECTION_COHORT_ENV], '3') + } +}) + +test('starts an unstamped cohort at zero and rejects a cohort change without disabled-control proof', async () => { + const harness = directorHarness() + await assert.rejects(deployDirector({ 'region-correction-cohort-percent': '1' }, + 'candidate-new', harness.operations), /exact disabled regional-rehome generation/) + assert.equal(harness.state.activeRevision, 'relay-00001-old') + assert.equal(harness.state.nextRevision, 2) + await deployDirector({}, 'candidate-new', harness.operations) + assert.equal(harness.state.revisions.get('relay-00003-new').env[DIRECTOR_CORRECTION_COHORT_ENV], '0') +}) + +test('sets a reviewed cohort only behind repeated disabled-control verification', async () => { + const harness = directorHarness() + let verified = 0 + const config = { 'region-correction-cohort-percent': '1', 'expected-rehome-generation': '7' } + await deployDirector(config, 'candidate-new', { ...harness.operations, + assertRegionalRehomeDisabled: async () => { verified++ } }) + assert.ok(verified >= 2) + assert.equal(harness.state.revisions.get('relay-00003-new').env[DIRECTOR_CORRECTION_COHORT_ENV], '1') +}) diff --git a/cloud/dev/scripts/read-relay-serving-regional-placement-version.mjs b/cloud/dev/scripts/read-relay-serving-regional-placement-version.mjs index 80d40277e5a..796d0e9d91a 100644 --- a/cloud/dev/scripts/read-relay-serving-regional-placement-version.mjs +++ b/cloud/dev/scripts/read-relay-serving-regional-placement-version.mjs @@ -53,7 +53,7 @@ export function readRelayServingRegionalPlacementVersion(input, dependencies = { try { service = run(gcloudArguments('services', input)) } catch (error) { - if (error?.code === 'NOT_FOUND') return { version: input.bootstrap_version } + if (error?.code === 'NOT_FOUND') return { version: input.bootstrap_version, cohort_percent: '0' } throw error } const serving = (service.status?.traffic ?? []).filter( @@ -67,12 +67,22 @@ export function readRelayServingRegionalPlacementVersion(input, dependencies = { throw new Error('Relay director must have exactly one revision serving 100% traffic') } const revision = run(gcloudArguments('revisions', input, serving[0].revisionName)) + const cohortSettings = (revision.spec?.containers ?? []).flatMap((container) => + (container.env ?? []).filter((environment) => + environment.name === 'ORCA_RELAY_REGION_CORRECTION_COHORT_PERCENT') + ) + if (cohortSettings.length > 1 || (cohortSettings.length === 1 && + (typeof cohortSettings[0].value !== 'string' || + !/^(?:[0-9]|[1-9][0-9]|100)$/.test(cohortSettings[0].value)))) { + throw new Error('serving region correction cohort is invalid') + } + const cohort_percent = cohortSettings[0]?.value ?? '0' const references = (revision.spec?.containers ?? []).flatMap((container) => (container.env ?? []).filter( (environment) => environment.name === 'ORCA_RELAY_REGIONAL_PLACEMENT_ENABLED' ) ) - if (references.length === 0) return { version: input.bootstrap_version } + if (references.length === 0) return { version: input.bootstrap_version, cohort_percent } const reference = normalizeSecretReference(references[0]) if ( references.length !== 1 || @@ -81,7 +91,7 @@ export function readRelayServingRegionalPlacementVersion(input, dependencies = { ) { throw new Error('serving regional placement secret reference is invalid') } - return { version: reference.version } + return { version: reference.version, cohort_percent } } // Why: the v2 API reports `valueSource.secretKeyRef.{secret,version}`, but diff --git a/cloud/dev/scripts/read-relay-serving-regional-placement-version.test.mjs b/cloud/dev/scripts/read-relay-serving-regional-placement-version.test.mjs index 8043fc23e94..9c49d4127a5 100644 --- a/cloud/dev/scripts/read-relay-serving-regional-placement-version.test.mjs +++ b/cloud/dev/scripts/read-relay-serving-regional-placement-version.test.mjs @@ -67,7 +67,7 @@ test('reads the exact version from the sole traffic-serving revision', () => { } }) - assert.deepEqual(result, { version: '11' }) + assert.deepEqual(result, { version: '11', cohort_percent: '0' }) assert.equal(calls[1][3], 'relay-serving') }) @@ -78,7 +78,7 @@ test('reads the gcloud v1 secret reference shape by bare id and by full resource ]) { assert.deepEqual(readRelayServingRegionalPlacementVersion(input, { run: (args) => args[1] === 'services' ? serving() : v1Revision(name, '1') - }), { version: '1' }) + }), { version: '1', cohort_percent: '0' }) } }) @@ -100,12 +100,12 @@ test('falls back only when the service or setting is absent', () => { notFound.code = 'NOT_FOUND' assert.deepEqual(readRelayServingRegionalPlacementVersion(input, { run: () => { throw notFound } - }), { version: '7' }) + }), { version: '7', cohort_percent: '0' }) assert.deepEqual(readRelayServingRegionalPlacementVersion(input, { run: (args) => args[1] === 'services' ? { status: { traffic: [{ revisionName: 'relay-serving', percent: 100 }] } } : { spec: { containers: [{ env: [] }] } } - }), { version: '7' }) + }), { version: '7', cohort_percent: '0' }) }) test('classifies real absent-service stderr without weakening revision failures', () => { @@ -136,3 +136,31 @@ test('rejects ambiguous traffic, malformed references, and read failures', () => run: () => { throw denied } }), denied) }) + + +test('preserves the serving cohort including explicit disable across later Terraform plans', () => { + for (const value of ['0', '1', '17', '100']) { + const servingRevision = revision() + servingRevision.spec.containers[0].env.push({ name: 'ORCA_RELAY_REGION_CORRECTION_COHORT_PERCENT', value }) + assert.deepEqual(readRelayServingRegionalPlacementVersion(input, { + run: (args) => args[1] === 'services' ? serving() : servingRevision + }), { version: '11', cohort_percent: value }) + } +}) + +test('fails closed on malformed, secret-backed or duplicate cohorts rather than resetting them', () => { + const name = 'ORCA_RELAY_REGION_CORRECTION_COHORT_PERCENT' + const cases = [ + [{ name, value: '101' }], [{ name, value: '-1' }], [{ name, value: '1.5' }], + [{ name, value: '' }], [{ name, value: '01' }], [{ name, value: 1 }], + [{ name, valueFrom: { secretKeyRef: { name: 'unexpected', key: '1' } } }], + [{ name, value: '1' }, { name, value: '2' }] + ] + for (const settings of cases) { + const servingRevision = revision() + servingRevision.spec.containers[0].env.push(...settings) + assert.throws(() => readRelayServingRegionalPlacementVersion(input, { + run: (args) => args[1] === 'services' ? serving() : servingRevision + }), /cohort is invalid/) + } +}) diff --git a/cloud/docs/orca-relay-operations.md b/cloud/docs/orca-relay-operations.md index cd989b58e94..f27b437fff9 100644 --- a/cloud/docs/orca-relay-operations.md +++ b/cloud/docs/orca-relay-operations.md @@ -491,3 +491,64 @@ Run and record each scenario in staging before launch: - return a dormant host, overload a cell, kill a cell, evacuate active work, and exercise pre-registration rollback. The served black-box relay suite validates the protocol/state transitions used by these procedures. The physical-device and real-GFE canaries remain separate launch gates; unit/black-box success cannot replace them. + +## Optional measured region correction (deployment gated) + +New optimization claims require both the durable regional-rehome control and +`ORCA_RELAY_REGION_CORRECTION_COHORT_PERCENT` (integer 0–100, default **0**). +Turning either gate off stops new optional moves; ordinary migration cleanup and +recovery continue. Legacy preferred-region hints do not certify a correction. Both +cells must advertise regional protocol3 and the authenticated desktop control must +advertise idle-regional-rehome-v1. The source must have no actual client sockets or +pending admission/control work; a live control socket alone does not prevent a move. + +The monitor/deploy identity can read **GET `/v1/admin/regional-rehome-preview`**. +It returns full-population eligibility/exclusion counts, open-migration capacity, +process-safety gating and aggregate migration outcomes; it never claims a +host or changes the failure budget. This is advisory, with separately read state: +concurrent assignments, capacity changes, rate pauses and control changes can make +the next claim differ. Inspect the durable control separately before enabling. +Do not treat an unavailable/failed preview as zero eligible hosts. + +`orca_relay_region_correction_outcomes` reports attempts by source/target, +registration/completion/abort state, oldest open age and target +reservation units every five minutes. `orca_relay_region_comparison` samples a +stable 10% of accepted reports (including unchanged hosts), keyed by host digest, +assignment epoch and decision generation. Existing control RTT and client-accept +logs include assignment epoch, control generation and drain mode; join those for +matched before/after and unchanged-cohort comparisons. Client accept latency is +connection setup, not application command round trip. No application-latency +improvement has been demonstrated by probe differences alone. + +Quiet live connections count as work and defer optional correction indefinitely. +A returning client may race with the short admission gate and retry normally. No +optimization timer may close an established client. Investigate failed registration, +ambiguous authority, stuck reservations and reconnect/failure rates against agreed +limits. A database outage can keep the source fenced until locked reconciliation +establishes its authority; timeout alone is not permission to reopen admissions. + +All directors must run the reviewed idle worker before enabling. Record the tested +immutable source and rollback revisions, then verify the ordinary migration recovery +path before rollout. There is no retained-source table or renewal protocol. Deploying +supporting cells/desktops and enabling a cohort require separate rollout authorization +and explicit numerical stop criteria; this change enables neither. + +### Setting the correction cohort during a reviewed director rollout + +The existing **Deploy Relay Production Director** workflow accepts +`region-correction-cohort-percent`: `preserve` (default) or an integer0–100. +It carries the cohort onto both candidate and compatible rollback revisions and +verifies the environment before promotion. If the predecessor has no setting, +`preserve` stamps zero. An explicit change requires the exact disabled durable +rehome generation; configuring a nonzero cohort does not itself enable the sweep. +The usual image, identity, health and traffic checks remain in force. No workflow +was dispatched as part of implementation. + +Terraform reads the cohort from the same traffic-serving revision used to preserve +regional placement. A later apply therefore preserves a workflow-set cohort, +including explicit zero; only an absent service/setting bootstraps to0. Malformed +or ambiguous live settings fail the plan instead of silently resetting the cohort. +The audited director workflow owns subsequent changes. +Before the first nonzero cohort, verify compatible protocol2 cells, updated +cleanup workers, preview eligibility, both serving/rollback images and the +explicitly approved observation/stop criteria. diff --git a/cloud/infra/terraform/relay.tf b/cloud/infra/terraform/relay.tf index 7a5124a00b1..5df7372ff07 100644 --- a/cloud/infra/terraform/relay.tf +++ b/cloud/infra/terraform/relay.tf @@ -169,6 +169,11 @@ resource "google_cloud_run_v2_service" "relay" { } } + env { + name = "ORCA_RELAY_REGION_CORRECTION_COHORT_PERCENT" + value = data.external.relay_serving_regional_placement_version.result.cohort_percent + } + ports { container_port = 8080 } diff --git a/cloud/packages/relay-contract/src/control-messages.ts b/cloud/packages/relay-contract/src/control-messages.ts index 0daf21e7c28..ebe9586407e 100644 --- a/cloud/packages/relay-contract/src/control-messages.ts +++ b/cloud/packages/relay-contract/src/control-messages.ts @@ -50,6 +50,7 @@ export const RELAY_HOST_CAPABILITIES_HEADER = 'x-orca-host-capabilities' // The host accepts kind/relayDeviceId on a pendingConns entry. A host that does // not advertise this parses those entries strictly and would drop the whole ack. export const RELAY_HOST_CAPABILITY_PENDING_CONN_DETAILS = 'pending-conn-details' +export const RELAY_HOST_CAPABILITY_IDLE_REGIONAL_REHOME = 'idle-regional-rehome-v1' export function parseRelayHostCapabilities( header: string | string[] | undefined @@ -121,11 +122,22 @@ export const DeviceRevokeSchema = z .object({ reqId: OpaqueIdSchema, relayDeviceId: OpaqueIdSchema }) .strict() -export const AuthRefreshSchema = z.object({ relayJwt: z.string().min(1).max(8 * 1024) }).strict() +export const AuthRefreshSchema = z + .object({ + relayJwt: z + .string() + .min(1) + .max(8 * 1024) + }) + .strict() export const DrainSchema = z .object({ - graceMs: z.number().int().nonnegative().max(60 * 60 * 1000), + graceMs: z + .number() + .int() + .nonnegative() + .max(60 * 60 * 1000), recovery: z.literal('resolve-director') }) .strict() diff --git a/cloud/packages/relay-contract/src/director-messages.ts b/cloud/packages/relay-contract/src/director-messages.ts index e697135b68e..0f57081014a 100644 --- a/cloud/packages/relay-contract/src/director-messages.ts +++ b/cloud/packages/relay-contract/src/director-messages.ts @@ -7,8 +7,15 @@ import { RelayHostIdSchema } from './wire-scalars.js' import { RelayRegionSchema } from './relay-regions.js' +import { + RegionCorrectionRequestSchema, + RegionCorrectionResponseSchema +} from './region-correction.js' -const SignedAssignmentLeaseSchema = z.string().min(1).max(8 * 1024) +const SignedAssignmentLeaseSchema = z + .string() + .min(1) + .max(8 * 1024) export const AssignmentRequestSchema = z .object({ @@ -17,7 +24,8 @@ export const AssignmentRequestSchema = z // Client-declared reconnection; the director verifies it against the // durable assignment before granting fast-lane admission. reconnect: z.boolean().optional(), - preferredRegion: RelayRegionSchema.optional() + preferredRegion: RelayRegionSchema.optional(), + regionCorrection: RegionCorrectionRequestSchema.optional() }) .strict() @@ -26,7 +34,8 @@ export const AssignmentResponseSchema = z v: z.literal(1), cellUrl: CanonicalHttpsOriginSchema, assignmentEpoch: GenerationSchema, - lease: SignedAssignmentLeaseSchema + lease: SignedAssignmentLeaseSchema, + regionCorrection: RegionCorrectionResponseSchema.optional() }) .strict() diff --git a/cloud/packages/relay-contract/src/idle-regional-rehome.ts b/cloud/packages/relay-contract/src/idle-regional-rehome.ts new file mode 100644 index 00000000000..9f9eff36593 --- /dev/null +++ b/cloud/packages/relay-contract/src/idle-regional-rehome.ts @@ -0,0 +1,26 @@ +import { z } from 'zod' +import { GenerationSchema, RelayHostIdSchema } from './wire-scalars.js' + +export const IdleRegionalRehomeRequestSchema = z + .object({ + v: z.literal(1), + attemptId: z.string().uuid(), + userId: z.string().min(1).max(256), + relayHostId: RelayHostIdSchema, + sourceCellId: z.string().min(1).max(128), + sourceCellIncarnation: z.string().uuid(), + sourceAssignmentEpoch: GenerationSchema.refine((value) => value > 0), + sourceGeneration: GenerationSchema.refine((value) => value > 0), + targetCellId: z.string().min(1).max(128) + }) + .strict() + +export const IdleRegionalRehomeResponseSchema = z + .object({ + v: z.literal(1), + outcome: z.enum(['busy', 'committed', 'deferred', 'stale']) + }) + .strict() + +export type IdleRegionalRehomeRequest = z.infer +export type IdleRegionalRehomeOutcome = z.infer['outcome'] diff --git a/cloud/packages/relay-contract/src/index.ts b/cloud/packages/relay-contract/src/index.ts index aab3b53b5f3..3b52ec503a1 100644 --- a/cloud/packages/relay-contract/src/index.ts +++ b/cloud/packages/relay-contract/src/index.ts @@ -13,3 +13,5 @@ export * from './resume-confirmation-contract.js' export * from './relay-regions.js' export * from './splice-state-machine.js' export * from './wire-scalars.js' +export * from './region-correction.js' +export * from './idle-regional-rehome.js' diff --git a/cloud/packages/relay-contract/src/region-correction.test.ts b/cloud/packages/relay-contract/src/region-correction.test.ts new file mode 100644 index 00000000000..8c0122341d4 --- /dev/null +++ b/cloud/packages/relay-contract/src/region-correction.test.ts @@ -0,0 +1,76 @@ +import { describe, expect, it } from 'vitest' +import { AssignmentRequestSchema, AssignmentResponseSchema } from './director-messages.js' +import { DrainSchema } from './control-messages.js' +import { RegionCorrectionRequestSchema } from './region-correction.js' + +const report = { + v: 1, + action: 'report', + generation: 3, + assignmentEpoch: 7, + policyVersion: 1, + outcome: 'conclusive', + measurements: { 'us-central1': 40, 'asia-east2': 180 } +} +const retention = { + mode: 'finish-existing', + attemptId: '11111111-1111-4111-8111-111111111111', + sourceGeneration: 3, + sourceAssignmentEpoch: 7 +} + +describe('region correction wire boundaries', () => { + it('keeps legacy assignment shapes readable without negotiated fields', () => { + expect( + AssignmentRequestSchema.parse({ v: 1, relayHostId: 'abcdefghijklmnop' }) + ).not.toHaveProperty('regionCorrection') + expect( + AssignmentResponseSchema.parse({ + v: 1, + cellUrl: 'https://cell.example', + assignmentEpoch: 1, + lease: 'synthetic-lease' + }) + ).not.toHaveProperty('regionCorrection') + }) + + it('accepts complete comparison evidence and explicit inconclusive reports', () => { + expect(RegionCorrectionRequestSchema.safeParse(report).success).toBe(true) + const { measurements: _measurements, ...basis } = report + expect( + RegionCorrectionRequestSchema.safeParse({ + ...basis, + outcome: 'inconclusive', + reason: 'probe-unavailable' + }).success + ).toBe(true) + }) + + it.each([ + { measurements: { 'us-central1': 40 } }, + { measurements: { 'us-central1': -1, 'asia-east2': 10 } }, + { measurements: { 'us-central1': Infinity, 'asia-east2': 10 } }, + { measurements: { 'us-central1': 120_001, 'asia-east2': 10 } }, + { generation: Number.MAX_SAFE_INTEGER + 1 }, + { assignmentEpoch: 1.2 }, + { policyVersion: 2 }, + { outcome: 'inconclusive', reason: 'timeout' } + ])('rejects ambiguous or unbounded evidence: %j', (override) => { + expect(RegionCorrectionRequestSchema.safeParse({ ...report, ...override }).success).toBe(false) + }) + + it('rejects reporting and issuing a window in the same request', () => { + expect( + RegionCorrectionRequestSchema.safeParse({ + ...report, + action: 'issue-window' + }).success + ).toBe(false) + }) + + it('uses ordinary drain and rejects the superseded retention extension', () => { + const ordinary = { graceMs: 0, recovery: 'resolve-director' } + expect(DrainSchema.parse(ordinary)).toEqual(ordinary) + expect(DrainSchema.safeParse({ ...ordinary, retention }).success).toBe(false) + }) +}) diff --git a/cloud/packages/relay-contract/src/region-correction.ts b/cloud/packages/relay-contract/src/region-correction.ts new file mode 100644 index 00000000000..5fffca0ad8c --- /dev/null +++ b/cloud/packages/relay-contract/src/region-correction.ts @@ -0,0 +1,60 @@ +import { z } from 'zod' +import { EpochMsSchema, GenerationSchema } from './wire-scalars.js' +import { RelayRegionSchema } from './relay-regions.js' + +const RttSchema = z.number().finite().nonnegative().max(120_000) +export const RegionMeasurementsSchema = z + .object({ + 'us-central1': RttSchema, + 'asia-east2': RttSchema + }) + .strict() + +export const RegionMeasurementWindowSchema = z + .object({ + generation: GenerationSchema, + expiresAt: EpochMsSchema, + assignmentEpoch: GenerationSchema, + incumbentRegion: RelayRegionSchema, + policyVersion: z.literal(1) + }) + .strict() + +const ReportBasis = { + v: z.literal(1), + action: z.literal('report'), + generation: GenerationSchema, + assignmentEpoch: GenerationSchema, + policyVersion: z.literal(1) +} + +export const RegionCorrectionRequestSchema = z.union([ + z.object({ v: z.literal(1), action: z.literal('issue-window') }).strict(), + z + .object({ + ...ReportBasis, + outcome: z.literal('conclusive'), + measurements: RegionMeasurementsSchema + }) + .strict(), + z + .object({ + ...ReportBasis, + outcome: z.literal('inconclusive'), + reason: z.string().min(1).max(64) + }) + .strict() +]) + +export const RegionCorrectionResponseSchema = z + .object({ + v: z.literal(1), + window: RegionMeasurementWindowSchema.optional(), + reportStatus: z.enum(['accepted', 'duplicate', 'stale', 'expired', 'basis-changed']).optional() + }) + .strict() + +export type RegionMeasurements = z.infer +export type RegionMeasurementWindow = z.infer +export type RegionCorrectionRequest = z.infer +export type RegionCorrectionResponse = z.infer