diff --git a/cloud/apps/relay/src/app.ts b/cloud/apps/relay/src/app.ts index f6e29dd11fb..70467351dcc 100644 --- a/cloud/apps/relay/src/app.ts +++ b/cloud/apps/relay/src/app.ts @@ -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 diff --git a/cloud/apps/relay/src/assignment-isolated-cell-replacement.test.ts b/cloud/apps/relay/src/assignment-isolated-cell-replacement.test.ts index b76013ce7a8..17815d2aad4 100644 --- a/cloud/apps/relay/src/assignment-isolated-cell-replacement.test.ts +++ b/cloud/apps/relay/src/assignment-isolated-cell-replacement.test.ts @@ -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' } ]) diff --git a/cloud/apps/relay/src/assignment-placed-marker.test.ts b/cloud/apps/relay/src/assignment-placed-marker.test.ts new file mode 100644 index 00000000000..f50fcaa8682 --- /dev/null +++ b/cloud/apps/relay/src/assignment-placed-marker.test.ts @@ -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 { + 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 }) + }) +}) diff --git a/cloud/apps/relay/src/assignment-store.test.ts b/cloud/apps/relay/src/assignment-store.test.ts index 39eeaa0f924..b5ed192760c 100644 --- a/cloud/apps/relay/src/assignment-store.test.ts +++ b/cloud/apps/relay/src/assignment-store.test.ts @@ -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 }) }) diff --git a/cloud/apps/relay/src/assignment-store.ts b/cloud/apps/relay/src/assignment-store.ts index 21a9eb8013b..1d2cb007060 100644 --- a/cloud/apps/relay/src/assignment-store.ts +++ b/cloud/apps/relay/src/assignment-store.ts @@ -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 { + ): Promise { 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 { + ): Promise { 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 { + ): Promise { 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) => { diff --git a/cloud/apps/relay/src/relay-region-placement.test.ts b/cloud/apps/relay/src/relay-region-placement.test.ts index e2fc0c4be07..502119ae28b 100644 --- a/cloud/apps/relay/src/relay-region-placement.test.ts +++ b/cloud/apps/relay/src/relay-region-placement.test.ts @@ -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([]) diff --git a/cloud/apps/relay/src/relay-region-selection-app.test.ts b/cloud/apps/relay/src/relay-region-selection-app.test.ts new file mode 100644 index 00000000000..fc195fc167c --- /dev/null +++ b/cloud/apps/relay/src/relay-region-selection-app.test.ts @@ -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 = {}) { + 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 = {} +): Promise { + 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' + } +}