diff --git a/cloud/apps/relay/src/app.ts b/cloud/apps/relay/src/app.ts index bdbd3f03d39..e0f2da12dc0 100644 --- a/cloud/apps/relay/src/app.ts +++ b/cloud/apps/relay/src/app.ts @@ -390,6 +390,8 @@ export function createRelayApp( context.json({ ok: true, connectionCapacityProtocol: 2, + // Placement takes min(configured, cell-reported) cap, so a fleet raise may precede the rolls. + ...(config.role === 'cell' ? {} : { capacityReportedCapMin: true }), drainPaceWindowMaxMs: DRAIN_PACE_WINDOW_MAX_MS }) ) diff --git a/cloud/apps/relay/src/assignment-connection-headroom-query.ts b/cloud/apps/relay/src/assignment-connection-headroom-query.ts index c0237aecf26..d6b66a4c527 100644 --- a/cloud/apps/relay/src/assignment-connection-headroom-query.ts +++ b/cloud/apps/relay/src/assignment-connection-headroom-query.ts @@ -1,5 +1,7 @@ +import { EFFECTIVE_HARD_CAP_SQL } from './cell-connection-cap.js' + export const ASSIGNMENT_CONNECTION_HEADROOM_QUERY = - `SELECT limits.cell_id, limits.hard_cap, limits.unobserved_bound, + `SELECT limits.cell_id, ${EFFECTIVE_HARD_CAP_SQL} AS effective_hard_cap, limits.unobserved_bound, connection_snapshot.enforced_connection_units, connection_snapshot.snapshot_at AS last_heartbeat_at, connection_snapshot.cell_incarnation AS connection_incarnation, diff --git a/cloud/apps/relay/src/assignment-connection-headroom.test.ts b/cloud/apps/relay/src/assignment-connection-headroom.test.ts index 21e1d6aa110..9bfeb10749a 100644 --- a/cloud/apps/relay/src/assignment-connection-headroom.test.ts +++ b/cloud/apps/relay/src/assignment-connection-headroom.test.ts @@ -164,6 +164,8 @@ describe('relay assignment connection headroom', () => { expect((await store.cellDeploymentStatus(LIMITED_CELL.id)).connectionCapacity).toEqual({ hardCap: 600, + reportedHardCap: 600, + effectiveHardCap: 600, controlRebindReserve: 100, ordinaryConnectionLimit: 500, unobservedBound: 50, @@ -177,7 +179,7 @@ describe('relay assignment connection headroom', () => { }) }) - it('fails placement closed while a cell changes connection limits', async () => { + it('keeps placing against the old cap until the cell reports a raised one', async () => { const { store } = await setup(449) const expanded = { ...LIMITED_CELL, @@ -185,12 +187,15 @@ describe('relay assignment connection headroom', () => { } await store.reconcileCells([expanded], true) + await expect(store.cellDeploymentStatus(expanded.id)).resolves.toMatchObject({ + connectionCapacity: { hardCap: 1_000, effectiveHardCap: 600, heartbeatFresh: true } + }) await expect( store.assign({ userId: 'transition-user', relayHostId: 'host000000000099' }) + ).resolves.toMatchObject({ cellId: expanded.id }) + await expect( + store.assign({ userId: 'transition-user-2', relayHostId: 'host000000000096' }) ).rejects.toThrow('relay_capacity_exhausted') - await expect(store.cellDeploymentStatus(expanded.id)).resolves.toMatchObject({ - connectionCapacity: { hardCap: 1_000, heartbeatFresh: false } - }) await store.recordCellHeartbeat({ cellId: expanded.id, @@ -216,38 +221,38 @@ describe('relay assignment connection headroom', () => { }) }) - it('fails placement closed while the admin path changes connection limits', async () => { - const { store } = await setup(449) - const expanded = { + it('fails placement closed while the admin path changes the unobserved bound', async () => { + const { store } = await setup(0) + const lowered = { ...LIMITED_CELL, - connectionHardCap: 1_000 as const + connectionUnobservedBound: 49 } - await store.configureCell(expanded, 'general') + await store.configureCell(lowered, 'general') await expect( store.assign({ userId: 'admin-transition', relayHostId: 'host000000000098' }) ).rejects.toThrow('relay_capacity_exhausted') - await expect(store.cellDeploymentStatus(expanded.id)).resolves.toMatchObject({ - connectionCapacity: { hardCap: 1_000, heartbeatFresh: false } + await expect(store.cellDeploymentStatus(lowered.id)).resolves.toMatchObject({ + connectionCapacity: { hardCap: 600, heartbeatFresh: false } }) await store.recordCellHeartbeat({ - cellId: expanded.id, - cellUrl: expanded.url, + cellId: lowered.id, + cellUrl: lowered.url, cellIncarnation: '22222222-2222-4222-8222-222222222222', startedAt: 200, ready: true, observedRequests: 0, - totalConnections: 449, + totalConnections: 0, inFlightConnections: 0, reservedConnectionUnits: 0, - enforcedConnectionUnits: 449, - connectionHardCap: 1_000, - connectionUnobservedBound: 50 + enforcedConnectionUnits: 0, + connectionHardCap: 600, + connectionUnobservedBound: 49 }) await expect( store.assign({ userId: 'admin-restored', relayHostId: 'host000000000097' }) - ).resolves.toMatchObject({ cellId: expanded.id }) + ).resolves.toMatchObject({ cellId: lowered.id }) }) it('admits exactly one reservation at the admission boundary', async () => { diff --git a/cloud/apps/relay/src/assignment-store-lock-order.test.ts b/cloud/apps/relay/src/assignment-store-lock-order.test.ts index 352b42305b7..fa670b0ae7f 100644 --- a/cloud/apps/relay/src/assignment-store-lock-order.test.ts +++ b/cloud/apps/relay/src/assignment-store-lock-order.test.ts @@ -273,7 +273,7 @@ class NewAssignmentLockDatabase implements RelayDatabase { return [ { cell_id: 'cell-general', - hard_cap: 600, + effective_hard_cap: 600, unobserved_bound: 99, enforced_connection_units: 0, outstanding_reservations: 0, diff --git a/cloud/apps/relay/src/assignment-store.ts b/cloud/apps/relay/src/assignment-store.ts index 8427b0a3e72..bffaf69c5af 100644 --- a/cloud/apps/relay/src/assignment-store.ts +++ b/cloud/apps/relay/src/assignment-store.ts @@ -58,6 +58,16 @@ import { import { ASSIGNMENT_CONNECTION_HEADROOM_QUERY } from './assignment-connection-headroom-query.js' +import { + CONNECTION_CAP_COLUMNS_SQL, + EFFECTIVE_HARD_CAP_SQL, + lockCellConnectionLimit, + readCellConnectionCap, + recordReportedHardCap, + reportedHardCapWithinLimit, + writeCellConnectionLimit, + type CellConnectionCap +} from './cell-connection-cap.js' import { AssignmentIdentityQueue } from './assignment/assignment-identity-queue.js' import { CONTROL_RENEWAL_BATCH_SQL, @@ -304,19 +314,14 @@ export type CellDeploymentBaseStatus = { enabled: boolean admissionState: CellAdmissionState capacityRequests: number - connectionCapacity: null | { - hardCap: number - controlRebindReserve: number - ordinaryConnectionLimit: number - unobservedBound: number - normalAdmissionPause: number + connectionCapacity: null | (CellConnectionCap & { observedConnections: number inFlightConnections: number reservedConnectionUnits: number enforcedConnectionUnits: number pendingControlReservations: number heartbeatFresh: boolean - } + }) runtime: null | { cellUrl: string cellIncarnation: string @@ -960,48 +965,7 @@ export class RelayAssignmentStore { stateFromEnabled(cell.initiallyEnabled !== false), now ) - if ( - cell.connectionHardCap !== undefined && - cell.connectionUnobservedBound !== undefined - ) { - const currentLimit = ( - await transaction.queryLocked( - `SELECT hard_cap, unobserved_bound FROM relay_cell_connection_limits - WHERE cell_id = ?`, - [cell.id] - ) - )[0] - const limitChanged = - currentLimit !== undefined && - (integer(currentLimit, 'hard_cap') !== cell.connectionHardCap || - integer(currentLimit, 'unobserved_bound') !== - cell.connectionUnobservedBound) - await transaction.query( - `INSERT INTO relay_cell_connection_limits - (cell_id, hard_cap, unobserved_bound, updated_at) - VALUES (?, ?, ?, ?) - ON CONFLICT (cell_id) DO UPDATE SET - hard_cap = excluded.hard_cap, - unobserved_bound = excluded.unobserved_bound, - updated_at = excluded.updated_at`, - [ - cell.id, - cell.connectionHardCap, - cell.connectionUnobservedBound, - now - ] - ) - if (limitChanged) { - await transaction.query( - `DELETE FROM relay_cell_connection_snapshots WHERE cell_id = ?`, - [cell.id] - ) - await transaction.query( - `DELETE FROM relay_cell_connection_runtime WHERE cell_id = ?`, - [cell.id] - ) - } - } + await writeCellConnectionLimit(transaction, cell, now) } const current = await transaction.query( `SELECT cell_id, enabled FROM relay_cells ORDER BY cell_id ASC` @@ -1975,13 +1939,9 @@ export class RelayAssignmentStore { ) { throw new Error('stale_cell_incarnation') } - const connectionLimit = ( - await transaction.queryLocked( - `SELECT * FROM relay_cell_connection_limits WHERE cell_id = ?`, - [input.cellId] - ) - )[0] + const connectionLimit = await lockCellConnectionLimit(transaction, input.cellId) if (connectionLimit) { + const reportedHardCap = reportedHardCapWithinLimit(connectionLimit, input) if ( input.totalConnections === undefined || input.inFlightConnections === undefined || @@ -1991,12 +1951,11 @@ export class RelayAssignmentStore { input.totalConnections + input.inFlightConnections + input.reservedConnectionUnits || - input.connectionHardCap !== integer(connectionLimit, 'hard_cap') || - input.connectionUnobservedBound !== - integer(connectionLimit, 'unobserved_bound') + reportedHardCap === undefined ) { throw new Error('cell_connection_telemetry_mismatch') } + await recordReportedHardCap(transaction, connectionLimit, input.cellId, reportedHardCap) } else if ( input.totalConnections !== undefined || input.inFlightConnections !== undefined || @@ -3638,43 +3597,7 @@ export class RelayAssignmentStore { ) await ensureCellAdmission(transaction, cell.id, state, now) await setCellAdmissionBeforeBoundary(transaction, cell.id, state, now) - if ( - cell.connectionHardCap !== undefined && - cell.connectionUnobservedBound !== undefined - ) { - const currentLimit = ( - await transaction.queryLocked( - `SELECT hard_cap, unobserved_bound FROM relay_cell_connection_limits - WHERE cell_id = ?`, - [cell.id] - ) - )[0] - const limitChanged = - currentLimit !== undefined && - (integer(currentLimit, 'hard_cap') !== cell.connectionHardCap || - integer(currentLimit, 'unobserved_bound') !== - cell.connectionUnobservedBound) - await transaction.query( - `INSERT INTO relay_cell_connection_limits - (cell_id, hard_cap, unobserved_bound, updated_at) - VALUES (?, ?, ?, ?) - ON CONFLICT (cell_id) DO UPDATE SET - hard_cap = excluded.hard_cap, - unobserved_bound = excluded.unobserved_bound, - updated_at = excluded.updated_at`, - [cell.id, cell.connectionHardCap, cell.connectionUnobservedBound, now] - ) - if (limitChanged) { - await transaction.query( - `DELETE FROM relay_cell_connection_snapshots WHERE cell_id = ?`, - [cell.id] - ) - await transaction.query( - `DELETE FROM relay_cell_connection_runtime WHERE cell_id = ?`, - [cell.id] - ) - } - } + await writeCellConnectionLimit(transaction, cell, now) }) } @@ -3983,7 +3906,7 @@ export class RelayAssignmentStore { if (!cellRow) throw new Error('cell_not_found') const connectionRow = ( await this.database.query( - `SELECT limits.hard_cap, limits.unobserved_bound, + `SELECT ${CONNECTION_CAP_COLUMNS_SQL}, connection_runtime.total_connections, connection_runtime.in_flight_connections, connection_runtime.reserved_connection_units, connection_runtime.enforced_connection_units, @@ -4016,16 +3939,7 @@ export class RelayAssignmentStore { capacityRequests: integer(cellRow, 'capacity_requests'), connectionCapacity: connectionRow ? { - hardCap: integer(connectionRow, 'hard_cap'), - controlRebindReserve: RELAY_ADMISSION_BUDGETS.reservedHostControls, - ordinaryConnectionLimit: - integer(connectionRow, 'hard_cap') - - RELAY_ADMISSION_BUDGETS.reservedHostControls, - unobservedBound: integer(connectionRow, 'unobserved_bound'), - normalAdmissionPause: - integer(connectionRow, 'hard_cap') - - RELAY_ADMISSION_BUDGETS.reservedHostControls - - integer(connectionRow, 'unobserved_bound'), + ...readCellConnectionCap(connectionRow), observedConnections: optionalInteger(connectionRow, 'total_connections') ?? 0, inFlightConnections: @@ -6913,7 +6827,7 @@ export class RelayAssignmentStore { AND NOT (reservation.user_id = ? AND reservation.relay_host_id = ? AND reservation.assignment_epoch = ?)) + limits.unobserved_bound - < limits.hard_cap - ? + < ${EFFECTIVE_HARD_CAP_SQL} - ? ) ) RETURNING cell_id`, @@ -8573,7 +8487,7 @@ export class RelayAssignmentStore { integer(row, 'enforced_connection_units') + integer(row, 'outstanding_reservations') + integer(row, 'unobserved_bound') < - integer(row, 'hard_cap') - + integer(row, 'effective_hard_cap') - RELAY_ADMISSION_BUDGETS.reservedHostControls return [text(row, 'cell_id'), hasHeadroom] as const }) diff --git a/cloud/apps/relay/src/cell-admission-startup.ts b/cloud/apps/relay/src/cell-admission-startup.ts index 677980f27b5..2a4bb24509a 100644 --- a/cloud/apps/relay/src/cell-admission-startup.ts +++ b/cloud/apps/relay/src/cell-admission-startup.ts @@ -28,8 +28,9 @@ export async function reconcileCellAdmissionAtStartup( // Admission is operator/director state. A new worker must not enable itself // before its distinct candidate has passed production preflight. if (config.role === 'cell') return - // Retries every transient failure, a lost reply included: the reconcile is upserts and - // limit-change deletes that converge on the same rows if a lost COMMIT had landed. + // Retries every transient failure, a lost reply included: the reconcile is upserts, a cap-only + // update and bound-change deletes, all of which converge on the same rows if a lost COMMIT had + // landed. await retryTransientDatabaseStartup( async () => await assignments.reconcileCellsAtStartup(config.cells), STARTUP_RECONCILE_RETRY, diff --git a/cloud/apps/relay/src/cell-connection-cap-postgres.test.ts b/cloud/apps/relay/src/cell-connection-cap-postgres.test.ts new file mode 100644 index 00000000000..324484a7d49 --- /dev/null +++ b/cloud/apps/relay/src/cell-connection-cap-postgres.test.ts @@ -0,0 +1,157 @@ +import { afterAll, beforeAll, describe, expect, it } from 'vitest' +import { ASSIGNMENT_CONNECTION_HEADROOM_QUERY } from './assignment-connection-headroom-query.js' +import { RelayAssignmentStore } from './assignment-store.js' +import type { RelayCellConfig } from './config.js' +import { openRelayDatabase, type RelayDatabase } from './database.js' + +const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL +const describePostgres = databaseUrl ? describe : describe.skip +const CELL_PREFIX = 'connection-cap-min-postgres' +const CELL_ID = `${CELL_PREFIX}-single` +const FLEET_SIZE = 16 + +function cellConfig(id: string, connectionHardCap: 1_000 | 3_000): RelayCellConfig { + return { + id, + url: `https://${id}.example.com`, + capacityRequests: 1_000, + connectionHardCap, + connectionUnobservedBound: 60 + } +} +const configured = (connectionHardCap: 1_000 | 3_000): RelayCellConfig => + cellConfig(CELL_ID, connectionHardCap) + +async function removeCells(database: RelayDatabase): Promise { + for (const table of [ + 'relay_cell_connection_snapshots', + 'relay_cell_connection_runtime', + 'relay_cell_connection_limits', + 'relay_cell_capabilities', + 'relay_cell_rehome_safety', + 'relay_cell_runtime', + 'relay_cell_admission', + 'relay_cell_regions', + 'relay_cells' + ]) { + await database.query(`DELETE FROM ${table} WHERE cell_id LIKE ?`, [`${CELL_PREFIX}-%`]) + } +} + +function storeOn(database: RelayDatabase): RelayAssignmentStore { + return new RelayAssignmentStore(database, () => 100, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) +} + +async function heartbeat( + store: RelayAssignmentStore, + cell: RelayCellConfig, + connectionHardCap: number, + units: number +): Promise { + await store.recordCellHeartbeat({ + cellId: cell.id, + cellUrl: cell.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: units, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: units, + connectionHardCap, + connectionUnobservedBound: 60 + }) +} + +describePostgres('PostgreSQL effective connection cap', () => { + let database: RelayDatabase + let racer: RelayDatabase + + beforeAll(async () => { + database = await openRelayDatabase({ databaseUrl, dataDir: '' }) + racer = await openRelayDatabase({ databaseUrl, dataDir: '' }) + await removeCells(database) + }) + + afterAll(async () => { + await removeCells(database) + await racer.close() + await database.close() + }) + + it('records the previous cap on a raise and places against the cell-reported cap', async () => { + const store = storeOn(database) + const beat = async (connectionHardCap: number, units: number): Promise => + await heartbeat(store, configured(1_000), connectionHardCap, units) + await store.reconcileCells([configured(1_000)], true) + await beat(1_000, 10) + // A pre-change director's row: no report stamped. + await database.query( + `UPDATE relay_cell_connection_limits SET reported_hard_cap = NULL WHERE cell_id = ?`, + [CELL_ID] + ) + + await store.reconcileCells([configured(3_000)], true) + const status = await store.cellDeploymentStatus(CELL_ID) + expect(status.connectionCapacity).toMatchObject({ + hardCap: 3_000, + reportedHardCap: 1_000, + effectiveHardCap: 1_000, + normalAdmissionPause: 840, + heartbeatFresh: true + }) + + await beat(1_000, 839) + const headroom = await database.query( + `${ASSIGNMENT_CONNECTION_HEADROOM_QUERY} WHERE limits.cell_id = ?`, + [CELL_ID] + ) + expect(Number(headroom[0]?.effective_hard_cap)).toBe(1_000) + await expect(beat(3_000, 0)).resolves.toBeUndefined() + await expect(store.cellDeploymentStatus(CELL_ID)).resolves.toMatchObject({ + connectionCapacity: { reportedHardCap: 3_000, effectiveHardCap: 3_000 } + }) + }) + + it('keeps 16 raised cells placeable while a heartbeat races the raise', async () => { + const ids = Array.from({ length: FLEET_SIZE }, (_, index) => `${CELL_PREFIX}-fleet-${index}`) + const lower = ids.map((id) => cellConfig(id, 1_000)) + const director = storeOn(database) + const otherInstance = storeOn(racer) + await director.reconcileCells(lower, true) + for (const cell of lower) await heartbeat(director, cell, 1_000, 10) + + // The heartbeat lands before or after the raise; either order must leave the same state. + await Promise.all([ + director.reconcileCells( + ids.map((id) => cellConfig(id, 3_000)), + true + ), + heartbeat(otherInstance, lower[0]!, 1_000, 11) + ]) + + for (const id of ids) { + await expect(director.cellDeploymentStatus(id)).resolves.toMatchObject({ + connectionCapacity: { + hardCap: 3_000, + reportedHardCap: 1_000, + effectiveHardCap: 1_000, + heartbeatFresh: true + } + }) + const rows = await database.query( + `SELECT cell_id FROM relay_cell_connection_snapshots WHERE cell_id = ?`, + [id] + ) + expect(rows).toHaveLength(1) + } + // The old processes keep heartbeating their 1,000 under the raised config. + for (const cell of lower) { + await expect(heartbeat(otherInstance, cell, 1_000, 12)).resolves.toBeUndefined() + } + }) +}) diff --git a/cloud/apps/relay/src/cell-connection-cap.test.ts b/cloud/apps/relay/src/cell-connection-cap.test.ts new file mode 100644 index 00000000000..e486288ad7a --- /dev/null +++ b/cloud/apps/relay/src/cell-connection-cap.test.ts @@ -0,0 +1,443 @@ +import { createHash } from 'node:crypto' +import { readdirSync, readFileSync } from 'node:fs' +import { join } from 'node:path' +import { fileURLToPath } from 'node:url' +import { afterEach, describe, expect, it, vi } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.js' +import { encodeMembership, type CellAdmissionMembership } from './cell-admission-selector.js' +import type { RelayCellConfig } from './config.js' +import { openInMemoryRelayDatabase, type RelayDatabase } from './database.js' + +const HEARTBEAT_TTL_MS = 45_000 +const BOUND = 60 +const RELAY_CELL_COUNT = 16 +const OLD_INCARNATION = '11111111-1111-4111-8111-111111111111' +const NEW_INCARNATION = '22222222-2222-4222-8222-222222222222' + +function cell(index: number, connectionHardCap: 1_000 | 3_000): RelayCellConfig { + const id = `us-c${String(index).padStart(2, '0')}` + return { + id, + url: `https://${id}.example.com`, + region: 'us-central1', + capacityRequests: 10_000, + connectionHardCap, + connectionUnobservedBound: BOUND + } +} + +const databases: RelayDatabase[] = [] + +afterEach(async () => { + vi.restoreAllMocks() + for (const database of databases.splice(0)) await database.close() +}) + +type Report = { + connectionHardCap: number + units?: number + incarnation?: string + startedAt?: number + connectionUnobservedBound?: number +} + +async function setup(cells: RelayCellConfig[]) { + const database = await openInMemoryRelayDatabase() + databases.push(database) + let now = 100 + const store = new RelayAssignmentStore(database, () => now, { + requireLiveCells: true, + heartbeatTtlMs: HEARTBEAT_TTL_MS + }) + await store.reconcileCells(cells, true) + const heartbeat = async (target: RelayCellConfig, report: Report): Promise => { + const units = report.units ?? 0 + await store.recordCellHeartbeat({ + cellId: target.id, + cellUrl: target.url, + cellIncarnation: report.incarnation ?? OLD_INCARNATION, + startedAt: report.startedAt ?? 50, + ready: true, + observedRequests: 0, + region: 'us-central1', + totalConnections: units, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: units, + connectionHardCap: report.connectionHardCap, + connectionUnobservedBound: report.connectionUnobservedBound ?? BOUND + }) + } + const capacity = async (cellId: string) => + (await store.cellDeploymentStatus(cellId)).connectionCapacity + const reportedHardCap = async (cellId: string): Promise => + ( + await database.query( + `SELECT reported_hard_cap FROM relay_cell_connection_limits WHERE cell_id = ?`, + [cellId] + ) + )[0]?.reported_hard_cap + const telemetryRows = async (cellId: string): Promise => { + let rows = 0 + for (const table of ['relay_cell_connection_snapshots', 'relay_cell_connection_runtime']) { + rows += (await database.query(`SELECT cell_id FROM ${table} WHERE cell_id = ?`, [cellId])) + .length + } + return rows + } + const isolateForRoll = async (cellId: string): Promise => { + const current = await store.inspectCellAdmissionSelector() + const membership: CellAdmissionMembership = { + existingOnly: [], + migrationOnly: [cellId], + general: cells.map(({ id }) => id).filter((id) => id !== cellId) + } + await store.applyCellAdmissionSelector({ + attemptId: `isolate-${cellId}`, + expectedGeneration: current.selector.generation, + expectedMembershipSha256: createHash('sha256') + .update(encodeMembership(current.selector.membership)) + .digest('hex'), + membership, + rollIsolatedCells: [cellId] + }) + } + return { + database, + store, + heartbeat, + capacity, + reportedHardCap, + telemetryRows, + isolateForRoll, + setNow: (value: number) => (now = value) + } +} + +describe('effective connection cap = min(configured, reported)', () => { + it('accepts a report below the configured cap and places against the lower one', async () => { + const configured = cell(1, 3_000) + const harness = await setup([configured]) + // 1,000 - 100 rebind reserve - 60 unobserved = 840 pause; 839 used leaves one seat. + await harness.heartbeat(configured, { connectionHardCap: 1_000, units: 839 }) + + expect(await harness.capacity(configured.id)).toMatchObject({ + hardCap: 3_000, + reportedHardCap: 1_000, + effectiveHardCap: 1_000, + ordinaryConnectionLimit: 900, + normalAdmissionPause: 840, + heartbeatFresh: true + }) + await expect( + harness.store.assign({ userId: 'user-a', relayHostId: 'host000000000001' }) + ).resolves.toMatchObject({ cellId: configured.id }) + await expect( + harness.store.assign({ userId: 'user-b', relayHostId: 'host000000000002' }) + ).rejects.toThrow('relay_capacity_exhausted') + }) + + it('still refuses a report above the configured cap or with a different bound', async () => { + const configured = cell(1, 1_000) + const harness = await setup([configured]) + + await expect(harness.heartbeat(configured, { connectionHardCap: 3_000 })).rejects.toThrow( + 'cell_connection_telemetry_mismatch' + ) + await expect( + harness.heartbeat(configured, { connectionHardCap: 1_000, connectionUnobservedBound: 59 }) + ).rejects.toThrow('cell_connection_telemetry_mismatch') + expect(await harness.reportedHardCap(configured.id)).toBeNull() + }) + + it.each(['reconcile', 'admin'] as const)( + 'keeps the previous cap and the telemetry rows across a %s raise', + async (path) => { + const before = cell(1, 1_000) + const harness = await setup([before]) + await harness.heartbeat(before, { connectionHardCap: 1_000, units: 839 }) + // A row a pre-change director wrote: it never stamped a report. + await harness.database.query( + `UPDATE relay_cell_connection_limits SET reported_hard_cap = NULL WHERE cell_id = ?`, + [before.id] + ) + const raised = cell(1, 3_000) + if (path === 'reconcile') await harness.store.reconcileCells([raised], true) + else await harness.store.configureCell(raised, 'general') + + expect(await harness.reportedHardCap(before.id)).toBe(1_000) + expect(await harness.telemetryRows(before.id)).toBe(2) + expect(await harness.capacity(before.id)).toMatchObject({ + hardCap: 3_000, + effectiveHardCap: 1_000, + normalAdmissionPause: 840, + heartbeatFresh: true + }) + // No heartbeat yet: one seat under the old cap, and not one more. + await expect( + harness.store.assign({ userId: 'user-a', relayHostId: 'host000000000001' }) + ).resolves.toMatchObject({ cellId: before.id }) + await expect( + harness.store.assign({ userId: 'user-b', relayHostId: 'host000000000002' }) + ).rejects.toThrow('relay_capacity_exhausted') + } + ) + + it('keeps a stored report across a lower-then-raise', async () => { + const before = cell(1, 3_000) + const harness = await setup([before]) + await harness.heartbeat(before, { connectionHardCap: 3_000 }) + await harness.store.reconcileCells([cell(1, 1_000)], true) + expect(await harness.capacity(before.id)).toMatchObject({ + reportedHardCap: 3_000, + effectiveHardCap: 1_000 + }) + await harness.store.reconcileCells([cell(1, 3_000)], true) + + expect(await harness.capacity(before.id)).toMatchObject({ effectiveHardCap: 3_000 }) + }) + + it('keeps telemetry on a lowering, places against it and refuses the higher report', async () => { + const before = cell(1, 3_000) + const harness = await setup([before]) + await harness.heartbeat(before, { connectionHardCap: 3_000, units: 839 }) + await harness.store.reconcileCells([cell(1, 1_000)], true) + + expect(await harness.telemetryRows(before.id)).toBe(2) + expect(await harness.capacity(before.id)).toMatchObject({ + hardCap: 1_000, + effectiveHardCap: 1_000, + normalAdmissionPause: 840, + heartbeatFresh: true + }) + await expect( + harness.store.assign({ userId: 'user-a', relayHostId: 'host000000000001' }) + ).resolves.toMatchObject({ cellId: before.id }) + await expect( + harness.store.assign({ userId: 'user-b', relayHostId: 'host000000000002' }) + ).rejects.toThrow('relay_capacity_exhausted') + await expect(harness.heartbeat(before, { connectionHardCap: 3_000 })).rejects.toThrow( + 'cell_connection_telemetry_mismatch' + ) + }) + + it('still resets telemetry and fails placement closed on a bound change', async () => { + const before = cell(1, 3_000) + const harness = await setup([before]) + await harness.heartbeat(before, { connectionHardCap: 3_000, units: 10 }) + await harness.store.reconcileCells([{ ...before, connectionUnobservedBound: 50 }], true) + + expect(await harness.telemetryRows(before.id)).toBe(0) + expect(await harness.capacity(before.id)).toMatchObject({ heartbeatFresh: false }) + await expect( + harness.store.assign({ userId: 'user-a', relayHostId: 'host000000000001' }) + ).rejects.toThrow('relay_capacity_exhausted') + }) + + it('survives revisions with different caps reconciling in turn during a deploy', async () => { + // A new revision raises 16 cells; an outgoing instance of the same code, still configured at + // 1,000, then boots and reconciles them back down; the new revision boots again. + const lower = Array.from({ length: RELAY_CELL_COUNT }, (_, index) => cell(index + 1, 1_000)) + const higher = lower.map((_, index) => cell(index + 1, 3_000)) + const harness = await setup(lower) + for (const target of lower) await harness.heartbeat(target, { connectionHardCap: 1_000 }) + let host = 0 + for (const config of [higher, lower, higher]) { + await harness.store.reconcileCells(config, true) + for (const target of lower) { + expect(await harness.telemetryRows(target.id)).toBe(2) + expect(await harness.capacity(target.id)).toMatchObject({ + reportedHardCap: 1_000, + effectiveHardCap: 1_000, + heartbeatFresh: true + }) + } + await expect( + harness.store.assign({ userId: `user-${host}`, relayHostId: `interleave${host++}` }) + ).resolves.toMatchObject({ placed: true }) + } + }) + + it('logs once when a cell reports a cap below its configured one', async () => { + const configured = cell(1, 3_000) + const harness = await setup([configured]) + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + await harness.heartbeat(configured, { connectionHardCap: 1_000 }) + await harness.heartbeat(configured, { connectionHardCap: 1_000 }) + await harness.heartbeat(configured, { connectionHardCap: 3_000 }) + + const events = warn.mock.calls + .map((call) => String(call[0])) + .filter((line) => line.includes('orca_relay_cell_reported_cap_below_configured')) + .map((line) => JSON.parse(line) as unknown) + expect(events).toEqual([ + { + event: 'orca_relay_cell_reported_cap_below_configured', + cellId: configured.id, + reportedHardCap: 1_000, + configuredHardCap: 3_000 + } + ]) + }) + + it('takes a raise only once the cell itself reports it', async () => { + const before = cell(1, 1_000) + const harness = await setup([before]) + await harness.heartbeat(before, { connectionHardCap: 1_000 }) + await harness.store.reconcileCells([cell(1, 3_000)], true) + await harness.heartbeat(before, { + connectionHardCap: 3_000, + incarnation: NEW_INCARNATION, + startedAt: 90 + }) + + expect(await harness.capacity(before.id)).toMatchObject({ + hardCap: 3_000, + reportedHardCap: 3_000, + effectiveHardCap: 3_000, + normalAdmissionPause: 2_840 + }) + }) + + it('places a mixed fleet against each cell’s own reported cap', async () => { + const small = cell(1, 3_000) + const large = cell(2, 3_000) + const harness = await setup([small, large]) + await harness.heartbeat(small, { connectionHardCap: 1_000, units: 840 }) + await harness.heartbeat(large, { connectionHardCap: 3_000, units: 840 }) + + expect(await harness.capacity(small.id)).toMatchObject({ normalAdmissionPause: 840 }) + expect(await harness.capacity(large.id)).toMatchObject({ normalAdmissionPause: 2_840 }) + for (let index = 0; index < 4; index++) { + await expect( + harness.store.assign({ userId: `user-${index}`, relayHostId: `host00000000000${index}` }) + ).resolves.toMatchObject({ cellId: large.id }) + } + }) + + it('has no unavailable window while a cell rolls 1,000 -> 3,000', async () => { + const cells = [cell(1, 1_000), cell(2, 1_000)] + const [rolled, neighbour] = cells + const harness = await setup(cells) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + await harness.heartbeat(rolled!, { connectionHardCap: 1_000 }) + await harness.heartbeat(neighbour!, { connectionHardCap: 1_000, units: 500 }) + const home = await harness.store.assign(identity) + expect(home.cellId).toBe(rolled!.id) + await harness.store.reconcileCells( + cells.map((_, index) => cell(index + 1, 3_000)), + true + ) + + // The old process keeps heartbeating its 1,000 well past the TTL; the host stays sticky. + for (const at of [30_000, 60_000, 90_000]) { + harness.setNow(at) + await harness.heartbeat(rolled!, { connectionHardCap: 1_000 }) + await harness.heartbeat(neighbour!, { connectionHardCap: 1_000, units: 500 }) + await expect(harness.store.assign(identity)).resolves.toMatchObject({ + cellId: rolled!.id, + assignmentEpoch: home.assignmentEpoch + }) + } + // The roll isolates the live old process: the isolated-incumbent lane re-places the host. + await harness.isolateForRoll(rolled!.id) + const moved = await harness.store.assign(identity) + expect(moved).toMatchObject({ cellId: neighbour!.id, placed: true }) + + // The replacement boots and reports the new cap. + harness.setNow(120_000) + await harness.heartbeat(rolled!, { + connectionHardCap: 3_000, + incarnation: NEW_INCARNATION, + startedAt: 110_000 + }) + expect(await harness.capacity(rolled!.id)).toMatchObject({ + effectiveHardCap: 3_000, + heartbeatFresh: true + }) + }) + + it('changes no placement, sticky or status answer on a 16-cell raise', async () => { + const before = Array.from({ length: RELAY_CELL_COUNT }, (_, index) => cell(index + 1, 1_000)) + const harness = await setup(before) + for (const [index, target] of before.entries()) { + await harness.heartbeat(target, { connectionHardCap: 1_000, units: 100 + index }) + } + const pinned = [] + for (let index = 0; index < RELAY_CELL_COUNT; index++) { + pinned.push( + await harness.store.assign({ userId: `user-${index}`, relayHostId: `pinned${index}` }) + ) + } + const statusBefore = await Promise.all( + before.map(async ({ id }) => await harness.store.reserveAwareCellDeploymentStatus(id)) + ) + + await harness.store.reconcileCells( + before.map((_, index) => cell(index + 1, 3_000)), + true + ) + + // No heartbeat has landed since the raise. + for (const [index, target] of before.entries()) { + const after = await harness.store.reserveAwareCellDeploymentStatus(target.id) + expect(after).toEqual({ + ...statusBefore[index], + status: { + ...statusBefore[index]!.status, + connectionCapacity: { + ...statusBefore[index]!.status.connectionCapacity, + hardCap: 3_000, + reportedHardCap: 1_000 + } + } + }) + } + for (const [index, assignment] of pinned.entries()) { + await expect( + harness.store.assign({ userId: `user-${index}`, relayHostId: `pinned${index}` }) + ).resolves.toMatchObject({ + cellId: assignment.cellId, + assignmentEpoch: assignment.assignmentEpoch + }) + } + for (let index = 0; index < RELAY_CELL_COUNT; index++) { + await expect( + harness.store.assign({ userId: `fresh-${index}`, relayHostId: `fresh${index}` }) + ).resolves.toMatchObject({ placed: true }) + } + }) +}) + +describe('connection cap readers', () => { + const source = fileURLToPath(new URL('.', import.meta.url)) + const allowed = new Set([ + 'cell-connection-cap.ts', + 'database.ts', + 'cell-admission-migration-registration.ts' + ]) + const sources = readdirSync(source, { recursive: true, encoding: 'utf8' }) + .filter((file) => file.endsWith('.ts') && !file.endsWith('.test.ts')) + .filter((file) => !file.startsWith('test-fixtures')) + .map((file) => ({ file, text: readFileSync(join(source, file), 'utf8') })) + const offenders = (pattern: RegExp): string[] => + sources + .filter(({ file, text }) => !allowed.has(file) && pattern.test(text)) + .map(({ file }) => file) + + // Placement, headroom and status must take the effective cap; a raw read would count seats a + // running cell does not have. Schema and new-cell registration (an exact config write) are the + // only other owners. + it('reads the cap columns only through cell-connection-cap.ts', () => { + expect(offenders(/\bhard_cap\b/)).toEqual([]) + expect(offenders(/\blimits\.\*/)).toEqual([]) + expect(offenders(/SELECT\s+\*\s+FROM\s+relay_cell_connection_limits/)).toEqual([]) + const owner = sources.find(({ file }) => file === 'cell-connection-cap.ts') + expect(owner?.text).toMatch(/\bhard_cap\b/) + }) + + it('has no in-process reader of the configured status cap', () => { + // connectionCapacity.hardCap stays configured for old ops scripts; nothing here may plan on it. + expect(offenders(/connectionCapacity\??\.hardCap/)).toEqual([]) + }) +}) diff --git a/cloud/apps/relay/src/cell-connection-cap.ts b/cloud/apps/relay/src/cell-connection-cap.ts new file mode 100644 index 00000000000..4636b510470 --- /dev/null +++ b/cloud/apps/relay/src/cell-connection-cap.ts @@ -0,0 +1,166 @@ +import { RELAY_ADMISSION_BUDGETS } from '@orca-cloud/relay-contract' +import type { RelayCellConfig } from './config.js' +import type { RelayDatabase, SqlRow } from './database.js' + +// The one owner of relay_cell_connection_limits.hard_cap. The director's configured cap and the +// cap the cell last reported enforcing can differ while a cell is being resized; placement must +// never count seats a running process does not have, so every capacity reader takes the lower. +// A NULL report (a row no heartbeat has stamped yet) compares unknown and falls to the ELSE arm. +export const EFFECTIVE_HARD_CAP_SQL = `(CASE WHEN limits.reported_hard_cap < limits.hard_cap + THEN limits.reported_hard_cap ELSE limits.hard_cap END)` + +// Column list for readers that report the cap; requires the table aliased as `limits`. +export const CONNECTION_CAP_COLUMNS_SQL = `limits.hard_cap AS configured_hard_cap, + limits.reported_hard_cap, ${EFFECTIVE_HARD_CAP_SQL} AS effective_hard_cap, + limits.unobserved_bound` + +export type CellConnectionCap = { + // Configured on the director; old scripts read this field, so its meaning is fixed. + hardCap: number + reportedHardCap: number | null + effectiveHardCap: number + controlRebindReserve: number + ordinaryConnectionLimit: number + unobservedBound: number + normalAdmissionPause: number +} + +// Derived limits follow the effective cap: a same-cap job counting free slots from them must not +// count a raise the cell has not taken yet. +export function readCellConnectionCap(row: SqlRow): CellConnectionCap { + const effectiveHardCap = capInteger(row, 'effective_hard_cap') + const unobservedBound = capInteger(row, 'unobserved_bound') + const ordinaryConnectionLimit = effectiveHardCap - RELAY_ADMISSION_BUDGETS.reservedHostControls + return { + hardCap: capInteger(row, 'configured_hard_cap'), + reportedHardCap: + row.reported_hard_cap === null || row.reported_hard_cap === undefined + ? null + : capInteger(row, 'reported_hard_cap'), + effectiveHardCap, + controlRebindReserve: RELAY_ADMISSION_BUDGETS.reservedHostControls, + ordinaryConnectionLimit, + unobservedBound, + normalAdmissionPause: ordinaryConnectionLimit - unobservedBound + } +} + +// A cell may run below the configured cap (it has not been rolled onto a raise yet), never above +// it, and its unobserved allowance is part of the placement ceiling so it must match exactly. +// Returns the cap to record, or undefined when the report does not fit. +export function reportedHardCapWithinLimit( + limit: SqlRow, + report: { connectionHardCap?: number; connectionUnobservedBound?: number } +): number | undefined { + if ( + report.connectionHardCap === undefined || + report.connectionHardCap > capInteger(limit, 'hard_cap') || + report.connectionUnobservedBound !== capInteger(limit, 'unobserved_bound') + ) { + return undefined + } + return report.connectionHardCap +} + +// The heartbeat's row read, locked so the report it records cannot race a configuration write. +export async function lockCellConnectionLimit( + transaction: RelayDatabase, + cellId: string +): Promise { + return ( + await transaction.queryLocked( + `SELECT limits.hard_cap, limits.unobserved_bound, limits.reported_hard_cap + FROM relay_cell_connection_limits limits WHERE limits.cell_id = ?`, + [cellId] + ) + )[0] +} + +// Written only on change: steady heartbeats leave the row untouched. +export async function recordReportedHardCap( + transaction: RelayDatabase, + limit: SqlRow, + cellId: string, + reportedHardCap: number +): Promise { + const stored = limit.reported_hard_cap + if (stored !== null && stored !== undefined) { + if (capInteger(limit, 'reported_hard_cap') === reportedHardCap) return + } + await transaction.query( + `UPDATE relay_cell_connection_limits SET reported_hard_cap = ? WHERE cell_id = ?`, + [reportedHardCap, cellId] + ) + const configuredHardCap = capInteger(limit, 'hard_cap') + // A cell left below its configured cap (wrong image or env) is otherwise silent. + if (reportedHardCap < configuredHardCap) { + console.warn( + JSON.stringify({ + event: 'orca_relay_cell_reported_cap_below_configured', + cellId, + reportedHardCap, + configuredHardCap + }) + ) + } +} + +// A cap-only change (same unobserved bound, either direction) changes nothing placement reads +// beyond min(): the stored report stands, or failing one (a row an older director wrote) the cap +// the cell was configured with until now is recorded as its report, so a raise waits for the cell +// to report it. Telemetry rows survive, so a deploy that interleaves revisions with different caps +// cannot wipe the fleet's headroom. A bound change keeps the old fail-closed reset. +export async function writeCellConnectionLimit( + transaction: RelayDatabase, + cell: Pick, + now: number +): Promise { + const { connectionHardCap, connectionUnobservedBound } = cell + // An uncapped config leaves any existing row alone, as before. + if (connectionHardCap === undefined || connectionUnobservedBound === undefined) return + const current = ( + await transaction.queryLocked( + `SELECT hard_cap, unobserved_bound FROM relay_cell_connection_limits WHERE cell_id = ?`, + [cell.id] + ) + )[0] + const capOnlyChange = + current !== undefined && + capInteger(current, 'unobserved_bound') === connectionUnobservedBound && + capInteger(current, 'hard_cap') !== connectionHardCap + if (capOnlyChange) { + await transaction.query( + `UPDATE relay_cell_connection_limits + SET hard_cap = ?, reported_hard_cap = COALESCE(reported_hard_cap, hard_cap), updated_at = ? + WHERE cell_id = ?`, + [connectionHardCap, now, cell.id] + ) + return + } + await transaction.query( + `INSERT INTO relay_cell_connection_limits + (cell_id, hard_cap, unobserved_bound, updated_at) + VALUES (?, ?, ?, ?) + ON CONFLICT (cell_id) DO UPDATE SET + hard_cap = excluded.hard_cap, + unobserved_bound = excluded.unobserved_bound, + updated_at = excluded.updated_at`, + [cell.id, connectionHardCap, connectionUnobservedBound, now] + ) + const boundChanged = + current !== undefined && capInteger(current, 'unobserved_bound') !== connectionUnobservedBound + if (boundChanged) { + await transaction.query(`DELETE FROM relay_cell_connection_snapshots WHERE cell_id = ?`, [ + cell.id + ]) + await transaction.query(`DELETE FROM relay_cell_connection_runtime WHERE cell_id = ?`, [ + cell.id + ]) + } +} + +function capInteger(row: SqlRow, field: string): number { + const value = Number(row[field]) + if (!Number.isSafeInteger(value)) throw new Error(`invalid_${field}`) + return value +} diff --git a/cloud/apps/relay/src/database.ts b/cloud/apps/relay/src/database.ts index 49e9e1ecfa6..ef294ed2f12 100644 --- a/cloud/apps/relay/src/database.ts +++ b/cloud/apps/relay/src/database.ts @@ -435,7 +435,8 @@ CREATE TABLE IF NOT EXISTS relay_cell_connection_limits ( cell_id TEXT PRIMARY KEY, hard_cap BIGINT NOT NULL, unobserved_bound BIGINT NOT NULL, - updated_at BIGINT NOT NULL + updated_at BIGINT NOT NULL, + reported_hard_cap BIGINT ); CREATE TABLE IF NOT EXISTS relay_cell_connection_runtime ( @@ -712,6 +713,9 @@ export const POSTGRES_SCHEMA_MIGRATIONS = [ // the rewrite is catalog-only; a cell isolated before this column existed reads as unmarked and // keeps its hosts pinned, which is the pre-existing behaviour. `ALTER TABLE relay_cell_admission ADD COLUMN IF NOT EXISTS roll_isolated_at BIGINT`, + // The cap the cell last reported enforcing; see cell-connection-cap.ts. Nullable with no default, + // so the rewrite is catalog-only, and a director that predates it neither reads nor clears it. + `ALTER TABLE relay_cell_connection_limits ADD COLUMN IF NOT EXISTS reported_hard_cap BIGINT`, // Dropped, not created: see the comment on relay_assignment_activity_leases. Deferrable because // this is the one boot where it has to take ACCESS EXCLUSIVE on a table under continuous write, // and all 28 directors reach it at once; a lock timeout here must not restart the instance, which @@ -1466,13 +1470,14 @@ async function applySchema(database: RelayDatabase): Promise { for (const statement of SCHEMA.split(';')) { if (statement.trim()) await database.query(statement) } - for (const [table, column] of [ - ['relay_control_capabilities', 'idle_regional_rehome'], - ['relay_region_rehome_attempts', 'source_generation'] + for (const [table, column, definition] of [ + ['relay_control_capabilities', 'idle_regional_rehome', 'BIGINT NOT NULL DEFAULT 0'], + ['relay_region_rehome_attempts', 'source_generation', 'BIGINT NOT NULL DEFAULT 0'], + ['relay_cell_connection_limits', 'reported_hard_cap', 'BIGINT'] ]) { const columns = await database.query('SELECT name FROM pragma_table_info(?)', [table]) if (!columns.some((existing) => existing.name === column)) { - await database.query(`ALTER TABLE ${table} ADD COLUMN ${column} BIGINT NOT NULL DEFAULT 0`) + await database.query(`ALTER TABLE ${table} ADD COLUMN ${column} ${definition}`) } } } diff --git a/cloud/apps/relay/src/relay.blackbox.test.ts b/cloud/apps/relay/src/relay.blackbox.test.ts index 5b24b53022b..d2a7d950f31 100644 --- a/cloud/apps/relay/src/relay.blackbox.test.ts +++ b/cloud/apps/relay/src/relay.blackbox.test.ts @@ -474,6 +474,7 @@ describe('served relay URL', () => { expect(await (await fetch(`${relayUrl}/health`)).json()).toEqual({ ok: true, connectionCapacityProtocol: 2, + capacityReportedCapMin: true, drainPaceWindowMaxMs: 1_200_000 }) expect(await (await fetch(`${relayUrl}/ready`)).json()).toEqual({ ok: true })