fix(relay): count region selection on every assign lane, fallback only for new placements (#32186)

* fix(relay): count region selection on every assign lane, fallback only for new placements

Co-Authored-By: Claude <noreply@anthropic.com>

* test(relay): pin placed marker on isolated re-place and kept pins; drain-return fallback

Co-Authored-By: Claude <noreply@anthropic.com>

---------

Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Jinwoo Hong
2026-10-23 18:35:53 -04:00
committed by GitHub
co-authored by Claude
parent 2c53a550ff
commit 6b24bf596c
7 changed files with 303 additions and 25 deletions
+13 -6
View File
@@ -31,6 +31,7 @@ import {
RelayAssignmentRowBusyError,
RelayCellInReserveModeError,
RelayHomeCellUnavailableError,
type AssignedRelayAssignment,
type AssignOptions,
type CellFenceAttemptEvidence,
type RelayAssignment,
@@ -445,6 +446,13 @@ export function createRelayApp(
? requestedRegion
: RELAY_DEFAULT_REGION
operations.recordRegionRequest?.(requestedRegion)
// Only an answer this call minted can be a fallback: a kept assignment was placed earlier.
const recordSelection = (assignment: RelayAssignment, placed: boolean) =>
operations.recordRegionSelection?.({
targetRegion,
selectedRegion: assignment.region,
fallback: placed && assignment.region !== targetRegion
})
const signAssignment = async (assignment: RelayAssignment) =>
await new SignJWT({
purpose: 'cell-assignment',
@@ -476,6 +484,7 @@ export function createRelayApp(
}
if (plan.kind === 'answer') {
operations.recordAssignmentAdmission?.(plan.lane === 'sticky' ? 'sticky' : 'placement')
recordSelection(plan.assignment, plan.lane !== 'sticky')
return context.json({
v: 1,
cellUrl: plan.assignment.cellUrl,
@@ -484,6 +493,7 @@ export function createRelayApp(
})
}
if (plan.kind === 'retry') {
operations.recordRegionSelection?.({ targetRegion, fallback: false })
operations.recordAssignmentUnavailable?.(
plan.reason === 'map-incomplete'
? 'reserve-map-incomplete'
@@ -592,7 +602,7 @@ export function createRelayApp(
return rejectPublicAssignment(context)
}
}
let assignment: RelayAssignment
let assignment: AssignedRelayAssignment
let regionCorrection: RegionCorrectionResponse | undefined
try {
if (body.data.regionCorrection?.action === 'report') {
@@ -626,6 +636,7 @@ export function createRelayApp(
}
} catch (error) {
if (error instanceof ReservePaceError) {
operations.recordRegionSelection?.({ targetRegion, fallback: false })
operations.recordAssignmentUnavailable?.('reserve-paced')
context.header('Retry-After', String(error.retryAfterSeconds))
return context.json({ error: 'assignment_temporarily_unavailable' }, 503)
@@ -671,11 +682,7 @@ export function createRelayApp(
} finally {
admission.release()
}
operations.recordRegionSelection?.({
targetRegion,
selectedRegion: assignment.region,
fallback: lane === 'placement' && assignment.region !== targetRegion
})
recordSelection(assignment, assignment.placed === true)
// Grant-side counterpart of the rejection log: reconnect grants are rare
// enough to log and make "which cell is this host on" answerable. The
// placement-lane ones matter most — they are the only record that a host
@@ -297,8 +297,8 @@ describe('re-placing a host off a cell isolated for a roll', () => {
const second = await store.assign(IDENTITY, 'us-central1')
// Identical, not merely same-celled: the clock has not moved, so every
// field including the lease deadline must match.
expect(second).toEqual(first)
// field including the lease deadline must match; only the first was placed.
expect(second).toEqual({ ...first, placed: undefined })
// One row read answers both the stranded rule and the roll-stamp check.
expect(counter.count('relay_cell_admission')).toBe(1)
expect(counter.count('relay_assignment_migrations')).toBe(0)
@@ -315,10 +315,9 @@ describe('re-placing a host off a cell isolated for a roll', () => {
await park(first.cellId, 'migration-only')
expect(await rollIsolatedAt(first.cellId)).toBeNull()
expect(await store.assign(IDENTITY, 'us-central1')).toMatchObject({
cellId: first.cellId,
assignmentEpoch: first.assignmentEpoch
})
const pinned = await store.assign(IDENTITY, 'us-central1')
expect(pinned).toMatchObject({ cellId: first.cellId, assignmentEpoch: first.assignmentEpoch })
expect(pinned.placed).toBeUndefined()
})
it('re-places a host once the roll stamp is set', async () => {
@@ -332,13 +331,13 @@ describe('re-placing a host off a cell isolated for a roll', () => {
expect(moved.cellId).not.toBe(first.cellId)
expect(moved.assignmentEpoch).toBe(first.assignmentEpoch + 1)
expect(moved.placed).toBe(true)
expect(counter.count('relay_assignment_migrations')).toBeGreaterThan(0)
// Durable: the next dial does not bounce back to the isolated cell.
expect(await store.assign(IDENTITY, 'us-central1')).toMatchObject({
cellId: moved.cellId,
assignmentEpoch: moved.assignmentEpoch
})
const kept = await store.assign(IDENTITY, 'us-central1')
expect(kept).toMatchObject({ cellId: moved.cellId, assignmentEpoch: moved.assignmentEpoch })
expect(kept.placed).toBeUndefined()
})
it.each([
@@ -584,7 +583,7 @@ describe('re-placing a host off a cell isolated for a roll', () => {
const regionReserve = vi.fn(async () => onReserve)
await expect(
store.assign(IDENTITY, 'asia-east2', 'asia-east2', undefined, { regionReserve })
).resolves.toEqual(onReserve)
).resolves.toEqual({ ...onReserve, placed: true })
expect(regionReserve).toHaveBeenCalledOnce()
expect(jsonEvents(warn, 'orca_relay_sticky_replaced_off_isolated_cell')).toEqual([
{
@@ -639,7 +638,7 @@ describe('re-placing a host off a cell isolated for a roll', () => {
warn.mockClear()
await expect(
store.assign(IDENTITY, 'asia-east2', 'asia-east2', undefined, { placeFresh: async () => onReserve })
).resolves.toEqual(onReserve)
).resolves.toEqual({ ...onReserve, placed: true })
expect(jsonEvents(warn, 'orca_relay_sticky_replaced_off_isolated_cell')).toMatchObject([
{ fromCellId: first.cellId, toCellId: 'asia-reserve', via: 'reserve' }
])
@@ -0,0 +1,62 @@
import { ASSIGNMENT_LIMITS } from '@orca-cloud/relay-contract'
import { afterEach, describe, expect, it } from 'vitest'
import { RelayAssignmentStore } from './assignment-store.js'
import { openInMemoryRelayDatabase, type RelayDatabase } from './database.js'
const CELL = { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }
const IDENTITY = { userId: 'user-a', relayHostId: 'host000000000001' }
describe('assign marks only the answers it minted as placed', () => {
let database: RelayDatabase | undefined
afterEach(async () => await database?.close())
async function setup(now: () => number): Promise<RelayAssignmentStore> {
database = await openInMemoryRelayDatabase()
const store = new RelayAssignmentStore(database, now)
await store.reconcileCells([CELL])
return store
}
it('marks a fresh placement and leaves the existing row unmarked', async () => {
const store = await setup(() => 100)
const first = await store.assign(IDENTITY)
const again = await store.assign(IDENTITY)
expect(first.placed).toBe(true)
expect(again).toMatchObject({ cellId: first.cellId, assignmentEpoch: first.assignmentEpoch })
expect(again.placed).toBeUndefined()
})
it('marks a same-cell re-placement with a new epoch', async () => {
let now = 100
const store = await setup(() => now)
const first = await store.assign(IDENTITY)
await store.changeActivity(IDENTITY, 'control', -1)
now += ASSIGNMENT_LIMITS.activityLeaseMs + 1
await store.releaseExpiredActivityLeases()
await store.releaseExpiredActivity()
now += ASSIGNMENT_LIMITS.dormantTtlMs
const replaced = await store.assign(IDENTITY)
expect(replaced).toMatchObject({ cellId: first.cellId, assignmentEpoch: first.assignmentEpoch + 1 })
expect(replaced.placed).toBe(true)
})
it('marks a reserve booking and keeps the returned assignment otherwise unchanged', async () => {
const store = await setup(() => 100)
const booking = {
...IDENTITY,
cellId: 'reserve-c1',
cellUrl: 'https://reserve-c1.example.com',
assignmentEpoch: 9,
leaseExpiresAt: 400_000,
region: 'us-central1' as const
}
const booked = await store.assign(IDENTITY, undefined, undefined, undefined, {
placeFresh: async () => booking
})
expect(booked).toEqual({ ...booking, placed: true })
})
})
@@ -1750,7 +1750,8 @@ describe('RelayAssignmentStore', () => {
const first = await store.assign(identity)
const second = await store.assign(identity)
expect(second).toEqual(first)
expect(first.placed).toBe(true)
expect(second).toEqual({ ...first, placed: undefined })
expect(await cellReservations(database!)).toEqual({ 'cell-a': 1, 'cell-b': 0 })
})
+8 -5
View File
@@ -189,6 +189,9 @@ export type RelayAssignment = AssignmentIdentity & {
region?: RelayRegion
}
// `placed`: this assign minted the answer (a new epoch or a reserve booking); absent for the existing row.
export type AssignedRelayAssignment = RelayAssignment & { placed?: true }
// Set only by reconnect verification, for a host whose home cell is isolated for
// a roll right now: its next assign re-places it rather than re-granting the pin.
export type ResolvedRelayAssignment = RelayAssignment & { homeCellRollIsolated?: true }
@@ -1015,7 +1018,7 @@ export class RelayAssignmentStore {
// Step 5: a fresh placement may first be booked on a reserve-mode cell, outside any
// transaction; and a database mint stays above every epoch the director's map knew.
step5: AssignOptions = {}
): Promise<RelayAssignment> {
): Promise<AssignedRelayAssignment> {
const sticky = await this.assignStickyWithLockRetry(identity, lockMode, preferredRegion)
if (sticky && !('isolatedFrom' in sticky)) return sticky
let isolatedFrom = sticky?.isolatedFrom
@@ -1042,7 +1045,7 @@ export class RelayAssignmentStore {
})
)
}
return reserved
return reserved && { ...reserved, placed: true as const }
}
const booked = step5.placeFresh ? await bookReserve(step5.placeFresh) : null
if (booked) return booked
@@ -1092,7 +1095,7 @@ export class RelayAssignmentStore {
epochFloor = 0,
hostRequested = false,
keepRegion = false
): Promise<RelayAssignment> {
): Promise<AssignedRelayAssignment> {
const deadline = Date.now() + ASSIGNMENT_LOCK_RETRY_DEADLINE_MS
let inventoryScope: AssignmentInventoryScope = 'none'
while (true) {
@@ -1386,7 +1389,7 @@ export class RelayAssignmentStore {
hostRequested = false,
// Refuse to spill out of placementRegion: AssignmentRegionFull instead.
keepRegion = false
): Promise<RelayAssignment> {
): Promise<AssignedRelayAssignment> {
const now = this.now()
// An isolated retry takes its tier's rows before anything reassigns this,
// so a lock timeout there must retry the same tier, not fall to 'all'.
@@ -1771,7 +1774,7 @@ export class RelayAssignmentStore {
now
)
return {
assignment: { ...identity, ...target, assignmentEpoch, leaseExpiresAt },
assignment: { ...identity, ...target, assignmentEpoch, leaseExpiresAt, placed: true as const },
events
}
}).catch((error: unknown) => {
@@ -206,7 +206,7 @@ describe('Relay regional placement', () => {
it('books the reserve cell rather than spill out of a full region', async () => {
await fillAsia()
const place = vi.fn(async () => onReserve)
await expect(assignAsia(place)).resolves.toEqual(onReserve)
await expect(assignAsia(place)).resolves.toEqual({ ...onReserve, placed: true })
expect(place).toHaveBeenCalledOnce()
// The refused in-region attempt wrote nothing.
expect(await database.query(`SELECT cell_id FROM relay_assignments`)).toEqual([])
@@ -0,0 +1,206 @@
import { describe, expect, it, vi, type Mock } from 'vitest'
import type { RelayConfig } from './config.js'
import { ReservePaceError } from './reserve-assignment.js'
const fakes = vi.hoisted(() => ({
verifyRelayToken: vi.fn(async (token: string) => ({ sub: 'user-1', relayHostId: token }))
}))
vi.mock('./relay-token-verifier.js', () => ({
createRelayTokenVerifier: () => fakes.verifyRelayToken,
readBearer: (value: string | undefined) => value?.replace(/^Bearer /, '') ?? null
}))
import { createRelayApp } from './app.js'
const HOST = 'selecthost000001'
function usAssignment(extra: Record<string, unknown> = {}) {
return {
userId: 'user-1',
relayHostId: HOST,
cellId: 'us-c1',
cellUrl: 'https://us-c1.relay.example.test',
region: 'us-central1' as const,
assignmentEpoch: 3,
leaseExpiresAt: Date.now() + 300_000,
...extra
}
}
describe('Relay region selection counters', () => {
it('counts a reserve sticky answer as selected, never as a fallback', async () => {
const recordRegionSelection = vi.fn()
const plan = vi.fn(async () => ({ kind: 'answer', lane: 'sticky', assignment: usAssignment() }))
const response = await assign({ recordRegionSelection, plan }, { reconnect: true })
expect(response.status).toBe(200)
expect(recordRegionSelection.mock.calls).toEqual([
[{ targetRegion: 'asia-east2', selectedRegion: 'us-central1', fallback: false }]
])
})
it('counts a reserve placement out of the target region as a fallback', async () => {
const recordRegionSelection = vi.fn()
const plan = vi.fn(async () => ({ kind: 'answer', lane: 'placement', assignment: usAssignment() }))
await assign({ recordRegionSelection, plan }, { reconnect: true })
expect(recordRegionSelection.mock.calls).toEqual([
[{ targetRegion: 'asia-east2', selectedRegion: 'us-central1', fallback: true }]
])
})
it('counts a reserve-lane retry as an unavailable region', async () => {
const recordRegionSelection = vi.fn()
const plan = vi.fn(async () => ({ kind: 'retry', retryAfterSeconds: 2, reason: 'paced' }))
const response = await assign({ recordRegionSelection, plan }, { reconnect: true })
expect(response.status).toBe(503)
expect(recordRegionSelection.mock.calls).toEqual([[{ targetRegion: 'asia-east2', fallback: false }]])
})
it('counts a reserve booking paced inside the database lane as an unavailable region', async () => {
const recordRegionSelection = vi.fn()
const assignFn = vi.fn(async () => {
throw new ReservePaceError(3)
})
const response = await assign({ recordRegionSelection, assign: assignFn })
expect(response.status).toBe(503)
expect(recordRegionSelection.mock.calls).toEqual([[{ targetRegion: 'asia-east2', fallback: false }]])
})
it('does not count the existing row returned on the placement lane as a fallback', async () => {
const recordRegionSelection = vi.fn()
await assign({ recordRegionSelection, assign: vi.fn(async () => usAssignment()) })
expect(recordRegionSelection.mock.calls).toEqual([
[{ targetRegion: 'asia-east2', selectedRegion: 'us-central1', fallback: false }]
])
})
it('counts a placement out of the target region on any lane as a fallback', async () => {
const recordRegionSelection = vi.fn()
// The reconnect verifies on the sticky lane; the store then re-places the host.
await assign(
{
recordRegionSelection,
assign: vi.fn(async () => usAssignment({ placed: true })),
resolve: vi.fn(async () => usAssignment())
},
{ reconnect: true }
)
expect(recordRegionSelection.mock.calls).toEqual([
[{ targetRegion: 'asia-east2', selectedRegion: 'us-central1', fallback: true }]
])
})
it('counts a drain-return re-placement out of the target region as a fallback', async () => {
const recordRegionSelection = vi.fn()
const recordAssignmentAdmission = vi.fn()
await assign(
{
recordRegionSelection,
recordAssignmentAdmission,
assign: vi.fn(async () => usAssignment({ placed: true })),
resolve: vi.fn(async () => usAssignment({ cellId: 'asia-c1', region: 'asia-east2', homeCellRollIsolated: true }))
},
{ reconnect: true }
)
expect(recordAssignmentAdmission).toHaveBeenCalledWith('drain-return')
expect(recordRegionSelection.mock.calls).toEqual([
[{ targetRegion: 'asia-east2', selectedRegion: 'us-central1', fallback: true }]
])
})
it('does not count a migration report as a fallback', async () => {
const recordRegionSelection = vi.fn()
await assign(
{
recordRegionSelection,
// A report carries no preference, so its target is the default region.
resolve: vi.fn(async () => usAssignment({ cellId: 'asia-c1', region: 'asia-east2' })),
exchangeRegionCorrection: vi.fn(async () => ({ v: 1, reportStatus: 'accepted' }))
},
{
regionCorrection: {
v: 1,
action: 'report',
generation: 2,
assignmentEpoch: 3,
policyVersion: 1,
outcome: 'inconclusive',
reason: 'timeout'
}
}
)
expect(recordRegionSelection).toHaveBeenCalledTimes(1)
expect(recordRegionSelection.mock.calls[0]![0]).toMatchObject({ fallback: false })
})
})
async function assign(
operations: {
recordRegionSelection: Mock
recordAssignmentAdmission?: Mock
plan?: Mock
assign?: Mock
resolve?: Mock
exchangeRegionCorrection?: Mock
},
extra: Record<string, unknown> = {}
): Promise<Response> {
const app = createRelayApp(config(), {
store: {} as never,
assignments: {
assign: operations.assign ?? vi.fn(),
resolve: operations.resolve ?? vi.fn(async () => null),
exchangeRegionCorrection: operations.exchangeRegionCorrection ?? vi.fn()
} as never,
...(operations.plan ? { reservePlacement: { plan: operations.plan } as never } : {}),
drain: vi.fn(),
ready: vi.fn(async () => true),
recordRegionSelection: operations.recordRegionSelection,
...(operations.recordAssignmentAdmission
? { recordAssignmentAdmission: operations.recordAssignmentAdmission }
: {})
})
return await app.request('/v1/assign', {
method: 'POST',
headers: { authorization: `Bearer ${HOST}`, 'content-type': 'application/json' },
body: JSON.stringify({ v: 1, relayHostId: HOST, preferredRegion: 'asia-east2', ...extra })
})
}
function config(): RelayConfig {
return {
port: 8080,
publicUrl: 'https://relay.example.test',
cellUrl: 'https://relay.example.test',
region: 'us-central1',
authIssuer: 'https://auth.example.test',
authAudience: 'orca-relay',
jwksUrl: 'https://auth.example.test/jwks',
assignmentSigningKey: new TextEncoder().encode('assignment-key-with-at-least-32-bytes'),
role: 'director',
cellId: 'director',
cells: [],
adminAudience: 'https://relay.example.test/v1/admin/drain',
deployServiceAccount: 'deploy@example.test',
runtimeServiceAccount: 'runtime@example.test',
adminJwksUrl: 'https://auth.example.test/jwks',
databasePoolMax: 3,
publicAssignmentsEnabled: true,
regionalPlacementEnabled: true,
publicAssignmentConcurrency: 2,
publicAssignmentQueueMax: 128,
publicAssignmentWaitMs: 4_000,
publicResolveConcurrency: 1,
publicResolveWaitMs: 5_000,
publicAssignmentRetryAfterSeconds: 5,
dataDir: './data'
}
}