diff --git a/cloud/apps/relay/src/app.ts b/cloud/apps/relay/src/app.ts index 7786432d585..db8586719d3 100644 --- a/cloud/apps/relay/src/app.ts +++ b/cloud/apps/relay/src/app.ts @@ -39,6 +39,7 @@ import { ASSIGNMENT_LEASE_AUDIENCE, assignmentLeaseKeyId } from './assignment-le import { AssignmentRejectionLogWindow } from './assignment-rejection-log-window.js' import { CELL_ADMISSION_STATES } from './cell-admission-selector.js' import { registerCellSeatFeedRoute } from './cell-seat-feed-route.js' +import { ReservePaceError, type ReserveAssignment, type ReservePlan } from './reserve-assignment.js' import { registerCellReserveRoutes } from './cell-reserve-routes.js' import type { DemoteRequest, ReserveOutcome, ReserveRequest } from './cell-reserve-contract.js' import type { CellSeatFeedPage } from './cell-seat-log.js' @@ -128,6 +129,8 @@ export function createRelayApp( cellIncarnation?: string cellSeatFeed?: (sinceSeq: number | null) => CellSeatFeedPage cellFlags?: () => AppliedControlFlags + // Step 5 on directors; absent keeps today's path for every request. + reservePlacement?: ReserveAssignment cellReserve?: (request: ReserveRequest) => ReserveOutcome[] cellDemote?: (request: DemoteRequest) => string cellReserverPoll?: () => void @@ -419,6 +422,56 @@ export function createRelayApp( ? requestedRegion : RELAY_DEFAULT_REGION operations.recordRegionRequest?.(requestedRegion) + const signAssignment = async (assignment: RelayAssignment) => + await new SignJWT({ + purpose: 'cell-assignment', + cellId: assignment.cellId, + cellUrl: assignment.cellUrl, + assignmentEpoch: assignment.assignmentEpoch, + relayHostId: claims.relayHostId + }) + .setProtectedHeader({ alg: 'HS256', kid: assignmentLeaseKid }) + .setIssuer(config.publicUrl) + .setAudience(ASSIGNMENT_LEASE_AUDIENCE) + .setSubject(claims.sub) + .setIssuedAt() + .setExpirationTime('5m') + .sign(config.assignmentSigningKey) + // Step 5: switched-on cells answer from the map; anything else keeps today's path below. + let step5: { epochFloor?: number; placeFresh?: () => Promise } = {} + if (operations.reservePlacement && !body.data.regionCorrection) { + let plan: ReservePlan + try { + plan = await operations.reservePlacement.plan(identity, { + reconnect: Boolean(body.data.reconnect), + region: targetRegion + }) + } catch (error) { + if (!(error instanceof ReservePaceError)) throw error + plan = { kind: 'retry', retryAfterSeconds: error.retryAfterSeconds, reason: 'paced' } + } + if (plan.kind === 'answer') { + operations.recordAssignmentAdmission?.(plan.lane === 'sticky' ? 'sticky' : 'placement') + return context.json({ + v: 1, + cellUrl: plan.assignment.cellUrl, + assignmentEpoch: plan.assignment.assignmentEpoch, + lease: await signAssignment(plan.assignment) + }) + } + if (plan.kind === 'retry') { + operations.recordAssignmentUnavailable?.( + plan.reason === 'map-incomplete' + ? 'reserve-map-incomplete' + : plan.reason === 'database' + ? 'reserve-database' + : 'reserve-paced' + ) + context.header('Retry-After', String(plan.retryAfterSeconds)) + return context.json({ error: 'assignment_temporarily_unavailable' }, 503) + } + step5 = { epochFloor: plan.epochFloor, placeFresh: plan.placeFresh } + } let admission: { release(): void } | null = null let lane: AssignmentAdmissionLane = 'placement' if (body.data.reconnect) { @@ -519,9 +572,19 @@ export function createRelayApp( 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) + // Without step 5 the call is exactly today's. + const step5Active = step5.epochFloor !== undefined || step5.placeFresh !== undefined + assignment = step5Active + ? await operations.assignments.assign( + identity, + requestedRegion, + requestedRegion ? targetRegion : undefined, + undefined, + step5 + ) + : requestedRegion + ? await operations.assignments.assign(identity, requestedRegion, targetRegion) + : await operations.assignments.assign(identity) } if (body.data.regionCorrection) { try { @@ -537,6 +600,11 @@ export function createRelayApp( } } } catch (error) { + if (error instanceof ReservePaceError) { + operations.recordAssignmentUnavailable?.('reserve-paced') + context.header('Retry-After', String(error.retryAfterSeconds)) + return context.json({ error: 'assignment_temporarily_unavailable' }, 503) + } if (error instanceof RelayAssignmentRowBusyError) { logAssignmentRejection({ route: 'assign', @@ -593,20 +661,7 @@ export function createRelayApp( ` host=${relayHostLogDigest(claims.relayHostId)} cell=${assignment.cellId}` ) } - const lease = await new SignJWT({ - purpose: 'cell-assignment', - cellId: assignment.cellId, - cellUrl: assignment.cellUrl, - assignmentEpoch: assignment.assignmentEpoch, - relayHostId: claims.relayHostId - }) - .setProtectedHeader({ alg: 'HS256', kid: assignmentLeaseKid }) - .setIssuer(config.publicUrl) - .setAudience(ASSIGNMENT_LEASE_AUDIENCE) - .setSubject(claims.sub) - .setIssuedAt() - .setExpirationTime('5m') - .sign(config.assignmentSigningKey) + const lease = await signAssignment(assignment) return context.json({ v: 1, cellUrl: assignment.cellUrl, @@ -814,6 +869,10 @@ export function createRelayApp( ...(operations.cellAdmitModeEffective ? { admitModeEffective: operations.cellAdmitModeEffective() } : {}), + // Placement as built, not as configured: `on` only when this director books reserve cells. + ...(config.role === 'director' + ? { reservePlacement: operations.reservePlacement?.placementMode ?? 'off' } + : {}), connectionCapacity: config.connectionHardCap === undefined ? null diff --git a/cloud/apps/relay/src/assignment-store.ts b/cloud/apps/relay/src/assignment-store.ts index 2541348d128..d085b55dae2 100644 --- a/cloud/apps/relay/src/assignment-store.ts +++ b/cloud/apps/relay/src/assignment-store.ts @@ -751,15 +751,26 @@ export class RelayAssignmentStore { placementRegion: RelayRegion = preferredRegion ?? RELAY_DEFAULT_REGION, // evacuateDeadCells re-enters placement from a sweep; it must not take the // bounded wait, whose 55P03 would surface as a terminal sweep failure. - lockMode: CellInventoryLockMode = 'request' + lockMode: CellInventoryLockMode = 'request', + // 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: { epochFloor?: number; placeFresh?: () => Promise } = {} ): Promise { const sticky = await this.assignStickyWithLockRetry(identity, lockMode, preferredRegion) if (sticky) return sticky + const booked = await step5.placeFresh?.() + if (booked) return booked // Only placement needs the global inventory critical section; queueing those // attempts locally avoids turning true placement bursts into NOWAIT storms. return await this.serializeAssignment( async () => - await this.assignWithLockRetry(identity, lockMode, preferredRegion, placementRegion) + await this.assignWithLockRetry( + identity, + lockMode, + preferredRegion, + placementRegion, + step5.epochFloor ?? 0 + ) ) } @@ -778,7 +789,8 @@ export class RelayAssignmentStore { identity: AssignmentIdentity, lockMode: CellInventoryLockMode, preferredRegion?: RelayRegion, - placementRegion: RelayRegion = preferredRegion ?? RELAY_DEFAULT_REGION + placementRegion: RelayRegion = preferredRegion ?? RELAY_DEFAULT_REGION, + epochFloor = 0 ): Promise { const deadline = Date.now() + ASSIGNMENT_LOCK_RETRY_DEADLINE_MS let inventoryScope: AssignmentInventoryScope = 'none' @@ -789,7 +801,8 @@ export class RelayAssignmentStore { inventoryScope, lockMode, preferredRegion, - placementRegion + placementRegion, + epochFloor ) } catch (error) { if (error instanceof AssignmentInventoryScopeChanged) { @@ -1038,7 +1051,8 @@ export class RelayAssignmentStore { inventoryScope: AssignmentInventoryScope, lockMode: CellInventoryLockMode, preferredRegion?: RelayRegion, - placementRegion: RelayRegion = preferredRegion ?? RELAY_DEFAULT_REGION + placementRegion: RelayRegion = preferredRegion ?? RELAY_DEFAULT_REGION, + epochFloor = 0 ): Promise { const now = this.now() // An isolated retry takes its tier's rows before anything reassigns this, @@ -1313,7 +1327,9 @@ export class RelayAssignmentStore { } } await this.adjustCellReservation(transaction, target.cellId, 1) - const assignmentEpoch = existing ? integer(existing, 'assignment_epoch') + 1 : 1 + // A host that left a reserve-mode cell inside the ledger's lag never reuses an epoch. + const assignmentEpoch = + Math.max(existing ? integer(existing, 'assignment_epoch') : 0, epochFloor) + 1 // Counters follow the leases the host still holds, plus the new control. const counts = activityCounts(keptLeases) counts.control += 1 @@ -7879,7 +7895,8 @@ export class RelayAssignmentStore { const now = this.now() const rows = await this.database.query( `SELECT cell.cell_id, cell.cell_url, cell.capacity_requests, region.region, - admission.roll_isolated_at, runtime.ready, runtime.last_heartbeat_at + admission.admission_state, admission.roll_isolated_at, runtime.ready, + runtime.last_heartbeat_at 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 @@ -7902,7 +7919,12 @@ export class RelayAssignmentStore { heartbeatExpiresAt !== null && heartbeatExpiresAt > now && integer(row, 'capacity_requests') > 0 && - optionalInteger(row, 'roll_isolated_at') === undefined + optionalInteger(row, 'roll_isolated_at') === undefined, + // Step 5 places only on general cells that are not isolated for a roll. + general: + optionalText(row, 'admission_state') === 'general' && + optionalInteger(row, 'roll_isolated_at') === undefined && + integer(row, 'capacity_requests') > 0 } }) } diff --git a/cloud/apps/relay/src/cell-reserve-client.test.ts b/cloud/apps/relay/src/cell-reserve-client.test.ts new file mode 100644 index 00000000000..16a400a693d --- /dev/null +++ b/cloud/apps/relay/src/cell-reserve-client.test.ts @@ -0,0 +1,55 @@ +import { describe, expect, it, vi } from 'vitest' +import { CellReserveClient } from './cell-reserve-client.js' + +const target = { cellId: 'c1', cellUrl: 'https://c1.example' } +const item = (userId: string, epoch = 3) => ({ + userId, + relayHostId: 'abcdefghijklmnop', + epoch, + ttlMs: 30_000 +}) + +describe('cell reserve client', () => { + it('batches a cell’s bookings into one request and maps each answer back', async () => { + const fetchImpl = vi.fn(async (_url: string | URL | Request, init?: RequestInit) => { + const body = JSON.parse(String(init?.body)) as { items: Array<{ userId: string }> } + return Response.json({ + v: 1, + results: body.items.map((entry) => + entry.userId === 'b' ? { outcome: 'seated-newer', epoch: 9 } : { outcome: 'ok' } + ) + }) + }) + const client = new CellReserveClient({ + directorId: 'director-a', + identityToken: async () => 'token', + fetch: fetchImpl + }) + const answers = await Promise.all([ + client.reserve(target, item('a')), + client.reserve(target, item('b')), + client.reserve(target, item('c'), true) + ]) + expect(answers).toEqual([{ outcome: 'ok' }, { outcome: 'seated-newer', epoch: 9 }, { outcome: 'ok' }]) + // One booking batch and one dry-run batch. + expect(fetchImpl).toHaveBeenCalledTimes(2) + const first = JSON.parse(String(fetchImpl.mock.calls[0]![1]?.body)) as Record + expect(first).toMatchObject({ v: 1, directorId: 'director-a', items: [{ userId: 'a' }, { userId: 'b' }] }) + expect(String(fetchImpl.mock.calls[0]![0])).toBe('https://c1.example/internal/reserve') + }) + + it('reads every item as unreachable when the cell fails or answers out of shape', async () => { + for (const fetchImpl of [ + async () => Promise.reject(new Error('timeout')), + async () => new Response('nope', { status: 503 }), + async () => Response.json({ v: 1, results: [] }) + ]) { + const client = new CellReserveClient({ + directorId: 'director-a', + identityToken: async () => 'token', + fetch: fetchImpl + }) + expect(await client.reserve(target, item('a'))).toEqual({ outcome: 'unreachable' }) + } + }) +}) diff --git a/cloud/apps/relay/src/cell-reserve-client.ts b/cloud/apps/relay/src/cell-reserve-client.ts new file mode 100644 index 00000000000..4daaa0beece --- /dev/null +++ b/cloud/apps/relay/src/cell-reserve-client.ts @@ -0,0 +1,108 @@ +import { + CELL_DEMOTE_PATH, + CELL_RESERVE_MAX_ITEMS, + CELL_RESERVE_PATH, + ReserveResponseSchema, + type DemoteRequest, + type ReserveItem +} from './cell-reserve-contract.js' +import type { ReserveAttempt } from './reserve-placement.js' + +// The director's calls to cells. Bookings for one cell batch for up to 5 ms or 64 items; +// a cell that does not answer within the timeout is unverifiable, and every item in the +// batch reads `unreachable` (one of the request's three tries). +export const CELL_RESERVE_FLUSH_MS = 5 +export const CELL_RESERVE_TIMEOUT_MS = 1_000 + +export type ReserveTarget = { cellId: string; cellUrl: string } + +type Pending = { item: ReserveItem; resolve: (attempt: ReserveAttempt) => void } + +export class CellReserveClient { + private readonly batches = new Map() + + constructor( + private readonly input: { + directorId: string + identityToken: () => Promise + fetch?: typeof fetch + flushMs?: number + timeoutMs?: number + } + ) {} + + reserve(target: ReserveTarget, item: ReserveItem, dryRun = false): Promise { + const key = `${target.cellId}\u0000${dryRun ? 'dry' : 'book'}` + return new Promise((resolve) => { + let batch = this.batches.get(key) + if (!batch) { + batch = { target, dryRun, items: [] } + this.batches.set(key, batch) + const timer = setTimeout(() => this.flush(key), this.input.flushMs ?? CELL_RESERVE_FLUSH_MS) + timer.unref?.() + } + batch.items.push({ item, resolve }) + if (batch.items.length >= CELL_RESERVE_MAX_ITEMS) this.flush(key) + }) + } + + // The cell treats repeats as one; a failure is only logged, and the next poll retries. + async demote(target: ReserveTarget, request: DemoteRequest): Promise { + try { + const response = await this.post(target, CELL_DEMOTE_PATH, request) + await response.body?.cancel().catch(() => undefined) + return response.ok + } catch { + return false + } + } + + private flush(key: string): void { + const batch = this.batches.get(key) + if (!batch) return + this.batches.delete(key) + void this.send(batch.target, batch.dryRun, batch.items) + } + + private async send(target: ReserveTarget, dryRun: boolean, items: Pending[]): Promise { + let results: Array<{ outcome: string; epoch?: number }> | null = null + try { + const response = await this.post(target, CELL_RESERVE_PATH, { + v: 1, + directorId: this.input.directorId, + ...(dryRun ? { dryRun: true } : {}), + items: items.map((pending) => pending.item) + }) + if (response.ok) { + const body = ReserveResponseSchema.safeParse(await response.json().catch(() => null)) + if (body.success && body.data.results.length === items.length) results = body.data.results + } else { + await response.body?.cancel().catch(() => undefined) + } + } catch { + results = null + } + items.forEach((pending, index) => { + const result = results?.[index] + pending.resolve( + !result + ? { outcome: 'unreachable' } + : result.outcome === 'seated-newer' && typeof result.epoch === 'number' + ? { outcome: 'seated-newer', epoch: result.epoch } + : { outcome: result.outcome } + ) + }) + } + + private async post(target: ReserveTarget, path: string, body: unknown): Promise { + return await (this.input.fetch ?? fetch)(new URL(path, target.cellUrl), { + method: 'POST', + headers: { + authorization: `Bearer ${await this.input.identityToken()}`, + 'content-type': 'application/json' + }, + body: JSON.stringify(body), + signal: AbortSignal.timeout(this.input.timeoutMs ?? CELL_RESERVE_TIMEOUT_MS) + }) + } +} diff --git a/cloud/apps/relay/src/config.ts b/cloud/apps/relay/src/config.ts index cd8f88eef5b..7e0406f80f4 100644 --- a/cloud/apps/relay/src/config.ts +++ b/cloud/apps/relay/src/config.ts @@ -99,6 +99,9 @@ const EnvSchema = z.object({ 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), + // Step 5 on directors: `off` (default) is today's path; `dry-run` asks the picked cell to + // check a booking without making it, and logs the answer; `on` places on switched-on cells. + ORCA_RELAY_RESERVE_PLACEMENT: z.enum(['off', 'dry-run', 'on']).default('off'), // Step 3 shadow directory: unset or empty polls no cell; `all` or a comma list of cell ids. ORCA_RELAY_SHADOW_SEAT_FEED_CELLS: z .string() @@ -218,6 +221,7 @@ export type RelayConfig = { rehomeDirectorServiceAccount?: string rehomeAudience?: string shadowSeatFeedCells?: 'all' | string[] + reservePlacement?: 'off' | 'dry-run' | 'on' runtimeServiceAccount: string directorUrl?: string heartbeatAudience?: string @@ -368,6 +372,7 @@ export function loadRelayConfig(env: NodeJS.ProcessEnv = process.env): RelayConf rehomeDirectorServiceAccount: parsed.ORCA_RELAY_REHOME_DIRECTOR_SERVICE_ACCOUNT, rehomeAudience: parsed.ORCA_RELAY_REHOME_AUDIENCE, shadowSeatFeedCells: parsed.ORCA_RELAY_SHADOW_SEAT_FEED_CELLS, + reservePlacement: parsed.ORCA_RELAY_RESERVE_PLACEMENT, runtimeServiceAccount: parsed.ORCA_RELAY_RUNTIME_SERVICE_ACCOUNT ?? parsed.ORCA_RELAY_DEPLOY_SERVICE_ACCOUNT, directorUrl, diff --git a/cloud/apps/relay/src/director-reserve-placement-status.test.ts b/cloud/apps/relay/src/director-reserve-placement-status.test.ts new file mode 100644 index 00000000000..190ceca4ffb --- /dev/null +++ b/cloud/apps/relay/src/director-reserve-placement-status.test.ts @@ -0,0 +1,81 @@ +import { describe, expect, it, vi } from 'vitest' +import type { RelayConfig } from './config.js' + +vi.mock('./admin-token-verifier.js', () => ({ + createAdminTokenVerifier: () => async (token: string) => token === 'deploy-token', + createReadOnlyAdminTokenVerifier: () => async () => false, + createRegionalRehomeControlApplyTokenVerifier: () => async () => false, + createRegionalRehomeRuntimeTokenVerifier: () => async () => false, + createRegionalRehomeTokenVerifier: () => async () => false, + createRuntimeTokenVerifier: () => async () => false +})) + +vi.mock('./relay-token-verifier.js', () => ({ + createRelayTokenVerifier: () => async () => null, + readBearer: (value: string | undefined) => value?.replace(/^Bearer /, '') ?? null +})) + +import { createRelayApp } from './app.js' + +// The deploy guard reads a director's placement here, so it must be what the director does. +describe('a director reports its reserve placement as built', () => { + const status = async (operations: Record, overrides: Partial = {}) => { + const app = createRelayApp(config(overrides), { + store: {} as never, + assignments: {} as never, + drain: vi.fn(), + ready: vi.fn(async () => true), + ...operations + }) + const response = await app.request('/v1/admin/runtime-status', { + method: 'POST', + headers: { authorization: 'Bearer deploy-token', 'content-type': 'application/json' }, + body: JSON.stringify({ v: 1 }) + }) + return (await response.json()) as Record + } + + it('says off when placement is configured on but was never built', async () => { + expect(await status({}, { reservePlacement: 'on' })).toMatchObject({ role: 'director', reservePlacement: 'off' }) + }) + + it('says what the built placement does', async () => { + expect(await status({ reservePlacement: { placementMode: 'on' } }, { reservePlacement: 'on' })).toMatchObject({ + reservePlacement: 'on' + }) + }) + + it('reports nothing for a cell', async () => { + expect(await status({}, { role: 'cell', cellId: 'production-gce-c1' })).not.toHaveProperty('reservePlacement') + }) +}) + +function config(overrides: Partial = {}): 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 Uint8Array(32), + role: 'director', + cellId: 'director', + cells: [], + adminAudience: 'https://relay.example.test/v1/admin/drain', + deployServiceAccount: 'deploy@example.test', + runtimeServiceAccount: 'relay-cell@example.test', + adminJwksUrl: 'https://auth.example.test/jwks', + databasePoolMax: 10, + publicAssignmentsEnabled: true, + publicAssignmentConcurrency: 2, + publicAssignmentQueueMax: 128, + publicAssignmentWaitMs: 4_000, + publicResolveConcurrency: 1, + publicResolveWaitMs: 5_000, + publicAssignmentRetryAfterSeconds: 5, + dataDir: './data', + ...overrides + } +} diff --git a/cloud/apps/relay/src/floating-promise-census.test.ts b/cloud/apps/relay/src/floating-promise-census.test.ts index 0a647836a51..19fd403a8db 100644 --- a/cloud/apps/relay/src/floating-promise-census.test.ts +++ b/cloud/apps/relay/src/floating-promise-census.test.ts @@ -12,6 +12,8 @@ const NEVER_REJECTS = new Map([ ['relay-background-operation.ts#runRelayBackgroundOperation', 'catches and logs every failure'], ['assignment-cleanup-steps.ts#runAssignmentCleanup', 'runs each step through the above'], ['cell-heartbeat-client.ts#send', 'one try/catch around the whole send'], + ['cell-reserve-client.ts#send', 'one try/catch around the request; resolves every item'], + ['cell-reserve-client.ts#demote', 'one try/catch around the request'], ['control-renewal-batch.ts#flush', 'rejects the waiters, never itself'], ['regional-rehome-worker.ts#run', 'one try/catch around the whole poll'], ['relay-control-flag-channel.ts#poll', 'one catch on the whole read'] diff --git a/cloud/apps/relay/src/google-metadata-identity-token.ts b/cloud/apps/relay/src/google-metadata-identity-token.ts index f51f6282a65..9b93ad6f7c1 100644 --- a/cloud/apps/relay/src/google-metadata-identity-token.ts +++ b/cloud/apps/relay/src/google-metadata-identity-token.ts @@ -20,3 +20,24 @@ export async function googleMetadataIdentityToken( } return token } + +// Google identity tokens live an hour; minting one per request would be a metadata read each. +// A failed mint is never reused. +export function reusedIdentityToken( + mint: () => Promise, + now: () => number = Date.now, + reuseMs = 10 * 60_000 +): () => Promise { + let token: { value: Promise; at: number } | undefined + return () => { + const at = now() + if (!token || at - token.at > reuseMs) { + const value = mint() + token = { value, at } + value.catch(() => { + if (token?.value === value) token = undefined + }) + } + return token.value + } +} diff --git a/cloud/apps/relay/src/relay-observability.ts b/cloud/apps/relay/src/relay-observability.ts index 612beb0fa1e..b24b0608ea6 100644 --- a/cloud/apps/relay/src/relay-observability.ts +++ b/cloud/apps/relay/src/relay-observability.ts @@ -63,6 +63,10 @@ export type AssignmentUnavailableCause = | 'relay_connection_headroom_exhausted' | 'relay_home_cell_unavailable' | 'relay_assignment_row_busy' + // Step 5: a starting director's map, a never-seen host with the database down, or intake. + | 'reserve-map-incomplete' + | 'reserve-database' + | 'reserve-paced' // Row-lock waiters seen in one pg_stat_activity sample, by who waits, on which // table, behind whom. Roles come from each pool's application_name. diff --git a/cloud/apps/relay/src/relay-server.ts b/cloud/apps/relay/src/relay-server.ts index 53151ffda87..d2f8e71ab56 100644 --- a/cloud/apps/relay/src/relay-server.ts +++ b/cloud/apps/relay/src/relay-server.ts @@ -47,6 +47,10 @@ import { createRelayTokenVerifier, readBearer } from './relay-token-verifier.js' import { closeRelayWebSocket } from './relay-websocket-close.js' import { ShadowDirectoryCompare } from './shadow-directory-compare.js' import { startShadowSeatPoller } from './shadow-seat-directory.js' +import { CellReserveClient } from './cell-reserve-client.js' +import { googleMetadataIdentityToken, reusedIdentityToken } from './google-metadata-identity-token.js' +import { ReserveAssignment } from './reserve-assignment.js' +import { ReservePlacer } from './reserve-placement.js' import { ProcessQueuedByteBudget } from './splice-forwarder.js' // A malformed percent-escape in the request target must be a client error, never a URIError @@ -168,12 +172,62 @@ export function createRelayServer( observeGrace: (event) => observability.recordReadinessGrace(event) }) const ready = readiness.check + // Set once placement is built: a director configured `on` that could not build it (no map or + // no rehome credential) books nothing, so it must not keep a reserve cell's dead-man fed. + const placing = { on: false } const shadowSeatPoller = startShadowSeatPoller(config, { - listCells: () => assignments.seatFeedCells() + listCells: () => assignments.seatFeedCells(), + reserver: () => placing.on }) const shadowCompare = shadowSeatPoller ? new ShadowDirectoryCompare(shadowSeatPoller.directory, options.now) : undefined + // Step 5 on directors: needs the map and the rehome credential the cells verify. + const directorId = randomUUID() + const rehomeAudience = config.rehomeAudience + const reservePlacement = + config.role === 'director' && + shadowSeatPoller && + rehomeAudience && + (config.reservePlacement ?? 'off') !== 'off' + ? new ReserveAssignment({ + mode: config.reservePlacement ?? 'off', + directory: shadowSeatPoller.directory, + cells: shadowSeatPoller.cells, + startedAt: shadowSeatPoller.startedAt, + placer: new ReservePlacer(options.now, options.random), + client: new CellReserveClient({ + directorId, + identityToken: reusedIdentityToken( + () => googleMetadataIdentityToken(rehomeAudience), + options.now + ) + }), + readRow: async (identity) => await assignments.hostWhereabouts(identity), + now: options.now, + random: options.random + }) + : undefined + placing.on = reservePlacement?.placementMode === 'on' + // A director that is not placing while a cell is in reserve mode strands that cell's hosts on + // its memory: its dead-man flips it back, and this says so loudly until then. + const offWithReserveTimer = + config.role === 'director' && shadowSeatPoller && !placing.on + ? setInterval(() => { + const { directory } = shadowSeatPoller + const cells = directory.cellIds().filter((cellId) => directory.admitModeOf(cellId) === 'reserve') + if (cells.length === 0) return + console.error( + JSON.stringify({ + event: 'orca_relay_reserve_placement_off_with_reserve_cells', + cells, + placement: reservePlacement?.placementMode ?? 'off', + configured: config.reservePlacement ?? 'off' + }) + ) + }, 60_000) + : null + offWithReserveTimer?.unref() const configuredConnectionLimits = config.connectionHardCap === undefined ? null @@ -328,6 +382,7 @@ export function createRelayServer( recordAssignmentUnavailable: (cause) => observability.recordAssignmentUnavailable?.(cause), recordRegionRequest: (region) => observability.recordRegionRequest?.(region), shadowSeats: shadowSeatPoller?.directory, + reservePlacement, compareShadowSeats: shadowCompare ? (route, identity, answer) => { shadowCompare.compare(route, identity, answer) @@ -754,6 +809,8 @@ export function createRelayServer( connectionSnapshot, ready, cellIncarnation, - shadowSeatPoller + shadowSeatPoller, + directorId, + reservePlacement } } diff --git a/cloud/apps/relay/src/reserve-assignment.test.ts b/cloud/apps/relay/src/reserve-assignment.test.ts new file mode 100644 index 00000000000..c48d8f3a4ac --- /dev/null +++ b/cloud/apps/relay/src/reserve-assignment.test.ts @@ -0,0 +1,283 @@ +import { describe, expect, it, vi } from 'vitest' +import type { CellReserveClient } from './cell-reserve-client.js' +import { ReserveAssignment, ReservePaceError, RESERVE_STARTUP_GATE_MS } from './reserve-assignment.js' +import { ReservePlacer, type ReserveAttempt } from './reserve-placement.js' +import { ShadowSeatDirectory, type SeatFeedCell, type SeatFeedResponse } from './shadow-seat-directory.js' + +const HOST = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } +const US = 'us-central1' as const + +function cellList(ids: string[], region = US): SeatFeedCell[] { + return ids.map((cellId) => ({ + cellId, + cellUrl: `https://${cellId}.example`, + region, + heartbeatExpiresAt: Number.MAX_SAFE_INTEGER, + requiredForComplete: true, + general: true + })) +} + +function feed( + cellId: string, + overrides: Partial & { admitMode?: 'db' | 'reserve' } = {} +): SeatFeedResponse { + const { admitMode = 'reserve', ...rest } = overrides + return { + v: 1, + cellId, + incarnation: `${cellId}-incarnation-1`, + seq: 1, + at: 0, + counts: { seats: 0, controls: 0, bookings: 0, units: 0, ceiling: 500 }, + intake: { perSec: 100, burst: 20, tokens: 20 }, + flagsApplied: { generation: 3, flags: { admitMode } }, + full: [], + ...rest + } +} + +function setup( + options: { + cells?: SeatFeedCell[] + feeds?: SeatFeedResponse[] + reserve?: (cellId: string, sticky: boolean) => ReserveAttempt + row?: (identity: { userId: string; relayHostId: string }) => Promise<{ cellId: string; assignmentEpoch: number } | null> + mode?: 'off' | 'dry-run' | 'on' + startedAt?: number + } = {} +) { + const now = { value: 100_000 } + const cells = options.cells ?? cellList(['c1', 'c2']) + const directory = new ShadowSeatDirectory() + directory.setCells(cells.map((cell) => cell.cellId)) + for (const response of options.feeds ?? cells.map((cell) => feed(cell.cellId))) { + directory.apply(response.cellId, response, now.value - 1, now.value - 1) + } + const reserved: Array<{ cellId: string; epoch: number; sticky: boolean; dryRun: boolean }> = [] + const demoted: Array<{ cellId: string; epoch: number; joinedAt: number }> = [] + const client = { + reserve: vi.fn(async (target: { cellId: string }, item: { epoch: number; sticky?: boolean }, dryRun = false) => { + reserved.push({ cellId: target.cellId, epoch: item.epoch, sticky: item.sticky ?? false, dryRun }) + return options.reserve?.(target.cellId, item.sticky ?? false) ?? { outcome: 'ok' } + }), + demote: vi.fn(async (target: { cellId: string }, request: { epoch: number; joinedAt: number }) => { + demoted.push({ cellId: target.cellId, epoch: request.epoch, joinedAt: request.joinedAt }) + return true + }) + } satisfies Pick + // Each director's estimate starts empty; a second of refill gives it its share. + const placer = new ReservePlacer(() => now.value, () => 0.25) + const assignment = new ReserveAssignment({ + mode: options.mode ?? 'on', + directory, + cells: () => cells, + startedAt: options.startedAt ?? 0, + placer, + // SAFETY: only reserve and demote are called, and both are implemented above. + // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: see above. + client: client as unknown as CellReserveClient, + readRow: options.row ?? (async () => null), + now: () => now.value, + random: () => 0.25, + log: () => undefined + }) + // Starts each cell's refill clock, as a director's first look does. + for (const cell of assignment.placementCells()) placer.hasBudget(cell) + now.value += 1_000 + for (const response of options.feeds ?? cells.map((cell) => feed(cell.cellId))) { + directory.apply(response.cellId, { ...response, full: response.full ?? [] }, now.value - 1, now.value - 1) + } + return { assignment, directory, reserved, demoted, now } +} + +function seated(epoch: number, joinedAt = 0) { + return { ...HOST, epoch, generation: 1, state: 'active', joinedAt } +} + +describe('reserve assignment on a director', () => { + it('keeps today’s path when off, or when no cell reports reserve mode', async () => { + expect((await setup({ mode: 'off' }).assignment.plan(HOST, { reconnect: true, region: US })).kind).toBe( + 'database' + ) + const databaseOnly = setup({ feeds: [feed('c1', { admitMode: 'db' }), feed('c2', { admitMode: 'db' })] }) + expect(await databaseOnly.assignment.plan(HOST, { reconnect: false, region: US })).toEqual({ + kind: 'database' + }) + }) + + it('reads what each cell acts on: a tripped dead-man is db, a flip back still leasing is reserve', async () => { + const tripped = setup({ + feeds: [feed('c1', { admitModeEffective: 'db' }), feed('c2', { admitModeEffective: 'db' })] + }) + expect(await tripped.assignment.plan(HOST, { reconnect: false, region: US })).toEqual({ kind: 'database' }) + const leasing = setup({ feeds: [feed('c1', { admitMode: 'db', admitModeEffective: 'reserve' }), feed('c2')] }) + expect(leasing.directory.admitModeOf('c1')).toBe('reserve') + }) + + it('before every live cell has answered, holds back only hosts that may be on an unheard cell', async () => { + const cells = cellList(['c1', 'c2']) + const onUnheard = { ...HOST, relayHostId: 'qrstuvwxyzabcdef' } + const gated = setup({ + cells, + feeds: [feed('c1')], + startedAt: 100_500, + row: async (identity) => + identity.relayHostId === onUnheard.relayHostId ? { cellId: 'c2', assignmentEpoch: 3 } : null + }) + // A fresh host goes to today's path, with no booking yet. + expect(await gated.assignment.plan(HOST, { reconnect: false, region: US })).toEqual({ + kind: 'database', + epochFloor: 0 + }) + expect(gated.reserved).toEqual([]) + expect(await gated.assignment.plan(onUnheard, { reconnect: true, region: US })).toMatchObject({ + kind: 'retry', + reason: 'map-incomplete' + }) + gated.now.value = 100_500 + RESERVE_STARTUP_GATE_MS + expect((await gated.assignment.plan(onUnheard, { reconnect: true, region: US })).kind).not.toBe('retry') + }) + + it('answers a reconnect from memory at the same epoch with no booking and no database read', async () => { + const row = vi.fn(async () => null) + const { assignment, reserved } = setup({ + feeds: [feed('c1', { full: [seated(7)] }), feed('c2')], + row + }) + const plan = await assignment.plan(HOST, { reconnect: true, region: US }) + expect(plan).toMatchObject({ kind: 'answer', lane: 'sticky', assignment: { cellId: 'c1', assignmentEpoch: 7 } }) + expect(reserved).toEqual([]) + expect(row).not.toHaveBeenCalled() + }) + + it('re-books a host on a cell that restarted since its seat, and only then answers it', async () => { + const { assignment, directory, reserved, now } = setup({ + feeds: [feed('c1', { full: [seated(7)] }), feed('c2')] + }) + // The cell restarted: its full snapshot no longer holds the seat. + directory.apply('c1', feed('c1', { incarnation: 'c1-incarnation-2' }), now.value, now.value) + const plan = await assignment.plan(HOST, { reconnect: true, region: US }) + expect(reserved[0]).toEqual({ cellId: 'c1', epoch: 7, sticky: true, dryRun: false }) + expect(plan).toMatchObject({ kind: 'answer', lane: 'sticky', assignment: { cellId: 'c1', assignmentEpoch: 7 } }) + }) + + it('never sends a drained or demoted host back, and books it elsewhere above every known epoch', async () => { + const { assignment, directory, reserved, now } = setup() + directory.apply( + 'c1', + feed('c1', { + full: undefined, + seq: 3, + changes: [ + { seq: 2, kind: 'join', ...HOST, epoch: 7, generation: 1, at: 1 }, + { seq: 3, kind: 'leave', ...HOST, epoch: 7, generation: 1, closeCode: 4503, at: 2 } + ] + }), + now.value, + now.value + ) + const plan = await assignment.plan(HOST, { reconnect: true, region: US }) + expect(plan).toMatchObject({ kind: 'answer', lane: 'placement', assignment: { assignmentEpoch: 8 } }) + expect(reserved.every((call) => !call.sticky && call.epoch === 8)).toBe(true) + }) + + it('leaves a host the map last saw on a database-mode cell to today’s path, with the floor', async () => { + const { assignment } = setup({ + feeds: [feed('c1', { admitMode: 'db', full: [seated(4)] }), feed('c2')] + }) + const plan = await assignment.plan(HOST, { reconnect: true, region: US }) + expect(plan).toMatchObject({ kind: 'database', epochFloor: 4 }) + expect(plan.kind === 'database' && typeof plan.placeFresh).toBe('function') + }) + + it('leaves a cell without a feed (c17/c18) on today’s path without starving the reserve pool', async () => { + const cells = cellList(['c1', 'c2', 'c17']) + let rowOnOld = false + const { assignment, directory, reserved, now } = setup({ + cells, + feeds: [feed('c1'), feed('c2', { admitMode: 'db' })], + row: async () => (rowOnOld ? { cellId: 'c17', assignmentEpoch: 5 } : null) + }) + directory.markNoFeed('c17', now.value) + expect(directory.isComplete()).toBe(true) + // An old image reports no ceiling; counting it as a database cell would zero the share. + expect(await assignment.plan(HOST, { reconnect: false, region: US })).toMatchObject({ + kind: 'answer', + lane: 'placement', + assignment: { cellId: 'c1' } + }) + expect(reserved.map((call) => call.cellId)).not.toContain('c17') + rowOnOld = true + expect( + await assignment.plan({ ...HOST, relayHostId: 'qrstuvwxyzabcdef' }, { reconnect: true, region: US }) + ).toMatchObject({ kind: 'database', epochFloor: 5 }) + // A cell that rolled back to an image without the feed takes no placement. + directory.markNoFeed('c1', now.value) + expect(directory.cellState('c1')?.ceiling).toBeUndefined() + expect(assignment.placementCells().find((cell) => cell.cellId === 'c1')?.general).toBe(false) + }) + + it('reads the row once for a host the map never saw, and refuses it while the database is down', async () => { + const down = setup({ row: async () => Promise.reject(new Error('timeout')) }) + expect(await down.assignment.plan(HOST, { reconnect: false, region: US })).toMatchObject({ + kind: 'retry', + reason: 'database' + }) + const known = setup({ row: async () => ({ cellId: 'c1', assignmentEpoch: 11 }) }) + expect(await known.assignment.plan(HOST, { reconnect: false, region: US })).toMatchObject({ + kind: 'answer', + assignment: { assignmentEpoch: 12 } + }) + }) + + it('paces with Retry-After when every reserve cell is out of budget and no database cell exists', async () => { + const { assignment } = setup({ reserve: () => ({ outcome: 'intake' }) }) + await expect(assignment.plan(HOST, { reconnect: false, region: US })).rejects.toBeInstanceOf( + ReservePaceError + ) + }) + + it('tells the old reserve seat to go when it books the host elsewhere, naming that seat exactly', async () => { + const { assignment, directory, demoted, now } = setup({ + feeds: [feed('c1', { full: [seated(5, 1_000)] }), feed('c2')], + // c1 is full, so the host lands on c2. + reserve: (cellId) => (cellId === 'c1' ? { outcome: 'full' } : { outcome: 'ok' }) + }) + // The host left c1 for a reason that moves it: its next placement is fresh, above 5. + directory.apply( + 'c1', + feed('c1', { + full: undefined, + seq: 2, + changes: [{ seq: 2, kind: 'drain-only', ...HOST, epoch: 5, generation: 1, at: 2 }] + }), + now.value, + now.value + ) + const plan = await assignment.plan(HOST, { reconnect: true, region: US }) + expect(plan).toMatchObject({ kind: 'answer', lane: 'placement', assignment: { assignmentEpoch: 6 } }) + expect(demoted).toEqual([{ cellId: 'c1', epoch: 5, joinedAt: 1_000 }]) + }) + + it('answers a reconnect only at the newest epoch, newest join first', async () => { + const { assignment, reserved } = setup({ + feeds: [feed('c1', { full: [seated(6, 1_000)] }), feed('c2', { full: [seated(5, 3_000)] })] + }) + // c2's seat joined later but at an older epoch: a stale rebind, never the answer. + expect(await assignment.plan(HOST, { reconnect: true, region: US })).toMatchObject({ + kind: 'answer', + lane: 'sticky', + assignment: { cellId: 'c1', assignmentEpoch: 6 } + }) + expect(reserved).toEqual([]) + }) + + it('in dry run only asks a cell to check a booking, and never changes the answer', async () => { + const { assignment, reserved } = setup({ mode: 'dry-run' }) + const plan = await assignment.plan(HOST, { reconnect: false, region: US }) + expect(plan.kind).toBe('database') + expect(plan.kind === 'database' && (await plan.placeFresh?.())).toBeNull() + expect(reserved).toMatchObject([{ dryRun: true }]) + }) +}) diff --git a/cloud/apps/relay/src/reserve-assignment.ts b/cloud/apps/relay/src/reserve-assignment.ts new file mode 100644 index 00000000000..701bdc1e0d4 --- /dev/null +++ b/cloud/apps/relay/src/reserve-assignment.ts @@ -0,0 +1,386 @@ +import { + ASSIGNMENT_LIMITS, + RELAY_CLOSE_CODE, + RELAY_DEFAULT_REGION, + type RelayRegion +} from '@orca-cloud/relay-contract' +import type { RelayAssignment } from './assignment-store.js' +import { CELL_RESERVE_TTL_MS } from './cell-reserve-contract.js' +import type { CellReserveClient } from './cell-reserve-client.js' +import { relayHostLogDigest } from './relay-host-log-digest.js' +import { + mintEpoch, + type PlacementCell, + type ReservePlacer +} from './reserve-placement.js' +import type { SeatFeedCell, ShadowSeat, ShadowSeatDirectory } from './shadow-seat-directory.js' + +// Step 5 on a director: answer a desktop from the in-memory map for cells that report +// admitMode=reserve, and fall back to today's database path for everything else. The cell +// decides every booking; this only estimates, picks and asks. + +// A leaver may rejoin its cell from memory this long, as the cell's own rule 2 allows. +export const RESERVE_STICKY_RECENT_MS = 10 * 60_000 +// A starting director waits this long for every live cell before it answers. +export const RESERVE_STARTUP_GATE_MS = 30_000 +const SUMMARY_INTERVAL_MS = 60_000 +// A leaver for these reasons was moved off its cell; it is not sent back there. +const NO_REJOIN_CLOSE_CODES: ReadonlySet = new Set([ + RELAY_CLOSE_CODE.DRAINING, + RELAY_CLOSE_CODE.WRONG_CELL +]) + +export type ReservePlacementMode = 'off' | 'dry-run' | 'on' + +export class ReservePaceError extends Error { + constructor(readonly retryAfterSeconds: number) { + super('reserve_cells_paced') + } +} + +type Identity = { userId: string; relayHostId: string } + +export type ReservePlan = + | { kind: 'answer'; assignment: RelayAssignment; lane: 'sticky' | 'placement' } + | { kind: 'retry'; retryAfterSeconds: number; reason: 'map-incomplete' | 'database' | 'paced' } + | { + kind: 'database' + epochFloor?: number + placeFresh?: () => Promise + } + +export class ReserveAssignment { + private summary: PlacementSummary = emptySummary() + private summaryAt: number + + constructor( + private readonly input: { + mode: ReservePlacementMode + directory: ShadowSeatDirectory + cells: () => readonly SeatFeedCell[] + startedAt: number + placer: ReservePlacer + client: CellReserveClient + // One primary-key read of the host's row: only for a host the map has never seen. + readRow: (identity: Identity) => Promise<{ cellId: string; assignmentEpoch: number } | null> + now?: () => number + random?: () => number + log?: (line: string) => void + } + ) { + this.summaryAt = this.now() + } + + // What this director actually does; the deploy guard reads it through runtime-status. + get placementMode(): ReservePlacementMode { + return this.input.mode + } + + private now(): number { + return (this.input.now ?? Date.now)() + } + + private random(): number { + return (this.input.random ?? Math.random)() + } + + placementCells(dryRun = false): PlacementCell[] { + const { directory } = this.input + return this.input.cells().flatMap((cell) => { + const state = directory.cellState(cell.cellId) + if (!state) return [] + const reserve = directory.admitModeOf(cell.cellId) === 'reserve' + return [ + { + cellId: cell.cellId, + region: cell.region, + // A dry run asks every step-5 cell, so the shadow covers cells still in database mode. + reserve: dryRun ? state.ceiling !== undefined : reserve, + general: cell.general === true && state.status === 'live', + draining: state.draining === true, + polledAt: state.status === 'live' ? (state.polledAt ?? null) : null, + seats: Math.max(0, (state.units ?? state.reportedSeats ?? 0) - (state.bookings ?? 0)), + bookings: state.bookings ?? 0, + ...(state.ceiling === undefined ? {} : { ceiling: state.ceiling }), + intakePerSec: state.intake?.perSec ?? 0, + intakeTokens: state.intake?.tokens ?? 0 + } + ] + }) + } + + private hasReserveCells(): boolean { + return this.input.directory.cellIds().some((cellId) => this.input.directory.admitModeOf(cellId) === 'reserve') + } + + async plan(identity: Identity, request: { reconnect: boolean; region: RelayRegion }): Promise { + this.flushSummary() + if (this.input.mode === 'off') return { kind: 'database' } + if (this.input.mode === 'dry-run') { + return { kind: 'database', placeFresh: async () => await this.dryRun(identity, request.region) } + } + if (!this.hasReserveCells()) return { kind: 'database' } + const now = this.now() + const { directory } = this.input + // Until every live cell has answered, only a host that may be on an unheard cell waits. + const incomplete = !directory.isComplete() && now - this.input.startedAt < RESERVE_STARTUP_GATE_MS + const seats = directory.seatsOf(identity.userId, identity.relayHostId) + const left = directory + .recentlyLeftOf(identity.userId, identity.relayHostId, now) + .filter((entry) => now - entry.at <= RESERVE_STICKY_RECENT_MS) + .at(-1) + const booking = directory.bookingOf(identity.userId, identity.relayHostId) + const known = [ + ...seats.map((seat) => seat.epoch), + ...(left ? [left.epoch] : []), + ...(booking ? [booking.epoch] : []) + ] + let row: { cellId: string; assignmentEpoch: number } | null = null + if (known.length === 0) { + try { + row = await this.input.readRow(identity) + } catch { + // A host the map never saw, with the database down: its epoch cannot be minted safely. + this.summary.retries.database += 1 + return { kind: 'retry', retryAfterSeconds: 2, reason: 'database' } + } + if (row) known.push(row.assignmentEpoch) + } + const epochFloor = known.length === 0 ? 0 : Math.max(...known) + const lastCellId = seats[0]?.cellId ?? left?.cellId ?? row?.cellId + if (incomplete) { + if (lastCellId !== undefined && directory.cellState(lastCellId)?.status === 'pending') { + this.summary.retries.mapIncomplete += 1 + return { kind: 'retry', retryAfterSeconds: 1, reason: 'map-incomplete' } + } + // Fresh placement waits for the full map; today's path places meanwhile. + if (lastCellId === undefined) return { kind: 'database', epochFloor } + } + + if (request.reconnect) { + const sticky = await this.stickyFromMemory(identity, seats, left, now) + if (sticky) { + this.summary.sticky += 1 + return { kind: 'answer', assignment: sticky, lane: 'sticky' } + } + } + const placeFresh = async (): Promise => + await this.placeFresh(identity, request.region, epochFloor, seats) + // Where the host was last seen decides who owns it: a database-mode cell keeps today's + // path, sticky included; only a fresh placement may be booked on a reserve-mode cell. + const lastOnReserve = lastCellId !== undefined && directory.admitModeOf(lastCellId) === 'reserve' + if (!lastOnReserve && (seats.length > 0 || left !== undefined || row !== null)) { + return { kind: 'database', epochFloor, placeFresh } + } + const placed = await placeFresh() + if (placed) return { kind: 'answer', assignment: placed, lane: 'placement' } + return { kind: 'database', epochFloor } + } + + // The same epoch on the same cell: no booking, no database, no row lock. A cell that + // restarted since forgot the seat, so it is asked to hold it first. + private async stickyFromMemory( + identity: Identity, + seats: ShadowSeat[], + left: { cellId: string; epoch: number; incarnation?: string; closeCode?: number } | undefined, + now: number + ): Promise { + // Highest epoch first, then the newest join: only the newest grant may be answered from + // memory, so a stale seat (or a leaver at a lower epoch) is never sent back. + const newest = Math.max(0, ...seats.map((seat) => seat.epoch), left?.epoch ?? 0) + const candidates: Array<{ + cellId: string + epoch: number + incarnation?: string + seated: boolean + joinedAt: number + }> = [ + ...seats + .filter((seat) => seat.state === 'active') + .map((seat) => ({ + cellId: seat.cellId, + epoch: seat.epoch, + incarnation: seat.incarnation, + seated: true, + joinedAt: seat.joinedAt + })), + ...(left && !NO_REJOIN_CLOSE_CODES.has(left.closeCode ?? 0) ? [{ ...left, seated: false, joinedAt: -1 }] : []) + ] + .filter((candidate) => candidate.epoch === newest) + .sort((first, second) => second.joinedAt - first.joinedAt) + for (const candidate of candidates) { + const cell = this.input.cells().find((entry) => entry.cellId === candidate.cellId) + const state = this.input.directory.cellState(candidate.cellId) + const placement = this.placementCells().find((entry) => entry.cellId === candidate.cellId) + if (!cell || !state || !placement || !this.input.placer.isCandidate(placement, cell.region, now)) { + continue + } + if (!candidate.seated && !this.input.placer.hasSeat(placement)) continue + if (candidate.incarnation !== undefined && candidate.incarnation !== state.incarnation) { + const answer = await this.input.client.reserve(cell, { + ...identity, + epoch: candidate.epoch, + ttlMs: CELL_RESERVE_TTL_MS, + sticky: true + }) + this.summary.stickyReserves += 1 + if (answer.outcome !== 'ok') continue + } + return { + ...identity, + cellId: cell.cellId, + cellUrl: cell.cellUrl, + assignmentEpoch: candidate.epoch, + leaseExpiresAt: now + ASSIGNMENT_LIMITS.activityLeaseMs, + region: cell.region + } + } + return null + } + + private async placeFresh( + identity: Identity, + region: RelayRegion, + epochFloor: number, + seats: readonly ShadowSeat[] = [] + ): Promise { + const placements = this.placementCells() + if (this.random() >= this.input.placer.reservePoolShare(placements, region)) return null + const regions = region === RELAY_DEFAULT_REGION ? [region] : [region, RELAY_DEFAULT_REGION] + for (const target of regions) { + const result = await this.input.placer.place({ + cells: placements, + region: target, + epoch: mintEpoch([epochFloor]), + reserve: async (cellId, epoch) => { + const cell = this.input.cells().find((entry) => entry.cellId === cellId) + if (!cell) return { outcome: 'unreachable' } + const answer = await this.input.client.reserve(cell, { + ...identity, + epoch, + ttlMs: CELL_RESERVE_TTL_MS + }) + this.summary.outcomes[answer.outcome] = (this.summary.outcomes[answer.outcome] ?? 0) + 1 + return answer + } + }) + this.summary.calls += result.calls + if (result.kind === 'placed') { + this.summary.placed += 1 + this.supersede(identity, seats, result.cellId, result.epoch) + const cell = this.input.cells().find((entry) => entry.cellId === result.cellId)! + return { + ...identity, + cellId: cell.cellId, + cellUrl: cell.cellUrl, + assignmentEpoch: result.epoch, + leaseExpiresAt: this.now() + ASSIGNMENT_LIMITS.activityLeaseMs, + region: cell.region + } + } + if (result.kind === 'pace') { + this.summary.retries.paced += 1 + throw new ReservePaceError(result.retryAfterSeconds) + } + // Database cells in this region take it; a region with none spills as today. + const hasDatabaseCells = placements.some( + (cell) => !cell.reserve && cell.region === target && cell.general && !cell.draining + ) + if (hasDatabaseCells) break + } + this.summary.database += 1 + return null + } + + // Never changes an answer: books nothing, and returns null so today's path places. + private async dryRun(identity: Identity, region: RelayRegion): Promise { + const placements = this.placementCells(true) + const cell = this.input.placer.pick(placements, region, new Set()) + const target = cell && this.input.cells().find((entry) => entry.cellId === cell.cellId) + if (!target) return null + const startedAt = performance.now() + void this.input.client + .reserve(target, { ...identity, epoch: 1, ttlMs: CELL_RESERVE_TTL_MS }, true) + .then((answer) => { + const elapsedMs = performance.now() - startedAt + this.summary.dryRun.outcomes[answer.outcome] = (this.summary.dryRun.outcomes[answer.outcome] ?? 0) + 1 + this.summary.dryRun.latenciesMs.push(elapsedMs) + }) + .catch(() => undefined) + return null + } + + // Supersede at placement: the old seats this director's map holds on reserve cells are told + // to go. Each cell checks the seat's epoch and join itself, so a stale map only misses. + private supersede(identity: Identity, seats: readonly ShadowSeat[], cellId: string, epoch: number): void { + for (const seat of seats) { + if (seat.cellId === cellId || seat.epoch >= epoch) continue + if (this.input.directory.admitModeOf(seat.cellId) !== 'reserve') continue + const cell = this.input.cells().find((entry) => entry.cellId === seat.cellId) + if (!cell) continue + this.summary.supersedes += 1 + void this.input.client.demote(cell, { + v: 1, + userId: identity.userId, + relayHostId: identity.relayHostId, + epoch: seat.epoch, + joinedAt: seat.joinedAt + }) + this.input.log?.( + JSON.stringify({ + event: 'orca_relay_reserve_supersede', + relayHostIdDigest: relayHostLogDigest(identity.relayHostId), + fromCellId: seat.cellId, + toCellId: cellId + }) + ) + } + } + + private flushSummary(): void { + const now = this.now() + if (now - this.summaryAt < SUMMARY_INTERVAL_MS) return + this.summaryAt = now + const { latenciesMs, outcomes } = this.summary.dryRun + latenciesMs.sort((left, right) => left - right) + const percentile = (share: number) => + latenciesMs.length === 0 + ? null + : Math.round(latenciesMs[Math.min(latenciesMs.length - 1, Math.floor(share * latenciesMs.length))]!) + ;(this.input.log ?? ((line: string) => console.log(line)))( + JSON.stringify({ + event: 'orca_relay_reserve_placement_summary', + mode: this.input.mode, + ...this.summary, + dryRun: { outcomes, count: latenciesMs.length, p50Ms: percentile(0.5), p99Ms: percentile(0.99) } + }) + ) + this.summary = emptySummary() + } +} + +type PlacementSummary = { + sticky: number + stickyReserves: number + placed: number + calls: number + database: number + supersedes: number + outcomes: Record + retries: { mapIncomplete: number; database: number; paced: number } + dryRun: { outcomes: Record; latenciesMs: number[] } +} + +function emptySummary(): PlacementSummary { + return { + sticky: 0, + stickyReserves: 0, + placed: 0, + calls: 0, + database: 0, + supersedes: 0, + outcomes: {}, + retries: { mapIncomplete: 0, database: 0, paced: 0 }, + dryRun: { outcomes: {}, latenciesMs: [] } + } +} diff --git a/cloud/apps/relay/src/shadow-seat-directory.test.ts b/cloud/apps/relay/src/shadow-seat-directory.test.ts index 8e90f442b66..83822afa8fd 100644 --- a/cloud/apps/relay/src/shadow-seat-directory.test.ts +++ b/cloud/apps/relay/src/shadow-seat-directory.test.ts @@ -139,7 +139,7 @@ describe('startShadowSeatPoller', () => { expect(cell.requests[1]?.searchParams.get('since')).toBe(`${INCARNATION}:2`) expect(cell.directory.seatsOf('user-a', 'host-1')).toEqual([]) expect(cell.directory.recentlyLeftOf('user-a', 'host-1', 1_000)).toEqual([ - { cellId: 'cell-a', epoch: 3, generation: 1, closeCode: 4001, at: 53 } + { cellId: 'cell-a', epoch: 3, generation: 1, incarnation: INCARNATION, closeCode: 4001, at: 53 } ]) expect(cell.directory.seatsOf('user-a', 'host-2')).toMatchObject([{ state: 'drain-only' }]) expect(cell.directory.seatsOf('user-a', 'host-3')).toMatchObject([ diff --git a/cloud/apps/relay/src/shadow-seat-directory.ts b/cloud/apps/relay/src/shadow-seat-directory.ts index 4785dbe4766..8cd0f7814db 100644 --- a/cloud/apps/relay/src/shadow-seat-directory.ts +++ b/cloud/apps/relay/src/shadow-seat-directory.ts @@ -1,7 +1,7 @@ import type { RelayRegion } from '@orca-cloud/relay-contract' import { z } from 'zod' import type { RelayConfig } from './config.js' -import { googleMetadataIdentityToken } from './google-metadata-identity-token.js' +import { googleMetadataIdentityToken, reusedIdentityToken } from './google-metadata-identity-token.js' // Step 3: each director polls every cell's seat feed and keeps who is seated // where in memory. It decides nothing; readers only compare and report. @@ -30,6 +30,15 @@ const SeatChangeSchema = z.object({ generation: z.number().int(), state: z.string().optional(), closeCode: z.number().int().optional(), + // Step 5: on `reserve`, and on the join a booking admitted. + reservedBy: z.string().optional(), + at: z.number() +}) +const RecentlyLeftSchema = z.object({ + userId: z.string(), + relayHostId: z.string(), + epoch: z.number().int(), + closeCode: z.number().int().optional(), at: z.number() }) const SeatSchema = z.object({ @@ -52,17 +61,29 @@ const SeatFeedSchema = z.object({ // Seats in the cell's own log, under the same generation rule: equals the map exactly. seats: z.number().int().nonnegative(), // Drops at the start of a close handshake, up to ~30 s before the leave lands. - controls: z.number().int().nonnegative() + controls: z.number().int().nonnegative(), + // Step 5: live bookings, every unit the cell's ledger counts, and its placement ceiling. + bookings: z.number().int().nonnegative(), + units: z.number().int().nonnegative(), + ceiling: z.number().int().nonnegative() }) .partial() .optional(), + intake: z + .object({ perSec: z.number().nonnegative(), burst: z.number(), tokens: z.number() }) + .optional(), flagsApplied: z .object({ generation: z.union([z.string(), z.number()]), flags: z.record(z.unknown()) }) .optional(), + // What the cell acts on: reserve also while a flip back is still leasing its controls, or + // db once its dead-man has flipped it back. Absent from older images. + admitModeEffective: z.enum(['db', 'reserve']).optional(), changes: z.array(SeatChangeSchema).optional(), // The cell cut the page (2,000 changes); poll again at once. more: z.boolean().optional(), - full: z.array(SeatSchema).optional() + full: z.array(SeatSchema).optional(), + // Step 5, with `full`: the cell's leavers of the last 10 minutes. + recentlyLeft: z.array(RecentlyLeftSchema).optional() }) export type SeatFeedResponse = z.infer @@ -74,6 +95,8 @@ export type SeatFeedCell = { heartbeatExpiresAt: number | null // Live, with capacity, and not roll-isolated: completeness waits only on these. requiredForComplete: boolean + // Admission state general and not roll-isolated: step 5 places only on these. + general?: boolean } export type ShadowSeat = { @@ -84,12 +107,15 @@ export type ShadowSeat = { state: string joinedAt: number observedAt: number + // The cell incarnation that reported it: a restart since means the cell forgot it. + incarnation?: string } export type RecentlyLeftSeat = { cellId: string epoch: number generation: number + incarnation?: string closeCode?: number // A full resync dropped the seat, so no close code was seen. resync?: true @@ -111,6 +137,24 @@ export type SeatFeedCellState = { reportedSeats?: number reportedControls?: number flagsApplied?: SeatFeedResponse['flagsApplied'] + admitModeEffective?: 'db' | 'reserve' + // Step 5: when the answered poll was sent, and the cell's own reserve numbers. + polledAt?: number + bookings?: number + units?: number + ceiling?: number + intake?: SeatFeedResponse['intake'] +} + +// A join a booking admitted, for the director that booked it (its after-the-fact row). +export type BookedJoin = { + cellId: string + userId: string + relayHostId: string + epoch: number + generation: number + joinedAt: number + reservedBy: string } type CellCursor = SeatFeedCellState & { seats: Map } @@ -131,6 +175,9 @@ export class ShadowSeatDirectory { private heartbeats: HeartbeatSnapshot | undefined // Per host, the newest database epoch seen and when: stands in for the grant time. private readonly databaseEpochs = new Map() + // Per host, the newest booking any cell reported (step 5 epoch mint). Age-ordered. + private readonly bookings = new Map() + onBookedJoin?: (join: BookedJoin) => void // Cells no longer listed are forgotten with their seats; new ones start pending. // `required` defaults to every listed cell. @@ -161,7 +208,8 @@ export class ShadowSeatDirectory { return `${cursor.incarnation}:${cursor.seq}` } - apply(cellId: string, response: SeatFeedResponse, now: number): void { + // `polledAt` is when the request left; the reply holds every change before it. + apply(cellId: string, response: SeatFeedResponse, now: number, polledAt = now): void { const cursor = this.cells.get(cellId) if (!cursor) return if (response.cellId !== cellId) { @@ -170,7 +218,19 @@ export class ShadowSeatDirectory { } let seq: number | null = response.seq if (response.full) { - this.replaceSeats(cursor, response.full, now) + this.replaceSeats(cursor, response.full, now, response.incarnation) + for (const left of response.recentlyLeft ?? []) { + const key = hostKey(left.userId, left.relayHostId) + if (cursor.seats.has(key)) continue + this.rememberLeft(key, { + cellId, + epoch: left.epoch, + generation: 0, + incarnation: response.incarnation, + ...(left.closeCode === undefined ? {} : { closeCode: left.closeCode }), + at: left.at + }) + } } else { seq = cursor.incarnation === response.incarnation @@ -195,8 +255,31 @@ export class ShadowSeatDirectory { cursor.reportedSeats = response.counts?.seats cursor.reportedControls = response.counts?.controls cursor.flagsApplied = response.flagsApplied + cursor.admitModeEffective = response.admitModeEffective + cursor.polledAt = polledAt + cursor.bookings = response.counts?.bookings + cursor.units = response.counts?.units + cursor.ceiling = response.counts?.ceiling + cursor.intake = response.intake } + // The newest booking any cell reported for this host since this director started. + bookingOf(userId: string, relayHostId: string): { cellId: string; epoch: number } | undefined { + return this.bookings.get(hostKey(userId, relayHostId)) + } + + // Applied state only: what the cell last said its switch is, never the desired object. + admitModeOf(cellId: string): 'db' | 'reserve' { + const cursor = this.cells.get(cellId) + if (cursor?.admitModeEffective) return cursor.admitModeEffective + return cursor?.flagsApplied?.flags.admitMode === 'reserve' ? 'reserve' : 'db' + } + + cellIds(): string[] { + return [...this.cells.keys()] + } + + // An old cell (404) has no feed; its last known seats stay, marked by status. markNoFeed(cellId: string, now: number): void { const cursor = this.cells.get(cellId) @@ -205,6 +288,12 @@ export class ShadowSeatDirectory { cursor.lastAnsweredAt = now cursor.incarnation = undefined cursor.seq = undefined + // An image without the feed takes no booking. Its reserve switch is #27058's to keep. + cursor.polledAt = undefined + cursor.bookings = undefined + cursor.units = undefined + cursor.ceiling = undefined + cursor.intake = undefined } fail(cellId: string, reason: string): void { @@ -332,6 +421,11 @@ export class ShadowSeatDirectory { if (change.seq !== seq + 1) return null seq = change.seq const key = hostKey(change.userId, change.relayHostId) + // A booking is not a seat and carries no generation. + if (change.kind === 'reserve') { + this.rememberBooking(key, cursor.cellId, change.epoch, change.at) + continue + } const seat = cursor.seats.get(key) // A superseded generation's leave can arrive after its successor's join. if (seat && change.generation < seat.generation) continue @@ -342,8 +436,20 @@ export class ShadowSeatDirectory { generation: change.generation, state: change.state ?? 'active', joinedAt: change.at, - observedAt: now + observedAt: now, + incarnation: cursor.incarnation }) + if (change.reservedBy !== undefined) { + this.onBookedJoin?.({ + cellId: cursor.cellId, + userId: change.userId, + relayHostId: change.relayHostId, + epoch: change.epoch, + generation: change.generation, + joinedAt: change.at, + reservedBy: change.reservedBy + }) + } } else if (change.kind === 'leave') { if (seat) { this.unseat(cursor, key) @@ -351,6 +457,7 @@ export class ShadowSeatDirectory { cellId: cursor.cellId, epoch: change.epoch, generation: change.generation, + incarnation: cursor.incarnation, closeCode: change.closeCode, at: change.at }) @@ -366,7 +473,12 @@ export class ShadowSeatDirectory { return seq } - private replaceSeats(cursor: CellCursor, full: SeatFeedResponse['full'] & object, now: number) { + private replaceSeats( + cursor: CellCursor, + full: SeatFeedResponse['full'] & object, + now: number, + incarnation: string + ) { const next = new Map() for (const seat of full) { next.set(hostKey(seat.userId, seat.relayHostId), { @@ -375,7 +487,8 @@ export class ShadowSeatDirectory { generation: seat.generation, state: seat.state, joinedAt: seat.joinedAt, - observedAt: now + observedAt: now, + incarnation }) } for (const [key, seat] of cursor.seats) { @@ -385,6 +498,7 @@ export class ShadowSeatDirectory { cellId: cursor.cellId, epoch: seat.epoch, generation: seat.generation, + ...(seat.incarnation === undefined ? {} : { incarnation: seat.incarnation }), resync: true, at: now }) @@ -414,6 +528,17 @@ export class ShadowSeatDirectory { if (seats.size === 0) this.hosts.delete(key) } + private rememberBooking(key: string, cellId: string, epoch: number, at: number): void { + const known = this.bookings.get(key) + if (known && known.epoch > epoch) return + this.bookings.delete(key) + this.bookings.set(key, { cellId, epoch, at }) + if (this.bookings.size > SHADOW_SEAT_RECENTLY_LEFT_MAX_HOSTS) { + const oldest = this.bookings.keys().next().value + if (oldest !== undefined) this.bookings.delete(oldest) + } + } + private rememberLeft(key: string, entry: RecentlyLeftSeat): void { const entries = this.recentlyLeft.get(key) ?? [] this.recentlyLeft.delete(key) @@ -439,10 +564,16 @@ export type ShadowSeatPollerOptions = { now?: () => number pollMs?: number log?: (line: string) => void + // True while this director places on reserve cells (ORCA_RELAY_RESERVE_PLACEMENT=on). + reserver?: () => boolean } export type ShadowSeatPoller = { directory: ShadowSeatDirectory + // The cell list as last read: url, region, liveness and admission. + cells: () => readonly SeatFeedCell[] + // When the poller started, for the startup gate. + startedAt: number // One round: refreshes the cell list when due, then polls every idle cell. tick: () => Promise stop: () => void @@ -470,28 +601,17 @@ export function startShadowSeatPoller( options.identityToken ?? ((tokenAudience: string) => googleMetadataIdentityToken(tokenAudience, fetchImpl)) const directory = new ShadowSeatDirectory() + const startedAt = now() const inFlight = new Set() let cells: SeatFeedCell[] = [] let cellsReadAt: number | undefined let cellListInFlight = false let cellListFailures = 0 let cellListRetryAt = 0 - let token: { value: Promise; at: number } | undefined let lastSummaryAt = now() let stopped = false - const identity = (): Promise => { - const at = now() - if (!token || at - token.at > IDENTITY_TOKEN_REUSE_MS) { - const value = tokenProvider(audience) - token = { value, at } - // A failed fetch must not be reused for ten minutes. - value.catch(() => { - if (token?.value === value) token = undefined - }) - } - return token.value - } + const identity = reusedIdentityToken(() => tokenProvider(audience), now, IDENTITY_TOKEN_REUSE_MS) const refreshCells = async (): Promise => { if (cellListInFlight) return @@ -538,8 +658,11 @@ export function startShadowSeatPoller( // Returns true when the cell cut the page and the next one should follow at once. const fetchPage = async (cell: SeatFeedCell): Promise => { const url = new URL('/v1/admin/cell-seats', cell.cellUrl) + const polledAt = now() const since = directory.since(cell.cellId) if (since !== undefined) url.searchParams.set('since', since) + // Placing on reserve cells: keeps each cell's dead-man from flipping it back to db. + if (options.reserver?.()) url.searchParams.set('reserver', '1') const response = await fetchImpl(url, { headers: { authorization: `Bearer ${await identity()}` }, signal: AbortSignal.timeout(SHADOW_SEAT_POLL_TIMEOUT_MS) @@ -560,7 +683,7 @@ export function startShadowSeatPoller( directory.fail(cell.cellId, 'malformed') return false } - directory.apply(cell.cellId, body.data, now()) + directory.apply(cell.cellId, body.data, now(), polledAt) return body.data.more === true && directory.cellState(cell.cellId)?.status === 'live' } @@ -605,6 +728,8 @@ export function startShadowSeatPoller( timer.unref() return { directory, + cells: () => cells, + startedAt, tick: () => round(true), stop: () => { stopped = true