diff --git a/cloud/apps/relay/src/app.ts b/cloud/apps/relay/src/app.ts index 8dd8deeb4e4..91fe5fcf3f6 100644 --- a/cloud/apps/relay/src/app.ts +++ b/cloud/apps/relay/src/app.ts @@ -39,6 +39,8 @@ 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 { registerCellReserveRoutes } from './cell-reserve-routes.js' +import type { DemoteRequest, ReserveOutcome, ReserveRequest } from './cell-reserve-contract.js' import type { CellSeatFeedPage } from './cell-seat-log.js' import { type CellFlags, supportedCellFlags } from './cell-flags.js' import type { AppliedControlFlags } from './relay-control-flag-channel.js' @@ -126,6 +128,13 @@ export function createRelayApp( cellIncarnation?: string cellSeatFeed?: (sinceSeq: number | null) => CellSeatFeedPage cellFlags?: () => AppliedControlFlags + cellReserve?: (request: ReserveRequest) => ReserveOutcome[] + cellDemote?: (request: DemoteRequest) => string + cellReserveCounts?: () => { + bookings: number + intake: { perSec: number; burst: number; tokens: number } + } | null + cellPlacementCeiling?: number isDraining?: () => boolean regionalRehomeSafetySnapshot?: () => RegionalRehomeSafetySnapshot runtimeCounts?: () => RelayRuntimeCounts @@ -333,7 +342,14 @@ export function createRelayApp( seatFeed: operations.cellSeatFeed, isDraining: operations.isDraining, runtimeCounts: operations.runtimeCounts, - cellFlags: operations.cellFlags + cellFlags: operations.cellFlags, + reserveCounts: operations.cellReserveCounts, + placementCeiling: operations.cellPlacementCeiling + }) + registerCellReserveRoutes(app, config, { + verifyRegionalRehomeToken, + reserve: operations.cellReserve, + demote: operations.cellDemote }) // Not /healthz: Google Front End reserves that path before the container. diff --git a/cloud/apps/relay/src/cell-flags.ts b/cloud/apps/relay/src/cell-flags.ts index f9e01b967f7..6543a826681 100644 --- a/cloud/apps/relay/src/cell-flags.ts +++ b/cloud/apps/relay/src/cell-flags.ts @@ -5,34 +5,51 @@ import type { ControlFlagParse } from './relay-control-flag-channel.js' // keys are dropped; a known key with a bad value voids the object. export type CellFlags = { readinessLocal: boolean - ticketCheck: 'off' | 'shadow' + // `enforce` admits a hello on a valid lease alone; only reserve mode reads it. + ticketCheck: 'off' | 'shadow' | 'enforce' // Off restores Node's default: any unhandled rejection ends the process (#26817 reverted). rejectionFence: boolean // Past statement_timeout before a lost reply ends its connection; absent keeps the boot value. readTimeoutMarginMs?: number + // Step 5: `reserve` admits booked and recently seated hosts without the database. + admitMode: 'db' | 'reserve' + // Bookings per second this cell accepts; absent means its region's default. + intakePerSec?: number + // Answers a director's dry-run booking check; off, the reserve endpoint is inert in db mode. + reserveDryRun: boolean } export const CELL_FLAG_DEFAULTS: CellFlags = { readinessLocal: false, ticketCheck: 'off', - rejectionFence: true + rejectionFence: true, + admitMode: 'db', + reserveDryRun: false } export const READ_TIMEOUT_MARGIN_MIN_MS = 1_000 export const READ_TIMEOUT_MARGIN_MAX_MS = 60_000 +// Until E-pre measures each cell's accept rate (S3 §4.2). +export function defaultIntakePerSec(region: string | undefined): number { + return region === 'asia-east2' ? 2 : 5 +} + export const CELL_FLAGS_APPLIED_EVENT = 'orca_relay_cell_flags_applied' const CellFlagsSchema = z.object({ readinessLocal: z.boolean().default(CELL_FLAG_DEFAULTS.readinessLocal), - ticketCheck: z.enum(['off', 'shadow']).default(CELL_FLAG_DEFAULTS.ticketCheck), + ticketCheck: z.enum(['off', 'shadow', 'enforce']).default(CELL_FLAG_DEFAULTS.ticketCheck), rejectionFence: z.boolean().default(CELL_FLAG_DEFAULTS.rejectionFence), readTimeoutMarginMs: z .number() .int() .min(READ_TIMEOUT_MARGIN_MIN_MS) .max(READ_TIMEOUT_MARGIN_MAX_MS) - .optional() + .optional(), + admitMode: z.enum(['db', 'reserve']).default(CELL_FLAG_DEFAULTS.admitMode), + intakePerSec: z.number().positive().max(1_000).optional(), + reserveDryRun: z.boolean().default(CELL_FLAG_DEFAULTS.reserveDryRun) }) const KNOWN_FLAG_KEYS = new Set(Object.keys(CellFlagsSchema.shape)) @@ -41,6 +58,9 @@ export type SupportedCellFlag = | { type: 'enum'; values: string[] } | { type: 'number'; min?: number; max?: number; integer?: boolean } +// Parsed, never offered: ticketCheck=enforce has no step that turns it on yet. +const UNOFFERED_FLAG_VALUES: Record = { ticketCheck: ['enforce'] } + // What this image accepts, read from the schema itself: the flag tool refuses a key or value // the cell would drop or void before it writes. export function supportedCellFlags(): Record { @@ -51,7 +71,10 @@ export function supportedCellFlags(): Record { inner = inner instanceof z.ZodDefault ? inner.removeDefault() : inner.unwrap() } if (inner instanceof z.ZodBoolean) supported[key] = { type: 'boolean' } - else if (inner instanceof z.ZodEnum) supported[key] = { type: 'enum', values: [...inner.options] } + else if (inner instanceof z.ZodEnum) { + const unoffered = UNOFFERED_FLAG_VALUES[key] ?? [] + supported[key] = { type: 'enum', values: inner.options.filter((value: string) => !unoffered.includes(value)) } + } else if (inner instanceof z.ZodNumber) { supported[key] = { type: 'number', diff --git a/cloud/apps/relay/src/cell-reserve-admission.blackbox.test.ts b/cloud/apps/relay/src/cell-reserve-admission.blackbox.test.ts new file mode 100644 index 00000000000..a952679ff53 --- /dev/null +++ b/cloud/apps/relay/src/cell-reserve-admission.blackbox.test.ts @@ -0,0 +1,444 @@ +import { createHash, createHmac } from 'node:crypto' +import { createServer, type Server } from 'node:http' +import { createServer as createNetServer } from 'node:net' +import { buildHostProofMacInput, HOST_CHALLENGE_PLAINTEXT_DOMAIN } from '@orca-cloud/relay-contract' +import { exportJWK, generateKeyPair, SignJWT } from 'jose' +import nacl from 'tweetnacl' +import { afterEach, describe, expect, it, vi } from 'vitest' +import WebSocket from 'ws' +import type { RawData } from 'ws' +import { CELL_FLAG_DEFAULTS, type CellFlags } from './cell-flags.js' +import type { RelayConfig } from './config.js' +import { PostgresDatabase } from './database.js' +import type { AppliedControlFlags } from './relay-control-flag-channel.js' +import { createRelayServer, HOST_HELLO_SHED_OLDEST_WAIT_MS } from './relay-server.js' + +// A step-5 cell whose database pool is wedged: every query waits forever and the hello shed +// is armed. Booked and recently seated hosts must still connect; everything else must not. + +async function unusedPort(): Promise { + const server = createNetServer() + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)) + const address = server.address() + if (!address || typeof address === 'string') throw new Error('missing test port') + await new Promise((resolve) => server.close(() => resolve())) + return address.port +} + +function stalledPool() { + const pool = { + options: { max: 2, connectionTimeoutMillis: 0 }, + totalCount: 2, + idleCount: 0, + waitingCount: 0, + connect: () => { + pool.waitingCount++ + return new Promise(() => {}) + }, + end: async () => undefined + } + return pool +} + +function cellConfig(port: number, issuer: string): RelayConfig { + const relayUrl = `http://127.0.0.1:${port}` + return { + port, + publicUrl: relayUrl, + cellUrl: relayUrl, + authIssuer: issuer, + authAudience: 'orca-relay', + jwksUrl: issuer, + assignmentSigningKey: new Uint8Array(32), + role: 'cell', + cellId: 'production-gce-c3', + cells: [{ id: 'production-gce-c3', url: relayUrl, capacityRequests: 4_000 }], + adminAudience: `${relayUrl}/admin`, + deployServiceAccount: 'deploy@example.com', + runtimeServiceAccount: 'runtime@example.com', + connectionHardCap: 600, + connectionUnobservedBound: 60, + adminJwksUrl: `${issuer}/admin-jwks`, + databasePoolMax: 2, + publicAssignmentsEnabled: true, + publicAssignmentConcurrency: 1, + publicAssignmentQueueMax: 128, + publicAssignmentWaitMs: 4_000, + publicResolveConcurrency: 1, + publicResolveWaitMs: 5_000, + publicAssignmentRetryAfterSeconds: 5, + dataDir: './test-data' + } +} + +function nextMessage(socket: WebSocket): Promise> { + return new Promise((resolve, reject) => { + socket.once('message', (data: RawData) => { + try { + resolve(JSON.parse(data.toString()) as Record) + } catch (error) { + reject(error) + } + }) + socket.once('error', reject) + }) +} + +describe('a reserve-mode cell admits from memory with its database wedged', () => { + const cleanup: Array<() => Promise | void> = [] + + afterEach(async () => { + for (const close of cleanup.splice(0).reverse()) await close() + vi.restoreAllMocks() + }) + + async function startCell( + flags: Partial, + connectionLedgerLimits?: { hardCap: number; controlReserve: number } + ) { + vi.spyOn(console, 'log').mockImplementation(() => undefined) + vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const keys = await generateKeyPair('ES256') + const publicJwk = await exportJWK(keys.publicKey) + const jwksServer: Server = createServer((_request, response) => { + response.setHeader('content-type', 'application/json') + response.end( + JSON.stringify({ keys: [{ ...publicJwk, kid: 'test-key', alg: 'ES256', use: 'sig' }] }) + ) + }) + await new Promise((resolve) => jwksServer.listen(0, '127.0.0.1', resolve)) + cleanup.push(() => new Promise((resolve) => jwksServer.close(() => resolve()))) + const jwksAddress = jwksServer.address() + if (!jwksAddress || typeof jwksAddress === 'string') throw new Error('missing JWKS address') + const issuer = `http://127.0.0.1:${jwksAddress.port}` + const realNow = Date.now() + let elapsedMs = 0 + vi.spyOn(Date, 'now').mockImplementation(() => realNow + elapsedMs) + const database = new PostgresDatabase(stalledPool() as never) + for (let index = 0; index < 2; index++) void database.query('SELECT 1').catch(() => {}) + elapsedMs += HOST_HELLO_SHED_OLDEST_WAIT_MS + let applied: AppliedControlFlags = { + generation: 3, + flags: { ...CELL_FLAG_DEFAULTS, ...flags } + } + const port = await unusedPort() + const relay = createRelayServer(cellConfig(port, issuer), database, { + cellFlags: () => applied, + ...(connectionLedgerLimits ? { connectionLedgerLimits } : {}) + }) + relay.server.listen(port, '127.0.0.1') + await new Promise((resolve) => relay.server.once('listening', resolve)) + cleanup.push(() => new Promise((resolve) => relay.server.close(() => resolve()))) + + const host = async () => { + const keyPair = nacl.box.keyPair() + const hostId = createHash('sha256').update(keyPair.publicKey).digest('base64url').slice(0, 16) + const token = await new SignJWT({ + prof: 'profile-1', + org: 'org-1', + purpose: 'host-control', + relayHostId: hostId + }) + .setProtectedHeader({ alg: 'ES256', kid: 'test-key' }) + .setIssuer(issuer) + .setAudience('orca-relay') + .setSubject('user-1') + .setIssuedAt() + .setExpirationTime('5m') + .sign(keys.privateKey) + return { keyPair, hostId, token } + } + + // The upgrade status, or the first control frame after the proof. + const connect = async ( + identity: Awaited>, + epoch: number + ): Promise<{ + status: number + ack?: Record + socket?: WebSocket + closeCode?: Promise + }> => { + const socket = new WebSocket(`ws://127.0.0.1:${port}/v1/host/control`, { + headers: { authorization: `Bearer ${identity.token}` }, + perMessageDeflate: false + }) + const opened = await new Promise<{ status: number }>((resolve, reject) => { + socket.once('unexpected-response', (request, response) => { + resolve({ status: response.statusCode ?? 0 }) + request.destroy() + }) + socket.once('open', () => resolve({ status: 101 })) + socket.once('error', reject) + }) + if (opened.status !== 101) return opened + cleanup.push(() => socket.terminate()) + const closeCode = new Promise((resolve) => socket.once('close', resolve)) + socket.send( + JSON.stringify({ + type: 'host-hello', + v: 1, + relayHostId: identity.hostId, + assignmentEpoch: epoch, + hostPublicKeyB64: Buffer.from(identity.keyPair.publicKey).toString('base64'), + appVersion: 'test' + }) + ) + const challenge = await Promise.race([ + nextMessage(socket), + closeCode.then(() => null), + new Promise((resolve) => setTimeout(() => resolve(null), 500)) + ]) + if (!challenge || challenge.type !== 'host-challenge') return { status: 101, socket, closeCode } + const plaintext = nacl.box.open( + Buffer.from(String(challenge.ciphertextB64), 'base64'), + Buffer.from(String(challenge.nonceB64), 'base64'), + Buffer.from(String(challenge.relayEphemeralPublicKeyB64), 'base64'), + identity.keyPair.secretKey + )! + const domainLength = new TextEncoder().encode(`${HOST_CHALLENGE_PLAINTEXT_DOMAIN}\0`).length + const transcriptLength = new DataView( + plaintext.buffer, + plaintext.byteOffset + domainLength, + 4 + ).getUint32(0, false) + const transcriptStart = domainLength + 4 + const transcript = plaintext.slice(transcriptStart, transcriptStart + transcriptLength) + const secret = plaintext.slice(transcriptStart + transcriptLength) + socket.send( + JSON.stringify({ + type: 'host-challenge-ack', + challengeId: challenge.challengeId, + proofB64: createHmac('sha256', secret) + .update(buildHostProofMacInput(transcript)) + .digest('base64') + }) + ) + return { status: 101, ack: await nextMessage(socket), socket } + } + + // An upgrade that never sends its hello, so it keeps its connection unit. + const upgradeOnly = async (identity: Awaited>): Promise => { + const socket = new WebSocket(`ws://127.0.0.1:${port}/v1/host/control`, { + headers: { authorization: `Bearer ${identity.token}` }, + perMessageDeflate: false + }) + cleanup.push(() => socket.terminate()) + return await new Promise((resolve, reject) => { + socket.once('unexpected-response', (request, response) => { + resolve(response.statusCode ?? 0) + request.destroy() + }) + socket.once('open', () => resolve(101)) + socket.once('error', reject) + }) + } + + const book = (identity: Awaited>, epoch: number) => + relay.sessions.reserve({ + v: 1, + directorId: 'director-a', + items: [{ userId: 'user-1', relayHostId: identity.hostId, epoch, ttlMs: 30_000 }] + }) + + return { + relay, + host, + connect, + upgradeOnly, + book, + setFlags: (next: Partial) => { + applied = { generation: applied.generation + 1, flags: { ...CELL_FLAG_DEFAULTS, ...next } } + } + } + } + + it('answers a director dry run only while reserveDryRun is on, and books nothing either way', async () => { + const cell = await startCell({}) + const identity = await cell.host() + const check = () => + cell.relay.sessions.reserve({ + v: 1, + directorId: 'director-a', + dryRun: true, + items: [{ userId: 'user-1', relayHostId: identity.hostId, epoch: 7, ttlMs: 30_000 }] + }) + expect(check()).toEqual([{ outcome: 'off' }]) + cell.setFlags({ reserveDryRun: true }) + expect(check()).toEqual([{ outcome: 'ok' }]) + expect(cell.relay.sessions.reserveCounts()?.bookings).toBe(0) + // A real booking on a database-mode cell stays off, whatever the dry-run switch says. + expect(cell.book(identity, 7)).toEqual([{ outcome: 'off' }]) + }) + + it('connects a booked host, and only at its booked epoch, without a database read', async () => { + const cell = await startCell({ admitMode: 'reserve' }) + const booked = await cell.host() + expect(cell.book(booked, 7)).toEqual([{ outcome: 'ok' }]) + const stranger = await cell.host() + // No booking and no seat: the wedged pool sheds it, as today. + expect((await cell.connect(stranger, 7)).status).toBe(503) + const reply = await cell.connect(booked, 7) + expect(reply.ack).toMatchObject({ type: 'host-hello-ack', generation: 1 }) + const page = cell.relay.sessions.seatFeed(null) + expect('full' in page && page.full).toMatchObject([{ relayHostId: booked.hostId, epoch: 7 }]) + const changes = cell.relay.sessions.seatFeed(0) + expect('changes' in changes && changes.changes.map((change) => [change.kind, change.reservedBy])).toEqual([ + ['reserve', 'director-a'], + ['join', 'director-a'] + ]) + }) + + it('lets a host that just left rejoin at the same epoch from memory', async () => { + const cell = await startCell({ admitMode: 'reserve' }) + const identity = await cell.host() + cell.book(identity, 4) + const first = await cell.connect(identity, 4) + expect(first.ack).toMatchObject({ type: 'host-hello-ack' }) + first.socket!.close() + await new Promise((resolve) => setTimeout(resolve, 50)) + expect((await cell.connect(identity, 4)).ack).toMatchObject({ type: 'host-hello-ack' }) + // A different epoch is not the seat it remembers: shed at the hello, not queued on the pool. + const other = await cell.connect(identity, 5) + expect(other.ack).toBeUndefined() + expect(await other.closeCode).toBe(4429) + }) + + it('keeps today’s path when the switch is off, even for a booking made while it was on', async () => { + const cell = await startCell({ admitMode: 'reserve' }) + const identity = await cell.host() + cell.book(identity, 2) + cell.setFlags({ admitMode: 'db' }) + expect((await cell.connect(identity, 2)).status).toBe(503) + expect(cell.book(identity, 3)).toEqual([{ outcome: 'off' }]) + }) + + it('demotes exactly the seat a director names, and repeats are one demotion', async () => { + const cell = await startCell({ admitMode: 'reserve' }) + const identity = await cell.host() + cell.book(identity, 9) + await cell.connect(identity, 9) + const page = cell.relay.sessions.seatFeed(null) + const joinedAt = 'full' in page ? page.full[0]!.joinedAt : -1 + const seat = { v: 1 as const, userId: 'user-1', relayHostId: identity.hostId, epoch: 9 } + expect(cell.relay.sessions.demote({ ...seat, joinedAt: joinedAt - 1 })).toBe('not-seated') + expect(cell.relay.sessions.demote({ ...seat, joinedAt })).toBe('demoted') + expect(cell.relay.sessions.demote({ ...seat, joinedAt })).toBe('already-demoted') + const changes = cell.relay.sessions.seatFeed(0) + expect('changes' in changes && changes.changes.at(-1)?.kind).toBe('drain-only') + // A demoted seat no longer holds its epoch against a new booking. + expect(cell.book(identity, 9)).toEqual([{ outcome: 'ok' }]) + }) + + it('gives every booking unit back when the cell starts to drain', async () => { + const cell = await startCell({ admitMode: 'reserve' }) + const identity = await cell.host() + cell.book(identity, 2) + expect(cell.relay.sessions.reserveCounts()?.bookings).toBe(1) + cell.relay.sessions.drain(30_000) + expect(cell.relay.sessions.reserveCounts()?.bookings).toBe(0) + expect(cell.book(identity, 3)).toEqual([{ outcome: 'draining' }]) + }) + + it('takes no demotion on a database-mode cell', async () => { + const cell = await startCell({ admitMode: 'reserve' }) + const identity = await cell.host() + cell.book(identity, 9) + await cell.connect(identity, 9) + const page = cell.relay.sessions.seatFeed(null) + const joinedAt = 'full' in page ? page.full[0]!.joinedAt : -1 + cell.setFlags({ admitMode: 'db' }) + const seat = { v: 1 as const, userId: 'user-1', relayHostId: identity.hostId, epoch: 9, joinedAt } + expect(cell.relay.sessions.demote(seat)).toBe('off') + const changes = cell.relay.sessions.seatFeed(0) + expect('changes' in changes && changes.changes.at(-1)?.kind).toBe('join') + }) + + it('ends a demotion with its socket, so the host booked back later is an ordinary seat', async () => { + const cell = await startCell({ admitMode: 'reserve' }) + const identity = await cell.host() + cell.book(identity, 9) + const loser = await cell.connect(identity, 9) + const page = cell.relay.sessions.seatFeed(null) + const joinedAt = 'full' in page ? page.full[0]!.joinedAt : -1 + const seat = { v: 1 as const, userId: 'user-1', relayHostId: identity.hostId, epoch: 9, joinedAt } + expect(cell.relay.sessions.demote(seat)).toBe('demoted') + // The loser goes on its own, well before the 60 s demotion close. + loser.socket!.close() + await new Promise((resolve) => setTimeout(resolve, 50)) + const left = cell.relay.sessions.seatFeed(0) + // Recorded as moved, so neither rule 2 nor a director's sticky sends it back at 9. + expect('changes' in left && left.changes.at(-1)).toMatchObject({ kind: 'leave', closeCode: 4409 }) + expect((await cell.connect(identity, 9)).ack).toBeUndefined() + // Booked back onto this cell at a newer epoch: admitted, and then remembered as seated. + expect(cell.book(identity, 11)).toEqual([{ outcome: 'ok' }]) + const back = await cell.connect(identity, 11) + expect(back.ack).toMatchObject({ type: 'host-hello-ack' }) + back.socket!.close() + await new Promise((resolve) => setTimeout(resolve, 50)) + expect((await cell.connect(identity, 11)).ack).toMatchObject({ type: 'host-hello-ack' }) + }) + + it('refuses a booked upgrade at the hard cap and leaves the feed units and bookings as they were', async () => { + const hardCap = 4 + const cell = await startCell({ admitMode: 'reserve' }, { hardCap, controlReserve: 1 }) + const [waiting, seated] = [await cell.host(), await cell.host()] + expect(cell.book(waiting, 5)[0]?.outcome).toBe('ok') + expect(cell.book(seated, 5)[0]?.outcome).toBe('ok') + expect((await cell.connect(seated, 5)).ack?.type).toBe('host-hello-ack') + // Rebinds over the live control may rise to the hard cap; unproven, they hold their units. + expect(await cell.upgradeOnly(seated)).toBe(101) + expect(await cell.upgradeOnly(seated)).toBe(101) + const before = { + units: cell.relay.connectionSnapshot()!.enforcedConnectionUnits, + bookings: cell.relay.sessions.reserveCounts()!.bookings + } + expect(before).toEqual({ units: hardCap, bookings: 1 }) + expect((await cell.connect(waiting, 5)).status).toBe(503) + expect({ + units: cell.relay.connectionSnapshot()!.enforcedConnectionUnits, + bookings: cell.relay.sessions.reserveCounts()!.bookings + }).toEqual(before) + }, 30_000) + + it('never crosses the hard cap under seeded bookings, hellos, rebooks and closes', async () => { + const hardCap = 6 + const cell = await startCell({ admitMode: 'reserve' }, { hardCap, controlReserve: 1 }) + const hosts = await Promise.all(Array.from({ length: 12 }, () => cell.host())) + const epochs = new Map() + const open: WebSocket[] = [] + let seed = 7 + const random = () => { + seed = (seed * 1_103_515_245 + 12_345) % 2_147_483_648 + return seed / 2_147_483_648 + } + let peak = 0 + for (let step = 0; step < 240; step += 1) { + const identity = hosts[Math.floor(random() * hosts.length)]! + const roll = random() + if (roll < 0.4) { + // A director's fresh booking, or a re-book at a newer epoch. + const epoch = (epochs.get(identity.hostId) ?? 1) + (random() < 0.5 ? 0 : 1) + if (cell.book(identity, epoch)[0]?.outcome === 'ok') epochs.set(identity.hostId, epoch) + } else if (roll < 0.8) { + const reply = await cell.connect(identity, epochs.get(identity.hostId) ?? 1) + if (reply.ack?.type === 'host-hello-ack') open.push(reply.socket!) + } else if (open.length > 0) { + const socket = open.splice(Math.floor(random() * open.length), 1)[0]! + // A newer control for the same host may already have replaced it. + if (socket.readyState !== WebSocket.CLOSED) { + const closed = new Promise((resolve) => socket.once('close', resolve)) + socket.close() + await closed + } + } + const counts = cell.relay.connectionSnapshot()! + peak = Math.max(peak, counts.physicalConnections) + expect(counts.physicalConnections).toBeLessThanOrEqual(hardCap) + expect(counts.enforcedConnectionUnits).toBeLessThanOrEqual(hardCap) + } + // The run did reach the cap, so the bound was exercised. + expect(peak).toBeGreaterThanOrEqual(hardCap - 1) + }, 60_000) +}) + diff --git a/cloud/apps/relay/src/cell-reserve-book.test.ts b/cloud/apps/relay/src/cell-reserve-book.test.ts index c408b51467c..1d5d051936a 100644 --- a/cloud/apps/relay/src/cell-reserve-book.test.ts +++ b/cloud/apps/relay/src/cell-reserve-book.test.ts @@ -110,4 +110,22 @@ describe('cell reserve book', () => { book.clear() expect(capacity.held()).toBe(0) }) + + it('hands a booking unit to the host upgrade, so the host counts once until its hello', () => { + const capacity = units(2) + const book = new CellReserveBook(capacity, () => 100) + book.reserve('d1', item('a', 2), reserveMode) + expect(capacity.held()).toBe(1) + book.handOff('a', 'abcdefghijklmnop') + book.handOff('a', 'abcdefghijklmnop') + expect(capacity.held()).toBe(0) + // Still a booking for the hello, at its own epoch only. + expect(book.has('a', 'abcdefghijklmnop')).toBe(true) + expect(book.take('a', 'abcdefghijklmnop', 3)).toBeNull() + // A newer booking for a handed-off host needs a unit of its own. + expect(book.reserve('d1', item('a', 3), reserveMode).outcome).toBe('ok') + expect(capacity.held()).toBe(1) + expect(book.take('a', 'abcdefghijklmnop', 3)).toMatchObject({ epoch: 3 }) + expect(capacity.held()).toBe(0) + }) }) diff --git a/cloud/apps/relay/src/cell-reserve-book.ts b/cloud/apps/relay/src/cell-reserve-book.ts index e67ed54161a..089f4727f89 100644 --- a/cloud/apps/relay/src/cell-reserve-book.ts +++ b/cloud/apps/relay/src/cell-reserve-book.ts @@ -33,7 +33,8 @@ export type CellBooking = { expiresAt: number } -type HeldBooking = CellBooking & { unit: BookingUnit } +// `handedOff`: the host's upgrade now counts this unit, so the booking holds none. +type HeldBooking = CellBooking & { unit: BookingUnit; handedOff: boolean } function hostKey(userId: string, relayHostId: string): string { return `${userId}\u0000${relayHostId}` @@ -104,11 +105,12 @@ export class CellReserveBook { return { outcome: 'seated-newer', epoch: known } } // A lower-epoch booking for this host hands its unit to the new one. - if (!booked && !this.capacity.canReserve()) return { outcome: 'full' } + const heldUnit = booked && !booked.handedOff ? booked.unit : undefined + if (!heldUnit && !this.capacity.canReserve()) return { outcome: 'full' } // The post-restart sticky re-booking books nothing new, so it takes no token. if (!item.sticky && this.intake.available(now) < 1) return { outcome: 'intake' } if (dryRun) return { outcome: 'ok' } - const unit = booked?.unit ?? this.capacity.tryReserve() + const unit = heldUnit ?? this.capacity.tryReserve() if (!unit) return { outcome: 'full' } if (!item.sticky) this.intake.take(now) this.bookings.set(key, { @@ -117,7 +119,8 @@ export class CellReserveBook { epoch: item.epoch, directorId, expiresAt: now + item.ttlMs, - unit + unit, + handedOff: false }) return { outcome: 'ok' } } @@ -127,6 +130,15 @@ export class CellReserveBook { return booking !== undefined && booking.expiresAt > this.now() } + // The booked host's upgrade takes over its unit, so the host counts once until its hello. + // The upgrade carries no epoch; the hello's take() still requires the booked one. + handOff(userId: string, relayHostId: string): void { + const booking = this.bookings.get(hostKey(userId, relayHostId)) + if (!booking || booking.handedOff || booking.expiresAt <= this.now()) return + booking.handedOff = true + booking.unit.release() + } + // The hello's rule 1: a booking for this host at exactly this epoch admits it once. take(userId: string, relayHostId: string, epoch: number): CellBooking | null { const key = hostKey(userId, relayHostId) diff --git a/cloud/apps/relay/src/cell-reserve-routes.test.ts b/cloud/apps/relay/src/cell-reserve-routes.test.ts new file mode 100644 index 00000000000..32390953df1 --- /dev/null +++ b/cloud/apps/relay/src/cell-reserve-routes.test.ts @@ -0,0 +1,62 @@ +import { Hono } from 'hono' +import { describe, expect, it, vi } from 'vitest' +import { registerCellReserveRoutes } from './cell-reserve-routes.js' +import type { RelayConfig } from './config.js' + +const cell = { + role: 'cell', + cellId: 'production-gce-c3', + rehomeAudience: 'https://relay.example/rehome', + rehomeDirectorServiceAccount: 'director@example.com' +} as RelayConfig + +function app(config: RelayConfig = cell) { + const reserve = vi.fn(() => [{ outcome: 'ok' as const }]) + const demote = vi.fn(() => 'demoted') + const routes = new Hono() + registerCellReserveRoutes(routes, config, { + verifyRegionalRehomeToken: async (token) => token === 'director-token', + reserve, + demote + }) + const post = (path: string, body: unknown, token = 'director-token') => + routes.request(path, { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: JSON.stringify(body) + }) + return { post, reserve, demote } +} + +const booking = { + v: 1, + directorId: 'director-a', + items: [{ userId: 'user-1', relayHostId: 'abcdefghijklmnop', epoch: 3, ttlMs: 30_000 }] +} + +describe('cell reserve routes', () => { + it('books for a director holding the rehome credential and refuses everyone else', async () => { + const { post, reserve } = app() + const ok = await post('/internal/reserve', booking) + expect(await ok.json()).toEqual({ v: 1, results: [{ outcome: 'ok' }] }) + expect((await post('/internal/reserve', booking, 'someone-else')).status).toBe(401) + expect((await post('/internal/reserve', { ...booking, items: [] })).status).toBe(400) + expect((await post('/internal/reserve', { ...booking, v: 2 })).status).toBe(400) + expect(reserve).toHaveBeenCalledTimes(1) + }) + + it('is absent on a director and on a cell without the rehome pair', async () => { + expect((await app({ ...cell, role: 'director' }).post('/internal/reserve', booking)).status).toBe(404) + const noPair = app({ ...cell, rehomeDirectorServiceAccount: undefined }) + expect((await noPair.post('/internal/demote', {})).status).toBe(404) + }) + + it('demotes only a fully named seat', async () => { + const { post, demote } = app() + const seat = { v: 1, userId: 'user-1', relayHostId: 'abcdefghijklmnop', epoch: 3 } + expect((await post('/internal/demote', seat)).status).toBe(400) + const reply = await post('/internal/demote', { ...seat, joinedAt: 1_000 }) + expect(await reply.json()).toEqual({ v: 1, outcome: 'demoted' }) + expect(demote).toHaveBeenCalledWith({ ...seat, joinedAt: 1_000 }) + }) +}) diff --git a/cloud/apps/relay/src/cell-reserve-routes.ts b/cloud/apps/relay/src/cell-reserve-routes.ts new file mode 100644 index 00000000000..46e103c1a95 --- /dev/null +++ b/cloud/apps/relay/src/cell-reserve-routes.ts @@ -0,0 +1,57 @@ +import type { Hono } from 'hono' +import { + CELL_DEMOTE_PATH, + CELL_RESERVE_PATH, + DemoteRequestSchema, + ReserveRequestSchema, + type DemoteRequest, + type ReserveOutcome, + type ReserveRequest +} from './cell-reserve-contract.js' +import type { RelayConfig } from './config.js' +import { readBearer } from './relay-token-verifier.js' + +const MAX_BODY_BYTES = 16 * 1024 + +// The directors' rehome credential, as the seat feed uses. Unlike the feed these write, but a +// director can already place any host through the database, so they grant nothing new. +export function registerCellReserveRoutes( + app: Hono, + config: RelayConfig, + input: { + verifyRegionalRehomeToken: (token: string) => Promise + reserve?: (request: ReserveRequest) => ReserveOutcome[] + demote?: (request: DemoteRequest) => string + } +): void { + const authorize = async (header: string | undefined): Promise<'ok' | 404 | 401> => { + if (config.role !== 'cell' || !input.reserve || !input.demote) return 404 + if (!config.rehomeAudience || !config.rehomeDirectorServiceAccount) return 404 + const bearer = readBearer(header) + return bearer && (await input.verifyRegionalRehomeToken(bearer)) ? 'ok' : 401 + } + const unauthorized = (status: 404 | 401) => + status === 404 ? { error: 'reserve_unavailable' } : { error: 'invalid_token' } + + app.post(CELL_RESERVE_PATH, async (context) => { + const auth = await authorize(context.req.header('authorization')) + if (auth !== 'ok') return context.json(unauthorized(auth), auth) + if (Number(context.req.header('content-length') ?? 0) > MAX_BODY_BYTES) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = ReserveRequestSchema.safeParse(await context.req.json().catch(() => null)) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + return context.json({ v: 1, results: input.reserve!(body.data) }) + }) + + app.post(CELL_DEMOTE_PATH, async (context) => { + const auth = await authorize(context.req.header('authorization')) + if (auth !== 'ok') return context.json(unauthorized(auth), auth) + if (Number(context.req.header('content-length') ?? 0) > MAX_BODY_BYTES) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = DemoteRequestSchema.safeParse(await context.req.json().catch(() => null)) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + return context.json({ v: 1, outcome: input.demote!(body.data) }) + }) +} diff --git a/cloud/apps/relay/src/cell-seat-feed-route.ts b/cloud/apps/relay/src/cell-seat-feed-route.ts index 178b88de369..3385353118c 100644 --- a/cloud/apps/relay/src/cell-seat-feed-route.ts +++ b/cloud/apps/relay/src/cell-seat-feed-route.ts @@ -19,6 +19,12 @@ export function registerCellSeatFeedRoute( isDraining?: () => boolean runtimeCounts?: () => RelayRuntimeCounts cellFlags?: () => AppliedControlFlags + // Step 5: what a director needs to estimate a seat and a booking here. + reserveCounts?: () => { + bookings: number + intake: { perSec: number; burst: number; tokens: number } + } | null + placementCeiling?: number now?: () => number } ): void { @@ -39,6 +45,7 @@ export function registerCellSeatFeedRoute( cursor !== null && cursor.incarnation === input.cellIncarnation ? cursor.seq : null const { seats, ...page } = input.seatFeed(sinceSeq) const counts = input.runtimeCounts?.() + const reserve = input.reserveCounts?.() ?? null return context.json({ v: 1, cellId: config.cellId, @@ -49,8 +56,17 @@ export function registerCellSeatFeedRoute( controls: counts?.controls ?? 0, // The feed's own view: a reader that applied every change holds exactly this many. // `controls` drops when a socket starts closing, `seats` when its close lands. - seats + seats, + ...(reserve && input.placementCeiling !== undefined + ? { + bookings: reserve.bookings, + // Every unit the ledger counts against the ceiling, bookings and phones included. + units: counts?.enforcedConnectionUnits ?? 0, + ceiling: input.placementCeiling + } + : {}) }, + ...(reserve ? { intake: reserve.intake } : {}), // Applied, never desired: generation 0 means no object has been read since boot. ...(input.cellFlags ? { flagsApplied: input.cellFlags() } : {}), ...page diff --git a/cloud/apps/relay/src/cell-seat-feed.test.ts b/cloud/apps/relay/src/cell-seat-feed.test.ts index 9bde68090a3..ecfbbd0c48c 100644 --- a/cloud/apps/relay/src/cell-seat-feed.test.ts +++ b/cloud/apps/relay/src/cell-seat-feed.test.ts @@ -69,8 +69,10 @@ const SeatFeedReplySchema = z.object({ generation: z.number(), flags: z.object({ readinessLocal: z.boolean(), - ticketCheck: z.enum(['off', 'shadow']), - rejectionFence: z.boolean() + ticketCheck: z.enum(['off', 'shadow', 'enforce']), + rejectionFence: z.boolean(), + admitMode: z.enum(['db', 'reserve']), + reserveDryRun: z.boolean() }), ignoredKeys: z.array(z.string()).optional() }) @@ -352,7 +354,10 @@ describe('cell seat feed', () => { readinessLocal: { type: 'boolean' }, ticketCheck: { type: 'enum', values: ['off', 'shadow'] }, rejectionFence: { type: 'boolean' }, - readTimeoutMarginMs: { type: 'number', min: 1_000, max: 60_000, integer: true } + readTimeoutMarginMs: { type: 'number', min: 1_000, max: 60_000, integer: true }, + admitMode: { type: 'enum', values: ['db', 'reserve'] }, + intakePerSec: { type: 'number', max: 1_000 }, + reserveDryRun: { type: 'boolean' } } }) }) diff --git a/cloud/apps/relay/src/cell-seat-log.test.ts b/cloud/apps/relay/src/cell-seat-log.test.ts index fcf8505c04b..d3854c8c8ab 100644 --- a/cloud/apps/relay/src/cell-seat-log.test.ts +++ b/cloud/apps/relay/src/cell-seat-log.test.ts @@ -104,4 +104,38 @@ describe('cell seat cursor', () => { expect(parseCellSeatCursor('a:-1')).toBe('invalid') expect(parseCellSeatCursor('a b:1')).toBe('invalid') }) + + it('never seats a booking, and remembers leavers for ten minutes for the full snapshot', () => { + let now = 0 + const log = new CellSeatLog(64, () => now) + const host = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + log.append({ kind: 'reserve', ...host, epoch: 4, generation: 0, reservedBy: 'd1', at: now }) + expect(log.seatCount()).toBe(0) + log.append({ kind: 'join', ...host, epoch: 4, generation: 1, reservedBy: 'd1', at: now }) + expect(log.seatOf(host.userId, host.relayHostId)).toMatchObject({ epoch: 4, state: 'active' }) + now = 1_000 + log.append({ kind: 'leave', ...host, epoch: 4, generation: 1, closeCode: 1006, at: now }) + expect(log.recentlyLeftOf(host.userId, host.relayHostId)).toEqual({ ...host, epoch: 4, closeCode: 1006, at: 1_000 }) + const page = log.read(null) + expect('full' in page && page.recentlyLeft).toEqual([{ ...host, epoch: 4, closeCode: 1006, at: 1_000 }]) + now = 1_000 + 10 * 60_000 + expect(log.recentlyLeftOf(host.userId, host.relayHostId)).toBeUndefined() + }) + + it('sends only the newest page of leavers in a full snapshot', () => { + const log = new CellSeatLog(CELL_SEAT_FEED_PAGE_MAX * 4, () => 0) + const leavers = CELL_SEAT_FEED_PAGE_MAX + 5 + for (let index = 0; index < leavers; index += 1) { + const host = { userId: `user-${index}`, relayHostId: 'abcdefghijklmnop', epoch: 1, generation: 1 } + log.append({ kind: 'join', ...host, at: 0 }) + log.append({ kind: 'leave', ...host, closeCode: 1006, at: 0 }) + } + const page = log.read(null) + const left = 'full' in page ? page.recentlyLeft : [] + expect(left).toHaveLength(CELL_SEAT_FEED_PAGE_MAX) + expect(left.at(-1)?.userId).toBe(`user-${leavers - 1}`) + // The cell itself still remembers the oldest. + expect(log.recentlyLeftOf('user-0', 'abcdefghijklmnop')).toBeDefined() + }) }) + diff --git a/cloud/apps/relay/src/cell-seat-log.ts b/cloud/apps/relay/src/cell-seat-log.ts index fd93b78e080..6524976911c 100644 --- a/cloud/apps/relay/src/cell-seat-log.ts +++ b/cloud/apps/relay/src/cell-seat-log.ts @@ -6,13 +6,15 @@ // Apply a change only when its generation is >= the seat's: a superseded generation's // socket can close after its successor joined. `full` is this log's own seats under the // same rules, so a reader that applies every change holds exactly `counts.seats`. +// Step 5 adds `reserve` (a booking, never a seat) and `reservedBy` on the join it admitted; +// a full snapshot also carries `recentlyLeft`, this cell's leavers of the last 10 minutes. export type CellSeatState = 'active' | 'drain-only' export type CellSeatChange = { seq: number // `active` lifts a drain-only the host cleared itself (a refreshed token after it expired). - kind: 'join' | 'leave' | 'drain-only' | 'active' + kind: 'join' | 'leave' | 'drain-only' | 'active' | 'reserve' userId: string relayHostId: string epoch: number @@ -21,6 +23,8 @@ export type CellSeatChange = { state?: CellSeatState // Only on leave. closeCode?: number + // On reserve, and on the join a booking admitted: the director that booked it. + reservedBy?: string at: number } @@ -33,9 +37,17 @@ export type CellSeat = { joinedAt: number } +export type RecentlyLeftCellSeat = { + userId: string + relayHostId: string + epoch: number + closeCode?: number + at: number +} + export type CellSeatPage = | { seq: number; changes: CellSeatChange[]; more: boolean } - | { seq: number; full: CellSeat[] } + | { seq: number; full: CellSeat[]; recentlyLeft: RecentlyLeftCellSeat[] } // What the registry hands the route: a page plus the log's own seat count. export type CellSeatFeedPage = CellSeatPage & { seats: number } @@ -43,13 +55,21 @@ export type CellSeatFeedPage = CellSeatPage & { seats: number } // ~10 minutes of a full-cell drain's joins and leaves (inferred), at ~130 B each. export const CELL_SEAT_LOG_CAPACITY = 20_000 export const CELL_SEAT_FEED_PAGE_MAX = 2_000 +// How long a leaver may rejoin here at the same epoch without the database (step 5 rule 2). +export const CELL_RECENT_SEAT_MS = 10 * 60_000 +const CELL_RECENT_SEAT_MAX = 20_000 export class CellSeatLog { private readonly ring: (CellSeatChange | undefined)[] private readonly seats = new Map() + // Insertion order is age order. + private readonly recent = new Map() private headSeq = 0 - constructor(private readonly capacity = CELL_SEAT_LOG_CAPACITY) { + constructor( + private readonly capacity = CELL_SEAT_LOG_CAPACITY, + private readonly now: () => number = Date.now + ) { this.ring = new Array(capacity) } @@ -67,7 +87,11 @@ export class CellSeatLog { read(sinceSeq: number | null): CellSeatPage { const oldestSeq = Math.max(1, this.headSeq - this.capacity + 1) if (sinceSeq === null || sinceSeq > this.headSeq || sinceSeq < oldestSeq - 1) { - return { seq: this.headSeq, full: [...this.seats.values()].map((seat) => ({ ...seat })) } + return { + seq: this.headSeq, + full: [...this.seats.values()].map((seat) => ({ ...seat })), + recentlyLeft: this.recentlyLeft() + } } const lastSeq = Math.min(this.headSeq, sinceSeq + CELL_SEAT_FEED_PAGE_MAX) const changes: CellSeatChange[] = [] @@ -77,11 +101,30 @@ export class CellSeatLog { return { seq: lastSeq, changes, more: lastSeq < this.headSeq } } + seatOf(userId: string, relayHostId: string): CellSeat | undefined { + return this.seats.get(JSON.stringify([userId, relayHostId])) + } + + recentlyLeftOf(userId: string, relayHostId: string): RecentlyLeftCellSeat | undefined { + const entry = this.recent.get(JSON.stringify([userId, relayHostId])) + return entry && entry.at > this.now() - CELL_RECENT_SEAT_MS ? entry : undefined + } + + // The newest page's worth: an older leaver the director misses only takes the database path. + private recentlyLeft(): RecentlyLeftCellSeat[] { + const cutoff = this.now() - CELL_RECENT_SEAT_MS + const live = [...this.recent.values()].filter((entry) => entry.at > cutoff) + return live.slice(-CELL_SEAT_FEED_PAGE_MAX).map((entry) => ({ ...entry })) + } + private applyToSeats(change: Omit): void { + // A booking is not a seat. + if (change.kind === 'reserve') return const key = JSON.stringify([change.userId, change.relayHostId]) const seat = this.seats.get(key) if (seat && change.generation < seat.generation) return if (change.kind === 'join') { + this.recent.delete(key) this.seats.set(key, { userId: change.userId, relayHostId: change.relayHostId, @@ -91,11 +134,29 @@ export class CellSeatLog { joinedAt: change.at }) } else if (change.kind === 'leave') { + if (!seat) return this.seats.delete(key) + this.rememberLeft(key, { + userId: change.userId, + relayHostId: change.relayHostId, + epoch: seat.epoch, + ...(change.closeCode === undefined ? {} : { closeCode: change.closeCode }), + at: change.at + }) } else if (seat) { seat.state = change.kind === 'active' ? 'active' : 'drain-only' } } + + private rememberLeft(key: string, entry: RecentlyLeftCellSeat): void { + this.recent.delete(key) + this.recent.set(key, entry) + const cutoff = entry.at - CELL_RECENT_SEAT_MS + for (const [oldestKey, oldest] of this.recent) { + if (this.recent.size <= CELL_RECENT_SEAT_MAX && oldest.at > cutoff) break + this.recent.delete(oldestKey) + } + } } // Malformed cursors are a caller bug, so they are refused rather than resynced. diff --git a/cloud/apps/relay/src/host-session-registry.ts b/cloud/apps/relay/src/host-session-registry.ts index 7de30853e84..5410addee44 100644 --- a/cloud/apps/relay/src/host-session-registry.ts +++ b/cloud/apps/relay/src/host-session-registry.ts @@ -28,8 +28,11 @@ import nacl from 'tweetnacl' import type WebSocket from 'ws' import type { RawData } from 'ws' import type { RelayConfig } from './config.js' +import type { AssignmentLeaseClass } from './assignment-lease.js' import type { AssignmentLeaseShadow } from './assignment-lease-shadow.js' import type { RelayAssignmentStore } from './assignment-store.js' +import type { CellAdmitMode, CellReserveBook } from './cell-reserve-book.js' +import type { DemoteRequest, ReserveOutcome, ReserveRequest } from './cell-reserve-contract.js' import { CellSeatLog, type CellSeatChange, type CellSeatFeedPage } from './cell-seat-log.js' import { ControlRenewalBatch } from './control-renewal-batch.js' import { RelayCredentialStore, type CredentialReservation } from './credential-store.js' @@ -205,6 +208,30 @@ const ACTIVATION_QUEUE_WAIT_MS = 30_000 export const CONTROL_LEASE_MS = 6 * 60 * 60 * 1000 export const CONTROL_LEASE_JITTER_MS = 30 * 60 * 1000 +// Step 5 on a cell: what the per-cell switch says, and the booking book it admits from. +export type CellReserveAdmission = { + mode: () => CellAdmitMode + ticketEnforce: () => boolean + dryRunEnabled: () => boolean + // The pool-pressure shed's condition: a hello memory cannot admit is shed, not queued. + databaseShedding: () => boolean + book: CellReserveBook + verifyLease: (input: { + lease: string | undefined + userId: string + relayHostId: string + helloEpoch: number + }) => Promise +} + +// A leaver that closed for these reasons was moved off this cell; it may not rejoin from memory. +const NO_REJOIN_CLOSE_CODES: ReadonlySet = new Set([ + RELAY_CLOSE_CODE.DRAINING, + RELAY_CLOSE_CODE.WRONG_CELL +]) +// A demoted duplicate keeps its splices this long, then closes so the desktop re-assigns. +export const DEMOTION_CLOSE_MS = 60_000 + export class HostSessionRegistry { private readonly sessions = new Map() private readonly activationQueues = new Map>() @@ -218,7 +245,10 @@ export class HostSessionRegistry { // Hosts whose drain has been sent. Paced sends land minutes apart, so "this cell is // draining" is not the same question as "this host has been told to leave". private readonly drainSentHosts = new Set() - private readonly seatLog = new CellSeatLog() + private readonly seatLog: CellSeatLog + // Seat key -> the demoted control socket. Its own set, not drainSentHosts: it ends with + // that socket (close, rebind or a new join), so a later booking here is a fresh seat. + private readonly demotedSeats = new Map() private readonly idleWork = new Map() private readonly idleAttempts = new Map< @@ -345,8 +375,115 @@ export class HostSessionRegistry { private readonly now: () => number = Date.now, private readonly random: () => number = Math.random, private readonly cellIncarnation?: string, - private readonly assignmentLeaseShadow?: AssignmentLeaseShadow - ) {} + private readonly assignmentLeaseShadow?: AssignmentLeaseShadow, + private readonly reserveAdmission?: CellReserveAdmission + ) { + this.seatLog = new CellSeatLog(undefined, now) + } + + // Reserve mode is only ever a cell's, and only while its applied switch says so. + private reserveMode(): boolean { + return this.config.role === 'cell' && this.reserveAdmission?.mode() === 'reserve' + } + + inReserveMode(): boolean { + return this.reserveMode() + } + + reserveCounts(): { bookings: number; intake: { perSec: number; burst: number; tokens: number } } | null { + const book = this.reserveAdmission?.book + return book ? { bookings: book.count(), intake: book.intakeSnapshot() } : null + } + + reserve(request: ReserveRequest): ReserveOutcome[] { + const admission = this.reserveAdmission + if (this.config.role !== 'cell' || !admission || (request.dryRun && !admission.dryRunEnabled())) { + return request.items.map(() => ({ outcome: 'off' })) + } + const context = { + mode: admission.mode(), + draining: this.draining, + seatedEpoch: (userId: string, relayHostId: string) => { + const seat = this.seatLog.seatOf(userId, relayHostId) + return seat && !this.demotedSeats.has(this.key(userId, relayHostId)) ? seat.epoch : undefined + } + } + return request.items.map((item) => { + const outcome = admission.book.reserve(request.directorId, item, context, request.dryRun) + if (outcome.outcome === 'ok' && !request.dryRun) { + this.seatLog.append({ + kind: 'reserve', + userId: item.userId, + relayHostId: item.relayHostId, + epoch: item.epoch, + generation: 0, + reservedBy: request.directorId, + at: this.now() + }) + } + return outcome + }) + } + + // The upgrade's view: a booked host claims its unit up to the hard cap, and neither a booked + // nor a recently seated host needs the database, so the pool-pressure shed skips them. + reserveAdmissionFor(identity: RelayIdentityKey): 'booked' | 'seated' | null { + if (!this.reserveMode()) return null + if (this.reserveAdmission!.book.has(identity.userId, identity.relayHostId)) return 'booked' + const seat = this.seatLog.seatOf(identity.userId, identity.relayHostId) + const left = this.seatLog.recentlyLeftOf(identity.userId, identity.relayHostId) + return seat || (left && !NO_REJOIN_CLOSE_CODES.has(left.closeCode ?? 0)) ? 'seated' : null + } + + // Rule 2: seated here at this epoch, or left here at it within 10 minutes for a reason that + // did not move the host away. + private seatedHereAt(userId: string, relayHostId: string, epoch: number): boolean { + if (this.drainSentHosts.has(relayHostId)) return false + if (this.demotedSeats.has(this.key(userId, relayHostId))) return false + const seat = this.seatLog.seatOf(userId, relayHostId) + if (seat) return seat.epoch === epoch + const left = this.seatLog.recentlyLeftOf(userId, relayHostId) + return left !== undefined && left.epoch === epoch && !NO_REJOIN_CLOSE_CODES.has(left.closeCode ?? 0) + } + + // A director names this seat the loser of a duplicate: stop taking phones for it now, and + // close it once its splices have had time to finish. Repeats are one demotion. + demote(request: DemoteRequest): 'demoted' | 'already-demoted' | 'not-seated' | 'off' { + // Only a reserve-mode cell takes demotions, so flipping admitMode back reverts the route. + if (!this.reserveMode()) return 'off' + const session = this.get(request) + const key = this.key(request.userId, request.relayHostId) + // The seat the director saw: the same host rejoining here later is a different seat. + const seat = this.seatLog.seatOf(request.userId, request.relayHostId) + if ( + !session || + session.state === 'closed' || + !session.socket || + session.assignmentEpoch !== request.epoch || + seat?.joinedAt !== request.joinedAt + ) { + return 'not-seated' + } + const socket = session.socket + if (this.demotedSeats.get(key) === socket) return 'already-demoted' + this.demotedSeats.set(key, socket) + this.markSeatDrainOnly(session) + const timer = setTimeout(() => { + if (this.demotedSeats.get(key) === socket) { + socket.close(RELAY_CLOSE_CODE.WRONG_CELL, 'seated on a newer cell') + } + }, DEMOTION_CLOSE_MS) + timer.unref?.() + console.warn( + JSON.stringify({ + event: 'orca_relay_cell_seat_demoted', + ...this.logIdentity(), + relayHostIdDigest: relayHostLogDigest(session.relayHostId), + assignmentEpoch: request.epoch + }) + ) + return 'demoted' + } // Renewals leave the heartbeat as an enqueue: one statement per cell per // window replaces one write transaction per host, which is what keeps the @@ -456,6 +593,12 @@ export class HostSessionRegistry { markStage('credential') const sessionKey = this.key(reservation.userId, hostId) const session = this.sessions.get(sessionKey) + if (session?.socket && this.demotedSeats.get(sessionKey) === session.socket) { + capacityReservation?.release() + await this.store.failReservation(reservation) + this.rejectClient(socket, RELAY_CLOSE_CODE.DRAINING) + return + } if ( !session || session.state !== 'active' || @@ -906,6 +1049,8 @@ export class HostSessionRegistry { drain(graceMs: number, options: { paceWindowMs?: number } = {}): void { this.draining = true + // A draining cell admits no booking, so the units they hold go back now. + this.reserveAdmission?.book.clear() // A later drain (an emergency one, or shutdown) owns every session again, so nothing // queued by an earlier paced drain may still fire: it would re-send and, worse, keep // the event loop alive for the rest of a window the operator just cut short. @@ -1016,19 +1161,50 @@ export class HostSessionRegistry { relayHostId: identity.relayHostId, helloEpoch: hello.data.assignmentEpoch }) ?? null + // Reserve mode admits, in order, a booking, a seat this cell remembers, and (dark) a valid + // lease; anything else falls to today's database check. The mode is read once per hello. + const reserveMode = this.reserveMode() + let reservedBy: string | undefined + let admittedFromMemory = false + if (reserveMode) { + const booking = this.reserveAdmission!.book.take( + identity.sub, + identity.relayHostId, + hello.data.assignmentEpoch + ) + reservedBy = booking?.directorId + admittedFromMemory = + booking !== null || + this.seatedHereAt(identity.sub, identity.relayHostId, hello.data.assignmentEpoch) || + (this.reserveAdmission!.ticketEnforce() && + (await this.reserveAdmission!.verifyLease({ + lease: assignmentLease, + userId: identity.sub, + relayHostId: identity.relayHostId, + helloEpoch: hello.data.assignmentEpoch + })) === 'valid') + } + if (reserveMode && !admittedFromMemory && this.reserveAdmission!.databaseShedding()) { + // Memory cannot admit it (e.g. seated here, but at another epoch), and the wedged pool + // would only time out its database check: shed it now, as the upgrade shed does. + this.observer.recordAuth(false) + socket.close(RELAY_CLOSE_CODE.LIMIT_EXCEEDED, 'relay temporarily unavailable') + return + } const assignmentReadStartedAt = performance.now() // Combined is staging-only compatibility; stamped cells require the durable director epoch. const assignmentValid = - this.config.role === 'combined' + admittedFromMemory || + (this.config.role === 'combined' ? hello.data.assignmentEpoch === 1 : await this.assignments.verifyCellAssignment({ userId: identity.sub, relayHostId: identity.relayHostId, cellId: this.config.cellId, assignmentEpoch: hello.data.assignmentEpoch - }) + })) // The database answer decides; the lease is only compared with it. - if (leaseShadow && leaseCheck) { + if (leaseShadow && leaseCheck && !admittedFromMemory) { leaseShadow.record( leaseCheck, assignmentValid, @@ -1112,7 +1288,8 @@ export class HostSessionRegistry { rebind, hello.data.assignmentEpoch, hello.data.appVersion, - connectionInclusionWatermark + connectionInclusionWatermark, + admittedFromMemory ? { reservedBy } : undefined ), socket, 'host activation' @@ -1128,7 +1305,9 @@ export class HostSessionRegistry { rebind: boolean, assignmentEpoch: number, appVersion: string, - connectionInclusionWatermark?: number + connectionInclusionWatermark?: number, + // Present when the hello was admitted from memory: no database registration. + reserved?: { reservedBy?: string } ): Promise { const release = this.beginIdleWork(identity.relayHostId) if (!release) { @@ -1162,7 +1341,8 @@ export class HostSessionRegistry { rebind, assignmentEpoch, appVersion, - connectionInclusionWatermark + connectionInclusionWatermark, + reserved ) }) this.activationQueues.set(key, activation) @@ -1182,10 +1362,11 @@ export class HostSessionRegistry { rebind: boolean, assignmentEpoch: number, appVersion: string, - connectionInclusionWatermark?: number + connectionInclusionWatermark?: number, + reserved?: { reservedBy?: string } ): Promise { let controlActivityId: string | null = null - if (this.config.role === 'cell') { + if (this.config.role === 'cell' && !reserved) { try { controlActivityId = await this.assignments.activateControl( { userId: identity.sub, relayHostId: identity.relayHostId }, @@ -1245,7 +1426,8 @@ export class HostSessionRegistry { existing.activityRenewalDueAt = this.now() + RELAY_PROTOCOL_LIMITS.controlPingIntervalMs this.wireActiveControl(existing) this.sendHelloAck(existing) - this.appendSeatChange(existing, 'join') + this.demotedSeats.delete(this.key(identity.sub, identity.relayHostId)) + this.appendSeatChange(existing, 'join', undefined, reserved?.reservedBy) if (existing.regionalDrainAttemptId) this.reassertRegionalDrain(existing) previousSocket?.close(RELAY_CLOSE_CODE.PEER_DROPPED, 'control rebound') return @@ -1315,9 +1497,10 @@ export class HostSessionRegistry { // A host that proved itself again is not signed out, whatever it said last. this.hostCloseReasons.forget(sessionKey) this.sessions.set(sessionKey, session) + this.demotedSeats.delete(sessionKey) this.wireActiveControl(session) this.sendHelloAck(session) - this.appendSeatChange(session, 'join') + this.appendSeatChange(session, 'join', undefined, reserved?.reservedBy) } seatFeed(sinceSeq: number | null): CellSeatFeedPage { @@ -1327,7 +1510,8 @@ export class HostSessionRegistry { private appendSeatChange( session: HostSession, kind: CellSeatChange['kind'], - closeCode?: number + closeCode?: number, + reservedBy?: string ): void { this.seatLog.append({ kind, @@ -1339,6 +1523,7 @@ export class HostSessionRegistry { ? { state: session.state } : {}), ...(closeCode === undefined ? {} : { closeCode }), + ...(reservedBy === undefined ? {} : { reservedBy }), at: this.now() }) } @@ -1376,8 +1561,12 @@ export class HostSessionRegistry { // Guarded on identity: a predecessor retired by a rebind must not stamp a // cause onto the live session that replaced it. if (session.socket === socket) { - this.hostCloseReasons.record(this.key(session.identity.sub, session.relayHostId), reason) - this.appendSeatChange(session, 'leave', code) + const seatKey = this.key(session.identity.sub, session.relayHostId) + this.hostCloseReasons.record(seatKey, reason) + // However it closed, a demoted seat was moved: it may not rejoin here from memory. + const demoted = this.demotedSeats.get(seatKey) === socket + if (demoted) this.demotedSeats.delete(seatKey) + this.appendSeatChange(session, 'leave', demoted ? RELAY_CLOSE_CODE.WRONG_CELL : code) } // One line per control close makes reconnect churners attributable by // host digest without exposing the raw relay host id. diff --git a/cloud/apps/relay/src/relay-connection-ledger.test.ts b/cloud/apps/relay/src/relay-connection-ledger.test.ts index f807c34d20d..d04a80bf665 100644 --- a/cloud/apps/relay/src/relay-connection-ledger.test.ts +++ b/cloud/apps/relay/src/relay-connection-ledger.test.ts @@ -169,4 +169,22 @@ describe('RelayConnectionLedger', () => { enforcedConnectionUnits: 0 }) }) + + it('books units only below the placement ceiling and releases each once', () => { + const ledger = new RelayConnectionLedger(6, 2) + const first = ledger.tryReserveBooking(2)! + expect(ledger.tryReserveBooking(2)).not.toBeNull() + expect(ledger.tryReserveBooking(2)).toBeNull() + // A host control may still rise to the socket ceiling past the bookings. + expect(ledger.tryReserveControl(false)).not.toBeNull() + first.release() + first.release() + expect(ledger.counts().reservedConnectionUnits).toBe(1) + // Never above the socket ceiling, whatever the caller's ceiling says. + expect(ledger.canReserveBooking(100)).toBe(true) + ledger.tryReserveControl(false) + ledger.tryReserveControl(false) + expect(ledger.canReserveBooking(100)).toBe(false) + }) }) + diff --git a/cloud/apps/relay/src/relay-connection-ledger.ts b/cloud/apps/relay/src/relay-connection-ledger.ts index 9c678695fa8..8ffc6356b94 100644 --- a/cloud/apps/relay/src/relay-connection-ledger.ts +++ b/cloud/apps/relay/src/relay-connection-ledger.ts @@ -146,6 +146,27 @@ export class RelayConnectionLedger { return this.reserveUpgrade(this.normalAdmissionLimit) } + // A step-5 booking holds one reserved unit, but only below the placement ceiling the + // director would use, so phones keep the headroom between it and the socket ceiling. + tryReserveBooking(limit: number): { release(): void } | null { + if (!this.canReserveBooking(limit)) return null + this.advanceWatermark() + this.reservedConnectionUnits++ + let released = false + return { + release: () => { + if (released) return + released = true + this.advanceWatermark() + this.reservedConnectionUnits-- + } + } + } + + canReserveBooking(limit: number): boolean { + return this.enforcedConnectionUnits + 1 <= Math.min(limit, this.normalAdmissionLimit) + } + counts(): RelayConnectionLedgerCounts { return { physicalConnections: this.physicalConnections, diff --git a/cloud/apps/relay/src/relay-control-flag-channel.test.ts b/cloud/apps/relay/src/relay-control-flag-channel.test.ts index c99f230124c..1245bd1f4fe 100644 --- a/cloud/apps/relay/src/relay-control-flag-channel.test.ts +++ b/cloud/apps/relay/src/relay-control-flag-channel.test.ts @@ -134,6 +134,20 @@ describe('control flag channel', () => { expect(output).not.toContain(ACCESS_TOKEN) }) + it('applies the step-5 switches: admit mode, intake rate and enforced tickets', async () => { + const google = fakeGoogle( + cellObject(5, { admitMode: 'reserve', intakePerSec: 3.5, ticketCheck: 'enforce' }) + ) + const channel = cellChannel(google.fetchImpl) + await channel.poll() + expect(channel.applied().flags).toEqual({ + ...CELL_FLAG_DEFAULTS, + admitMode: 'reserve', + intakePerSec: 3.5, + ticketCheck: 'enforce' + }) + }) + it('applies the fence and read-timeout switches within their bounds', async () => { const google = fakeGoogle(cellObject(5, { rejectionFence: false, readTimeoutMarginMs: 4_000 })) const channel = cellChannel(google.fetchImpl) @@ -157,22 +171,24 @@ describe('control flag channel', () => { }) it('ignores unknown keys but voids the object on a bad known value or another cell', async () => { - const google = fakeGoogle(cellObject(5, { readinessLocal: true, admitMode: 'memory' })) + const google = fakeGoogle(cellObject(5, { readinessLocal: true, placer: 'memory' })) const channel = cellChannel(google.fetchImpl) await channel.poll() expect(channel.applied()).toEqual({ generation: 5, flags: { ...CELL_FLAG_DEFAULTS, readinessLocal: true }, - ignoredKeys: ['admitMode'] + ignoredKeys: ['placer'] }) - expect(appliedLines().at(-1)).toMatchObject({ ignoredKeys: ['admitMode'] }) + expect(appliedLines().at(-1)).toMatchObject({ ignoredKeys: ['placer'] }) for (const [generation, object] of [ - [6, cellObject(6, { readinessLocal: false, ticketCheck: 'enforce' })], - [7, cellObject(7, { readinessLocal: 'no' })], - [8, cellObject(8, { readTimeoutMarginMs: 500 })], - [9, cellObject(9, { readTimeoutMarginMs: 61_000 })], - [10, cellObject(10, { rejectionFence: 'off' })], - [11, cellObject(11, { readinessLocal: false }, 'production-gce-c8')] + [6, cellObject(6, { readinessLocal: false, ticketCheck: 'strict' })], + [7, cellObject(7, { readinessLocal: false, admitMode: 'memory' })], + [8, cellObject(8, { readinessLocal: false, intakePerSec: -1 })], + [9, cellObject(9, { readinessLocal: 'no' })], + [10, cellObject(10, { readTimeoutMarginMs: 500 })], + [11, cellObject(11, { readTimeoutMarginMs: 61_000 })], + [12, cellObject(12, { rejectionFence: 'off' })], + [13, cellObject(13, { readinessLocal: false }, 'production-gce-c8')] ] as const) { google.state.object = object await channel.poll() diff --git a/cloud/apps/relay/src/relay-server.ts b/cloud/apps/relay/src/relay-server.ts index 821ac12df6d..525cd6b977a 100644 --- a/cloud/apps/relay/src/relay-server.ts +++ b/cloud/apps/relay/src/relay-server.ts @@ -9,7 +9,8 @@ import { RELAY_CLOSE_CODE, RELAY_DEFAULT_REGION, RELAY_PROTOCOL_LIMITS, - RelayAuthSchema + RelayAuthSchema, + cellPlacementCeiling } from '@orca-cloud/relay-contract' import type { IncomingMessage } from 'node:http' import { randomUUID } from 'node:crypto' @@ -19,7 +20,10 @@ import { WebSocketServer } from 'ws' import type WebSocket from 'ws' import type { RawData } from 'ws' import { createRelayApp } from './app.js' -import type { CellFlags } from './cell-flags.js' +import { classifyAssignmentLease } from './assignment-lease.js' +import { defaultIntakePerSec, type CellFlags } from './cell-flags.js' +import { CellReserveBook } from './cell-reserve-book.js' +import type { DemoteRequest, ReserveRequest } from './cell-reserve-contract.js' import { AssignmentLeaseShadow } from './assignment-lease-shadow.js' import { RelayAssignmentStore } from './assignment-store.js' import type { RelayConfig } from './config.js' @@ -60,6 +64,8 @@ function decodePathSegment(value: string): string | null { // tripping it. On 10-07 every window that met both also logged SQL failures. export const HOST_HELLO_SHED_OLDEST_WAIT_MS = 1_000 const HOST_HELLO_SHED_RETRY_AFTER_SECONDS = 2 +// Shipped desktops ignore it; it is for the ones that learn to read it. +const CELL_FULL_RETRY_AFTER_SECONDS = 2 function rejectUpgrade( socket: NodeJS.WritableStream, @@ -165,6 +171,41 @@ export function createRelayServer( const shadowCompare = shadowSeatPoller ? new ShadowDirectoryCompare(shadowSeatPoller.directory, options.now) : undefined + const configuredConnectionLimits = + config.connectionHardCap === undefined + ? null + : (options.connectionLedgerLimits ?? { + hardCap: config.connectionHardCap, + controlReserve: RELAY_ADMISSION_BUDGETS.reservedHostControls + }) + const connectionLedger = + configuredConnectionLimits === null + ? null + : new RelayConnectionLedger( + configuredConnectionLimits.hardCap, + configuredConnectionLimits.controlReserve + ) + // Step 5: a booking may fill the cell only to the ceiling the directors place against. + const placementCeiling = + config.connectionHardCap === undefined || config.connectionUnobservedBound === undefined + ? null + : cellPlacementCeiling(config.connectionHardCap, config.connectionUnobservedBound) + const reserveBook = + config.role === 'cell' && connectionLedger && placementCeiling !== null + ? new CellReserveBook( + { + tryReserve: () => connectionLedger.tryReserveBooking(placementCeiling), + canReserve: () => connectionLedger.canReserveBooking(placementCeiling) + }, + () => + options.cellFlags?.().flags.intakePerSec ?? defaultIntakePerSec(config.region), + options.now + ) + : null + // The hello shed's test: the pool is already timing queries out. + const databaseShedding = (): boolean => + readRelayDatabasePoolPressure(database).databasePoolWaiting >= config.databasePoolMax && + readRelayDatabasePoolOldestWaitMs(database) >= HOST_HELLO_SHED_OLDEST_WAIT_MS const queuedBytes = new ProcessQueuedByteBudget() const sessions = new HostSessionRegistry( config, @@ -180,8 +221,26 @@ export function createRelayServer( enabled: () => options.cellFlags?.().flags.ticketCheck === 'shadow', key: config.assignmentSigningKey, cellId: config.cellId - }) + }), + reserveBook + ? { + mode: () => options.cellFlags?.().flags.admitMode ?? 'db', + ticketEnforce: () => options.cellFlags?.().flags.ticketCheck === 'enforce', + dryRunEnabled: () => options.cellFlags?.().flags.reserveDryRun === true, + databaseShedding: () => databaseShedding(), + book: reserveBook, + verifyLease: (input) => + classifyAssignmentLease({ + ...input, + key: config.assignmentSigningKey, + cellId: config.cellId + }) + } + : undefined ) + // Expired bookings give their units back even while no director is reserving. + const reserveSweepTimer = reserveBook ? setInterval(() => reserveBook.sweep(), 1_000) : null + reserveSweepTimer?.unref() const app = createRelayApp(config, { store, assignments, @@ -204,6 +263,14 @@ export function createRelayServer( cellIncarnation, cellSeatFeed: (sinceSeq) => sessions.seatFeed(sinceSeq), cellFlags: options.cellFlags, + ...(reserveBook && placementCeiling !== null + ? { + cellReserve: (request: ReserveRequest) => sessions.reserve(request), + cellDemote: (request: DemoteRequest) => sessions.demote(request), + cellReserveCounts: () => sessions.reserveCounts(), + cellPlacementCeiling: placementCeiling + } + : {}), isDraining: () => sessions.isDraining(), runtimeCounts: () => runtimeCounts(), regionalRehomeSafetySnapshot: () => ({ @@ -258,20 +325,6 @@ export function createRelayServer( }) let preAuthConnections = 0 let totalConnections = 0 - const configuredConnectionLimits = - config.connectionHardCap === undefined - ? null - : (options.connectionLedgerLimits ?? { - hardCap: config.connectionHardCap, - controlReserve: RELAY_ADMISSION_BUDGETS.reservedHostControls - }) - const connectionLedger = - configuredConnectionLimits === null - ? null - : new RelayConnectionLedger( - configuredConnectionLimits.hardCap, - configuredConnectionLimits.controlReserve - ) const preAuthBySource = new Map() const preAuthAttemptsBySource = new Map() @@ -567,21 +620,28 @@ export function createRelayServer( userId: identity.sub, relayHostId: identity.relayHostId }) + // Null outside reserve mode, so today's path is untouched there. + const reserveAdmission = sessions.reserveAdmissionFor({ + userId: identity.sub, + relayHostId: identity.relayHostId + }) // A hello costs several pooled queries, each failing after a 2 s wait. Shed // only while the pool is already timing them out: the refused desktop gets // the same connect error and backoff a timed-out hello gives it today, 2 s // sooner (shipped desktops ignore Retry-After). A rebind over a live control // is a lease rotation, not a reconnect, so it is never refused here. - if ( - !isRebind && - readRelayDatabasePoolPressure(database).databasePoolWaiting >= config.databasePoolMax && - readRelayDatabasePoolOldestWaitMs(database) >= HOST_HELLO_SHED_OLDEST_WAIT_MS - ) { + if (!isRebind && reserveAdmission === null && databaseShedding()) { observability.recordHostHelloShed() rejectUpgrade(socket, 503, 'Service Unavailable', HOST_HELLO_SHED_RETRY_AFTER_SECONDS) return } - const controlUpgrade = connectionLedger?.tryReserveControl(isRebind) ?? null + // A booked host may rise to the hard cap. Once its upgrade holds a unit, the booking's + // goes back, so the host counts once; a refused upgrade leaves the booking whole. + const controlUpgrade = + connectionLedger?.tryReserveControl(isRebind || reserveAdmission === 'booked') ?? null + if (controlUpgrade && reserveAdmission === 'booked') { + reserveBook?.handOff(identity.sub, identity.relayHostId) + } if ( (connectionLedger && !controlUpgrade) || (!connectionLedger && totalConnections >= RELAY_ADMISSION_BUDGETS.cloudRunConcurrency) @@ -589,7 +649,8 @@ export function createRelayServer( rejectUpgrade( socket, connectionLedger ? 503 : 429, - connectionLedger ? 'Service Unavailable' : 'Too Many Requests' + connectionLedger ? 'Service Unavailable' : 'Too Many Requests', + connectionLedger && sessions.inReserveMode() ? CELL_FULL_RETRY_AFTER_SECONDS : undefined ) return }