diff --git a/.github/scripts/check-root-directory-entries.mjs b/.github/scripts/check-root-directory-entries.mjs index 5e7dcf3f0ff..5b63326691d 100644 --- a/.github/scripts/check-root-directory-entries.mjs +++ b/.github/scripts/check-root-directory-entries.mjs @@ -13,6 +13,10 @@ function readRootEntries(sha) { return stdout.split('\0').filter(Boolean) } +// Why: the Cloud workspace import is the one reviewed root addition; it stays +// listed until it lands on main, after which the base tree carries it. +const REVIEWED_ROOT_ENTRIES = new Set(['cloud']) + function checkRootDirectoryEntries(argv) { if (argv.length !== 2) { console.error(`Usage: ${process.argv[1]} `) @@ -21,7 +25,9 @@ function checkRootDirectoryEntries(argv) { const [baseSha, headSha] = argv const baseEntries = new Set(readRootEntries(baseSha)) - const blockedEntries = readRootEntries(headSha).filter((entry) => !baseEntries.has(entry)) + const blockedEntries = readRootEntries(headSha).filter( + (entry) => !baseEntries.has(entry) && !REVIEWED_ROOT_ENTRIES.has(entry) + ) if (blockedEntries.length === 0) { console.log('Root directory guard passed: no new root-level files or folders.') diff --git a/.github/workflows/cloud-verify.yml b/.github/workflows/cloud-verify.yml new file mode 100644 index 00000000000..22724342aab --- /dev/null +++ b/.github/workflows/cloud-verify.yml @@ -0,0 +1,103 @@ +name: Cloud Verify + +on: + pull_request: + paths: + - cloud/** + - .github/workflows/cloud-*.yml + push: + branches: [main] + paths: + - cloud/** + - .github/workflows/cloud-*.yml + +permissions: + contents: read + +concurrency: + group: cloud-verify-${{ github.ref }} + cancel-in-progress: true + +defaults: + run: + working-directory: cloud + +jobs: + security: + name: Secret scan + runs-on: blacksmith-2vcpu-ubuntu-2204 + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Scan Cloud history with Gitleaks + run: >- + docker run --rm + --volume "${GITHUB_WORKSPACE}:/repo:ro" + zricethezav/gitleaks@sha256:cdbb7c955abce02001a9f6c9f602fb195b7fadc1e812065883f695d1eeaba854 + git /repo --config /repo/cloud/.gitleaks.toml + --log-opts="--all -- cloud :(glob).github/workflows/cloud-*.yml" + + - name: Scan the single-commit Cloud snapshot with TruffleHog + run: >- + docker run --rm + --volume "${GITHUB_WORKSPACE}:/repo:ro" + trufflesecurity/trufflehog@sha256:5dc064868ba7933601b5cbaea6954954d524ddd5dc6222a9667acea70068bf7d + filesystem /repo --no-verification --fail + --include-paths=/repo/cloud/.trufflehog-include-paths.txt + --exclude-paths=/repo/cloud/.trufflehog-exclude-paths.txt + + # Compiles the workspace. No Postgres service: nothing here reaches a + # database, and the service container costs ~13s of startup. + build: + runs-on: blacksmith-4vcpu-ubuntu-2204 + steps: + - uses: actions/checkout@v4 + + - uses: pnpm/action-setup@v4 + with: + package_json_file: cloud/package.json + + - uses: actions/setup-node@v4 + with: + node-version: 24 + + - run: pnpm install --frozen-lockfile + - run: pnpm build + - run: pnpm typecheck + + # Runs in parallel with build. `pnpm test` compiles the one workspace + # package it needs through the relay pretest hook, so it does not depend on + # `pnpm build` having run. + test: + runs-on: blacksmith-4vcpu-ubuntu-2204 + services: + postgres: + image: postgres:16-alpine + env: + POSTGRES_DB: orca_relay_test + POSTGRES_PASSWORD: relay_test + POSTGRES_USER: relay_test + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U relay_test -d orca_relay_test" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + env: + ORCA_RELAY_TEST_POSTGRES_URL: postgres://relay_test:relay_test@127.0.0.1:5432/orca_relay_test + steps: + - uses: actions/checkout@v4 + + - uses: pnpm/action-setup@v4 + with: + package_json_file: cloud/package.json + + - uses: actions/setup-node@v4 + with: + node-version: 24 + + - run: pnpm install --frozen-lockfile + - run: pnpm test diff --git a/.oxfmtrc.json b/.oxfmtrc.json index b2aa0deabfa..629c7a8e411 100644 --- a/.oxfmtrc.json +++ b/.oxfmtrc.json @@ -3,5 +3,6 @@ "singleQuote": true, "semi": false, "printWidth": 100, - "trailingComma": "none" + "trailingComma": "none", + "ignorePatterns": ["cloud/**"] } diff --git a/.oxlintrc.json b/.oxlintrc.json index cae0d09dd58..47ea817ebe4 100644 --- a/.oxlintrc.json +++ b/.oxlintrc.json @@ -174,5 +174,11 @@ } } ], - "ignorePatterns": ["**/node_modules", "**/dist", "**/out", "tests/e2e/.cross-version-checkouts"] + "ignorePatterns": [ + "**/node_modules", + "**/dist", + "**/out", + "cloud/**", + "tests/e2e/.cross-version-checkouts" + ] } diff --git a/README.md b/README.md index 0f99bfc877f..7a3cbe2360c 100644 --- a/README.md +++ b/README.md @@ -252,6 +252,9 @@ Pair with your desktop app to monitor and steer your agents from your phone. Want to contribute or run locally? See our [CONTRIBUTING.md](.github/CONTRIBUTING.md) guide. +The relay that pairs the mobile app with a desktop host is also in this repository under +[`cloud/`](cloud/README.md), with a separate pnpm workspace and setup guide. + Orca contributors diff --git a/cloud/.dockerignore b/cloud/.dockerignore new file mode 100644 index 00000000000..bc93f9a9f1b --- /dev/null +++ b/cloud/.dockerignore @@ -0,0 +1,15 @@ +.git/ +.github/ +.local/ +.terraform/ +node_modules/ +dist/ +coverage/ +.env +.env.local +.env.local.generated +*.log +*.tfplan +*.tfstate +*.tfstate.* + diff --git a/cloud/.editorconfig b/cloud/.editorconfig new file mode 100644 index 00000000000..c814d3002a4 --- /dev/null +++ b/cloud/.editorconfig @@ -0,0 +1,10 @@ +root = true + +[*] +charset = utf-8 +end_of_line = lf +insert_final_newline = true +indent_style = space +indent_size = 2 +trim_trailing_whitespace = true + diff --git a/cloud/.gitignore b/cloud/.gitignore new file mode 100644 index 00000000000..8105ed487f9 --- /dev/null +++ b/cloud/.gitignore @@ -0,0 +1,24 @@ +node_modules/ +dist/ +coverage/ +.turbo/ +.local/ +.env +.env.local +.env.local.generated +*.log + +# Terraform/OpenTofu local state and plans must stay out of git. +**/.terraform/ +*.tfstate +*.tfstate.* +*.tfplan +crash.log +override.tf +override.tf.json +*_override.tf +*_override.tf.json + +# Local dev signing key + SQLite live under data/ — never commit (contains a +# private key and dev PII). Prod supplies the key via ORCA_CLOUD_SIGNING_KEY_PEM. +data/ diff --git a/cloud/.gitleaks.toml b/cloud/.gitleaks.toml new file mode 100644 index 00000000000..430d17622b8 --- /dev/null +++ b/cloud/.gitleaks.toml @@ -0,0 +1,7 @@ +[extend] +useDefault = true + +[[allowlists]] +description = "Explicit Relay test signing key" +regexTarget = "secret" +regexes = ['''^test-assignment-key-with-at-least-32-bytes$'''] diff --git a/cloud/.node-version b/cloud/.node-version new file mode 100644 index 00000000000..14db15d775c --- /dev/null +++ b/cloud/.node-version @@ -0,0 +1,2 @@ +24 + diff --git a/cloud/.npmrc b/cloud/.npmrc new file mode 100644 index 00000000000..e05d6820ed5 --- /dev/null +++ b/cloud/.npmrc @@ -0,0 +1,3 @@ +engine-strict=false +package-manager-strict=true + diff --git a/cloud/.trufflehog-exclude-paths.txt b/cloud/.trufflehog-exclude-paths.txt new file mode 100644 index 00000000000..ea20399af05 --- /dev/null +++ b/cloud/.trufflehog-exclude-paths.txt @@ -0,0 +1 @@ +(^|/)cloud/apps/relay/src/postgres-idle-client-error\.test\.ts$ diff --git a/cloud/.trufflehog-include-paths.txt b/cloud/.trufflehog-include-paths.txt new file mode 100644 index 00000000000..abd82ac6354 --- /dev/null +++ b/cloud/.trufflehog-include-paths.txt @@ -0,0 +1,2 @@ +(^|/)cloud/ +(^|/)\.github/workflows/cloud-[^/]+\.yml$ diff --git a/cloud/README.md b/cloud/README.md new file mode 100644 index 00000000000..e65b41a1952 --- /dev/null +++ b/cloud/README.md @@ -0,0 +1,50 @@ +# Orca Relay + +The relay that connects the Orca mobile app to a desktop host. Phones and +desktops never talk to each other directly: each opens an outbound WebSocket +to a relay cell, the relay pairs the two sessions, and it splices frames +between them. A director assigns hosts to cells and coordinates migrations; +cells carry the user connections. + +This directory is an independent pnpm workspace inside the Orca monorepo. Run +its commands from `cloud/`, not the repository root. The source is covered by +the repository's root [MIT license](../LICENSE). + +## Packages + +- `packages/relay-contract`: the wire contract shared by the relay, the + desktop app, and the mobile app (frame shapes, close codes, admission budgets, + splice state machine). +- `apps/relay`: the relay server. The same image runs as a director or a cell + depending on `ORCA_RELAY_ROLE`. + +## Local development + +```sh +cd cloud +pnpm install +pnpm build +pnpm test +``` + +`pnpm test` runs the SQLite-backed suites. Tests that need PostgreSQL run only +when `ORCA_RELAY_TEST_POSTGRES_URL` points at a disposable PostgreSQL 16 or 17 +database, for example: + +```sh +docker run --rm -d --name orca-relay-pg -e POSTGRES_HOST_AUTH_METHOD=trust \ + -e POSTGRES_DB=orca_relay_test -p 55440:5432 postgres:16-alpine +ORCA_RELAY_TEST_POSTGRES_URL=postgres://postgres@127.0.0.1:55440/orca_relay_test \ + pnpm --filter @orca-cloud/relay test +docker rm -f orca-relay-pg +``` + +Configuration is read from environment variables validated in +`apps/relay/src/config.ts`. `ORCA_RELAY_ASSIGNMENT_SIGNING_KEY` (at least 32 +bytes) is the only required value; everything else has a local default. + +## Deployment + +The hosted relay's infrastructure, deployment workflows, and operational tooling +are maintained privately. `Cloud Verify` in this repository builds, typechecks, +lints, tests, and secret-scans the relay on every change under `cloud/`. diff --git a/cloud/apps/relay/Dockerfile b/cloud/apps/relay/Dockerfile new file mode 100644 index 00000000000..12516cbf749 --- /dev/null +++ b/cloud/apps/relay/Dockerfile @@ -0,0 +1,25 @@ +FROM node:24-alpine AS build +WORKDIR /app +RUN corepack enable +COPY package.json pnpm-lock.yaml pnpm-workspace.yaml tsconfig.base.json ./ +COPY packages/relay-contract/package.json packages/relay-contract/package.json +COPY apps/relay/package.json apps/relay/package.json +RUN pnpm install --frozen-lockfile +COPY packages/relay-contract packages/relay-contract +COPY apps/relay apps/relay +RUN pnpm --filter @orca-cloud/relay-contract build && pnpm --filter @orca-cloud/relay build + +FROM node:24-alpine AS runtime +ENV NODE_ENV=production +ENV PORT=8080 +WORKDIR /app +RUN corepack enable +COPY package.json pnpm-lock.yaml pnpm-workspace.yaml ./ +COPY packages/relay-contract/package.json packages/relay-contract/package.json +COPY apps/relay/package.json apps/relay/package.json +COPY --from=build /app/packages/relay-contract/dist packages/relay-contract/dist +COPY --from=build /app/apps/relay/dist apps/relay/dist +RUN pnpm install --prod --frozen-lockfile --filter @orca-cloud/relay... +USER node +EXPOSE 8080 +CMD ["node", "apps/relay/dist/index.js"] diff --git a/cloud/apps/relay/package.json b/cloud/apps/relay/package.json new file mode 100644 index 00000000000..4c2b2e4269c --- /dev/null +++ b/cloud/apps/relay/package.json @@ -0,0 +1,35 @@ +{ + "name": "@orca-cloud/relay", + "private": true, + "version": "0.0.0", + "type": "module", + "main": "dist/index.js", + "scripts": { + "build": "pnpm clean && tsc -p tsconfig.build.json", + "clean": "node -e \"require('fs').rmSync('dist', { recursive: true, force: true })\"", + "dev": "tsx watch src/index.ts", + "lint": "tsc -p tsconfig.json --noEmit", + "pretest": "pnpm --filter @orca-cloud/relay-contract build", + "start": "node dist/index.js", + "test": "vitest run", + "typecheck": "tsc -p tsconfig.json --noEmit" + }, + "dependencies": { + "@hono/node-server": "^1.19.14", + "@orca-cloud/relay-contract": "workspace:*", + "hono": "^4.12.27", + "jose": "^6.1.3", + "pg": "^8.22.0", + "tweetnacl": "^1.0.3", + "ws": "^8.18.3", + "zod": "^3.25.76" + }, + "devDependencies": { + "@types/node": "^24.10.0", + "@types/pg": "^8.20.0", + "@types/ws": "^8.18.1", + "tsx": "^4.21.0", + "typescript": "^5.9.3", + "vitest": "^4.0.8" + } +} diff --git a/cloud/apps/relay/src/admin-token-verifier.test.ts b/cloud/apps/relay/src/admin-token-verifier.test.ts new file mode 100644 index 00000000000..f0870dbb09c --- /dev/null +++ b/cloud/apps/relay/src/admin-token-verifier.test.ts @@ -0,0 +1,98 @@ +import { describe, expect, it } from 'vitest' +import { + RELAY_ASIA_PROOF_ADMIN_ROUTES, + RELAY_CAPACITY_ADMIN_ROUTES, + RELAY_FENCE_BROKER_ADMIN_ROUTES, + RELAY_FENCE_ADMIN_ROUTES, + RELAY_MONITOR_ADMIN_ROUTES, + relayAdminIdentityMayAccess +} from './admin-token-verifier.js' + +const mutationRoutes = [ + '/v1/admin/drain', + '/v1/admin/evacuate', + '/v1/admin/migration-complete', + '/v1/admin/migration-supersede-cell', + '/v1/admin/rebalance-dormant', + '/v1/admin/admission-selector/apply', + '/v1/admin/admission-selector/add-migration-cells', + '/v1/admin/cell-state', + '/v1/admin/cell-fence-adopt-legacy', + '/v1/admin/cell-fence-commit-legacy-adoption', + '/v1/admin/cell-fence-attest', + '/v1/admin/cell-fence-attempt-prepare', + '/v1/admin/cell-fence-attempt-start', + '/v1/admin/cell-fence-attempt-operation', + '/v1/admin/cell-fence-attempt-abort', + '/v1/admin/drain-attempt-prepare', + '/v1/admin/drain-attempt-send', + '/v1/admin/drain-attempt-receipt', + '/v1/admin/drain-attempt-recover-forward', + '/v1/admin/cell-config', + '/v1/admin/evacuate-cell', + '/v1/admin/cell-heartbeat', + '/v1/admin/regional-rehome-control', + '/v1/admin/regional-rehome-trust-probe' +] as const + +describe('Relay admin route authorization', () => { + it('allows the staging capacity identity only its transition routes', () => { + for (const route of RELAY_CAPACITY_ADMIN_ROUTES) { + expect(relayAdminIdentityMayAccess('capacity', route)).toBe(true) + } + for (const route of mutationRoutes) { + if ((RELAY_CAPACITY_ADMIN_ROUTES as readonly string[]).includes(route)) continue + expect(relayAdminIdentityMayAccess('capacity', route)).toBe(false) + } + expect(RELAY_CAPACITY_ADMIN_ROUTES).toContain('/v1/admin/cell-state') + expect(relayAdminIdentityMayAccess('capacity', '/v1/admin/evacuation-status')).toBe(false) + }) + + it('allows the Asia proof identity only its selector and status routes', () => { + for (const route of RELAY_ASIA_PROOF_ADMIN_ROUTES) { + expect(relayAdminIdentityMayAccess('asia-proof', route)).toBe(true) + } + expect(relayAdminIdentityMayAccess('asia-proof', '/v1/admin/drain')).toBe(false) + expect(relayAdminIdentityMayAccess('asia-proof', '/v1/admin/cell-state')).toBe(false) + expect(relayAdminIdentityMayAccess('asia-proof', '/v1/admin/admission-selector/apply')).toBe(false) + expect(relayAdminIdentityMayAccess('asia-proof', '/v1/admin/add-migration-cells')).toBe(false) + }) + + it('keeps the monitor identity on exact aggregate read routes', () => { + for (const route of RELAY_MONITOR_ADMIN_ROUTES) { + expect(relayAdminIdentityMayAccess('monitor', route)).toBe(true) + } + for (const route of [...mutationRoutes, ...RELAY_FENCE_ADMIN_ROUTES]) { + if ((RELAY_MONITOR_ADMIN_ROUTES as readonly string[]).includes(route)) continue + expect(relayAdminIdentityMayAccess('monitor', route)).toBe(false) + } + expect(RELAY_MONITOR_ADMIN_ROUTES).toContain('/v1/admin/evacuation-status') + }) + + it('allows only reviewed fence evidence mutations beyond aggregate reads', () => { + for (const route of RELAY_FENCE_ADMIN_ROUTES) { + expect(relayAdminIdentityMayAccess('fence', route)).toBe(true) + } + for (const route of mutationRoutes) { + if ((RELAY_FENCE_ADMIN_ROUTES as readonly string[]).includes(route)) continue + expect(relayAdminIdentityMayAccess('fence', route)).toBe(false) + } + }) + + it('allows the broker only exact fence inspection and mutation routes', () => { + for (const route of RELAY_FENCE_BROKER_ADMIN_ROUTES) { + expect(relayAdminIdentityMayAccess('fence-broker', route)).toBe(true) + } + for (const route of [...mutationRoutes, ...RELAY_FENCE_ADMIN_ROUTES]) { + if ((RELAY_FENCE_BROKER_ADMIN_ROUTES as readonly string[]).includes(route)) continue + expect(relayAdminIdentityMayAccess('fence-broker', route)).toBe(false) + } + }) + + it('rejects unknown routes for dedicated identities', () => { + for (const identity of ['capacity', 'asia-proof', 'monitor', 'fence', 'fence-broker'] as const) { + expect(relayAdminIdentityMayAccess(identity, '/v1/admin/future-mutation')).toBe(false) + expect(relayAdminIdentityMayAccess(identity, '/health')).toBe(false) + } + }) +}) diff --git a/cloud/apps/relay/src/admin-token-verifier.ts b/cloud/apps/relay/src/admin-token-verifier.ts new file mode 100644 index 00000000000..4b8ad26e695 --- /dev/null +++ b/cloud/apps/relay/src/admin-token-verifier.ts @@ -0,0 +1,202 @@ +import { createRemoteJWKSet, jwtVerify } from 'jose' +import type { RelayConfig } from './config.js' + +export const RELAY_MONITOR_ADMIN_ROUTES = [ + '/v1/admin/admission-selector/status', + '/v1/admin/cell-status', + '/v1/admin/evacuation-status', + '/v1/admin/regional-rehome-control', + '/v1/admin/runtime-status' +] as const + +export const RELAY_FENCE_ADMIN_ROUTES = [ + ...RELAY_MONITOR_ADMIN_ROUTES, + '/v1/admin/evacuation-capacity', + '/v1/admin/cell-fence-attempt-status' +] as const + +export const RELAY_CAPACITY_ADMIN_ROUTES = [ + '/v1/admin/admission-selector/status', + '/v1/admin/admission-selector/apply', + '/v1/admin/cell-state', + '/v1/admin/cell-status', + '/v1/admin/runtime-status', + '/v1/admin/drain' +] as const + +export const RELAY_ASIA_PROOF_ADMIN_ROUTES = [ + '/v1/admin/admission-selector/status', + '/v1/admin/admission-selector/apply-staging-asia-proof', + '/v1/admin/cell-status', + '/v1/admin/runtime-status' +] as const + +export const RELAY_FENCE_BROKER_ADMIN_ROUTES = [ + ...RELAY_FENCE_ADMIN_ROUTES, + '/v1/admin/cell-fence-adopt-legacy', + '/v1/admin/cell-fence-commit-legacy-adoption', + '/v1/admin/cell-fence-attest', + '/v1/admin/cell-fence-attempt-prepare', + '/v1/admin/cell-fence-attempt-start', + '/v1/admin/cell-fence-attempt-plan', + '/v1/admin/cell-fence-attempt-operation', + '/v1/admin/cell-fence-attempt-abort', + '/v1/admin/migration-supersede-cell' +] as const + +type RelayAdminIdentity = + | 'deploy' + | 'capacity' + | 'asia-proof' + | 'monitor' + | 'fence' + | 'fence-broker' + +function createGoogleServiceTokenVerifier(input: { + jwksUrl: string + audience: string + serviceAccount: string +}): (token: string) => Promise { + const jwks = createRemoteJWKSet(new URL(input.jwksUrl)) + return async (token) => { + try { + const { payload } = await jwtVerify(token, jwks, { + issuer: ['https://accounts.google.com', 'accounts.google.com'], + audience: input.audience, + algorithms: ['RS256'] + }) + return payload.email === input.serviceAccount && payload.email_verified === true + } catch { + return false + } + } +} + +function createGoogleServiceTokenIdentityVerifier(input: { + jwksUrl: string + audience: string + serviceAccounts: ReadonlyMap +}): (token: string) => Promise { + const jwks = createRemoteJWKSet(new URL(input.jwksUrl)) + return async (token) => { + try { + const { payload } = await jwtVerify(token, jwks, { + issuer: ['https://accounts.google.com', 'accounts.google.com'], + audience: input.audience, + algorithms: ['RS256'] + }) + if (payload.email_verified !== true || typeof payload.email !== 'string') return null + return input.serviceAccounts.get(payload.email) ?? null + } catch { + return null + } + } +} + +export function relayAdminIdentityMayAccess( + identity: RelayAdminIdentity, + route: string +): boolean { + if (identity === 'deploy') return route.startsWith('/v1/admin/') + if (identity === 'capacity') { + return (RELAY_CAPACITY_ADMIN_ROUTES as readonly string[]).includes(route) + } + if (identity === 'asia-proof') { + return (RELAY_ASIA_PROOF_ADMIN_ROUTES as readonly string[]).includes(route) + } + if (identity === 'monitor') { + return (RELAY_MONITOR_ADMIN_ROUTES as readonly string[]).includes(route) + } + if (identity === 'fence') { + return (RELAY_FENCE_ADMIN_ROUTES as readonly string[]).includes(route) + } + return (RELAY_FENCE_BROKER_ADMIN_ROUTES as readonly string[]).includes(route) +} + +export function createReadOnlyAdminTokenVerifier( + config: RelayConfig +): (token: string) => Promise { + const serviceAccounts = new Map() + if (config.monitorServiceAccount) serviceAccounts.set(config.monitorServiceAccount, 'monitor') + if (config.fenceServiceAccount) serviceAccounts.set(config.fenceServiceAccount, 'fence') + if (serviceAccounts.size === 0) return async () => false + const verifyIdentity = createGoogleServiceTokenIdentityVerifier({ + jwksUrl: config.adminJwksUrl, + audience: config.adminAudience, + serviceAccounts + }) + return async (token) => (await verifyIdentity(token)) !== null +} + +export function createAdminTokenVerifier( + config: RelayConfig +): (token: string, route?: string) => Promise { + const serviceAccounts = new Map([ + [config.deployServiceAccount, 'deploy'] + ]) + if (config.capacityServiceAccount) { + serviceAccounts.set(config.capacityServiceAccount, 'capacity') + } + if (config.asiaProofServiceAccount) { + serviceAccounts.set(config.asiaProofServiceAccount, 'asia-proof') + } + if (config.monitorServiceAccount) serviceAccounts.set(config.monitorServiceAccount, 'monitor') + if (config.fenceServiceAccount) serviceAccounts.set(config.fenceServiceAccount, 'fence') + if (config.fenceBrokerServiceAccount) { + serviceAccounts.set(config.fenceBrokerServiceAccount, 'fence-broker') + } + const verifyIdentity = createGoogleServiceTokenIdentityVerifier({ + jwksUrl: config.adminJwksUrl, + audience: config.adminAudience, + serviceAccounts + }) + return async (token, route) => { + const identity = await verifyIdentity(token) + return identity !== null && (!route || relayAdminIdentityMayAccess(identity, route)) + } +} + +export function createRegionalRehomeControlApplyTokenVerifier( + config: RelayConfig +): (token: string) => Promise { + return createGoogleServiceTokenVerifier({ + jwksUrl: config.adminJwksUrl, + audience: config.adminAudience, + serviceAccount: config.deployServiceAccount + }) +} + +export function createRuntimeTokenVerifier( + config: RelayConfig +): (token: string) => Promise { + if (!config.heartbeatAudience) return async () => false + return createGoogleServiceTokenVerifier({ + jwksUrl: config.adminJwksUrl, + audience: config.heartbeatAudience, + serviceAccount: config.runtimeServiceAccount + }) +} + +export function createRegionalRehomeTokenVerifier( + config: RelayConfig +): (token: string) => Promise { + if (!config.rehomeAudience || !config.rehomeDirectorServiceAccount) { + return async () => false + } + return createGoogleServiceTokenVerifier({ + jwksUrl: config.adminJwksUrl, + audience: config.rehomeAudience, + serviceAccount: config.rehomeDirectorServiceAccount + }) +} + +export function createRegionalRehomeRuntimeTokenVerifier( + config: RelayConfig +): (token: string) => Promise { + if (!config.rehomeAudience) return async () => false + return createGoogleServiceTokenVerifier({ + jwksUrl: config.adminJwksUrl, + audience: config.rehomeAudience, + serviceAccount: config.runtimeServiceAccount + }) +} diff --git a/cloud/apps/relay/src/app.ts b/cloud/apps/relay/src/app.ts new file mode 100644 index 00000000000..3df01d9e9ce --- /dev/null +++ b/cloud/apps/relay/src/app.ts @@ -0,0 +1,1850 @@ +import { + AssignmentRequestSchema, + isRelayCellConnectionHardCap, + RELAY_ADMISSION_BUDGETS, + RELAY_DEFAULT_REGION, + RELAY_MAX_CELL_CONNECTION_UNOBSERVED_BOUND, + RELAY_PROTOCOL_LIMITS, + RelayRegionSchema, + ResolveRequestSchema, + cellPlacementCeiling, + relayCellAdmissionBounds, + type RelayCellConnectionHardCap, + type RelayRegion +} from '@orca-cloud/relay-contract' +import { Hono, type Context } from 'hono' +import { SignJWT } from 'jose' +import { z } from 'zod' +import { + createAdminTokenVerifier, + createReadOnlyAdminTokenVerifier, + createRegionalRehomeControlApplyTokenVerifier, + createRegionalRehomeRuntimeTokenVerifier, + createRegionalRehomeTokenVerifier, + createRuntimeTokenVerifier +} from './admin-token-verifier.js' +import type { + CellFenceAttemptEvidence, + RelayAssignment, + RelayAssignmentStore +} from './assignment-store.js' +import { AssignmentRejectionLogWindow } from './assignment-rejection-log-window.js' +import { CELL_ADMISSION_STATES } from './cell-admission-selector.js' +import { RELAY_MAX_CELL_CAPACITY_REQUESTS, type RelayConfig } from './config.js' +import type { RelayCredentialStore } from './credential-store.js' +import { isRelayDatabaseTransientError } from './database.js' +import { googleMetadataIdentityToken } from './google-metadata-identity-token.js' +import { + RelayPublicAssignmentAdmission, + type AssignmentAdmissionRejection +} from './public-assignment-admission.js' +import { relayHostLogDigest } from './relay-host-log-digest.js' +import type { RelayRuntimeCounts } from './relay-observability.js' +import { + isRegionalRehomeTrustProbe, + probeRegionalRehomeTrust +} from './regional-rehome-trust-probe.js' +import { createRelayTokenVerifier, readBearer } from './relay-token-verifier.js' +import { stagingAsiaProofMembership } from './staging-asia-proof-admission.js' + +const RelayCellConnectionHardCapSchema = z.custom( + isRelayCellConnectionHardCap +) + +const ASSIGNMENT_REJECTION_LOG_WINDOW_MS = 10_000 +const REGION_CATALOG_CACHE_MS = 30_000 + +type AdmissionRejectionLogEntry = { + route: 'assign' | 'resolve' + lane: 'sticky' | 'placement' + hinted: boolean + relayHostId: string + reason: AssignmentAdmissionRejection +} + +export function createRelayApp( + config: RelayConfig, + operations: { + store: RelayCredentialStore + assignments: RelayAssignmentStore + drain: (graceMs: number) => void + drainHost?: (input: { + attemptId: string + userId: string + relayHostId: string + sourceAssignmentEpoch: number + graceMs: number + }) => 'accepted' | 'already-accepted' | 'host-not-connected' + regionalRehomeIdentityToken?: (audience: string) => Promise + regionalRehomeFetch?: typeof fetch + regionalRehomeTrustProbeHostExists?: (input: { + userId: string + relayHostId: string + }) => boolean + cellIncarnation?: string + isDraining?: () => boolean + runtimeCounts?: () => RelayRuntimeCounts + ready: () => Promise + recordAssignmentAdmission?: ( + outcome: 'sticky' | 'sticky-rejected' | 'placement' | 'placement-rejected' + ) => void + recordAssignmentRejectionReason?: ( + lane: 'sticky' | 'placement', + reason: AssignmentAdmissionRejection + ) => void + recordRegionRequest?: (region: RelayRegion | undefined) => void + recordRegionSelection?: (input: { + targetRegion: RelayRegion + selectedRegion?: RelayRegion + fallback: boolean + }) => void + } +): Hono { + const app = new Hono() + let regionCatalogCache: + | { expiresAt: number; value: Awaited> } + | undefined + let regionCatalogRefresh: + | Promise>> + | undefined + const regionCatalog = async () => { + const now = Date.now() + if (regionCatalogCache && regionCatalogCache.expiresAt > now) { + return regionCatalogCache.value + } + regionCatalogRefresh ??= operations.assignments.regionCatalog().then((value) => { + regionCatalogCache = { expiresAt: Date.now() + REGION_CATALOG_CACHE_MS, value } + return value + }) + try { + return await regionCatalogRefresh + } finally { + regionCatalogRefresh = undefined + } + } + const verifyRelayToken = createRelayTokenVerifier(config) + const verifyAdminToken = createAdminTokenVerifier(config) + const verifyReadOnlyAdminToken = createReadOnlyAdminTokenVerifier(config) + const verifyRegionalRehomeControlApplyToken = + createRegionalRehomeControlApplyTokenVerifier(config) + const verifyRuntimeToken = createRuntimeTokenVerifier(config) + const verifyRegionalRehomeToken = createRegionalRehomeTokenVerifier(config) + const verifyRegionalRehomeRuntimeToken = + createRegionalRehomeRuntimeTokenVerifier(config) + const regionalRehomeFetch = operations.regionalRehomeFetch ?? fetch + const regionalRehomeIdentityToken = + operations.regionalRehomeIdentityToken ?? + ((audience: string) => googleMetadataIdentityToken(audience, regionalRehomeFetch)) + const publicAssignmentAdmission = new RelayPublicAssignmentAdmission({ + maxConcurrent: config.publicAssignmentConcurrency, + maxQueued: config.publicAssignmentQueueMax, + waitMs: config.publicAssignmentWaitMs, + maxReservedConcurrent: config.publicResolveConcurrency, + reservedWaitMs: config.publicResolveWaitMs, + minIntervalMs: config.publicAssignmentRetryAfterSeconds * 1_000, + onRejected: (reason) => operations.recordAssignmentRejectionReason?.('placement', reason) + }) + const rejectPublicAssignment = (context: Context): Response => { + context.header('Retry-After', String(config.publicAssignmentRetryAfterSeconds)) + return context.json({ error: 'assignments_temporarily_unavailable' }, 503) + } + // Reconnecting hosts declare themselves and are verified against the durable + // assignment inside this bounded lane, so a placement backlog can never + // starve session recovery. Unhinted traffic never touches this lane. + const stickyRetryAfterSeconds = config.publicStickyRetryAfterSeconds ?? 2 + const stickyAssignmentAdmission = new RelayPublicAssignmentAdmission({ + maxConcurrent: config.publicStickyConcurrency ?? 1, + maxQueued: config.publicStickyQueueMax ?? 64, + waitMs: config.publicStickyWaitMs ?? 2_000, + minIntervalMs: stickyRetryAfterSeconds * 1_000, + onRejected: (reason) => operations.recordAssignmentRejectionReason?.('sticky', reason) + }) + const rejectStickyAssignment = (context: Context): Response => { + context.header('Retry-After', String(stickyRetryAfterSeconds)) + return context.json({ error: 'assignments_temporarily_unavailable' }, 503) + } + // Aggregate counters cannot separate a handful of pathological hosts from a broad + // population, so every admission rejection names its host and reason. Keyed on + // route:lane:reason rather than host, the log stays bounded under load. + const rejectionLogWindow = new AssignmentRejectionLogWindow({ + windowMs: ASSIGNMENT_REJECTION_LOG_WINDOW_MS, + onWindowClosed: ({ suppressed, sample }) => logAssignmentRejection({ ...sample, suppressed }) + }) + const logAdmissionRejection = (input: { + route: 'assign' | 'resolve' + lane: 'sticky' | 'placement' + hinted: boolean + relayHostId: string + reason: AssignmentAdmissionRejection | undefined + }): void => { + if (!input.reason) return + const entry: AdmissionRejectionLogEntry = { ...input, reason: input.reason } + if (!rejectionLogWindow.admit(`${entry.route}:${entry.lane}:${entry.reason}`, entry)) return + logAssignmentRejection(entry) + } + app.use('/v1/admin/*', async (context, next) => { + if ( + context.req.path === '/v1/admin/cell-heartbeat' || + context.req.path === '/v1/admin/cell-rehome-status' + ) { + return await next() + } + const bearer = readBearer(context.req.header('authorization')) + if (!bearer || !(await verifyAdminToken(bearer))) return await next() + if (!(await verifyAdminToken(bearer, context.req.path))) { + return context.json({ error: 'invalid_token' }, 401) + } + return await next() + }) + + // Not /healthz: Google Front End reserves that path before the container. + app.get('/health', (context) => + context.json({ ok: true, connectionCapacityProtocol: 2 }) + ) + app.get('/ready', async (context) => + (await operations.ready()) + ? context.json({ ok: true }) + : context.json({ error: 'dependency_unavailable' }, 503) + ) + app.get('/v1/regions', async (context) => { + if (config.role === 'cell') return context.json({ error: 'director_only' }, 404) + return context.json({ v: 1, regions: await regionCatalog() }) + }) + app.post('/v1/assign', async (context) => { + if (config.role === 'cell') return context.json({ error: 'director_only' }, 404) + if (!config.publicAssignmentsEnabled) return rejectPublicAssignment(context) + const bearer = readBearer(context.req.header('authorization')) + if (!bearer) return context.json({ error: 'invalid_token' }, 401) + const claims = await verifyRelayToken(bearer) + if (!claims) return context.json({ error: 'invalid_token' }, 401) + if (Number(context.req.header('content-length') ?? 0) > 4 * 1024) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AssignmentRequestSchema.safeParse(await context.req.json().catch(() => null)) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + if (body.data.relayHostId !== claims.relayHostId) { + return context.json({ error: 'host_identity_mismatch' }, 403) + } + const identity = { userId: claims.sub, relayHostId: claims.relayHostId } + const requestedRegion = body.data.preferredRegion + const targetRegion = + config.regionalPlacementEnabled !== false && requestedRegion + ? requestedRegion + : RELAY_DEFAULT_REGION + operations.recordRegionRequest?.(requestedRegion) + let admission: { release(): void } | null = null + let lane: 'sticky' | 'placement' = 'placement' + if (body.data.reconnect) { + let rejection: AssignmentAdmissionRejection | undefined + const fastLane = await stickyAssignmentAdmission.acquire(claims.relayHostId, (reason) => { + rejection = reason + }) + if (!fastLane) { + operations.recordAssignmentAdmission?.('sticky-rejected') + logAdmissionRejection({ + route: 'assign', + lane: 'sticky', + hinted: true, + relayHostId: claims.relayHostId, + reason: rejection + }) + return rejectStickyAssignment(context) + } + let verified = false + try { + verified = (await operations.assignments.resolve(identity)) !== null + } catch (error) { + fastLane.release() + operations.recordAssignmentAdmission?.('sticky-rejected') + if (isRelayDatabaseTransientError(error)) { + logAssignmentRejection({ + route: 'assign-verify', + lane: 'none', + hinted: true, + relayHostId: claims.relayHostId, + reason: operationError(error) + }) + return rejectStickyAssignment(context) + } + throw error + } + if (verified) { + lane = 'sticky' + admission = fastLane + operations.recordAssignmentAdmission?.('sticky') + } else { + // An unverified hint joins the placement lane with today's semantics. + fastLane.release() + } + } + if (!admission) { + let rejection: AssignmentAdmissionRejection | undefined + admission = await publicAssignmentAdmission.acquire(claims.relayHostId, (reason) => { + rejection = reason + }) + operations.recordAssignmentAdmission?.(admission ? 'placement' : 'placement-rejected') + if (!admission) { + logAdmissionRejection({ + route: 'assign', + lane: 'placement', + hinted: Boolean(body.data.reconnect), + relayHostId: claims.relayHostId, + reason: rejection + }) + return rejectPublicAssignment(context) + } + } + let assignment: RelayAssignment + try { + assignment = requestedRegion + ? await operations.assignments.assign(identity, requestedRegion, targetRegion) + : await operations.assignments.assign(identity) + } catch (error) { + if (isRelayAssignmentCapacityError(error) || isRelayDatabaseTransientError(error)) { + logAssignmentRejection({ + route: 'assign', + lane, + hinted: Boolean(body.data.reconnect), + relayHostId: claims.relayHostId, + reason: operationError(error) + }) + } + if (isRelayAssignmentCapacityError(error)) { + if (lane === 'placement') { + operations.recordRegionSelection?.({ targetRegion, fallback: false }) + } + return context.json({ error: operationError(error) }, 503) + } + if (isRelayDatabaseTransientError(error)) { + return lane === 'sticky' ? rejectStickyAssignment(context) : rejectPublicAssignment(context) + } + throw error + } finally { + admission.release() + } + operations.recordRegionSelection?.({ + targetRegion, + selectedRegion: assignment.region, + fallback: lane === 'placement' && assignment.region !== targetRegion + }) + // Grant-side counterpart of the rejection log: reconnect grants are rare + // enough to log and make "which cell is this host on" answerable. + if (lane === 'sticky') { + console.warn( + `[orca-relay] assignment granted lane=sticky host=${relayHostLogDigest(claims.relayHostId)}` + + ` cell=${assignment.cellId}` + ) + } + const lease = await new SignJWT({ + purpose: 'cell-assignment', + cellId: assignment.cellId, + cellUrl: assignment.cellUrl, + assignmentEpoch: assignment.assignmentEpoch, + relayHostId: claims.relayHostId + }) + .setProtectedHeader({ alg: 'HS256' }) + .setIssuer(config.publicUrl) + .setAudience('orca-relay-cell') + .setSubject(claims.sub) + .setIssuedAt() + .setExpirationTime('5m') + .sign(config.assignmentSigningKey) + return context.json({ + v: 1, + cellUrl: assignment.cellUrl, + assignmentEpoch: assignment.assignmentEpoch, + lease + }) + }) + app.post('/v1/resolve', async (context) => { + if (config.role === 'cell') return context.json({ error: 'director_only' }, 404) + if (!config.publicAssignmentsEnabled) return rejectPublicAssignment(context) + if (Number(context.req.header('content-length') ?? 0) > RELAY_PROTOCOL_LIMITS.maxHttpBodyBytes) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = ResolveRequestSchema.safeParse(await context.req.json().catch(() => null)) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + let rejection: AssignmentAdmissionRejection | undefined + const admission = await publicAssignmentAdmission.acquireReserved( + body.data.relayHostId, + (reason) => { + rejection = reason + } + ) + if (!admission) { + logAdmissionRejection({ + route: 'resolve', + lane: 'placement', + hinted: false, + relayHostId: body.data.relayHostId, + reason: rejection + }) + return rejectPublicAssignment(context) + } + try { + const resolved = await operations.store.resolveResume( + body.data.relayHostId, + body.data.resumeToken + ) + if (!resolved) return context.json({ error: 'invalid_credential' }, 401) + const identity = { + userId: resolved.userId, + relayHostId: body.data.relayHostId + } + // This also migrates credentials created by the staging-only combined service. + const assignment = + (await operations.assignments.resolve(identity)) ?? + (await operations.assignments.assign(identity)) + return context.json({ + v: 1, + cellUrl: assignment.cellUrl, + assignmentEpoch: assignment.assignmentEpoch, + leaseExpiresAt: assignment.leaseExpiresAt + }) + } catch (error) { + if (isRelayAssignmentCapacityError(error) || isRelayDatabaseTransientError(error)) { + logAssignmentRejection({ + route: 'resolve', + lane: 'none', + hinted: false, + relayHostId: body.data.relayHostId, + reason: operationError(error) + }) + } + if (isRelayAssignmentCapacityError(error)) { + return context.json({ error: operationError(error) }, 503) + } + if (isRelayDatabaseTransientError(error)) return rejectPublicAssignment(context) + throw error + } finally { + admission.release() + } + }) + app.post('/v1/admin/drain', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + const body = z + .object({ v: z.literal(1), graceMs: z.number().int().nonnegative().max(60 * 60 * 1000) }) + .strict() + .safeParse(await context.req.json().catch(() => null)) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + operations.drain(body.data.graceMs) + return context.json({ ok: true }) + }) + app.post('/v1/admin/host-drain', async (context) => { + if (config.role !== 'cell' || !operations.drainHost) { + return context.json({ error: 'cell_only' }, 404) + } + const bearer = readBearer(context.req.header('authorization')) + if (!bearer || !(await verifyRegionalRehomeToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = RegionalHostDrainSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + if ( + body.data.sourceCellId !== config.cellId || + !operations.cellIncarnation || + body.data.sourceCellIncarnation !== operations.cellIncarnation + ) { + return context.json({ error: 'regional_rehome_source_generation_mismatch' }, 409) + } + try { + const trustProbe = isRegionalRehomeTrustProbe(body.data) + let sharedRuntimeIdentityRejected: true | undefined + if (trustProbe) { + const runtimeToken = await regionalRehomeIdentityToken(config.rehomeAudience!) + if ( + !(await verifyRegionalRehomeRuntimeToken(runtimeToken)) || + (await verifyRegionalRehomeToken(runtimeToken)) + ) { + throw new Error('regional_rehome_shared_runtime_rejection_not_proven') + } + if ( + !operations.regionalRehomeTrustProbeHostExists || + operations.regionalRehomeTrustProbeHostExists(body.data) + ) { + throw new Error('regional_rehome_trust_probe_host_not_absent') + } + sharedRuntimeIdentityRejected = true + } + const outcome = operations.drainHost(body.data) + return context.json({ + v: 1, + outcome, + ...(sharedRuntimeIdentityRejected ? { sharedRuntimeIdentityRejected } : {}) + }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/runtime-status', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = RuntimeStatusSchema.safeParse(await context.req.json().catch(() => null)) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + return context.json({ + v: 1, + role: config.role, + cellId: config.cellId, + cellUrl: config.cellUrl, + region: config.region ?? RELAY_DEFAULT_REGION, + imageDigest: config.imageDigest ?? null, + draining: operations.isDraining?.() ?? false, + regionalRehomeProtocol: + config.rehomeAudience && config.rehomeDirectorServiceAccount ? 1 : 0, + connectionCapacity: + config.connectionHardCap === undefined + ? null + : { + hardCap: config.connectionHardCap, + controlRebindReserve: RELAY_ADMISSION_BUDGETS.reservedHostControls, + ordinaryConnectionLimit: relayCellAdmissionBounds(config.connectionHardCap) + .socketAdmissionCeiling, + unobservedBound: config.connectionUnobservedBound!, + normalAdmissionPause: cellPlacementCeiling( + config.connectionHardCap, + config.connectionUnobservedBound! + ) + }, + runtime: operations.runtimeCounts?.() ?? null + }) + }) + app.post('/v1/admin/cell-heartbeat', async (context) => { + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + const bearer = readBearer(context.req.header('authorization')) + if (!bearer || !(await verifyRuntimeToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = CellHeartbeatSchema.safeParse(await context.req.json().catch(() => null)) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + await operations.assignments.recordCellHeartbeat(body.data) + return context.json({ ok: true }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/cell-rehome-status', async (context) => { + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + const bearer = readBearer(context.req.header('authorization')) + if (!bearer || !(await verifyRuntimeToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = CellRegionalRehomeStatusSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + await operations.assignments.recordCellRegionalRehomeStatus(body.data) + return context.json({ ok: true }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/regional-rehome-control', async (context) => { + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + const bearer = readBearer(context.req.header('authorization')) + if (!bearer || !(await verifyAdminToken(bearer, context.req.path))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = RegionalRehomeControlSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + if (body.data.action === 'inspect') { + const control = await operations.assignments.inspectRegionalRehomeControl() + return context.json({ v: 1, control }) + } + if (!(await verifyRegionalRehomeControlApplyToken(bearer))) { + return context.json({ error: 'insufficient_permission' }, 403) + } + try { + const control = await operations.assignments.applyRegionalRehomeControl(body.data) + return context.json({ v: 1, control }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/regional-rehome-trust-probe', async (context) => { + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + const bearer = readBearer(context.req.header('authorization')) + if (!bearer || !(await verifyAdminToken(bearer, context.req.path))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = RegionalRehomeTrustProbeSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + if (!config.rehomeAudience || !config.rehomeDirectorServiceAccount) { + return context.json({ error: 'regional_rehome_trust_not_configured' }, 409) + } + try { + const source = await operations.assignments.cellDeploymentStatus( + body.data.sourceCellId + ) + if ( + source.region !== RELAY_DEFAULT_REGION || + !source.runtime || + source.runtime.cellIncarnation !== body.data.sourceCellIncarnation || + !source.runtime.ready || + !source.runtime.heartbeatFresh || + source.runtime.regionalRehomeProtocol < 1 + ) { + throw new Error('regional_rehome_trust_probe_source_unavailable') + } + const result = await probeRegionalRehomeTrust({ + sourceCellUrl: source.cellUrl, + sourceCellId: body.data.sourceCellId, + sourceCellIncarnation: body.data.sourceCellIncarnation, + audience: config.rehomeAudience, + identityToken: regionalRehomeIdentityToken, + fetch: regionalRehomeFetch + }) + return context.json(result) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/evacuate', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminAssignmentMoveSchema.safeParse(await context.req.json().catch(() => null)) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const migration = await operations.assignments.startEvacuation( + { userId: body.data.userId, relayHostId: body.data.relayHostId }, + body.data.targetCellId + ) + return context.json({ v: 1, migration }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/migration-complete', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminMigrationCompleteSchema.safeParse(await context.req.json().catch(() => null)) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + await operations.assignments.completeEvacuation( + { userId: body.data.userId, relayHostId: body.data.relayHostId }, + body.data.assignmentEpoch + ) + return context.json({ ok: true }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/migration-supersede-cell', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminRegisteredCellMigrationSupersedeSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const superseded = await operations.assignments.supersedeRegisteredCellEvacuations( + body.data.sourceCellId, + body.data.currentTargetCellId, + body.data.replacementTargetCellId, + body.data.limit + ) + return context.json({ v: 1, superseded }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/rebalance-dormant', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminAssignmentMoveSchema.safeParse(await context.req.json().catch(() => null)) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const assignment = await operations.assignments.rebalanceDormant( + { userId: body.data.userId, relayHostId: body.data.relayHostId }, + body.data.targetCellId + ) + return context.json({ v: 1, assignment }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/admission-selector/apply', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminAdmissionSelectorApplySchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const result = await operations.assignments.applyCellAdmissionSelector(body.data) + return context.json({ v: 1, ...result }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/admission-selector/apply-staging-asia-proof', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminStagingAsiaProofAdmissionApplySchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const current = await operations.assignments.inspectCellAdmissionSelector() + const result = await operations.assignments.applyCellAdmissionSelector({ + attemptId: body.data.attemptId, + expectedGeneration: body.data.expectedGeneration, + membership: stagingAsiaProofMembership(current.selector.membership, body.data.state) + }) + return context.json({ v: 1, ...result }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/admission-selector/status', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminAdmissionSelectorStatusSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const result = await operations.assignments.inspectCellAdmissionSelector( + body.data.attemptId + ) + return context.json({ v: 1, ...result }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/admission-selector/add-migration-cells', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminAdmissionSelectorAddMigrationCellsSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const result = await operations.assignments.addMigrationCells({ + attemptId: body.data.attemptId, + expectedGeneration: body.data.expectedGeneration, + cells: body.data.cells.map((cell) => ({ + id: cell.cellId, + url: cell.cellUrl, + capacityRequests: cell.capacityRequests, + region: cell.region, + connectionHardCap: cell.connectionHardCap, + connectionUnobservedBound: cell.connectionUnobservedBound + })) + }) + return context.json({ v: 1, ...result }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/cell-state', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminCellStateSchema.safeParse(await context.req.json().catch(() => null)) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + await operations.assignments.setCellAdmissionState( + body.data.cellId, + 'state' in body.data + ? body.data.state + : body.data.enabled + ? 'general' + : 'existing-only' + ) + return context.json({ ok: true }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/cell-fence-adopt-legacy', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminCellFenceLegacyAdoptionSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const expiresAt = await operations.assignments.adoptLegacyCellFence( + body.data.cellId, + body.data.cellIncarnation + ) + return context.json({ v: 1, cellId: body.data.cellId, expiresAt }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/cell-fence-commit-legacy-adoption', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminCellFenceLegacyAdoptionCommitSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + await operations.assignments.commitLegacyCellFenceAdoption( + body.data.cellId, + body.data.cellIncarnation + ) + return context.json({ v: 1, cellId: body.data.cellId, committed: true }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/cell-fence-attest', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminCellFenceAttestSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const evidence = cellFenceAttemptEvidence(body.data) + const result = await operations.assignments.attestCellFenceAttempt( + evidence, + body.data.gceOperation + ) + return context.json({ + v: 1, + cellId: body.data.cellId, + expiresAt: result.expiresAt, + attempt: result.attempt + }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/cell-fence-attempt-prepare', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminCellFenceAttemptPrepareSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const evidence = cellFenceAttemptEvidence(body.data) + const attempt = await operations.assignments.prepareCellFenceAttempt(evidence) + return context.json({ v: 1, attempt }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/cell-fence-attempt-start', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminCellFenceAttemptUpdateSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const evidence = cellFenceAttemptEvidence(body.data) + const result = await operations.assignments.startCellFenceApply( + evidence, + body.data.invocationId, + body.data.invocationRequestReason + ) + return context.json({ v: 1, ...result }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/cell-fence-attempt-plan', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminCellFencePlanSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const evidence = cellFenceAttemptEvidence(body.data) + const attempt = await operations.assignments.bindCellFencePlanGeneration( + evidence, + body.data.planObjectGeneration + ) + return context.json({ v: 1, attempt }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/cell-fence-attempt-operation', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminCellFenceOperationSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const evidence = cellFenceAttemptEvidence(body.data) + const result = await operations.assignments.recordCellFenceOperation( + evidence, + body.data.invocationId, + body.data.invocationRequestReason, + body.data.gceOperation + ) + return context.json({ v: 1, ...result }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/cell-fence-attempt-status', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminCellFenceAttemptStatusSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const attempt = await operations.assignments.cellFenceAttempt(body.data.cellId) + return context.json({ v: 1, attempt }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/cell-fence-attempt-abort', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminCellFenceAttemptAbortSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const evidence = cellFenceAttemptEvidence(body.data) + const attempt = await operations.assignments.abortCellFenceAttempt(evidence) + return context.json({ v: 1, attempt }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/drain-attempt-prepare', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminDrainAttemptPrepareSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const result = await operations.assignments.prepareCellDrainAttempt({ + attemptId: body.data.attemptId, + cellId: body.data.cellId, + cellIncarnation: body.data.cellIncarnation, + traceValue: body.data.traceValue, + plannedGraceMs: body.data.graceMs + }) + return context.json({ v: 1, ...result }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/drain-attempt-send', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminDrainAttemptMutationSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const attempt = await operations.assignments.beginCellDrainSend(body.data) + return context.json({ v: 1, attempt }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/drain-attempt-receipt', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminDrainAttemptReceiptSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const attempt = await operations.assignments.recordCellDrainApplicationReceipt( + body.data + ) + return context.json({ v: 1, attempt }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/drain-attempt-recover-forward', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminDrainAttemptRecoverSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const result = await operations.assignments.prepareCellDrainRecovery(body.data) + return context.json({ v: 1, ...result }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/cell-config', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminCellConfigSchema.safeParse(await context.req.json().catch(() => null)) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + await operations.assignments.configureCell( + { + id: body.data.cellId, + url: body.data.cellUrl, + capacityRequests: body.data.capacityRequests, + connectionHardCap: body.data.connectionHardCap, + connectionUnobservedBound: body.data.connectionUnobservedBound + }, + 'state' in body.data ? body.data.state : body.data.enabled + ) + return context.json({ ok: true }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/evacuate-cell', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminCellEvacuateSchema.safeParse(await context.req.json().catch(() => null)) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const started = await operations.assignments.startActiveCellEvacuations( + body.data.sourceCellId, + body.data.targetCellId, + body.data.limit + ) + return context.json({ v: 1, started }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/evacuation-capacity', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminCellEvacuationCapacitySchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const capacity = await operations.assignments.cellEvacuationCapacity( + body.data.sourceCellId, + body.data.targetCellId + ) + return context.json({ v: 1, ...capacity }) + } catch (error) { + return context.json({ error: operationError(error) }, 409) + } + }) + app.post('/v1/admin/evacuation-status', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminCellEvacuationStatusSchema.safeParse( + await context.req.json().catch(() => null) + ) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + if (body.data.completeReady && (await verifyReadOnlyAdminToken(bearer))) { + return context.json({ error: 'insufficient_permission' }, 403) + } + const status = await operations.assignments.cellEvacuationStatus( + body.data.sourceCellId, + body.data.targetCellId, + body.data.completeReady + ) + return context.json({ v: 1, ...status }) + }) + app.post('/v1/admin/cell-status', async (context) => { + const bearer = readBearer(context.req.header('authorization')) + if (config.role !== 'director') return context.json({ error: 'director_only' }, 404) + if (!bearer || !(await verifyAdminToken(bearer))) { + return context.json({ error: 'invalid_token' }, 401) + } + if (requestTooLarge(context.req.header('content-length'))) { + return context.json({ error: 'request_too_large' }, 413) + } + const body = AdminCellStatusSchema.safeParse(await context.req.json().catch(() => null)) + if (!body.success) return context.json({ error: 'invalid_request' }, 400) + try { + const status = await operations.assignments.cellDeploymentStatus(body.data.cellId) + return context.json({ v: 1, status }) + } catch (error) { + return context.json({ error: operationError(error) }, 404) + } + }) + return app +} + +const AdminAssignmentMoveSchema = z + .object({ + v: z.literal(1), + userId: z.string().min(1).max(256), + relayHostId: z.string().regex(/^[A-Za-z0-9_-]{16}$/), + targetCellId: z.string().min(1).max(128) + }) + .strict() + +const RuntimeStatusSchema = z.object({ v: z.literal(1) }).strict() + +const RegionalRehomeSafetySchema = z + .object({ + observedAt: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER), + sqlFailures: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER), + reconnects: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER), + controlActivityRecoveryFailures: z + .number() + .int() + .nonnegative() + .max(Number.MAX_SAFE_INTEGER), + databasePoolWaiting: z.number().int().nonnegative().max(100), + databasePoolWaitersMax: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER), + databasePoolWaitMsMax: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER), + // Cells spread the full pool-pressure counts into the payload; rejecting + // the extra gauges 400'd every rehome-status heartbeat since 2026-08-15, + // so no cell ever recorded rehome protocol 1. + databasePoolTotal: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER).optional(), + databasePoolIdle: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER).optional(), + databasePoolOldestWaitMs: z + .number() + .int() + .nonnegative() + .max(Number.MAX_SAFE_INTEGER) + .optional() + }) + .strict() + +const CellHeartbeatSchema = z + .object({ + v: z.literal(1), + cellId: z.string().min(1).max(128), + cellUrl: z.string().url().max(2_048).refine(isCanonicalRelayOrigin), + region: RelayRegionSchema.default(RELAY_DEFAULT_REGION), + cellIncarnation: z.string().uuid(), + startedAt: z.number().int().positive().max(Number.MAX_SAFE_INTEGER), + ready: z.boolean(), + observedRequests: z.number().int().nonnegative().max(RELAY_MAX_CELL_CAPACITY_REQUESTS), + totalConnections: z + .number() + .int() + .nonnegative() + .max(RELAY_MAX_CELL_CAPACITY_REQUESTS) + .optional(), + inFlightConnections: z + .number() + .int() + .nonnegative() + .max(RELAY_MAX_CELL_CAPACITY_REQUESTS) + .optional(), + reservedConnectionUnits: z + .number() + .int() + .nonnegative() + .max(RELAY_MAX_CELL_CAPACITY_REQUESTS) + .optional(), + enforcedConnectionUnits: z + .number() + .int() + .nonnegative() + .max(RELAY_MAX_CELL_CAPACITY_REQUESTS) + .optional(), + connectionInclusionWatermark: z + .number() + .int() + .nonnegative() + .max(Number.MAX_SAFE_INTEGER) + .optional(), + connectionHardCap: RelayCellConnectionHardCapSchema.optional(), + connectionUnobservedBound: z + .number() + .int() + .nonnegative() + .max(RELAY_MAX_CELL_CONNECTION_UNOBSERVED_BOUND) + .optional() + }) + .strict() + .superRefine((value, context) => { + const values = [ + value.totalConnections, + value.inFlightConnections, + value.reservedConnectionUnits, + value.enforcedConnectionUnits, + value.connectionHardCap, + value.connectionUnobservedBound + ] + if ( + values.some((candidate) => candidate !== undefined) && + values.some((candidate) => candidate === undefined) + ) { + context.addIssue({ + code: 'custom', + message: 'connection telemetry must be complete' + }) + } + if ( + value.enforcedConnectionUnits !== undefined && + value.totalConnections !== undefined && + value.inFlightConnections !== undefined && + value.reservedConnectionUnits !== undefined && + value.enforcedConnectionUnits !== + value.totalConnections + + value.inFlightConnections + + value.reservedConnectionUnits + ) { + context.addIssue({ + code: 'custom', + message: 'connection telemetry sum does not match' + }) + } + if ( + value.connectionHardCap !== undefined && + value.connectionUnobservedBound! > + relayCellAdmissionBounds(value.connectionHardCap).maxUnobservedBound + ) { + context.addIssue({ + code: 'custom', + message: 'connection unobserved bound must leave ordinary admission capacity' + }) + } + }) + +const CellRegionalRehomeStatusSchema = z + .object({ + v: z.literal(1), + cellId: z.string().min(1).max(128), + cellIncarnation: z.string().uuid(), + regionalRehomeProtocol: z.number().int().min(0).max(1), + safety: RegionalRehomeSafetySchema + }) + .strict() + +const RegionalRehomeControlSchema = z.discriminatedUnion('action', [ + z.object({ v: z.literal(1), action: z.literal('inspect') }).strict(), + z.object({ + v: z.literal(1), + action: z.literal('apply'), + expectedGeneration: z.number().int().nonnegative(), + enabled: z.boolean(), + notBefore: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER), + ratePerMinute: z.number().int().min(1).max(120), + preferenceMaxAgeMs: z + .number() + .int() + .min(60_000) + .max(30 * 24 * 60 * 60_000), + drainGraceMs: z.number().int().min(60_000).max(60 * 60_000), + confirmation: z.enum([ + 'ENABLE_REGIONAL_REHOMING', + 'DISABLE_REGIONAL_REHOMING' + ]) + }).strict() +]).superRefine((value, context) => { + if (value.action !== 'apply') return + const expected = value.enabled + ? 'ENABLE_REGIONAL_REHOMING' + : 'DISABLE_REGIONAL_REHOMING' + if (value.confirmation !== expected) { + context.addIssue({ code: 'custom', message: 'confirmation does not match state' }) + } +}) + +const RegionalRehomeTrustProbeSchema = z + .object({ + v: z.literal(1), + sourceCellId: z.string().min(1).max(128), + sourceCellIncarnation: z.string().uuid() + }) + .strict() + +const AdminMigrationCompleteSchema = z + .object({ + v: z.literal(1), + userId: z.string().min(1).max(256), + relayHostId: z.string().regex(/^[A-Za-z0-9_-]{16}$/), + assignmentEpoch: z.number().int().positive().max(Number.MAX_SAFE_INTEGER) + }) + .strict() + +const AdminRegisteredCellMigrationSupersedeSchema = z + .object({ + v: z.literal(1), + sourceCellId: z.string().min(1).max(128), + currentTargetCellId: z.string().min(1).max(128), + replacementTargetCellId: z.string().min(1).max(128), + limit: z.number().int().min(1).max(100), + confirmation: z.literal('SUPERSEDE_REGISTERED_CELL_MIGRATIONS') + }) + .strict() + .refine( + (value) => + new Set([ + value.sourceCellId, + value.currentTargetCellId, + value.replacementTargetCellId + ]).size === 3 + ) + +const CellIdSchema = z.string().min(1).max(128) +const CellAdmissionStateSchema = z.enum(CELL_ADMISSION_STATES) +const AdmissionSelectorAttemptIdSchema = z.string().regex(/^[A-Za-z0-9_-]{8,128}$/) +const AdmissionSelectorMembershipSchema = z + .object({ + existingOnly: z.array(CellIdSchema).max(256), + migrationOnly: z.array(CellIdSchema).max(256), + general: z.array(CellIdSchema).max(256) + }) + .strict() + +const AdminAdmissionSelectorApplySchema = z + .object({ + v: z.literal(1), + attemptId: AdmissionSelectorAttemptIdSchema, + expectedGeneration: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER), + expectedMembershipSha256: z.string().regex(/^[a-f0-9]{64}$/).optional(), + membership: AdmissionSelectorMembershipSchema + }) + .strict() + .refine( + (value) => value.expectedGeneration > 0 || value.expectedMembershipSha256 !== undefined + ) + +const AdminStagingAsiaProofAdmissionApplySchema = z + .object({ + v: z.literal(1), + attemptId: AdmissionSelectorAttemptIdSchema, + expectedGeneration: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER), + state: z.enum(['general', 'migration-only']) + }) + .strict() + +const AdminAdmissionSelectorStatusSchema = z + .object({ v: z.literal(1), attemptId: AdmissionSelectorAttemptIdSchema.optional() }) + .strict() + +const AdminAdmissionSelectorAddMigrationCellsSchema = z + .object({ + v: z.literal(1), + attemptId: AdmissionSelectorAttemptIdSchema, + expectedGeneration: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER), + cells: z + .array( + z + .object({ + cellId: CellIdSchema, + cellUrl: z.string().url().max(2_048).refine(isCanonicalRelayOrigin), + capacityRequests: z + .number() + .int() + .positive() + .max(RELAY_MAX_CELL_CAPACITY_REQUESTS), + region: RelayRegionSchema.default(RELAY_DEFAULT_REGION), + connectionHardCap: RelayCellConnectionHardCapSchema, + connectionUnobservedBound: z + .number() + .int() + .nonnegative() + .max(RELAY_MAX_CELL_CONNECTION_UNOBSERVED_BOUND) + }) + .strict() + .superRefine((value, context) => { + if ( + value.connectionUnobservedBound > + relayCellAdmissionBounds(value.connectionHardCap).maxUnobservedBound + ) { + context.addIssue({ + code: 'custom', + path: ['connectionUnobservedBound'], + message: 'connection unobserved bound must leave ordinary admission capacity' + }) + } + }) + ) + .min(1) + .max(128) + }) + .strict() + .refine( + ({ cells }) => + new Set(cells.map(({ cellId }) => cellId)).size === cells.length && + new Set(cells.map(({ cellUrl }) => cellUrl)).size === cells.length + ) + +const AdminCellStateSchema = z.union([ + z.object({ v: z.literal(1), cellId: CellIdSchema, enabled: z.boolean() }).strict(), + z.object({ v: z.literal(1), cellId: CellIdSchema, state: CellAdmissionStateSchema }).strict() +]) + +const CellConfigShape = { + v: z.literal(1), + cellId: CellIdSchema, + cellUrl: z.string().url().max(2_048).refine(isCanonicalRelayOrigin), + capacityRequests: z.number().int().positive().max(RELAY_MAX_CELL_CAPACITY_REQUESTS), + connectionHardCap: RelayCellConnectionHardCapSchema.optional(), + connectionUnobservedBound: z + .number() + .int() + .nonnegative() + .max(RELAY_MAX_CELL_CONNECTION_UNOBSERVED_BOUND) + .optional() +} as const + +const AdminCellConfigSchema = z + .union([ + z + .object({ + ...CellConfigShape, + enabled: z.boolean() + }) + .strict(), + z + .object({ + ...CellConfigShape, + state: CellAdmissionStateSchema + }) + .strict() + ]) + .refine( + (value) => + (value.connectionHardCap === undefined) === + (value.connectionUnobservedBound === undefined), + { message: 'connection hard cap and unobserved bound must be configured together' } + ) + .refine( + (value) => + value.connectionHardCap === undefined || + value.connectionUnobservedBound! <= + relayCellAdmissionBounds(value.connectionHardCap).maxUnobservedBound, + { message: 'connection unobserved bound must leave ordinary admission capacity' } + ) + +const CellPairShape = { + v: z.literal(1), + sourceCellId: z.string().min(1).max(128), + targetCellId: z.string().min(1).max(128) +} as const + +const AdminCellEvacuateSchema = z + .object({ ...CellPairShape, limit: z.number().int().min(1).max(100) }) + .strict() + .refine((value) => value.sourceCellId !== value.targetCellId) + +const AdminCellEvacuationCapacitySchema = z + .object(CellPairShape) + .strict() + .refine((value) => value.sourceCellId !== value.targetCellId) + +const AdminCellEvacuationStatusSchema = z + .object({ ...CellPairShape, completeReady: z.boolean().default(false) }) + .strict() + .refine((value) => value.sourceCellId !== value.targetCellId) + +const TerraformStateLineageSchema = z + .string() + .regex(/^[0-9a-f]{8}(?:-[0-9a-f]{4}){3}-[0-9a-f]{12}$/i) + +const CellFenceAttemptBaseShape = { + attemptId: z.string().uuid(), + environment: z.enum(['staging', 'production']), + cellId: z.string().min(1).max(128), + cellIncarnation: z.string().uuid(), + migName: z.string().min(1).max(128), + instanceGroup: z.string().url().max(2048), + generationIdentity: z.string().url().max(2048), + fenceCommit: z.string().regex(/^[a-f0-9]{40}$/), + planSha256: z.string().regex(/^[a-f0-9]{64}$/), + planObjectName: z + .string() + .regex(/^terraform\/state\/relay-fence-plans\/(?:staging|production)\/[0-9a-f-]{36}\.tfplan$/), + varFileSha256: z.string().regex(/^[a-f0-9]{64}$/), + terraformStateLineage: TerraformStateLineageSchema, + terraformStateSerial: z.number().int().nonnegative().safe(), + terraformStateObjectGeneration: z.string().regex(/^[1-9][0-9]{0,30}$/), + terraformStateObjectSha256: z.string().regex(/^[a-f0-9]{64}$/), + requestReason: z + .string() + .regex(/^orca-relay-fence\/[0-9a-f]{8}-[0-9a-f-]{27}$/) +} as const +const CellFenceAttemptEvidenceShape = { + ...CellFenceAttemptBaseShape, + planObjectGeneration: z.string().regex(/^[1-9][0-9]{0,30}$/) +} as const + +const AdminCellFenceAttemptPrepareSchema = z + .object({ + v: z.literal(1), + ...CellFenceAttemptBaseShape, + confirmation: z.literal('PREPARE_TERRAFORM_CELL_FENCE') + }) + .strict() + +const AdminCellFencePlanSchema = z + .object({ + v: z.literal(1), + ...CellFenceAttemptEvidenceShape, + confirmation: z.literal('BIND_TERRAFORM_CELL_FENCE_PLAN') + }) + .strict() + +const AdminCellFenceAttemptUpdateSchema = z + .object({ + v: z.literal(1), + ...CellFenceAttemptEvidenceShape, + invocationId: z.string().uuid(), + invocationRequestReason: z.string().min(1).max(256), + confirmation: z.literal('START_TERRAFORM_CELL_FENCE') + }) + .strict() + +const AdminCellFenceOperationSchema = z + .object({ + v: z.literal(1), + ...CellFenceAttemptEvidenceShape, + invocationId: z.string().uuid(), + invocationRequestReason: z.string().min(1).max(256), + gceOperation: z.string().min(1).max(256), + confirmation: z.literal('RECORD_TERRAFORM_CELL_FENCE_OPERATION') + }) + .strict() + +const AdminCellFenceAttemptStatusSchema = z + .object({ v: z.literal(1), cellId: z.string().min(1).max(128) }) + .strict() + +const AdminCellFenceLegacyAdoptionSchema = z + .object({ + v: z.literal(1), + cellId: z.string().min(1).max(128), + cellIncarnation: z.string().uuid(), + confirmation: z.literal('ADOPT_LEGACY_TERRAFORM_CELL_FENCE') + }) + .strict() + +const AdminCellFenceLegacyAdoptionCommitSchema = z + .object({ + v: z.literal(1), + cellId: z.string().min(1).max(128), + cellIncarnation: z.string().uuid(), + confirmation: z.literal('COMMIT_LEGACY_TERRAFORM_CELL_FENCE_ADOPTION') + }) + .strict() + +const AdminCellFenceAttemptAbortSchema = z + .object({ + v: z.literal(1), + ...CellFenceAttemptBaseShape, + confirmation: z.literal('ABORT_UNSTARTED_TERRAFORM_CELL_FENCE') + }) + .strict() + +const AdminCellFenceAttestSchema = z + .object({ + v: z.literal(1), + ...CellFenceAttemptEvidenceShape, + gceOperation: z.string().min(1).max(256), + confirmation: z.literal('ATTEST_TERRAFORM_FENCED_CELL') + }) + .strict() + +const AdminDrainAttemptPrepareSchema = z + .object({ + v: z.literal(1), + attemptId: z.string().uuid(), + cellId: z.string().min(1).max(128), + cellIncarnation: z.string().uuid(), + traceValue: z.string().uuid(), + graceMs: z.literal(120_000), + confirmation: z.literal('PREPARE_LEGACY_DRAIN') + }) + .strict() + +const AdminDrainAttemptMutationSchema = z + .object({ + v: z.literal(1), + attemptId: z.string().uuid(), + cellId: z.string().min(1).max(128), + cellIncarnation: z.string().uuid() + }) + .strict() + +const AdminDrainAttemptReceiptSchema = AdminDrainAttemptMutationSchema.extend({ + traceValue: z.string().uuid(), + backendStatus: z.number().int().min(200).max(299), + backendInstance: z.string().min(1).max(256).optional() +}).strict() + +const AdminDrainAttemptRecoverSchema = z + .object({ + v: z.literal(1), + cellId: z.string().min(1).max(128), + cellIncarnation: z.string().uuid(), + confirmation: z.literal('RECOVER_LEGACY_DRAIN') + }) + .strict() + +const RegionalHostDrainSchema = z + .object({ + v: z.literal(1), + attemptId: z.string().uuid(), + userId: z.string().min(1).max(256), + relayHostId: z.string().regex(/^[A-Za-z0-9_-]{16}$/), + sourceCellId: z.string().min(1).max(128), + sourceCellIncarnation: z.string().uuid(), + sourceAssignmentEpoch: z.number().int().positive(), + graceMs: z.number().int().nonnegative().max(60 * 60 * 1000) + }) + .strict() + +const AdminCellStatusSchema = z + .object({ v: z.literal(1), cellId: z.string().min(1).max(128) }) + .strict() + +function cellFenceAttemptEvidence( + value: CellFenceAttemptEvidence +): CellFenceAttemptEvidence { + return { + attemptId: value.attemptId, + environment: value.environment, + cellId: value.cellId, + cellIncarnation: value.cellIncarnation, + migName: value.migName, + instanceGroup: value.instanceGroup, + generationIdentity: value.generationIdentity, + fenceCommit: value.fenceCommit, + planSha256: value.planSha256, + planObjectName: value.planObjectName, + planObjectGeneration: value.planObjectGeneration, + varFileSha256: value.varFileSha256, + terraformStateLineage: value.terraformStateLineage, + terraformStateSerial: value.terraformStateSerial, + terraformStateObjectGeneration: value.terraformStateObjectGeneration, + terraformStateObjectSha256: value.terraformStateObjectSha256, + requestReason: value.requestReason + } +} + +function operationError(error: unknown): string { + return error instanceof Error ? error.message : 'operation_failed' +} + +export { relayHostLogDigest } + +// `suppressed` is present only on a window-closing line, and counts the rejections +// that line stands for beyond the one already logged when the window opened. +function logAssignmentRejection(input: { + route: 'assign' | 'assign-verify' | 'resolve' + lane: 'sticky' | 'placement' | 'none' + hinted: boolean + relayHostId: string + reason: string + suppressed?: number +}): void { + console.warn( + `[orca-relay] assignment rejected route=${input.route} lane=${input.lane}` + + ` hinted=${input.hinted} reason=${input.reason}` + + ` host=${relayHostLogDigest(input.relayHostId)}` + + (input.suppressed === undefined ? '' : ` suppressed=${input.suppressed}`) + ) +} + +function isRelayAssignmentCapacityError(error: unknown): boolean { + return ( + error instanceof Error && + ['relay_capacity_exhausted', 'relay_connection_headroom_exhausted'].includes( + error.message + ) + ) +} + +function isCanonicalRelayOrigin(value: string): boolean { + const url = new URL(value) + const loopback = ['127.0.0.1', 'localhost', '::1', '[::1]'].includes(url.hostname) + return ( + url.origin === value && + url.pathname === '/' && + (url.protocol === 'https:' || (loopback && url.protocol === 'http:')) + ) +} + +function requestTooLarge(contentLength: string | undefined): boolean { + return Number(contentLength ?? 0) > RELAY_PROTOCOL_LIMITS.maxHttpBodyBytes +} diff --git a/cloud/apps/relay/src/assignment-cleanup-steps.test.ts b/cloud/apps/relay/src/assignment-cleanup-steps.test.ts new file mode 100644 index 00000000000..02c21f76ada --- /dev/null +++ b/cloud/apps/relay/src/assignment-cleanup-steps.test.ts @@ -0,0 +1,67 @@ +import { describe, expect, it, vi } from 'vitest' +import { + assignmentCleanupSteps, + runAssignmentCleanup, + type AssignmentCleanupStore +} from './assignment-cleanup-steps.js' + +function stubStore(overrides: Partial = {}) { + const calls: string[] = [] + const method = (name: string) => + vi.fn(async () => { + calls.push(name) + }) + const store: AssignmentCleanupStore = { + refreshRegionalRehomeLeases: method('refreshRegionalRehomeLeases'), + completeReadyEvacuations: method('completeReadyEvacuations'), + completeReadyRegionalRehomes: method('completeReadyRegionalRehomes'), + abortExpiredEvacuations: method('abortExpiredEvacuations'), + abortExpiredRegionalRehomes: method('abortExpiredRegionalRehomes'), + reapRegionalRehomeAttempts: method('reapRegionalRehomeAttempts'), + releaseExpiredActivityLeases: method('releaseExpiredActivityLeases'), + releaseExpiredActivity: method('releaseExpiredActivity'), + releaseExpiredRegionPreferences: method('releaseExpiredRegionPreferences'), + evacuateDeadCells: method('evacuateDeadCells'), + ...overrides + } + return { store, calls } +} + +describe('assignment cleanup steps', () => { + it('runs every later sweep when an early one fails, naming the step', async () => { + const { store, calls } = stubStore({ + completeReadyRegionalRehomes: vi.fn(async () => { + throw new Error('regional_rehome_assignment_mismatch') + }) + }) + const warn = vi.fn() + + await runAssignmentCleanup(store, warn) + + expect(calls).toEqual([ + 'refreshRegionalRehomeLeases', + 'completeReadyEvacuations', + 'abortExpiredEvacuations', + 'abortExpiredRegionalRehomes', + 'reapRegionalRehomeAttempts', + 'releaseExpiredActivityLeases', + 'releaseExpiredActivity', + 'releaseExpiredRegionPreferences', + 'evacuateDeadCells' + ]) + expect(warn).toHaveBeenCalledTimes(1) + expect(String(warn.mock.calls[0]![0])).toContain( + '[orca-relay] assignment cleanup failed: complete-ready-regional-rehomes' + ) + }) + + it('covers all ten sweeps exactly once per run', async () => { + const { store, calls } = stubStore() + + await runAssignmentCleanup(store) + + expect(calls).toHaveLength(10) + expect(new Set(calls).size).toBe(10) + expect(assignmentCleanupSteps(store)).toHaveLength(10) + }) +}) diff --git a/cloud/apps/relay/src/assignment-cleanup-steps.ts b/cloud/apps/relay/src/assignment-cleanup-steps.ts new file mode 100644 index 00000000000..0b6547d2ae4 --- /dev/null +++ b/cloud/apps/relay/src/assignment-cleanup-steps.ts @@ -0,0 +1,52 @@ +import { runRelayBackgroundOperation } from './relay-background-operation.js' + +// The ten periodic assignment sweeps the director runs every 30s. Each step +// re-derives its state from the database and is idempotent, so they carry no +// intra-tick ordering dependency — which is what makes per-step isolation +// sound: one failing sweep costs one tick of itself, never the other nine. +// (A single poisoned rehome row once silenced the whole chained form +// fleet-wide.) Sweep failures are logged, never fed into the rehome worker's +// dispatch-failure budget: a sweep exception is not a dispatch failure and +// must not durably disable regional rehoming. +export type AssignmentCleanupStore = { + refreshRegionalRehomeLeases(): Promise + completeReadyEvacuations(): Promise + completeReadyRegionalRehomes(): Promise + abortExpiredEvacuations(): Promise + abortExpiredRegionalRehomes(): Promise + reapRegionalRehomeAttempts(): Promise + releaseExpiredActivityLeases(): Promise + releaseExpiredActivity(): Promise + releaseExpiredRegionPreferences(): Promise + evacuateDeadCells(): Promise +} + +export function assignmentCleanupSteps( + assignments: AssignmentCleanupStore +): ReadonlyArray Promise]> { + return [ + ['refresh-regional-rehome-leases', () => assignments.refreshRegionalRehomeLeases()], + ['complete-ready-evacuations', () => assignments.completeReadyEvacuations()], + ['complete-ready-regional-rehomes', () => assignments.completeReadyRegionalRehomes()], + ['abort-expired-evacuations', () => assignments.abortExpiredEvacuations()], + ['abort-expired-regional-rehomes', () => assignments.abortExpiredRegionalRehomes()], + ['reap-regional-rehome-attempts', () => assignments.reapRegionalRehomeAttempts()], + ['release-expired-activity-leases', () => assignments.releaseExpiredActivityLeases()], + ['release-expired-activity', () => assignments.releaseExpiredActivity()], + ['release-expired-region-preferences', () => assignments.releaseExpiredRegionPreferences()], + ['evacuate-dead-cells', () => assignments.evacuateDeadCells()] + ] +} + +export async function runAssignmentCleanup( + assignments: AssignmentCleanupStore, + warn?: (message: string) => void +): Promise { + for (const [step, operation] of assignmentCleanupSteps(assignments)) { + await runRelayBackgroundOperation( + operation, + `[orca-relay] assignment cleanup failed: ${step}`, + warn + ) + } +} diff --git a/cloud/apps/relay/src/assignment-connection-headroom-postgres.test.ts b/cloud/apps/relay/src/assignment-connection-headroom-postgres.test.ts new file mode 100644 index 00000000000..6ac9521c3d6 --- /dev/null +++ b/cloud/apps/relay/src/assignment-connection-headroom-postgres.test.ts @@ -0,0 +1,342 @@ +import pg from 'pg' +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 { + openRelayDatabase, + type RelayDatabase +} from './database.js' + +const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL +const describePostgres = databaseUrl ? describe : describe.skip +const headroomIndexName = 'relay_control_connection_reservation_headroom' +const cell = { + id: 'connection-headroom-postgres', + url: 'https://connection-headroom-postgres.example.com', + capacityRequests: 1_000, + connectionHardCap: 600 as const, + connectionUnobservedBound: 50 +} + +describePostgres('PostgreSQL assignment connection headroom', () => { + const databases: RelayDatabase[] = [] + + beforeAll(async () => { + databases.push( + await openRelayDatabase({ databaseUrl, dataDir: '' }), + await openRelayDatabase({ databaseUrl, dataDir: '' }), + await openRelayDatabase({ databaseUrl, dataDir: '' }) + ) + }) + + afterAll(async () => { + const database = databases[0] + if (database) { + await database.query( + `DELETE FROM relay_control_connection_reservations + WHERE user_id LIKE 'connection-headroom-postgres-%'` + ) + await database.query( + `DELETE FROM relay_assignment_activity_leases + WHERE user_id LIKE 'connection-headroom-postgres-%'` + ) + await database.query( + `DELETE FROM relay_assignments + WHERE user_id LIKE 'connection-headroom-postgres-%'` + ) + await database.query( + `DELETE FROM relay_cell_connection_runtime WHERE cell_id = ?`, + [cell.id] + ) + await database.query( + `DELETE FROM relay_cell_connection_limits WHERE cell_id = ?`, + [cell.id] + ) + await database.query(`DELETE FROM relay_cell_runtime WHERE cell_id = ?`, [cell.id]) + await database.query(`DELETE FROM relay_cells WHERE cell_id = ?`, [cell.id]) + } + for (const connection of databases) await connection.close() + }) + + it('commits only one parallel assignment at the admission boundary', async () => { + const stores = databases.map( + (database) => + new RelayAssignmentStore(database, () => 100, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + ) + await databases[0]!.query( + `DELETE FROM relay_control_connection_reservations + WHERE user_id LIKE 'connection-headroom-postgres-%'` + ) + await databases[0]!.query( + `DELETE FROM relay_assignment_activity_leases + WHERE user_id LIKE 'connection-headroom-postgres-%'` + ) + await databases[0]!.query( + `DELETE FROM relay_assignments + WHERE user_id LIKE 'connection-headroom-postgres-%'` + ) + await stores[0]!.reconcileCells([cell], false) + await stores[0]!.recordCellHeartbeat({ + cellId: cell.id, + cellUrl: cell.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: 449, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 449, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }) + for (const suffix of ['1', '2']) { + await databases[0]!.query( + `INSERT INTO relay_assignments + (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, + last_activity_at, reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + [ + `connection-headroom-postgres-${suffix}`, + `headroomhost000${suffix}`, + cell.id, + 1, + 90_100, + 100, + 0, + 0, + 0, + 0, + 0, + 0 + ] + ) + } + + const results = await Promise.allSettled([ + stores[1]!.assign({ + userId: 'connection-headroom-postgres-1', + relayHostId: 'headroomhost0001' + }), + stores[2]!.assign({ + userId: 'connection-headroom-postgres-2', + relayHostId: 'headroomhost0002' + }) + ]) + expect(results.filter(({ status }) => status === 'fulfilled')).toHaveLength(1) + expect(results.filter(({ status }) => status === 'rejected')).toHaveLength(1) + + const assignments = await databases[0]!.query( + `SELECT COALESCE(SUM(reserved_controls), 0) AS count FROM relay_assignments + WHERE user_id LIKE 'connection-headroom-postgres-%'` + ) + const pending = await databases[0]!.query( + `SELECT COUNT(*) AS count FROM relay_assignment_activity_leases + WHERE user_id LIKE 'connection-headroom-postgres-%' + AND activity_kind = 'control' + AND activity_id LIKE 'control-pending:%'` + ) + const reservations = await databases[0]!.query( + `SELECT COUNT(*) AS count FROM relay_control_connection_reservations + WHERE user_id LIKE 'connection-headroom-postgres-%' + AND state <> 'released'` + ) + expect(Number(assignments[0]!.count)).toBe(1) + expect(Number(pending[0]!.count)).toBe(1) + expect(Number(reservations[0]!.count)).toBe(1) + }, 15_000) + + it('uses the composite headroom index when released history dominates', async () => { + await databases[0]!.query( + `INSERT INTO relay_control_connection_reservations + (reservation_id, idempotency_key, user_id, relay_host_id, + assignment_epoch, cell_id, state, created_at, timeout_at, updated_at) + SELECT 'headroom-plan-' || value, 'headroom-plan-' || value, + 'connection-headroom-postgres-plan', 'plan-host-' || value, + 1, ?, 'released', 1, 2, 1 + FROM generate_series(1, 50000) AS value`, + [cell.id] + ) + await databases[0]!.query(`ANALYZE relay_control_connection_reservations`) + + const client = new pg.Client({ connectionString: databaseUrl }) + await client.connect() + let reservationIndexPlan: Record | undefined + try { + const plan = await client.query( + `EXPLAIN (FORMAT JSON) ${ASSIGNMENT_CONNECTION_HEADROOM_QUERY}` + ) + reservationIndexPlan = findReservationHeadroomIndexPlan( + plan.rows[0]?.['QUERY PLAN'] + ) + } finally { + await client.end() + } + + expect(reservationIndexPlan).toBeDefined() + expect(String(reservationIndexPlan?.['Index Cond'])).toContain('cell_id') + expect(String(reservationIndexPlan?.['Index Cond'])).toContain('state = ANY') + expect(reservationIndexPlan?.['Filter']).toBeUndefined() + }) + + it('deduplicates expired debt after a fresh PostgreSQL snapshot', async () => { + const now = 200 + const store = new RelayAssignmentStore(databases[0]!, () => now, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + const identity = { + userId: 'connection-headroom-postgres-debt', + relayHostId: 'headroomdebt0001' + } + await databases[0]!.query( + `INSERT INTO relay_control_connection_reservations + (reservation_id, idempotency_key, user_id, relay_host_id, + assignment_epoch, cell_id, state, inclusion_watermark, + claim_activity_id, created_at, timeout_at, claimed_at, released_at, + updated_at) + VALUES + (?, ?, ?, ?, 1, ?, 'late-arrival-debt', NULL, NULL, 100, 150, NULL, NULL, 150), + (?, ?, ?, ?, 1, ?, 'late-arrival-debt', NULL, NULL, 101, 150, NULL, NULL, 150)`, + [ + 'postgres-debt-1', + 'postgres-debt-1', + identity.userId, + identity.relayHostId, + cell.id, + 'postgres-debt-2', + 'postgres-debt-2', + identity.userId, + identity.relayHostId, + cell.id + ] + ) + + await store.recordCellHeartbeat({ + cellId: cell.id, + cellUrl: cell.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: 449, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 449, + connectionInclusionWatermark: 10, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }) + + expect( + await databases[0]!.query( + `SELECT state + FROM relay_control_connection_reservations + WHERE user_id = ? + ORDER BY created_at`, + [identity.userId] + ) + ).toEqual([ + { state: 'late-arrival-debt' }, + { state: 'released' } + ]) + }) + + it('releases an aborted older epoch after a fresh PostgreSQL snapshot', async () => { + const now = 300 + const store = new RelayAssignmentStore(databases[0]!, () => now, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + const identity = { + userId: 'connection-headroom-postgres-aborted', + relayHostId: 'headroomabort001' + } + await databases[0]!.query( + `INSERT INTO relay_assignments + (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, + last_activity_at, reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases) + VALUES (?, ?, ?, 3, ?, ?, 0, 0, 0, 0, 0, 0)`, + [identity.userId, identity.relayHostId, cell.id, now, now] + ) + await databases[0]!.query( + `INSERT INTO relay_assignment_migrations + (user_id, relay_host_id, source_cell_id, target_cell_id, + previous_epoch, assignment_epoch, source_request_units, + target_reserved_units, expires_at, target_registered_at, + completed_at, aborted_at, created_at, updated_at) + VALUES (?, ?, ?, ?, 1, 2, 0, 1, ?, ?, NULL, ?, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + 'source', + cell.id, + now - 2, + now - 3, + now - 1, + now - 4, + now - 1 + ] + ) + await databases[0]!.query( + `INSERT INTO relay_control_connection_reservations + (reservation_id, idempotency_key, user_id, relay_host_id, + assignment_epoch, cell_id, state, inclusion_watermark, + claim_activity_id, created_at, timeout_at, claimed_at, released_at, + updated_at) + VALUES (?, ?, ?, ?, 2, ?, 'late-arrival-debt', NULL, NULL, ?, ?, NULL, NULL, ?)`, + [ + 'postgres-aborted-reservation', + 'postgres-aborted-reservation', + identity.userId, + identity.relayHostId, + cell.id, + now - 4, + now - 2, + now - 1 + ] + ) + + await store.recordCellHeartbeat({ + cellId: cell.id, + cellUrl: cell.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: 449, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 449, + connectionInclusionWatermark: 20, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }) + + expect( + await databases[0]!.query( + `SELECT state FROM relay_control_connection_reservations + WHERE reservation_id = ?`, + ['postgres-aborted-reservation'] + ) + ).toEqual([{ state: 'released' }]) + }) +}) + +function findReservationHeadroomIndexPlan( + value: unknown +): Record | undefined { + if (value === null || typeof value !== 'object') return undefined + const record = value as Record + if (!Array.isArray(value) && record['Index Name'] === headroomIndexName) return record + for (const child of Object.values(value)) { + const match = findReservationHeadroomIndexPlan(child) + if (match) return match + } + return undefined +} diff --git a/cloud/apps/relay/src/assignment-connection-headroom-query.ts b/cloud/apps/relay/src/assignment-connection-headroom-query.ts new file mode 100644 index 00000000000..c0237aecf26 --- /dev/null +++ b/cloud/apps/relay/src/assignment-connection-headroom-query.ts @@ -0,0 +1,15 @@ +export const ASSIGNMENT_CONNECTION_HEADROOM_QUERY = + `SELECT limits.cell_id, limits.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, + current_runtime.cell_incarnation AS current_incarnation, + (SELECT COUNT(*) FROM relay_control_connection_reservations reservation + WHERE reservation.cell_id = limits.cell_id + AND reservation.state IN + ('reserved', 'late-arrival-debt', 'claimed')) AS outstanding_reservations + FROM relay_cell_connection_limits limits + LEFT JOIN relay_cell_connection_snapshots connection_snapshot + ON connection_snapshot.cell_id = limits.cell_id + LEFT JOIN relay_cell_runtime current_runtime + ON current_runtime.cell_id = limits.cell_id` diff --git a/cloud/apps/relay/src/assignment-connection-headroom.test.ts b/cloud/apps/relay/src/assignment-connection-headroom.test.ts new file mode 100644 index 00000000000..19f5d46aa33 --- /dev/null +++ b/cloud/apps/relay/src/assignment-connection-headroom.test.ts @@ -0,0 +1,1007 @@ +import { ASSIGNMENT_LIMITS } from '@orca-cloud/relay-contract' +import { afterEach, describe, expect, it } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.js' +import type { RelayCellConfig } from './config.js' +import { + openInMemoryRelayDatabase, + type RelayDatabase +} from './database.js' + +const LIMITED_CELL: RelayCellConfig = { + id: 'limited', + url: 'https://limited.example.com', + capacityRequests: 1_000, + connectionHardCap: 600, + connectionUnobservedBound: 50 +} + +const databases: RelayDatabase[] = [] + +afterEach(async () => { + for (const database of databases.splice(0)) await database.close() +}) + +async function setup( + connectionUnits: number, + now: () => number = () => 100 +): Promise<{ database: RelayDatabase; store: RelayAssignmentStore }> { + const database = await openInMemoryRelayDatabase() + databases.push(database) + const store = new RelayAssignmentStore(database, now, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + await store.reconcileCells([LIMITED_CELL], true) + await store.recordCellHeartbeat({ + cellId: LIMITED_CELL.id, + cellUrl: LIMITED_CELL.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: connectionUnits, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: connectionUnits, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }) + return { database, store } +} + +async function setupHeadroomReassignment(): Promise<{ + database: RelayDatabase + store: RelayAssignmentStore + source: RelayCellConfig + target: RelayCellConfig +}> { + const database = await openInMemoryRelayDatabase() + databases.push(database) + const store = new RelayAssignmentStore(database, () => 100, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + const source = { + ...LIMITED_CELL, + id: 'saturated-source', + url: 'https://saturated-source.example.com' + } + const target = { + ...LIMITED_CELL, + id: 'available-target', + url: 'https://available-target.example.com' + } + await store.reconcileCells([source, target], true) + await store.recordCellHeartbeat({ + cellId: source.id, + cellUrl: source.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: 450, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 450, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }) + await store.recordCellHeartbeat({ + cellId: target.id, + cellUrl: target.url, + cellIncarnation: '22222222-2222-4222-8222-222222222222', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }) + return { database, store, source, target } +} + +describe('relay assignment connection headroom', () => { + it('requires exact, internally consistent telemetry from limited cells', async () => { + const database = await openInMemoryRelayDatabase() + databases.push(database) + const store = new RelayAssignmentStore(database, () => 100) + await store.reconcileCells([LIMITED_CELL], true) + const heartbeat = { + cellId: LIMITED_CELL.id, + cellUrl: LIMITED_CELL.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0 + } + + await expect(store.recordCellHeartbeat(heartbeat)).rejects.toThrow( + 'cell_connection_telemetry_mismatch' + ) + await expect( + store.recordCellHeartbeat({ + ...heartbeat, + totalConnections: 550, + inFlightConnections: 2, + reservedConnectionUnits: 3, + enforcedConnectionUnits: 554, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }) + ).rejects.toThrow('cell_connection_telemetry_mismatch') + await expect( + store.recordCellHeartbeat({ + ...heartbeat, + totalConnections: 601, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 601, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }) + ).resolves.toBeUndefined() + }) + + it('reports the limited-cell capacity contract and telemetry components', async () => { + const { store } = await setup(0) + await store.recordCellHeartbeat({ + cellId: LIMITED_CELL.id, + cellUrl: LIMITED_CELL.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: 120, + inFlightConnections: 2, + reservedConnectionUnits: 3, + enforcedConnectionUnits: 125, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }) + + expect((await store.cellDeploymentStatus(LIMITED_CELL.id)).connectionCapacity).toEqual({ + hardCap: 600, + controlRebindReserve: 100, + ordinaryConnectionLimit: 500, + unobservedBound: 50, + normalAdmissionPause: 450, + observedConnections: 120, + inFlightConnections: 2, + reservedConnectionUnits: 3, + enforcedConnectionUnits: 125, + pendingControlReservations: 0, + heartbeatFresh: true + }) + }) + + it('fails placement closed while a cell changes connection limits', async () => { + const { store } = await setup(449) + const expanded = { + ...LIMITED_CELL, + connectionHardCap: 1_000 as const + } + + await store.reconcileCells([expanded], true) + await expect( + store.assign({ userId: 'transition-user', relayHostId: 'host000000000099' }) + ).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, + cellUrl: expanded.url, + cellIncarnation: '22222222-2222-4222-8222-222222222222', + startedAt: 200, + ready: true, + observedRequests: 0, + totalConnections: 449, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 449, + connectionHardCap: 1_000, + connectionUnobservedBound: 50 + }) + await expect(store.cellDeploymentStatus(expanded.id)).resolves.toMatchObject({ + connectionCapacity: { + hardCap: 1_000, + ordinaryConnectionLimit: 900, + normalAdmissionPause: 850, + heartbeatFresh: true + } + }) + }) + + it('fails placement closed while the admin path changes connection limits', async () => { + const { store } = await setup(449) + const expanded = { + ...LIMITED_CELL, + connectionHardCap: 1_000 as const + } + + await store.configureCell(expanded, '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 store.recordCellHeartbeat({ + cellId: expanded.id, + cellUrl: expanded.url, + cellIncarnation: '22222222-2222-4222-8222-222222222222', + startedAt: 200, + ready: true, + observedRequests: 0, + totalConnections: 449, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 449, + connectionHardCap: 1_000, + connectionUnobservedBound: 50 + }) + await expect( + store.assign({ userId: 'admin-restored', relayHostId: 'host000000000097' }) + ).resolves.toMatchObject({ cellId: expanded.id }) + }) + + it('admits exactly one reservation at the admission boundary', async () => { + const { database, store } = await setup(449) + const results = await Promise.allSettled([ + store.assign({ userId: 'user-1', relayHostId: 'host000000000001' }), + store.assign({ userId: 'user-2', relayHostId: 'host000000000002' }) + ]) + + expect(results.filter(({ status }) => status === 'fulfilled')).toHaveLength(1) + expect( + results.filter(({ status }) => status === 'rejected').map((result) => + result.status === 'rejected' ? result.reason : null + ) + ).toEqual([expect.objectContaining({ message: 'relay_capacity_exhausted' })]) + expect( + await database.query( + `SELECT activity_id FROM relay_assignment_activity_leases + WHERE activity_kind = 'control' AND activity_id LIKE 'control-pending:%'` + ) + ).toHaveLength(1) + }) + + it('rejects at the boundary before mutating durable assignment state', async () => { + const { database, store } = await setup(450) + + await expect( + store.assign({ userId: 'user-1', relayHostId: 'host000000000001' }) + ).rejects.toThrow('relay_capacity_exhausted') + expect(await database.query(`SELECT * FROM relay_assignments`)).toEqual([]) + expect(await database.query(`SELECT * FROM relay_assignment_activity_leases`)).toEqual([]) + }) + + it('charges every active reservation state at the admission boundary', async () => { + const { database, store } = await setup(447) + await database.query( + `INSERT INTO relay_control_connection_reservations + (reservation_id, idempotency_key, user_id, relay_host_id, + assignment_epoch, cell_id, state, created_at, timeout_at, updated_at) + VALUES + ('active-reserved', 'active-reserved', 'state-user', 'state-host-1', + 1, ?, 'reserved', 1, 2, 1), + ('active-debt', 'active-debt', 'state-user', 'state-host-2', + 1, ?, 'late-arrival-debt', 1, 2, 1), + ('active-claimed', 'active-claimed', 'state-user', 'state-host-3', + 1, ?, 'claimed', 1, 2, 1)`, + [LIMITED_CELL.id, LIMITED_CELL.id, LIMITED_CELL.id] + ) + + await expect( + store.assign({ userId: 'blocked-user', relayHostId: 'blockedhost00001' }) + ).rejects.toThrow('relay_capacity_exhausted') + expect(await database.query(`SELECT * FROM relay_assignment_activity_leases`)).toEqual([]) + }) + + it('does not charge released reservation history', async () => { + const { database, store } = await setup(449) + await database.query( + `INSERT INTO relay_control_connection_reservations + (reservation_id, idempotency_key, user_id, relay_host_id, + assignment_epoch, cell_id, state, created_at, timeout_at, updated_at) + VALUES ('released-history', 'released-history', 'state-user', 'state-host', + 1, ?, 'released', 1, 2, 1)`, + [LIMITED_CELL.id] + ) + + await expect( + store.assign({ userId: 'admitted-user', relayHostId: 'admittedhost0001' }) + ).resolves.toMatchObject({ cellId: LIMITED_CELL.id }) + }) + + it('rejects sticky control restoration before mutating assignment state', async () => { + const { database, store } = await setup(450) + const identity = { userId: 'sticky-user', relayHostId: 'stickyhost000001' } + await database.query( + `INSERT INTO relay_assignments + (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, + last_activity_at, reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + LIMITED_CELL.id, + 1, + 10_000, + 100, + 0, + 0, + 1, + 0, + 0, + 0 + ] + ) + + await expect(store.assign(identity)).rejects.toThrow( + 'relay_connection_headroom_exhausted' + ) + expect( + await database.query( + `SELECT cell_id, assignment_epoch, reserved_controls, reserved_invites + FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([ + { + cell_id: LIMITED_CELL.id, + assignment_epoch: 1, + reserved_controls: 0, + reserved_invites: 1 + } + ]) + expect(await database.query(`SELECT * FROM relay_assignment_activity_leases`)).toEqual([]) + }) + + it('moves a zero-activity assignment off a general cell without connection headroom', async () => { + const { database, store, source, target } = await setupHeadroomReassignment() + const identity = { userId: 'movable-user', relayHostId: 'movablehost000001' } + await database.query( + `INSERT INTO relay_assignments + (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, + last_activity_at, reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases) + VALUES (?, ?, ?, ?, ?, ?, 0, 0, 0, 0, 0, 0)`, + [identity.userId, identity.relayHostId, source.id, 7, 10_000, 100] + ) + await database.query( + `INSERT INTO relay_control_connection_reservations + (reservation_id, idempotency_key, user_id, relay_host_id, + assignment_epoch, cell_id, state, inclusion_watermark, + claim_activity_id, created_at, timeout_at, claimed_at, released_at, + updated_at) + VALUES (?, ?, ?, ?, ?, ?, 'late-arrival-debt', NULL, NULL, ?, ?, NULL, NULL, ?)`, + [ + 'superseded-reservation', + 'superseded-reservation', + identity.userId, + identity.relayHostId, + 7, + source.id, + 50, + 99, + 100 + ] + ) + + await expect(store.assign(identity)).resolves.toMatchObject({ + cellId: target.id, + assignmentEpoch: 8 + }) + expect( + await database.query( + `SELECT cell_id, assignment_epoch, reserved_controls + FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ cell_id: target.id, assignment_epoch: 8, reserved_controls: 1 }]) + expect( + await database.query( + `SELECT cell_id, activity_kind FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ cell_id: target.id, activity_kind: 'control' }]) + expect( + await database.query( + `SELECT cell_id, assignment_epoch, state + FROM relay_control_connection_reservations + WHERE user_id = ? AND relay_host_id = ? + ORDER BY assignment_epoch`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([ + { cell_id: source.id, assignment_epoch: 7, state: 'released' }, + { cell_id: target.id, assignment_epoch: 8, state: 'reserved' } + ]) + }) + + it('keeps non-control activity pinned when its cell has no connection headroom', async () => { + const { database, store, source } = await setupHeadroomReassignment() + const identity = { userId: 'pinned-user', relayHostId: 'pinnedhost0000001' } + await database.query( + `INSERT INTO relay_assignments + (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, + last_activity_at, reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases) + VALUES (?, ?, ?, ?, ?, ?, 0, 0, 1, 0, 0, 0)`, + [identity.userId, identity.relayHostId, source.id, 7, 10_000, 100] + ) + + await expect(store.assign(identity)).rejects.toThrow( + 'relay_connection_headroom_exhausted' + ) + expect(await store.resolve(identity)).toMatchObject({ + cellId: source.id, + assignmentEpoch: 7 + }) + }) + + it.each(['existing-only', 'migration-only'] as const)( + 'keeps a zero-activity assignment pinned on a %s cell without connection headroom', + async (admission) => { + const { database, store, source } = await setupHeadroomReassignment() + const identity = { + userId: `pinned-${admission}-user`, + relayHostId: `pinned${admission.replace('-', '')}` + } + await store.setCellAdmissionState(source.id, admission) + await database.query( + `INSERT INTO relay_assignments + (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, + last_activity_at, reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases) + VALUES (?, ?, ?, ?, ?, ?, 0, 0, 0, 0, 0, 0)`, + [identity.userId, identity.relayHostId, source.id, 7, 10_000, 100] + ) + + await expect(store.assign(identity)).rejects.toThrow( + 'relay_connection_headroom_exhausted' + ) + expect(await store.resolve(identity)).toMatchObject({ + cellId: source.id, + assignmentEpoch: 7 + }) + expect(await database.query(`SELECT * FROM relay_assignment_activity_leases`)).toEqual([]) + } + ) + + it('renames a pending reservation exactly once on control activation', async () => { + const { database, store } = await setup(449) + const identity = { userId: 'user-1', relayHostId: 'host000000000001' } + const assignment = await store.assign(identity) + + await store.activateControl(identity, { + cellId: LIMITED_CELL.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + await store.activateControl(identity, { + cellId: LIMITED_CELL.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + + expect( + await database.query( + `SELECT activity_id, request_units FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ activity_id: 'control:limited:1', request_units: 1 }]) + }) + + it('keeps expired pending reservations charged as late-arrival debt', async () => { + let now = 100 + const { database, store } = await setup(449, () => now) + await store.assign({ userId: 'user-1', relayHostId: 'host000000000001' }) + now += ASSIGNMENT_LIMITS.activityLeaseMs + 1 + await store.releaseExpiredActivityLeases() + await store.recordCellHeartbeat({ + cellId: LIMITED_CELL.id, + cellUrl: LIMITED_CELL.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: 449, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 449, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }) + + await expect( + store.assign({ userId: 'user-2', relayHostId: 'host000000000002' }) + ).rejects.toThrow('relay_capacity_exhausted') + expect( + await database.query( + `SELECT activity_id FROM relay_assignment_activity_leases + WHERE activity_kind = 'control' AND activity_id LIKE 'control-pending:%'` + ) + ).toHaveLength(0) + expect( + await database.query( + `SELECT state FROM relay_control_connection_reservations` + ) + ).toEqual([{ state: 'late-arrival-debt' }]) + }) + + it('reconciles late-arrival debt only after a covering absolute snapshot', async () => { + let now = 100 + const { database, store } = await setup(449, () => now) + const identity = { userId: 'late-user', relayHostId: 'latehost00000001' } + const assignment = await store.assign(identity) + now += ASSIGNMENT_LIMITS.activityLeaseMs + 1 + await store.releaseExpiredActivityLeases() + await store.recordCellHeartbeat({ + cellId: LIMITED_CELL.id, + cellUrl: LIMITED_CELL.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: 449, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 449, + connectionInclusionWatermark: 4, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }) + + await store.activateControl(identity, { + cellId: LIMITED_CELL.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1, + connectionInclusionWatermark: 5 + }) + expect( + await database.query( + `SELECT state, inclusion_watermark FROM relay_control_connection_reservations` + ) + ).toEqual([{ state: 'claimed', inclusion_watermark: 5 }]) + + await store.recordCellHeartbeat({ + cellId: LIMITED_CELL.id, + cellUrl: LIMITED_CELL.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: 450, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 450, + connectionInclusionWatermark: 5, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }) + expect( + await database.query( + `SELECT state, released_at FROM relay_control_connection_reservations` + ) + ).toEqual([{ state: 'released', released_at: now }]) + await expect( + store.assign({ userId: 'blocked-user', relayHostId: 'blockedhost00001' }) + ).rejects.toThrow('relay_capacity_exhausted') + }) + + it('rejects duplicate and out-of-order absolute snapshots', async () => { + const { database, store } = await setup(0) + const heartbeat = { + cellId: LIMITED_CELL.id, + cellUrl: LIMITED_CELL.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: 10, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 10, + connectionInclusionWatermark: 10, + connectionHardCap: 600 as const, + connectionUnobservedBound: 50 + } + await store.recordCellHeartbeat(heartbeat) + await expect(store.recordCellHeartbeat(heartbeat)).rejects.toThrow( + 'stale_connection_snapshot' + ) + await expect( + store.recordCellHeartbeat({ + ...heartbeat, + totalConnections: 9, + enforcedConnectionUnits: 9, + connectionInclusionWatermark: 9 + }) + ).rejects.toThrow('stale_connection_snapshot') + expect( + await database.query( + `SELECT inclusion_watermark, enforced_connection_units + FROM relay_cell_connection_snapshots` + ) + ).toEqual([{ inclusion_watermark: 10, enforced_connection_units: 10 }]) + }) + + it('keeps crash retries bound to one reservation claim', async () => { + let now = 100 + const { database, store } = await setup(448, () => now) + const identity = { userId: 'retry-user', relayHostId: 'retryhost0000001' } + const assignment = await store.assign(identity) + await new RelayAssignmentStore(database, () => now, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }).assign(identity) + expect( + await database.query( + `SELECT state FROM relay_control_connection_reservations` + ) + ).toEqual([{ state: 'reserved' }]) + now += ASSIGNMENT_LIMITS.activityLeaseMs + 1 + await store.releaseExpiredActivityLeases() + await store.recordCellHeartbeat({ + cellId: LIMITED_CELL.id, + cellUrl: LIMITED_CELL.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: 448, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 448, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }) + await store.assign(identity) + const restarted = new RelayAssignmentStore(database, () => now, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + + const activation = { + cellId: LIMITED_CELL.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1, + connectionInclusionWatermark: 5 + } + await restarted.activateControl(identity, activation) + await restarted.activateControl(identity, activation) + + expect( + await database.query( + `SELECT state, claim_activity_id + FROM relay_control_connection_reservations + ORDER BY created_at ASC, reservation_id ASC` + ) + ).toEqual([{ state: 'claimed', claim_activity_id: 'control:limited:1' }]) + }) + + it('releases only redundant expired debt after a fresh absolute snapshot', async () => { + let now = 100 + const { database, store } = await setup(449, () => now) + const identity = { userId: 'debt-user', relayHostId: 'debthost000000001' } + const assignment = await store.assign(identity) + now += ASSIGNMENT_LIMITS.activityLeaseMs + 1 + await store.releaseExpiredActivityLeases() + await database.query( + `INSERT INTO relay_control_connection_reservations + (reservation_id, idempotency_key, user_id, relay_host_id, + assignment_epoch, cell_id, state, inclusion_watermark, + claim_activity_id, created_at, timeout_at, claimed_at, released_at, + updated_at) + VALUES (?, ?, ?, ?, ?, ?, 'late-arrival-debt', NULL, NULL, ?, ?, NULL, NULL, ?)`, + [ + 'duplicate-reservation', + 'duplicate-reservation', + identity.userId, + identity.relayHostId, + assignment.assignmentEpoch, + LIMITED_CELL.id, + 101, + now - 1, + now + ] + ) + + await store.recordCellHeartbeat({ + cellId: LIMITED_CELL.id, + cellUrl: LIMITED_CELL.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: 449, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 449, + connectionInclusionWatermark: 5, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }) + + expect( + await database.query( + `SELECT state, COUNT(*) AS count + FROM relay_control_connection_reservations + GROUP BY state ORDER BY state` + ) + ).toEqual([ + { state: 'late-arrival-debt', count: 1 }, + { state: 'released', count: 1 } + ]) + }) + + it('releases expired debt for a snapshot-proven aborted older epoch', async () => { + const now = 1_000 + const { database, store } = await setup(0, () => now) + const identity = { userId: 'aborted-user', relayHostId: 'abortedhost00001' } + await database.query( + `INSERT INTO relay_assignments + (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, + last_activity_at, reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases) + VALUES (?, ?, ?, ?, ?, ?, 0, 0, 0, 0, 0, 0)`, + [identity.userId, identity.relayHostId, LIMITED_CELL.id, 3, now, now] + ) + await database.query( + `INSERT INTO relay_assignment_migrations + (user_id, relay_host_id, source_cell_id, target_cell_id, + previous_epoch, assignment_epoch, source_request_units, + target_reserved_units, expires_at, target_registered_at, + completed_at, aborted_at, created_at, updated_at) + VALUES (?, ?, ?, ?, 1, 2, 0, 1, ?, ?, NULL, ?, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + 'source', + LIMITED_CELL.id, + now - 2, + now - 3, + now - 1, + now - 4, + now - 1 + ] + ) + await database.query( + `INSERT INTO relay_control_connection_reservations + (reservation_id, idempotency_key, user_id, relay_host_id, + assignment_epoch, cell_id, state, inclusion_watermark, + claim_activity_id, created_at, timeout_at, claimed_at, released_at, + updated_at) + VALUES (?, ?, ?, ?, 2, ?, 'late-arrival-debt', NULL, NULL, ?, ?, NULL, NULL, ?)`, + [ + 'aborted-reservation', + 'aborted-reservation', + identity.userId, + identity.relayHostId, + LIMITED_CELL.id, + now - 4, + now - 2, + now - 1 + ] + ) + + await store.recordCellHeartbeat({ + cellId: LIMITED_CELL.id, + cellUrl: LIMITED_CELL.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0, + connectionInclusionWatermark: 5, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }) + + expect( + await database.query( + `SELECT state FROM relay_control_connection_reservations + WHERE reservation_id = ?`, + ['aborted-reservation'] + ) + ).toEqual([{ state: 'released' }]) + }) + + it('fails a limited cell closed on stale telemetry while leaving legacy cells unchanged', async () => { + let now = 100 + const { store } = await setup(0, () => now) + now += 45_001 + await expect( + store.assign({ userId: 'user-1', relayHostId: 'host000000000001' }) + ).rejects.toThrow('relay_capacity_exhausted') + + const legacyDatabase = await openInMemoryRelayDatabase() + databases.push(legacyDatabase) + const legacyStore = new RelayAssignmentStore(legacyDatabase, () => now, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + const legacy = { + id: 'legacy', + url: 'https://legacy.example.com', + capacityRequests: 10 + } + await legacyStore.reconcileCells([legacy], true) + await legacyStore.recordCellHeartbeat({ + cellId: legacy.id, + cellUrl: legacy.url, + cellIncarnation: '22222222-2222-4222-8222-222222222222', + startedAt: now, + ready: true, + observedRequests: 0 + }) + await expect( + legacyStore.assign({ userId: 'user-2', relayHostId: 'host000000000002' }) + ).resolves.toMatchObject({ cellId: 'legacy' }) + }) + + it('does not create connection debt for mixed-version uncapped cells', async () => { + let now = 100 + const legacy = { + id: 'legacy', + url: 'https://legacy.example.com', + capacityRequests: 10 + } + const { database, store } = await setup(0, () => now) + await store.reconcileCells([legacy], true) + await store.recordCellHeartbeat({ + cellId: legacy.id, + cellUrl: legacy.url, + cellIncarnation: '22222222-2222-4222-8222-222222222222', + startedAt: 50, + ready: true, + observedRequests: 0 + }) + + await store.assign({ userId: 'legacy-user', relayHostId: 'legacyhost000001' }) + expect( + await database.query(`SELECT * FROM relay_control_connection_reservations`) + ).toEqual([]) + expect(await store.releaseExpiredActivityLeases()).toBe(0) + now += ASSIGNMENT_LIMITS.activityLeaseMs + 1 + expect(await store.releaseExpiredActivityLeases()).toBe(1) + expect( + await database.query(`SELECT * FROM relay_control_connection_reservations`) + ).toEqual([]) + }) + + it('rejects dormant rebalance before changing the assignment epoch', async () => { + const now = ASSIGNMENT_LIMITS.dormantTtlMs + 100 + const database = await openInMemoryRelayDatabase() + databases.push(database) + const store = new RelayAssignmentStore(database, () => now, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + const source = { + id: 'source', + url: 'https://source.example.com', + capacityRequests: 10 + } + await store.reconcileCells([source, LIMITED_CELL], true) + await store.recordCellHeartbeat({ + cellId: source.id, + cellUrl: source.url, + cellIncarnation: '22222222-2222-4222-8222-222222222222', + startedAt: now - 50, + ready: true, + observedRequests: 0 + }) + await store.recordCellHeartbeat({ + cellId: LIMITED_CELL.id, + cellUrl: LIMITED_CELL.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: now - 50, + ready: true, + observedRequests: 0, + totalConnections: 450, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 450, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }) + const identity = { userId: 'dormant-user', relayHostId: 'dormanthost00001' } + await database.query( + `INSERT INTO relay_assignments + (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, + last_activity_at, reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + source.id, + 7, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0 + ] + ) + + await expect(store.rebalanceDormant(identity, LIMITED_CELL.id)).rejects.toThrow( + 'relay_connection_headroom_exhausted' + ) + expect(await store.resolve(identity)).toMatchObject({ + cellId: source.id, + assignmentEpoch: 7 + }) + expect(await database.query(`SELECT * FROM relay_assignment_activity_leases`)).toEqual([]) + }) + + it('rejects an evacuation target at its pause before migration mutation', async () => { + const database = await openInMemoryRelayDatabase() + databases.push(database) + const store = new RelayAssignmentStore(database, () => 100, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + const source = { + id: 'source', + url: 'https://source.example.com', + capacityRequests: 10 + } + await store.reconcileCells( + [source, { ...LIMITED_CELL, initiallyEnabled: false }], + true + ) + await store.recordCellHeartbeat({ + cellId: source.id, + cellUrl: source.url, + cellIncarnation: '22222222-2222-4222-8222-222222222222', + startedAt: 50, + ready: true, + observedRequests: 0 + }) + await store.recordCellHeartbeat({ + cellId: LIMITED_CELL.id, + cellUrl: LIMITED_CELL.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: 450, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 450, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }) + const identity = { userId: 'user-1', relayHostId: 'host000000000001' } + const assignment = await store.assign(identity) + expect(assignment.cellId).toBe(source.id) + await store.setCellEnabled(LIMITED_CELL.id, true) + + await expect( + store.startEvacuation(identity, LIMITED_CELL.id) + ).rejects.toThrow('relay_connection_headroom_exhausted') + expect(await database.query(`SELECT * FROM relay_assignment_migrations`)).toEqual([]) + expect(await store.resolve(identity)).toMatchObject({ + cellId: source.id, + assignmentEpoch: assignment.assignmentEpoch + }) + }) +}) diff --git a/cloud/apps/relay/src/assignment-control-supersession-postgres.test.ts b/cloud/apps/relay/src/assignment-control-supersession-postgres.test.ts new file mode 100644 index 00000000000..10193b78cc6 --- /dev/null +++ b/cloud/apps/relay/src/assignment-control-supersession-postgres.test.ts @@ -0,0 +1,130 @@ +import { afterAll, beforeAll, describe, expect, it } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.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 = { + id: 'control-supersession-postgres', + url: 'https://control-supersession-postgres.example.com', + capacityRequests: 1_000, + connectionHardCap: 600 as const, + connectionUnobservedBound: 50 +} +const identity = { + userId: 'control-supersession-postgres-user', + relayHostId: 'supersedehost001' +} + +describePostgres('PostgreSQL control supersession', () => { + const databases: RelayDatabase[] = [] + + beforeAll(async () => { + databases.push( + await openRelayDatabase({ databaseUrl, dataDir: '' }), + await openRelayDatabase({ databaseUrl, dataDir: '' }) + ) + }) + + afterAll(async () => { + const database = databases[0] + if (database) { + await database.query( + `DELETE FROM relay_control_connection_reservations WHERE user_id = ?`, + [identity.userId] + ) + await database.query( + `DELETE FROM relay_assignment_activity_leases WHERE user_id = ?`, + [identity.userId] + ) + await database.query(`DELETE FROM relay_assignments WHERE user_id = ?`, [identity.userId]) + await database.query(`DELETE FROM relay_cell_connection_runtime WHERE cell_id = ?`, [cell.id]) + await database.query(`DELETE FROM relay_cell_connection_limits WHERE cell_id = ?`, [cell.id]) + await database.query(`DELETE FROM relay_cell_runtime WHERE cell_id = ?`, [cell.id]) + await database.query(`DELETE FROM relay_cells WHERE cell_id = ?`, [cell.id]) + } + for (const connection of databases) await connection.close() + }) + + it('serializes parallel generations into one durable control', async () => { + const stores = databases.map((database) => new RelayAssignmentStore(database, () => 100)) + await stores[0]!.reconcileCells([cell]) + await stores[0]!.recordCellHeartbeat({ + cellId: cell.id, + cellUrl: cell.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0, + connectionInclusionWatermark: 1, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }) + const assignment = await stores[0]!.assign(identity) + + await Promise.all([ + stores[0]!.activateControl(identity, { + cellId: cell.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1, + connectionInclusionWatermark: 10 + }), + stores[1]!.activateControl(identity, { + cellId: cell.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 2, + connectionInclusionWatermark: 11 + }) + ]) + await databases[0]!.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, ?, 'control', ?, 1, 90100, 100), + (?, ?, ?, 'control', ?, 1, 90100, 100)`, + [ + identity.userId, + identity.relayHostId, + `control:${cell.id}:100`, + cell.id, + identity.userId, + identity.relayHostId, + `control:${cell.id}:101`, + cell.id + ] + ) + await stores[0]!.activateControl(identity, { + cellId: cell.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 102, + connectionInclusionWatermark: 12 + }) + + const controls = await databases[0]!.query( + `SELECT COUNT(*) AS count FROM relay_assignment_activity_leases + WHERE user_id = ? AND activity_kind = 'control'`, + [identity.userId] + ) + expect(Number(controls[0]!.count)).toBe(1) + const assignments = await databases[0]!.query( + `SELECT reserved_controls FROM relay_assignments WHERE user_id = ?`, + [identity.userId] + ) + expect(Number(assignments[0]!.reserved_controls)).toBe(1) + const cells = await databases[0]!.query( + `SELECT reserved_requests FROM relay_cells WHERE cell_id = ?`, + [cell.id] + ) + expect(Number(cells[0]!.reserved_requests)).toBe(1) + const claims = await databases[0]!.query( + `SELECT claim_activity_id FROM relay_control_connection_reservations + WHERE user_id = ? AND state = 'claimed'`, + [identity.userId] + ) + expect(claims).toEqual([{ claim_activity_id: `control:${cell.id}:102` }]) + }, 15_000) +}) diff --git a/cloud/apps/relay/src/assignment-deployment-status-postgres.test.ts b/cloud/apps/relay/src/assignment-deployment-status-postgres.test.ts new file mode 100644 index 00000000000..9e8e3b85bd9 --- /dev/null +++ b/cloud/apps/relay/src/assignment-deployment-status-postgres.test.ts @@ -0,0 +1,87 @@ +import pg from 'pg' +import { afterAll, beforeAll, describe, expect, it } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.js' +import { openRelayDatabase, type RelayDatabase } from './database.js' + +const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL +const describePostgres = databaseUrl ? describe : describe.skip +const schema = 'relay_deployment_status_test' +const cell = { + id: 'deployment-status-postgres', + url: 'https://deployment-status-postgres.example.com', + capacityRequests: 20 +} + +describePostgres('PostgreSQL deployment status', () => { + let database: RelayDatabase | undefined + + beforeAll(async () => { + const client = new pg.Client({ connectionString: databaseUrl }) + await client.connect() + try { + await client.query(`DROP SCHEMA IF EXISTS ${schema} CASCADE`) + await client.query(`CREATE SCHEMA ${schema}`) + } finally { + await client.end() + } + const url = new URL(databaseUrl!) + url.searchParams.set('options', `-c search_path=${schema}`) + database = await openRelayDatabase({ databaseUrl: url.toString(), dataDir: '' }) + }) + + afterAll(async () => { + await database?.close() + const client = new pg.Client({ connectionString: databaseUrl }) + await client.connect() + try { + await client.query(`DROP SCHEMA IF EXISTS ${schema} CASCADE`) + } finally { + await client.end() + } + }) + + it('classifies pending controls and restart blockers in one PostgreSQL query', async () => { + const store = new RelayAssignmentStore(database!, () => 100) + await store.reconcileCells([cell]) + await store.assign({ userId: 'status-user', relayHostId: 'a1b2c3d4e5f6' }) + + expect(await store.cellDeploymentStatus(cell.id)).toMatchObject({ + activityLeases: 1, + activityRequestUnits: 1, + reservedRequests: 1, + restartBlockingActivityLeases: 0, + restartBlockingActivityRequestUnits: 0, + restartBlockingReservedRequests: 0 + }) + + await database!.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES ('blocking-user', 'b1c2d3e4f5a6', 'invite:test', 'invite', ?, 1, 90100, 100), + ('malformed-user', 'c1d2e3f4a5b6', 'control-pending:bad', 'control', ?, 2, 90100, 100)`, + [cell.id, cell.id] + ) + await database!.query( + `UPDATE relay_cells SET reserved_requests = reserved_requests + 3 WHERE cell_id = ?`, + [cell.id] + ) + + expect(await store.cellDeploymentStatus(cell.id)).toMatchObject({ + activityLeases: 3, + activityRequestUnits: 4, + reservedRequests: 4, + restartBlockingActivityLeases: 2, + restartBlockingActivityRequestUnits: 3, + restartBlockingReservedRequests: 3 + }) + + await database!.query( + `UPDATE relay_cells SET reserved_requests = 0 WHERE cell_id = ?`, + [cell.id] + ) + expect(await store.cellDeploymentStatus(cell.id)).toMatchObject({ + restartBlockingReservedRequests: -1 + }) + }) +}) diff --git a/cloud/apps/relay/src/assignment-identity-queue.test.ts b/cloud/apps/relay/src/assignment-identity-queue.test.ts new file mode 100644 index 00000000000..930f0b692df --- /dev/null +++ b/cloud/apps/relay/src/assignment-identity-queue.test.ts @@ -0,0 +1,70 @@ +import { describe, expect, it } from 'vitest' +import { AssignmentIdentityQueue } from './assignment-identity-queue.js' + +function deferred(): { + promise: Promise + resolve: () => void +} { + let resolve!: () => void + const promise = new Promise((complete) => { + resolve = complete + }) + return { promise, resolve } +} + +describe('AssignmentIdentityQueue', () => { + it('serializes operations for the same assignment', async () => { + const queue = new AssignmentIdentityQueue() + const firstStarted = deferred() + const firstRelease = deferred() + const started: string[] = [] + const identity = { userId: 'user-1', relayHostId: 'host-1' } + + const first = queue.run(identity, async () => { + started.push('first') + firstStarted.resolve() + await firstRelease.promise + }) + const second = queue.run(identity, async () => { + started.push('second') + }) + + await firstStarted.promise + expect(started).toEqual(['first']) + firstRelease.resolve() + await Promise.all([first, second]) + expect(started).toEqual(['first', 'second']) + }) + + it('allows different assignments to run concurrently', async () => { + const queue = new AssignmentIdentityQueue() + const release = deferred() + const started: string[] = [] + + const first = queue.run({ userId: 'user-1', relayHostId: 'host-1' }, async () => { + started.push('first') + await release.promise + }) + const second = queue.run({ userId: 'user-2', relayHostId: 'host-1' }, async () => { + started.push('second') + }) + + await second + expect(started).toEqual(['first', 'second']) + release.resolve() + await first + }) + + it('continues after a rejected operation', async () => { + const queue = new AssignmentIdentityQueue() + const identity = { userId: 'user-1', relayHostId: 'host-1' } + + const failed = queue.run(identity, async () => { + throw new Error('failed') + }) + const recovered = queue.run(identity, async () => 'recovered') + + await expect(failed).rejects.toThrow('failed') + await expect(recovered).resolves.toBe('recovered') + }) +}) diff --git a/cloud/apps/relay/src/assignment-identity-queue.ts b/cloud/apps/relay/src/assignment-identity-queue.ts new file mode 100644 index 00000000000..f39d9c0a03c --- /dev/null +++ b/cloud/apps/relay/src/assignment-identity-queue.ts @@ -0,0 +1,24 @@ +export type AssignmentIdentity = { + userId: string + relayHostId: string +} + +export class AssignmentIdentityQueue { + private readonly tails = new Map>() + + async run(identity: AssignmentIdentity, operation: () => Promise): Promise { + const key = JSON.stringify([identity.userId, identity.relayHostId]) + const previous = this.tails.get(key) ?? Promise.resolve() + const result = previous.catch(() => undefined).then(operation) + const tail = result.then( + () => undefined, + () => undefined + ) + this.tails.set(key, tail) + try { + return await result + } finally { + if (this.tails.get(key) === tail) this.tails.delete(key) + } + } +} diff --git a/cloud/apps/relay/src/assignment-inventory-snapshot.test.ts b/cloud/apps/relay/src/assignment-inventory-snapshot.test.ts new file mode 100644 index 00000000000..afc35f7e601 --- /dev/null +++ b/cloud/apps/relay/src/assignment-inventory-snapshot.test.ts @@ -0,0 +1,107 @@ +import { describe, expect, it } from 'vitest' +import { + formatAssignmentInventorySnapshot, + readAssignmentInventorySnapshot +} from './assignment-inventory-snapshot.js' +import { openInMemoryRelayDatabase } from './database.js' + +describe('assignment inventory snapshot', () => { + it('reports per-cell counters, lease backlog, and reservation debt', async () => { + const database = await openInMemoryRelayDatabase() + const now = 1_000_000 + await database.query( + `INSERT INTO relay_cells + (cell_id, cell_url, enabled, capacity_requests, reserved_requests, + observed_requests, last_heartbeat_at, updated_at) + VALUES ('cell-a', 'https://a.example.test', 1, 4000, 3999, 5, ?, ?)`, + [now, now] + ) + await database.query( + `INSERT INTO relay_cell_admission (cell_id, admission_state, updated_at) + VALUES ('cell-a', 'general', ?)`, + [now] + ) + await database.query( + `INSERT INTO relay_cell_runtime + (cell_id, cell_url, cell_incarnation, started_at, ready, observed_requests, + last_heartbeat_at, updated_at) + VALUES ('cell-a', 'https://a.example.test', 'inc-1', ?, 1, 5, ?, ?)`, + [now - 60_000, now - 10_000, now] + ) + await database.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, request_units, + expires_at, updated_at) + VALUES + ('user-1', 'host-1', 'control:1', 'control', 'cell-a', 1, ?, ?), + ('user-1', 'host-2', 'control:1', 'control', 'cell-a', 3, ?, ?)`, + [now - 1, now, now + 90_000, now] + ) + await database.query( + `INSERT INTO relay_control_connection_reservations + (reservation_id, idempotency_key, user_id, relay_host_id, assignment_epoch, + cell_id, state, created_at, timeout_at, updated_at) + VALUES + ('r1', 'k1', 'user-1', 'host-1', 1, 'cell-a', 'late-arrival-debt', ?, ?, ?), + ('r2', 'k2', 'user-1', 'host-2', 1, 'cell-a', 'reserved', ?, ?, ?), + ('r3', 'k3', 'user-1', 'host-3', 1, 'cell-a', 'claimed', ?, ?, ?), + ('r4', 'k4', 'user-1', 'host-4', 1, 'cell-a', 'released', ?, ?, ?)`, + [now, now, now, now, now, now, now, now, now, now, now, now] + ) + + const snapshot = await readAssignmentInventorySnapshot(database, now) + + expect(snapshot.cells).toEqual([ + { + cellId: 'cell-a', + region: 'us-central1', + admissionState: 'general', + enabled: true, + capacityRequests: 4000, + reservedRequests: 3999, + runtimeReady: true, + heartbeatAgeMs: 10_000 + } + ]) + expect(snapshot.activityLeases).toEqual({ total: 2, expired: 1, requestUnits: 4 }) + expect(snapshot.connectionReservations).toEqual({ outstanding: 3, lateArrivalDebt: 1 }) + expect(snapshot.regionalRehomes).toEqual({ + active: 0, + awaitingReceipt: 0, + targetRegistered: 0, + completedLast24Hours: 0, + abortedLast24Hours: 0, + oldestActiveAgeMs: null + }) + + const lines = formatAssignmentInventorySnapshot(snapshot) + expect(lines).toHaveLength(3) + expect(lines[0]).toContain('cellId=cell-a') + expect(lines[0]).toContain('reserved=3999') + expect(lines[1]).toContain('expiredLeases=1') + expect(lines[1]).toContain('lateArrivalDebt=1') + expect(lines[2]).toContain('active=0') + }) + + it('reports cells missing runtime and admission rows without failing', async () => { + const database = await openInMemoryRelayDatabase() + await database.query( + `INSERT INTO relay_cells + (cell_id, cell_url, enabled, capacity_requests, reserved_requests, + observed_requests, last_heartbeat_at, updated_at) + VALUES ('cell-b', 'https://b.example.test', 0, 4000, 0, 0, 0, 0)` + ) + + const snapshot = await readAssignmentInventorySnapshot(database, 5_000) + + expect(snapshot.cells[0]).toMatchObject({ + cellId: 'cell-b', + admissionState: 'unset', + enabled: false, + runtimeReady: null, + heartbeatAgeMs: null + }) + expect(snapshot.activityLeases).toEqual({ total: 0, expired: 0, requestUnits: 0 }) + expect(snapshot.regionalRehomes.active).toBe(0) + }) +}) diff --git a/cloud/apps/relay/src/assignment-inventory-snapshot.ts b/cloud/apps/relay/src/assignment-inventory-snapshot.ts new file mode 100644 index 00000000000..0675bd49b94 --- /dev/null +++ b/cloud/apps/relay/src/assignment-inventory-snapshot.ts @@ -0,0 +1,170 @@ +import type { RelayDatabase, SqlRow } from './database.js' + +export type CellInventorySnapshotRow = { + cellId: string + region: string + admissionState: string + enabled: boolean + capacityRequests: number + reservedRequests: number + runtimeReady: boolean | null + heartbeatAgeMs: number | null +} + +export type AssignmentInventorySnapshot = { + cells: CellInventorySnapshotRow[] + activityLeases: { total: number; expired: number; requestUnits: number } + connectionReservations: { outstanding: number; lateArrivalDebt: number } + regionalRehomes: { + active: number + awaitingReceipt: number + targetRegistered: number + completedLast24Hours: number + abortedLast24Hours: number + oldestActiveAgeMs: number | null + } +} + +// Plain SELECTs only: this snapshot must never take the cell-inventory lock, +// or observability itself would add to the assign-path lock contention. +export async function readAssignmentInventorySnapshot( + database: RelayDatabase, + now: number +): Promise { + const cellRows = await database.query( + `SELECT cell.cell_id, cell.enabled, cell.capacity_requests, cell.reserved_requests, + admission.admission_state, region.region, + runtime.ready AS runtime_ready, runtime.last_heartbeat_at AS runtime_heartbeat_at + FROM relay_cells cell + LEFT JOIN relay_cell_admission admission ON admission.cell_id = cell.cell_id + LEFT JOIN relay_cell_regions region ON region.cell_id = cell.cell_id + LEFT JOIN relay_cell_runtime runtime ON runtime.cell_id = cell.cell_id + ORDER BY cell.cell_id ASC` + ) + const leaseRow = ( + await database.query( + `SELECT COUNT(*) AS total, + COALESCE(SUM(CASE WHEN expires_at <= ? THEN 1 ELSE 0 END), 0) AS expired, + COALESCE(SUM(request_units), 0) AS request_units + FROM relay_assignment_activity_leases`, + [now] + ) + )[0] + const reservationRow = ( + await database.query( + // relay_cells is append-only in production; joining it lets the composite + // index skip released history without hiding any production cell's debt. + `SELECT COUNT(reservation.reservation_id) AS outstanding, + COALESCE(SUM(CASE WHEN reservation.state = 'late-arrival-debt' + THEN 1 ELSE 0 END), 0) AS late_arrival_debt + FROM relay_cells cell + LEFT JOIN relay_control_connection_reservations reservation + ON reservation.cell_id = cell.cell_id + AND reservation.state IN ('reserved', 'late-arrival-debt', 'claimed')` + ) + )[0] + const regionalRehomeRow = ( + await database.query( + `SELECT + COALESCE(SUM(CASE WHEN attempt.completed_at IS NULL AND attempt.aborted_at IS NULL + THEN 1 ELSE 0 END), 0) AS active, + COALESCE(SUM(CASE WHEN attempt.completed_at IS NULL AND attempt.aborted_at IS NULL + AND attempt.drain_receipt_at IS NULL + THEN 1 ELSE 0 END), 0) AS awaiting_receipt, + COALESCE(SUM(CASE WHEN attempt.completed_at IS NULL AND attempt.aborted_at IS NULL + AND migration.target_registered_at IS NOT NULL + THEN 1 ELSE 0 END), 0) AS target_registered, + COALESCE(SUM(CASE WHEN attempt.completed_at >= ? THEN 1 ELSE 0 END), 0) + AS completed_last_24_hours, + COALESCE(SUM(CASE WHEN attempt.aborted_at >= ? THEN 1 ELSE 0 END), 0) + AS aborted_last_24_hours, + MIN(CASE WHEN attempt.completed_at IS NULL AND attempt.aborted_at IS NULL + THEN attempt.created_at END) AS oldest_active_at + FROM relay_region_rehome_attempts attempt + LEFT JOIN relay_assignment_migrations migration + ON migration.user_id = attempt.user_id + AND migration.relay_host_id = attempt.relay_host_id + AND migration.assignment_epoch = attempt.assignment_epoch`, + [now - 24 * 60 * 60_000, now - 24 * 60 * 60_000] + ) + )[0] + const oldestActiveAt = optionalInteger(regionalRehomeRow, 'oldest_active_at') + return { + cells: cellRows.map((row) => ({ + cellId: asText(row, 'cell_id'), + region: optionalText(row, 'region') ?? 'us-central1', + admissionState: optionalText(row, 'admission_state') ?? 'unset', + enabled: asInteger(row, 'enabled') === 1, + capacityRequests: asInteger(row, 'capacity_requests'), + reservedRequests: asInteger(row, 'reserved_requests'), + runtimeReady: row['runtime_ready'] == null ? null : asInteger(row, 'runtime_ready') === 1, + heartbeatAgeMs: + row['runtime_heartbeat_at'] == null ? null : now - asInteger(row, 'runtime_heartbeat_at') + })), + activityLeases: { + total: asInteger(leaseRow, 'total'), + expired: asInteger(leaseRow, 'expired'), + requestUnits: asInteger(leaseRow, 'request_units') + }, + connectionReservations: { + outstanding: asInteger(reservationRow, 'outstanding'), + lateArrivalDebt: asInteger(reservationRow, 'late_arrival_debt') + }, + regionalRehomes: { + active: asInteger(regionalRehomeRow, 'active'), + awaitingReceipt: asInteger(regionalRehomeRow, 'awaiting_receipt'), + targetRegistered: asInteger(regionalRehomeRow, 'target_registered'), + completedLast24Hours: asInteger(regionalRehomeRow, 'completed_last_24_hours'), + abortedLast24Hours: asInteger(regionalRehomeRow, 'aborted_last_24_hours'), + oldestActiveAgeMs: oldestActiveAt === null ? null : now - oldestActiveAt + } + } +} + +export function formatAssignmentInventorySnapshot( + snapshot: AssignmentInventorySnapshot +): string[] { + const lines = snapshot.cells.map( + (cell) => + `[orca-relay] cell inventory cellId=${cell.cellId}` + + ` region=${cell.region}` + + ` admission=${cell.admissionState} enabled=${cell.enabled}` + + ` capacity=${cell.capacityRequests} reserved=${cell.reservedRequests}` + + ` ready=${cell.runtimeReady ?? 'none'}` + + ` heartbeatAgeMs=${cell.heartbeatAgeMs ?? 'none'}` + ) + lines.push( + `[orca-relay] lease inventory leases=${snapshot.activityLeases.total}` + + ` expiredLeases=${snapshot.activityLeases.expired}` + + ` leaseRequestUnits=${snapshot.activityLeases.requestUnits}` + + ` outstandingReservations=${snapshot.connectionReservations.outstanding}` + + ` lateArrivalDebt=${snapshot.connectionReservations.lateArrivalDebt}` + ) + lines.push( + `[orca-relay] regional rehome inventory active=${snapshot.regionalRehomes.active}` + + ` awaitingReceipt=${snapshot.regionalRehomes.awaitingReceipt}` + + ` targetRegistered=${snapshot.regionalRehomes.targetRegistered}` + + ` completedLast24Hours=${snapshot.regionalRehomes.completedLast24Hours}` + + ` abortedLast24Hours=${snapshot.regionalRehomes.abortedLast24Hours}` + + ` oldestActiveAgeMs=${snapshot.regionalRehomes.oldestActiveAgeMs ?? 'none'}` + ) + return lines +} + +function asText(row: SqlRow | undefined, column: string): string { + return String(row?.[column] ?? '') +} + +function optionalText(row: SqlRow | undefined, column: string): string | null { + const value = row?.[column] + return value == null ? null : String(value) +} + +function asInteger(row: SqlRow | undefined, column: string): number { + return Number(row?.[column] ?? 0) +} + +function optionalInteger(row: SqlRow | undefined, column: string): number | null { + const value = row?.[column] + return value == null ? null : Number(value) +} diff --git a/cloud/apps/relay/src/assignment-rejection-log-window.test.ts b/cloud/apps/relay/src/assignment-rejection-log-window.test.ts new file mode 100644 index 00000000000..8ef078eff86 --- /dev/null +++ b/cloud/apps/relay/src/assignment-rejection-log-window.test.ts @@ -0,0 +1,100 @@ +import { describe, expect, it } from 'vitest' +import { AssignmentRejectionLogWindow } from './assignment-rejection-log-window.js' + +describe('assignment rejection log window', () => { + it('emits once per key per window and closes it with its own suppressed count', () => { + const timers = manualTimers() + const closed: { key: string; suppressed: number; sample: string }[] = [] + const window = new AssignmentRejectionLogWindow({ + windowMs: 10_000, + schedule: timers.schedule, + onWindowClosed: (input) => closed.push(input) + }) + const key = 'assign:placement:host-rate-limited' + + expect(window.admit(key, 'host-a')).toBe(true) + expect(window.admit(key, 'host-b')).toBe(false) + expect(window.admit(key, 'host-c')).toBe(false) + expect(closed).toEqual([]) + + timers.runPending() + expect(closed).toEqual([{ key, suppressed: 2, sample: 'host-c' }]) + + // The next window starts clean instead of inheriting the closed window's count. + expect(window.admit(key, 'host-d')).toBe(true) + timers.runPending() + expect(closed).toHaveLength(1) + }) + + it('reports the final count for a key that goes quiet', () => { + const timers = manualTimers() + const closed: { key: string; suppressed: number }[] = [] + const window = new AssignmentRejectionLogWindow({ + windowMs: 10_000, + schedule: timers.schedule, + onWindowClosed: ({ key, suppressed }) => closed.push({ key, suppressed }) + }) + + window.admit('assign:sticky:host-rate-limited', 'host-a') + window.admit('assign:sticky:host-rate-limited', 'host-a') + window.admit('assign:sticky:host-rate-limited', 'host-a') + // No further rejection ever arrives for this key. + timers.runPending() + + expect(closed).toEqual([{ key: 'assign:sticky:host-rate-limited', suppressed: 2 }]) + }) + + it('keeps distinct keys on independent windows', () => { + const timers = manualTimers() + const window = new AssignmentRejectionLogWindow({ + windowMs: 10_000, + schedule: timers.schedule, + onWindowClosed: () => undefined + }) + + expect(window.admit('assign:placement:host-rate-limited', 'host-a')).toBe(true) + expect(window.admit('assign:placement:queue-full', 'host-a')).toBe(true) + expect(window.admit('assign:sticky:host-rate-limited', 'host-a')).toBe(true) + expect(window.admit('assign:placement:queue-full', 'host-a')).toBe(false) + }) + + it('bounds the tracked keys and reports what an evicted window suppressed', () => { + const timers = manualTimers() + const closed: { key: string; suppressed: number }[] = [] + const window = new AssignmentRejectionLogWindow({ + windowMs: 10_000, + schedule: timers.schedule, + onWindowClosed: ({ key, suppressed }) => closed.push({ key, suppressed }) + }) + + window.admit('key-0', 'host-a') + window.admit('key-0', 'host-b') + for (let index = 1; index < 200; index++) window.admit(`key-${index}`, 'host-a') + + expect(closed[0]).toEqual({ key: 'key-0', suppressed: 1 }) + // Evicted keys emit again rather than staying silently suppressed. + expect(window.admit('key-0', 'host-a')).toBe(true) + expect(window.admit('key-199', 'host-a')).toBe(false) + }) +}) + +function manualTimers(): { + schedule: (callback: () => void) => () => void + runPending: () => void +} { + const pending = new Map void>() + let nextId = 0 + return { + schedule: (callback) => { + const id = nextId++ + pending.set(id, callback) + return () => pending.delete(id) + }, + runPending: () => { + for (const [id, callback] of [...pending]) { + pending.delete(id) + callback() + } + } + } +} diff --git a/cloud/apps/relay/src/assignment-rejection-log-window.ts b/cloud/apps/relay/src/assignment-rejection-log-window.ts new file mode 100644 index 00000000000..21ad4e72af9 --- /dev/null +++ b/cloud/apps/relay/src/assignment-rejection-log-window.ts @@ -0,0 +1,60 @@ +type CancelWindow = () => void + +// Admission rejections run at ~17/s per director instance, so the log summarizes +// instead of streaming: one line when a key's window opens, then a closing line +// carrying whatever that same window suppressed. The closing line is driven by the +// window's own timer, so a key that falls quiet still reports its final count +// instead of waiting for a rejection that may never arrive. +const MAX_TRACKED_KEYS = 64 + +type RejectionWindow = { + suppressed: number + sample: TSample + cancel: CancelWindow +} + +export class AssignmentRejectionLogWindow { + private readonly windows = new Map>() + + constructor( + private readonly options: { + windowMs: number + onWindowClosed: (input: { key: string; suppressed: number; sample: TSample }) => void + schedule?: (callback: () => void, delayMs: number) => CancelWindow + } + ) {} + + // Returns true when the caller should log this rejection immediately. + admit(key: string, sample: TSample): boolean { + const open = this.windows.get(key) + if (open) { + open.suppressed++ + // Keep the most recent host/reason so the closing line names a real rejection. + open.sample = sample + return false + } + const schedule = this.options.schedule ?? defaultSchedule + const window: RejectionWindow = { suppressed: 0, sample, cancel: () => undefined } + this.windows.set(key, window) + if (this.windows.size > MAX_TRACKED_KEYS) { + this.close(this.windows.keys().next().value!) + } + window.cancel = schedule(() => this.close(key), this.options.windowMs) + return true + } + + private close(key: string): void { + const window = this.windows.get(key) + if (!window) return + this.windows.delete(key) + window.cancel() + if (window.suppressed === 0) return + this.options.onWindowClosed({ key, suppressed: window.suppressed, sample: window.sample }) + } +} + +function defaultSchedule(callback: () => void, delayMs: number): CancelWindow { + const timer = setTimeout(callback, delayMs) + timer.unref?.() + return () => clearTimeout(timer) +} diff --git a/cloud/apps/relay/src/assignment-rejection-logging.test.ts b/cloud/apps/relay/src/assignment-rejection-logging.test.ts new file mode 100644 index 00000000000..b00d791fce9 --- /dev/null +++ b/cloud/apps/relay/src/assignment-rejection-logging.test.ts @@ -0,0 +1,244 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import type { RelayAssignment } from './assignment-store.js' +import type { RelayConfig } from './config.js' + +const fakes = vi.hoisted(() => ({ + verifyRelayToken: vi.fn(async (token: string) => ({ + sub: 'user-1', + relayHostId: token + })) +})) + +vi.mock('./relay-token-verifier.js', () => ({ + createRelayTokenVerifier: () => fakes.verifyRelayToken, + readBearer: (value: string | undefined) => value?.replace(/^Bearer /, '') ?? null +})) + +import { createRelayApp, relayHostLogDigest } from './app.js' + +describe('assignment rejection logging', () => { + afterEach(() => { + vi.useRealTimers() + vi.restoreAllMocks() + }) + + it('logs the store reason and host digest when assign() rejects on capacity', async () => { + const host = 'hhhhhhhhhhhhhhhh' + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const app = createRelayApp(config(), { + store: {} as never, + assignments: { + assign: vi.fn(async () => { + throw new Error('relay_capacity_exhausted') + }), + resolve: vi.fn(async () => assignment('cell-r', host)) + } as never, + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + + const hinted = await app.request('/v1/assign', assignmentRequest(host, { reconnect: true })) + + expect(hinted.status).toBe(503) + expect(await hinted.json()).toEqual({ error: 'relay_capacity_exhausted' }) + const line = warn.mock.calls.map((call) => String(call[0])).find((entry) => + entry.includes('assignment rejected') + ) + expect(line).toContain('route=assign') + expect(line).toContain('lane=sticky') + expect(line).toContain('hinted=true') + expect(line).toContain('reason=relay_capacity_exhausted') + expect(line).toContain(`host=${relayHostLogDigest(host)}`) + expect(line).not.toContain(host) + }) + + it('logs an unhinted placement rejection without the raw host id', async () => { + const host = 'gggggggggggggggg' + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const app = createRelayApp(config(), { + store: {} as never, + assignments: { + assign: vi.fn(async () => { + throw new Error('relay_connection_headroom_exhausted') + }), + resolve: vi.fn(async () => null) + } as never, + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + + const response = await app.request('/v1/assign', assignmentRequest(host)) + + expect(response.status).toBe(503) + const line = warn.mock.calls.map((call) => String(call[0])).find((entry) => + entry.includes('assignment rejected') + ) + expect(line).toContain('route=assign') + expect(line).toContain('lane=placement') + expect(line).toContain('hinted=false') + expect(line).toContain('reason=relay_connection_headroom_exhausted') + expect(line).not.toContain(host) + }) + + it('names the throttled host once per window and closes it with the suppressed count', async () => { + vi.useFakeTimers() + const host = 'mmmmmmmmmmmmmmmm' + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const app = createRelayApp(config(), { + store: {} as never, + assignments: { + assign: vi.fn(async () => assignment('cell-r', host)), + resolve: vi.fn(async () => null) + } as never, + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + + expect((await app.request('/v1/assign', assignmentRequest(host))).status).toBe(200) + expect((await app.request('/v1/assign', assignmentRequest(host))).status).toBe(503) + expect((await app.request('/v1/assign', assignmentRequest(host))).status).toBe(503) + + const throttled = (): string[] => + warn.mock.calls + .map((call) => String(call[0])) + .filter((entry) => entry.includes('reason=host-rate-limited')) + expect(throttled()).toHaveLength(1) + expect(throttled()[0]).toContain('route=assign') + expect(throttled()[0]).toContain('lane=placement') + expect(throttled()[0]).toContain(`host=${relayHostLogDigest(host)}`) + expect(throttled()[0]).not.toContain('suppressed=') + expect(throttled()[0]).not.toContain(host) + + await vi.advanceTimersByTimeAsync(10_000) + expect(throttled()).toHaveLength(2) + expect(throttled()[1]).toContain('suppressed=1') + expect(throttled()[1]).toContain(`host=${relayHostLogDigest(host)}`) + }) + + it('stays silent for successful assignments', async () => { + const host = 'ssssssssssssssss' + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const app = createRelayApp(config(), { + store: {} as never, + assignments: { + assign: vi.fn(async () => assignment('cell-r', host)), + resolve: vi.fn(async () => null) + } as never, + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + + expect((await app.request('/v1/assign', assignmentRequest(host))).status).toBe(200) + expect( + warn.mock.calls.map((call) => String(call[0])).filter((entry) => + entry.includes('assignment rejected') + ) + ).toEqual([]) + }) +}) + +function assignmentRequest( + relayHostId: string, + extra: Record = {} +): RequestInit & { headers: Record } { + return { + method: 'POST', + headers: { + authorization: `Bearer ${relayHostId}`, + 'content-type': 'application/json' + }, + body: JSON.stringify({ v: 1, relayHostId, ...extra }) + } +} + +function assignment(cellId: string, relayHostId: string): RelayAssignment { + return { + userId: 'user-1', + relayHostId, + cellId, + cellUrl: `https://${cellId}.relay.example.test`, + assignmentEpoch: 1, + leaseExpiresAt: Date.now() + 300_000 + } +} + +function config(overrides: Partial = {}): RelayConfig { + return { + port: 8080, + publicUrl: 'https://relay.example.test', + cellUrl: 'https://relay.example.test', + authIssuer: 'https://auth.example.test', + authAudience: 'orca-relay', + jwksUrl: 'https://auth.example.test/jwks', + assignmentSigningKey: new TextEncoder().encode('assignment-key-with-at-least-32-bytes'), + role: 'director', + cellId: 'director', + cells: [], + adminAudience: 'https://relay.example.test/v1/admin/drain', + deployServiceAccount: 'deploy@example.test', + runtimeServiceAccount: 'runtime@example.test', + adminJwksUrl: 'https://auth.example.test/jwks', + databasePoolMax: 3, + publicAssignmentsEnabled: true, + publicAssignmentConcurrency: 2, + publicAssignmentQueueMax: 128, + publicAssignmentWaitMs: 4_000, + publicResolveConcurrency: 1, + publicResolveWaitMs: 5_000, + publicAssignmentRetryAfterSeconds: 5, + dataDir: './data', + ...overrides + } +} + +describe('assignment grant logging', () => { + afterEach(() => { + vi.restoreAllMocks() + }) + + it('logs sticky grants with cell id and host digest only', async () => { + const host = 'gggggggggggggggg' + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const app = createRelayApp(config(), { + store: {} as never, + assignments: { + assign: vi.fn(async () => assignment('cell-r', host)), + resolve: vi.fn(async () => assignment('cell-r', host)) + } as never, + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + + const response = await app.request('/v1/assign', assignmentRequest(host, { reconnect: true })) + + expect(response.status).toBe(200) + const line = warn.mock.calls.map((call) => String(call[0])).find((entry) => + entry.includes('assignment granted') + ) + expect(line).toContain('lane=sticky') + expect(line).toContain('cell=cell-r') + expect(line).toContain(`host=${relayHostLogDigest(host)}`) + expect(line).not.toContain(host) + }) + + it('does not log unhinted placement grants', async () => { + const host = 'pppppppppppppppp' + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const app = createRelayApp(config(), { + store: {} as never, + assignments: { + assign: vi.fn(async () => assignment('cell-r', host)), + resolve: vi.fn(async () => null) + } as never, + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + + expect((await app.request('/v1/assign', assignmentRequest(host))).status).toBe(200) + expect( + warn.mock.calls.map((call) => String(call[0])).filter((entry) => + entry.includes('assignment granted') + ) + ).toEqual([]) + }) +}) diff --git a/cloud/apps/relay/src/assignment-store-lock-order.test.ts b/cloud/apps/relay/src/assignment-store-lock-order.test.ts new file mode 100644 index 00000000000..61baedfb1e2 --- /dev/null +++ b/cloud/apps/relay/src/assignment-store-lock-order.test.ts @@ -0,0 +1,486 @@ +import { describe, expect, it } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.js' +import type { RelayDatabase, RelayLockOptions, SqlRow } from './database.js' + +const identity = { userId: 'user-a', relayHostId: 'host000000000001' } +const activityId = 'splice:connection-1' + +class LockOrderDatabase implements RelayDatabase { + readonly lockedTables: string[] = [] + + constructor( + private readonly cleanupCandidate: boolean, + private readonly currentLeaseExpiresAt: number, + private readonly activityLeasePresent = true + ) {} + + async query(sql: string): Promise { + if (sql.includes('SELECT user_id, relay_host_id, activity_id')) { + return this.cleanupCandidate + ? [ + { + user_id: identity.userId, + relay_host_id: identity.relayHostId, + activity_id: activityId + } + ] + : [] + } + if ( + sql.includes('UPDATE relay_cells SET reserved_requests') && + sql.includes('RETURNING cell_id') + ) { + this.lockedTables.push('cell') + return [{ cell_id: 'cell-a' }] + } + return [{ changes: 1 }] + } + + async queryLocked(sql: string): Promise { + if (sql.includes('FROM relay_assignments ')) { + this.lockedTables.push('assignment') + return [{ cell_id: 'cell-a', assignment_epoch: 1 }] + } + if (sql.includes('FROM relay_assignment_activity_leases')) { + this.lockedTables.push('activity') + if (!this.activityLeasePresent) return [] + return [ + { + user_id: identity.userId, + relay_host_id: identity.relayHostId, + activity_id: activityId, + activity_kind: 'splice', + cell_id: 'cell-a', + request_units: 2, + expires_at: this.currentLeaseExpiresAt + } + ] + } + if (sql.trim() === 'SELECT * FROM relay_cells ORDER BY cell_id ASC') { + this.lockedTables.push('cell-inventory') + return [{ cell_id: 'cell-a', reserved_requests: 3, capacity_requests: 10 }] + } + if (sql.includes('FROM relay_cells')) { + this.lockedTables.push('cell') + return [{ reserved_requests: 3, capacity_requests: 10 }] + } + return [] + } + + async transaction(operation: (transaction: RelayDatabase) => Promise): Promise { + return await operation(this) + } + + async close(): Promise {} +} + +class ReassignmentLockOrderDatabase implements RelayDatabase { + readonly locks: string[] = [] + + async query(sql: string): Promise { + if (sql.includes('SELECT cell_id, region FROM relay_cell_regions')) { + return ['cell-a', 'cell-b'].map((cell_id) => ({ cell_id, region: 'us-central1' })) + } + if (sql.includes('SELECT region FROM relay_cell_regions')) { + return [{ region: 'us-central1' }] + } + if (sql.includes('SELECT * FROM relay_cells WHERE cell_id')) { + return [cellRow('cell-b', 1)] + } + if (sql.includes('JOIN relay_cell_runtime') && sql.includes('cell.cell_id = ?')) { + return [] + } + if (sql.includes('SELECT cell_id, observed_requests FROM relay_cell_runtime')) { + return [{ cell_id: 'cell-a', observed_requests: 0 }] + } + if (sql.includes('LEFT JOIN relay_cell_admission')) { + return [ + { cell_id: 'cell-a', admission_state: 'general' }, + { cell_id: 'cell-b', admission_state: 'general' } + ] + } + if (sql.includes('FROM relay_cell_connection_limits')) return [] + return [{ changes: 1 }] + } + + async queryLocked(sql: string, params: unknown[] = []): Promise { + if (sql.includes('FROM relay_assignments')) { + this.locks.push('assignment') + return [ + { + user_id: identity.userId, + relay_host_id: identity.relayHostId, + cell_id: 'cell-b', + assignment_epoch: 1, + lease_expires_at: 100, + last_activity_at: 100, + reserved_controls: 0, + reserved_splices: 0, + reserved_invites: 0, + pending_installs: 0, + pending_confirmations: 0, + migration_leases: 0 + } + ] + } + if (sql.trim() === 'SELECT * FROM relay_cells ORDER BY cell_id ASC') { + this.locks.push('cell-inventory') + return [cellRow('cell-a', 0), cellRow('cell-b', 1)] + } + if (sql.includes('SELECT * FROM relay_cells WHERE cell_id')) { + const cellId = String(params[0]) + this.locks.push(cellId) + return [cellRow(cellId, cellId === 'cell-b' ? 1 : 0)] + } + return [] + } + + async transaction(operation: (transaction: RelayDatabase) => Promise): Promise { + return await operation(this) + } + + async close(): Promise {} +} + +class AggregateCleanupDatabase implements RelayDatabase { + failIfUnavailable: boolean | null = null + + async query(): Promise { + return [{ changes: 1 }] + } + + async queryLocked( + sql: string, + _params: unknown[] = [], + options: RelayLockOptions = {} + ): Promise { + if (sql.includes('FROM relay_assignments WHERE lease_expires_at')) { + this.failIfUnavailable = options.failIfUnavailable ?? false + } + return [] + } + + async transaction(operation: (transaction: RelayDatabase) => Promise): Promise { + return await operation(this) + } + + async close(): Promise {} +} + +class HeartbeatLockDatabase implements RelayDatabase { + readonly locks: string[] = [] + readonly reservationCleanupTransactions: number[] = [] + legacyHeartbeatWritten = false + private transactionNumber = 0 + private activeTransaction = 0 + + constructor( + private readonly cleanupIncarnation = '11111111-1111-4111-8111-111111111111' + ) {} + + async query(sql: string, params: unknown[] = []): Promise { + if (sql.includes('SELECT region FROM relay_cell_regions')) { + return [{ cell_id: String(params[0]), region: 'us-central1' }] + } + if (sql.includes('SELECT * FROM relay_cells WHERE cell_id')) { + return [cellRow('cell-a', 0)] + } + if (sql.includes('UPDATE relay_cells SET observed_requests')) { + this.legacyHeartbeatWritten = true + } + if (sql.includes('UPDATE relay_control_connection_reservations')) { + this.reservationCleanupTransactions.push(this.activeTransaction) + } + return [{ changes: 1 }] + } + + async queryLocked(sql: string): Promise { + if (sql.includes('FROM relay_cells')) { + this.locks.push('cell') + return [cellRow('cell-a', 0)] + } + if (sql.includes('FROM relay_cell_runtime')) { + this.locks.push('runtime') + return this.activeTransaction === 2 + ? [{ cell_incarnation: this.cleanupIncarnation }] + : [] + } + if (sql.includes('FROM relay_cell_connection_limits')) { + this.locks.push('connection-limit') + return [{ hard_cap: 600, unobserved_bound: 99 }] + } + if (sql.includes('FROM relay_cell_connection_snapshots')) { + this.locks.push('snapshot') + return this.activeTransaction === 2 + ? [{ cell_incarnation: this.cleanupIncarnation, inclusion_watermark: 0 }] + : [] + } + return [] + } + + async transaction(operation: (transaction: RelayDatabase) => Promise): Promise { + this.activeTransaction = ++this.transactionNumber + const result = await operation(this) + this.activeTransaction = 0 + return result + } + + async close(): Promise {} +} + +class NewAssignmentLockDatabase implements RelayDatabase { + readonly inventoryLocks: string[] = [] + private generalLockFailed = false + + constructor( + private failGeneralOnce = false, + private readonly assignmentAppearsAfterFailure = false + ) {} + + async query(sql: string, params: unknown[] = []): Promise { + if (sql.includes('SELECT cell_id, region FROM relay_cell_regions')) { + return [ + { cell_id: 'cell-existing', region: 'us-central1' }, + { cell_id: 'cell-general', region: 'us-central1' } + ] + } + if (sql.includes('SELECT region FROM relay_cell_regions')) { + return [{ cell_id: String(params[0]), region: 'us-central1' }] + } + if (sql.includes('SELECT cell_id, observed_requests FROM relay_cell_runtime')) { + return [{ cell_id: 'cell-general', observed_requests: 0 }] + } + if (sql.includes('LEFT JOIN relay_cell_admission')) { + return [{ cell_id: 'cell-general', admission_state: 'general' }] + } + if (sql.includes('JOIN relay_cell_runtime') && sql.includes('cell.cell_id = ?')) { + return [{ cell_id: 'cell-existing' }] + } + if (sql.includes('FROM relay_cell_connection_limits')) { + return [ + { + cell_id: 'cell-general', + hard_cap: 600, + unobserved_bound: 99, + enforced_connection_units: 0, + outstanding_reservations: 0, + last_heartbeat_at: 100, + connection_incarnation: 'incarnation-a', + current_incarnation: 'incarnation-a' + } + ] + } + return [{ changes: 1 }] + } + + async queryLocked(sql: string): Promise { + if ( + sql.includes('FROM relay_assignments') && + this.assignmentAppearsAfterFailure && + this.generalLockFailed + ) { + return [ + { + user_id: identity.userId, + relay_host_id: identity.relayHostId, + cell_id: 'cell-existing', + assignment_epoch: 1, + lease_expires_at: 100, + last_activity_at: 100, + reserved_controls: 1, + reserved_splices: 0, + reserved_invites: 0, + pending_installs: 0, + pending_confirmations: 0, + migration_leases: 0 + } + ] + } + if (sql.includes('SELECT cell_id FROM relay_cell_admission')) { + this.inventoryLocks.push('general') + if (this.failGeneralOnce) { + this.failGeneralOnce = false + this.generalLockFailed = true + throw new Error('database_lock_unavailable') + } + return [cellRow('cell-general', 0)] + } + if (sql.trim() === 'SELECT * FROM relay_cells ORDER BY cell_id ASC') { + this.inventoryLocks.push('all') + return [cellRow('cell-existing', 0), cellRow('cell-general', 0)] + } + if (sql.includes('SELECT * FROM relay_cells WHERE cell_id')) { + return [cellRow('cell-general', 0)] + } + return [] + } + + async transaction(operation: (transaction: RelayDatabase) => Promise): Promise { + return await operation(this) + } + + async close(): Promise {} +} + +function cellRow(cellId: string, reservedRequests: number): SqlRow { + return { + cell_id: cellId, + cell_url: `https://${cellId}.example.com`, + enabled: 1, + capacity_requests: 10, + reserved_requests: reservedRequests, + observed_requests: 0 + } +} + +describe('RelayAssignmentStore activity lock order', () => { + it('updates the cell only after assignment activity during release', async () => { + const database = new LockOrderDatabase(false, 100) + const store = new RelayAssignmentStore(database, () => 100) + + await expect(store.releaseActivity(identity, activityId)).resolves.toBe(true) + expect(database.lockedTables).toEqual(['assignment', 'activity', 'cell']) + }) + + it('updates the cell only after inserting new activity', async () => { + const database = new LockOrderDatabase(false, 100, false) + const store = new RelayAssignmentStore(database, () => 100) + + await expect( + store.acquireActivity(identity, { activityId, kind: 'splice', cellId: 'cell-a' }) + ).resolves.toBeUndefined() + expect(database.lockedTables).toEqual(['assignment', 'activity', 'cell']) + }) + + it('rechecks an expired candidate after taking the assignment lock', async () => { + const database = new LockOrderDatabase(true, 101) + const store = new RelayAssignmentStore(database, () => 100) + + await expect(store.releaseExpiredActivityLeases()).resolves.toBe(0) + expect(database.lockedTables).toEqual(['assignment', 'activity']) + }) + + it('fails fast before aggregate cleanup waits on mixed-version assignment rows', async () => { + const database = new AggregateCleanupDatabase() + const store = new RelayAssignmentStore(database, () => 100) + + await expect(store.releaseExpiredActivity()).resolves.toBe(0) + expect(database.failIfUnavailable).toBe(true) + }) + + it('releases the placement capacity lock before heartbeat reservation cleanup', async () => { + const database = new HeartbeatLockDatabase() + const store = new RelayAssignmentStore(database, () => 100, { requireLiveCells: true }) + + await store.recordCellHeartbeat({ + cellId: 'cell-a', + cellUrl: 'https://cell-a.example.com', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0, + connectionHardCap: 600, + connectionUnobservedBound: 99 + }) + + expect(database.locks).toEqual([ + 'cell', + 'runtime', + 'connection-limit', + 'snapshot', + 'runtime', + 'snapshot' + ]) + expect(database.legacyHeartbeatWritten).toBe(false) + expect(database.reservationCleanupTransactions).toEqual([2, 2, 2]) + }) + + it('does not let an old heartbeat clean replacement-incarnation reservations', async () => { + const database = new HeartbeatLockDatabase('22222222-2222-4222-8222-222222222222') + const store = new RelayAssignmentStore(database, () => 100, { requireLiveCells: true }) + + await store.recordCellHeartbeat({ + cellId: 'cell-a', + cellUrl: 'https://cell-a.example.com', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0, + connectionHardCap: 600, + connectionUnobservedBound: 99 + }) + + expect(database.reservationCleanupTransactions).toEqual([]) + }) + + it('locks only general-admission inventory for a brand-new assignment', async () => { + const database = new NewAssignmentLockDatabase() + const store = new RelayAssignmentStore(database, () => 100, { + requireLiveCells: true, + heartbeatTtlMs: 45 + }) + + await expect(store.assign(identity)).resolves.toMatchObject({ + cellId: 'cell-general', + assignmentEpoch: 1 + }) + expect(database.inventoryLocks).toEqual(['general']) + }) + + it('keeps a brand-new assignment retry scoped to general admission', async () => { + const database = new NewAssignmentLockDatabase(true) + const store = new RelayAssignmentStore(database, () => 100, { + requireLiveCells: true, + heartbeatTtlMs: 45 + }) + + await expect(store.assign(identity)).resolves.toMatchObject({ + cellId: 'cell-general', + assignmentEpoch: 1 + }) + expect(database.inventoryLocks).toEqual(['general', 'general']) + }) + + it('restarts with full inventory if an assignment appears during a general retry', async () => { + const database = new NewAssignmentLockDatabase(true, true) + const store = new RelayAssignmentStore(database, () => 100, { + requireLiveCells: true, + heartbeatTtlMs: 45 + }) + + await expect(store.assign(identity)).resolves.toMatchObject({ + cellId: 'cell-existing', + assignmentEpoch: 1 + }) + expect(database.inventoryLocks).toEqual(['general', 'general', 'all']) + }) + + it('probes the sticky cell before locking full inventory for dead-cell reassignment', async () => { + const database = new ReassignmentLockOrderDatabase() + const store = new RelayAssignmentStore(database, () => 100, { + requireLiveCells: true, + heartbeatTtlMs: 45 + }) + + await expect(store.assign(identity)).resolves.toMatchObject({ + cellId: 'cell-a', + assignmentEpoch: 2 + }) + expect(database.locks).toEqual([ + 'assignment', + 'cell-b', + 'assignment', + 'cell-inventory', + 'cell-b', + 'cell-a' + ]) + }) +}) diff --git a/cloud/apps/relay/src/assignment-store.test.ts b/cloud/apps/relay/src/assignment-store.test.ts new file mode 100644 index 00000000000..e9d27a83db4 --- /dev/null +++ b/cloud/apps/relay/src/assignment-store.test.ts @@ -0,0 +1,4471 @@ +import { ASSIGNMENT_LIMITS } from '@orca-cloud/relay-contract' +import { createHash } from 'node:crypto' +import { afterEach, describe, expect, it, vi } from 'vitest' +import type { RelayCellConfig } from './config.js' +import { RelayAssignmentStore, STRANDED_MIGRATION_ABANDON_MS } from './assignment-store.js' +import { + encodeMembership, + type CellAdmissionMembership +} from './cell-admission-selector.js' +import { + openInMemoryRelayDatabase, + type RelayDatabase, + type RelayLockOptions, + type RelayTransactionOptions, + type SqlRow +} from './database.js' + +const CELLS: RelayCellConfig[] = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 2 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 2 } +] +const NO_EXPIRED_EVACUATION_DIAGNOSTICS = { + expiredUnregistered: 0, + repairableExpiredUnregistered: 0, + abortableExpiredUnregistered: 0, + blockedExpiredUnregistered: 0, + blockedExpiredOnNewerTargetAssignment: 0 +} +const NO_REGISTERED_EVACUATION_DIAGNOSTICS = { + registeredSourceActive: 0, + registeredCompletable: 0, + registeredTargetInactive: 0 +} +const NO_ACTIVE_MIGRATION_LEASE = { + oldestExpiresAt: null, + oldestRemainingMs: null +} +const FENCE_PLAN_BINDING = { + planObjectName: + 'terraform/state/relay-fence-plans/production/22222222-2222-4222-8222-222222222222.tfplan', + planObjectGeneration: '123456789', + varFileSha256: 'c'.repeat(64), + terraformStateLineage: '33333333-3333-4333-8333-333333333333', + terraformStateSerial: 7, + terraformStateObjectGeneration: '987654321', + terraformStateObjectSha256: 'd'.repeat(64), + requestReason: + 'orca-relay-fence/22222222-2222-4222-8222-222222222222' +} +const FENCE_INVOCATION_ID = '55555555-5555-4555-8555-555555555555' +const FENCE_INVOCATION_REASON = + `${FENCE_PLAN_BINDING.requestReason}/${FENCE_INVOCATION_ID}` + +async function applyGenerationZeroSelector( + store: RelayAssignmentStore, + input: { + attemptId: string + membership: CellAdmissionMembership + } +) { + const current = (await store.inspectCellAdmissionSelector()).selector.membership + return await store.applyCellAdmissionSelector({ + ...input, + expectedGeneration: 0, + expectedMembershipSha256: createHash('sha256') + .update(encodeMembership(current)) + .digest('hex') + }) +} + +function cellFenceEvidence(attemptId = '22222222-2222-4222-8222-222222222222') { + return { + attemptId, + environment: 'production' as const, + cellId: 'cell-a', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + migName: 'orca-relay-c1', + instanceGroup: 'https://compute.example/instanceGroups/orca-relay-c1', + generationIdentity: 'https://compute.example/instanceTemplates/orca-relay-c1-abc', + fenceCommit: 'a'.repeat(40), + planSha256: 'b'.repeat(64), + ...FENCE_PLAN_BINDING + } +} + +class OneShotInventoryFailureDatabase implements RelayDatabase { + readonly retryLocks: string[] = [] + private armed = false + private failed = false + + constructor(private readonly delegate: RelayDatabase) {} + + arm(): void { + this.armed = true + } + + async query(sql: string, params: unknown[] = []): Promise { + return await this.delegate.query(sql, params) + } + + async queryLocked( + sql: string, + params: unknown[] = [], + options: RelayLockOptions = {} + ): Promise { + return await this.lockedQuery(this.delegate, sql, params, options) + } + + async transaction(operation: (transaction: RelayDatabase) => Promise): Promise { + return await this.delegate.transaction( + async (transaction) => + await operation({ + query: async (sql, params = []) => await transaction.query(sql, params), + queryLocked: async (sql, params = [], options = {}) => + await this.lockedQuery(transaction, sql, params, options), + transaction: async (nested) => await transaction.transaction(nested), + close: async () => {} + }) + ) + } + + async close(): Promise { + await this.delegate.close() + } + + private async lockedQuery( + database: RelayDatabase, + sql: string, + params: unknown[], + options: RelayLockOptions + ): Promise { + const inventory = sql.trim() === 'SELECT * FROM relay_cells ORDER BY cell_id ASC' + if (this.armed && inventory && options.failIfUnavailable) { + this.armed = false + this.failed = true + this.retryLocks.push('inventory-nowait-failed') + throw new Error('database_lock_unavailable') + } + if (this.failed && inventory) this.retryLocks.push('inventory-first') + if (this.failed && sql.includes('FROM relay_assignments')) { + this.retryLocks.push(options.failIfUnavailable ? 'assignment-nowait' : 'assignment') + } + return await database.queryLocked(sql, params, options) + } +} + +class RepeatedDrainAccountingFailureDatabase implements RelayDatabase { + refreshAttempts = 0 + private armed = false + private failuresRemaining = 0 + + constructor(private readonly delegate: RelayDatabase) {} + + arm(failures: number): void { + this.armed = true + this.failuresRemaining = failures + this.refreshAttempts = 0 + } + + async query(sql: string, params: unknown[] = []): Promise { + return await this.delegate.query(sql, params) + } + + async queryLocked( + sql: string, + params: unknown[] = [], + options: RelayLockOptions = {} + ): Promise { + return await this.lockedQuery(this.delegate, sql, params, options) + } + + async transaction(operation: (transaction: RelayDatabase) => Promise): Promise { + return await this.delegate.transaction( + async (transaction) => + await operation({ + query: async (sql, params = []) => await transaction.query(sql, params), + queryLocked: async (sql, params = [], options = {}) => + await this.lockedQuery(transaction, sql, params, options), + transaction: async (nested) => await transaction.transaction(nested), + close: async () => {} + }) + ) + } + + async close(): Promise { + await this.delegate.close() + } + + private async lockedQuery( + database: RelayDatabase, + sql: string, + params: unknown[], + options: RelayLockOptions + ): Promise { + const drainRefresh = + sql.includes('SELECT migration.*') && + sql.includes('WHERE migration.source_cell_id = ?') && + !sql.includes('source_cell_incarnation') + if (this.armed && drainRefresh) { + this.refreshAttempts++ + if (this.failuresRemaining > 0) { + this.failuresRemaining-- + throw new Error('migration_activity_accounting_mismatch') + } + } + return await database.queryLocked(sql, params, options) + } +} + +// Runs a side effect between the sticky lane and the placement lane, the window +// in which a host can acquire control after the sticky read called it dormant. +class SeedBetweenAssignmentLanesDatabase implements RelayDatabase { + private transactions = 0 + private armed = false + + constructor( + private readonly delegate: RelayDatabase, + private readonly seed: () => Promise + ) {} + + arm(): void { + this.armed = true + this.transactions = 0 + } + + async query(sql: string, params: unknown[] = []): Promise { + return await this.delegate.query(sql, params) + } + + async queryLocked( + sql: string, + params: unknown[] = [], + options: RelayLockOptions = {} + ): Promise { + return await this.delegate.queryLocked(sql, params, options) + } + + async transaction( + operation: (transaction: RelayDatabase) => Promise, + options?: RelayTransactionOptions + ): Promise { + if (this.armed && ++this.transactions === 2) await this.seed() + return await this.delegate.transaction(operation, options) + } + + async close(): Promise { + await this.delegate.close() + } +} + +describe('RelayAssignmentStore', () => { + let database: RelayDatabase | undefined + + afterEach(async () => await database?.close()) + + async function setup(now: () => number, cells = CELLS): Promise { + database = await openInMemoryRelayDatabase() + const store = new RelayAssignmentStore(database, now) + await store.reconcileCells(cells) + return store + } + + async function setupWithHeartbeats( + now: () => number, + cells = CELLS + ): Promise { + database = await openInMemoryRelayDatabase() + const store = new RelayAssignmentStore(database, now, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + await store.reconcileCells(cells) + return store + } + + async function heartbeat( + store: RelayAssignmentStore, + cell = CELLS[0]!, + input: { + incarnation?: string + startedAt?: number + ready?: boolean + observedRequests?: number + } = {} + ): Promise { + await store.recordCellHeartbeat({ + cellId: cell.id, + cellUrl: cell.url, + cellIncarnation: input.incarnation ?? '11111111-1111-4111-8111-111111111111', + startedAt: input.startedAt ?? 50, + ready: input.ready ?? true, + observedRequests: input.observedRequests ?? 0, + ...(cell.connectionHardCap === undefined + ? {} + : { + totalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0, + connectionHardCap: cell.connectionHardCap, + connectionUnobservedBound: cell.connectionUnobservedBound + }) + }) + } + + it('admits only cells with a fresh ready heartbeat', async () => { + let now = 100 + const store = await setupWithHeartbeats(() => now, [CELLS[0]!]) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + + await expect(store.assign(identity)).rejects.toThrow('relay_capacity_exhausted') + await heartbeat(store, CELLS[0]!, { ready: false }) + await expect(store.assign(identity)).rejects.toThrow('relay_capacity_exhausted') + await heartbeat(store) + expect((await store.assign(identity)).cellId).toBe('cell-a') + now += 45_001 + expect(await store.resolve(identity)).toBeNull() + }) + + it('fences stale incarnations and origin mismatches', async () => { + const store = await setupWithHeartbeats(() => 100, [CELLS[0]!]) + await heartbeat(store, CELLS[0]!, { startedAt: 50 }) + await expect( + heartbeat(store, { ...CELLS[0]!, url: 'https://other.example.com' }) + ).rejects.toThrow('cell_origin_mismatch') + await expect( + heartbeat(store, CELLS[0]!, { + incarnation: '22222222-2222-4222-8222-222222222222', + startedAt: 50 + }) + ).rejects.toThrow('stale_cell_incarnation') + await heartbeat(store, CELLS[0]!, { + incarnation: '22222222-2222-4222-8222-222222222222', + startedAt: 51 + }) + await expect(heartbeat(store, CELLS[0]!, { startedAt: 50 })).rejects.toThrow( + 'stale_cell_incarnation' + ) + }) + + it('records receipt-relative legacy drain states and pins post-send migrations', async () => { + let now = 100 + const store = await setupWithHeartbeats(() => now) + await heartbeat(store, CELLS[0]!) + await heartbeat(store, CELLS[1]!, { + incarnation: '22222222-2222-4222-8222-222222222222' + }) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + await store.assign(identity) + await store.setCellEnabled('cell-a', false) + const migration = await store.startEvacuation(identity, 'cell-b') + await database!.query( + `UPDATE relay_assignments SET migration_leases = migration_leases + 1 + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + const incarnation = '11111111-1111-4111-8111-111111111111' + const attemptId = '33333333-3333-4333-8333-333333333333' + const traceValue = '44444444-4444-4444-8444-444444444444' + + await expect( + store.prepareCellDrainAttempt({ + attemptId, + cellId: 'cell-a', + cellIncarnation: incarnation, + traceValue, + plannedGraceMs: 120_000 + }) + ).resolves.toMatchObject({ state: 'prepared', shouldSend: false }) + expect( + await database!.query(`SELECT * FROM relay_post_drain_migration_pins`) + ).toEqual([]) + await expect( + store.prepareCellDrainRecovery({ + attemptId, + cellId: 'cell-a', + cellIncarnation: incarnation + }) + ).resolves.toMatchObject({ + shouldSend: false, + preparedAttempt: { + attemptId, + state: 'prepared', + traceValue, + plannedGraceMs: 120_000 + } + }) + + await expect( + store.beginCellDrainSend({ + attemptId, + cellId: 'cell-a', + cellIncarnation: incarnation + }) + ).resolves.toMatchObject({ + state: 'send-may-have-started', + shouldSend: true, + sendPermitExpiresAt: 30_100 + }) + await expect( + database!.query( + `SELECT migration_leases FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).resolves.toEqual([{ migration_leases: 1 }]) + await expect( + store.beginCellDrainSend({ + attemptId, + cellId: 'cell-a', + cellIncarnation: incarnation + }) + ).resolves.toMatchObject({ + state: 'send-may-have-started', + shouldSend: false + }) + await expect( + store.prepareCellDrainRecovery({ + attemptId, + cellId: 'cell-a', + cellIncarnation: incarnation + }) + ).rejects.toThrow('drain_application_receipt_missing') + expect( + await database!.query( + `SELECT drain_attempt_id, source_cell_id, source_cell_incarnation, + target_cell_id, assignment_epoch + FROM relay_post_drain_migration_pins` + ) + ).toEqual([ + { + drain_attempt_id: attemptId, + source_cell_id: 'cell-a', + source_cell_incarnation: incarnation, + target_cell_id: 'cell-b', + assignment_epoch: migration.assignmentEpoch + } + ]) + + now = migration.expiresAt + 1 + expect(await store.abortExpiredEvacuations()).toBe(0) + expect(await store.releaseExpiredActivityLeases()).toBe(1) + expect( + await database!.query( + `SELECT activity_kind, cell_id FROM relay_assignment_activity_leases + WHERE user_id = ? ORDER BY activity_kind`, + [identity.userId] + ) + ).toEqual([ + { activity_kind: 'control', cell_id: 'cell-b' }, + { activity_kind: 'migration', cell_id: 'cell-b' } + ]) + + now = 200_000 + await expect( + store.recordCellDrainApplicationReceipt({ + attemptId, + cellId: 'cell-a', + cellIncarnation: incarnation, + traceValue, + backendStatus: 200 + }) + ).resolves.toMatchObject({ + state: 'application-receipt', + applicationReceiptAt: 200_000, + retryAfter: 350_000 + }) + await expect( + store.prepareCellDrainRecovery({ attemptId, cellId: 'cell-a', cellIncarnation: incarnation }) + ).rejects.toThrow('drain_recovery_too_early') + now = 350_000 + await expect( + store.prepareCellDrainRecovery({ attemptId, cellId: 'cell-a', cellIncarnation: incarnation }) + ).resolves.toEqual({ + shouldSend: true, + retryAfter: 350_000 + }) + await expect( + store.prepareCellDrainRecovery({ attemptId, cellId: 'cell-a', cellIncarnation: incarnation }) + ).resolves.toEqual({ + shouldSend: false, + retryAfter: 350_000 + }) + }) + + it('recovers a proven drain once after the source cell is replaced', async () => { + let now = 100 + const store = await setupWithHeartbeats(() => now) + const oldIncarnation = '11111111-1111-4111-8111-111111111111' + const newIncarnation = '22222222-2222-4222-8222-222222222222' + const newerIncarnation = '55555555-5555-4555-8555-555555555555' + const attempt = { + attemptId: '33333333-3333-4333-8333-333333333333', + cellId: 'cell-a', + cellIncarnation: oldIncarnation, + traceValue: '44444444-4444-4444-8444-444444444444', + plannedGraceMs: 120_000 + } + await heartbeat(store) + await store.setCellEnabled('cell-a', false) + await store.prepareCellDrainAttempt(attempt) + await store.beginCellDrainSend(attempt) + now = 200 + await store.recordCellDrainApplicationReceipt({ + ...attempt, + backendStatus: 200 + }) + + now = 150_200 + await heartbeat(store, CELLS[0]!, { + incarnation: newIncarnation, + startedAt: 51 + }) + await expect( + store.prepareCellDrainRecovery({ + cellId: 'cell-a', + cellIncarnation: newIncarnation + }) + ).resolves.toEqual({ shouldSend: true, retryAfter: 150_200 }) + await expect( + store.prepareCellDrainRecovery({ + cellId: 'cell-a', + cellIncarnation: newIncarnation + }) + ).resolves.toEqual({ shouldSend: false, retryAfter: 150_200 }) + + now = 150_300 + await heartbeat(store, CELLS[0]!, { + incarnation: newerIncarnation, + startedAt: 52 + }) + const concurrentRecoveries = await Promise.all([ + store.prepareCellDrainRecovery({ + cellId: 'cell-a', + cellIncarnation: newerIncarnation + }), + store.prepareCellDrainRecovery({ + cellId: 'cell-a', + cellIncarnation: newerIncarnation + }) + ]) + expect(concurrentRecoveries.map((recovery) => recovery.shouldSend).sort()).toEqual([ + false, + true + ]) + await expect( + store.prepareCellDrainRecovery({ + cellId: 'cell-a', + cellIncarnation: newerIncarnation + }) + ).resolves.toEqual({ shouldSend: false, retryAfter: 150_200 }) + await expect( + database!.query( + `SELECT cell_incarnation FROM relay_cell_drain_recovery_attempts + WHERE drain_attempt_id = ? ORDER BY cell_incarnation`, + [attempt.attemptId] + ) + ).resolves.toEqual([ + { cell_incarnation: newIncarnation }, + { cell_incarnation: newerIncarnation } + ]) + }) + + it('rejects an invalid receipt and a prepared drain from an old incarnation', async () => { + let now = 100 + const store = await setupWithHeartbeats(() => now) + const oldIncarnation = '11111111-1111-4111-8111-111111111111' + const newIncarnation = '22222222-2222-4222-8222-222222222222' + const attempt = { + attemptId: '33333333-3333-4333-8333-333333333333', + cellId: 'cell-a', + cellIncarnation: oldIncarnation, + traceValue: '44444444-4444-4444-8444-444444444444', + plannedGraceMs: 120_000 + } + await heartbeat(store) + await store.setCellEnabled('cell-a', false) + await store.prepareCellDrainAttempt(attempt) + now = 150_200 + await heartbeat(store, CELLS[0]!, { + incarnation: newIncarnation, + startedAt: 51 + }) + await expect( + store.prepareCellDrainRecovery({ + cellId: 'cell-a', + cellIncarnation: newIncarnation + }) + ).rejects.toThrow('drain_application_receipt_missing') + + await database!.query( + `UPDATE relay_cell_drain_attempt_states + SET state = 'application-receipt', application_receipt_at = ?, + receipt_cell_incarnation = ?, retry_after = ? + WHERE attempt_id = ?`, + [200, oldIncarnation, 150_200, attempt.attemptId] + ) + await expect( + store.prepareCellDrainRecovery({ + cellId: 'cell-a', + cellIncarnation: newIncarnation + }) + ).rejects.toThrow('drain_application_receipt_missing') + }) + + it('restores an expired registered migration lease before a prepared drain send', async () => { + let now = 100 + const store = await setupWithHeartbeats(() => now) + await heartbeat(store, CELLS[0]!) + await heartbeat(store, CELLS[1]!, { + incarnation: '22222222-2222-4222-8222-222222222222' + }) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + await store.assign(identity) + await store.setCellEnabled('cell-a', false) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + const attempt = { + attemptId: '55555555-5555-4555-8555-555555555555', + cellId: 'cell-a', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + traceValue: '66666666-6666-4666-8666-666666666666', + plannedGraceMs: 120_000 + } + await store.prepareCellDrainAttempt(attempt) + + now = migration.expiresAt + 1 + expect(await store.releaseExpiredActivityLeases()).toBeGreaterThan(0) + await heartbeat(store, CELLS[0]!) + await heartbeat(store, CELLS[1]!, { + incarnation: '22222222-2222-4222-8222-222222222222' + }) + await expect( + store.beginCellDrainSend({ + attemptId: attempt.attemptId, + cellId: attempt.cellId, + cellIncarnation: attempt.cellIncarnation + }) + ).resolves.toMatchObject({ + state: 'send-may-have-started', + shouldSend: true + }) + await expect( + database!.query( + `SELECT activity_id, activity_kind, cell_id, request_units, expires_at + FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).resolves.toContainEqual({ + activity_id: `migration:${migration.assignmentEpoch}`, + activity_kind: 'migration', + cell_id: 'cell-b', + request_units: 1, + expires_at: now + ASSIGNMENT_LIMITS.migrationLeaseMs + }) + }) + + it('refreshes registered migration leases before prepared drain recovery', async () => { + let now = 100 + const cappedCells = CELLS.map((cell) => ({ + ...cell, + connectionHardCap: 600 as const, + connectionUnobservedBound: 50 + })) + const store = await setupWithHeartbeats(() => now, cappedCells) + await heartbeat(store, cappedCells[0]!) + await heartbeat(store, cappedCells[1]!, { + incarnation: '22222222-2222-4222-8222-222222222222' + }) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + await store.assign(identity) + await store.setCellEnabled('cell-a', false) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + const attempt = { + attemptId: '99999999-9999-4999-8999-999999999999', + cellId: 'cell-a', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + traceValue: 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa', + plannedGraceMs: 120_000 + } + await store.prepareCellDrainAttempt(attempt) + + now = migration.expiresAt - 500_000 + await heartbeat(store, cappedCells[0]!) + await heartbeat(store, cappedCells[1]!, { + incarnation: '22222222-2222-4222-8222-222222222222' + }) + await expect(store.prepareCellDrainRecovery(attempt)).resolves.toMatchObject({ + shouldSend: false, + preparedAttempt: { attemptId: attempt.attemptId } + }) + const refreshedExpiresAt = now + ASSIGNMENT_LIMITS.migrationLeaseMs + await expect( + database!.query( + `SELECT expires_at FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, migration.assignmentEpoch] + ) + ).resolves.toEqual([{ expires_at: refreshedExpiresAt }]) + await expect( + database!.query( + `SELECT state, timeout_at FROM relay_control_connection_reservations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, migration.assignmentEpoch] + ) + ).resolves.toEqual([{ state: 'reserved', timeout_at: refreshedExpiresAt }]) + }) + + it('bounds repeated accounting repair before prepared drain recovery', async () => { + const delegate = await openInMemoryRelayDatabase() + const retryDatabase = new RepeatedDrainAccountingFailureDatabase(delegate) + database = retryDatabase + const store = new RelayAssignmentStore(retryDatabase, () => 100, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + await store.reconcileCells(CELLS) + await heartbeat(store, CELLS[0]!) + await heartbeat(store, CELLS[1]!, { + incarnation: '22222222-2222-4222-8222-222222222222' + }) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + await store.assign(identity) + await store.setCellEnabled('cell-a', false) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + const attempt = { + attemptId: '99999999-9999-4999-8999-999999999999', + cellId: 'cell-a', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + traceValue: 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa', + plannedGraceMs: 120_000 + } + await store.prepareCellDrainAttempt(attempt) + retryDatabase.arm(2) + + await expect(store.prepareCellDrainRecovery(attempt)).resolves.toMatchObject({ + shouldSend: false, + preparedAttempt: { attemptId: attempt.attemptId } + }) + expect(retryDatabase.refreshAttempts).toBe(3) + + retryDatabase.arm(4) + await expect(store.prepareCellDrainRecovery(attempt)).rejects.toThrow( + 'migration_activity_accounting_mismatch' + ) + expect(retryDatabase.refreshAttempts).toBe(4) + }) + + it('retires a pinned migration superseded by a newer authoritative assignment', async () => { + let now = 100 + const cells = [ + ...CELLS, + { + id: 'cell-c', + url: 'https://relay-c.example.com', + capacityRequests: 2 + } + ] + const store = await setupWithHeartbeats(() => now, cells) + await heartbeat(store, cells[0]!) + await heartbeat(store, cells[1]!, { + incarnation: '22222222-2222-4222-8222-222222222222' + }) + await heartbeat(store, cells[2]!, { + incarnation: '33333333-3333-4333-8333-333333333333' + }) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + await store.assign(identity) + await store.setCellEnabled('cell-a', false) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + const attempt = { + attemptId: '99999999-9999-4999-8999-999999999999', + cellId: 'cell-a', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + traceValue: 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa', + plannedGraceMs: 120_000 + } + await store.prepareCellDrainAttempt(attempt) + await store.beginCellDrainSend(attempt) + const receipt = await store.recordCellDrainApplicationReceipt({ + ...attempt, + backendStatus: 200 + }) + + await database!.query( + `DELETE FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + await database!.query( + `UPDATE relay_assignments SET reserved_controls = 0, reserved_splices = 0, + reserved_invites = 0, pending_installs = 0, pending_confirmations = 0, + migration_leases = 0, lease_expires_at = 0, last_activity_at = 0 + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + await database!.query( + `UPDATE relay_cells SET reserved_requests = 0 WHERE cell_id IN (?, ?)`, + ['cell-a', 'cell-b'] + ) + now = ASSIGNMENT_LIMITS.dormantTtlMs + 1 + await heartbeat(store, cells[2]!, { + incarnation: '33333333-3333-4333-8333-333333333333' + }) + const current = await store.rebalanceDormant(identity, 'cell-c') + await store.activateControl(identity, { + cellId: 'cell-c', + assignmentEpoch: current.assignmentEpoch, + generation: 1 + }) + + expect(receipt.retryAfter).toBeLessThanOrEqual(now) + await heartbeat(store, cells[0]!) + await expect(store.prepareCellDrainRecovery(attempt)).resolves.toMatchObject({ + shouldSend: true + }) + await expect( + database!.query( + `SELECT aborted_at FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, migration.assignmentEpoch] + ) + ).resolves.toEqual([{ aborted_at: now }]) + await expect( + database!.query( + `SELECT cell_id, assignment_epoch FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).resolves.toEqual([ + { cell_id: 'cell-c', assignment_epoch: current.assignmentEpoch } + ]) + await expect( + database!.query( + `SELECT activity_id, cell_id FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).resolves.toEqual([{ activity_id: 'control:cell-c:1', cell_id: 'cell-c' }]) + }) + + it('refuses to restore a missing migration lease without durable target registration', async () => { + let now = 100 + const store = await setupWithHeartbeats(() => now) + await heartbeat(store, CELLS[0]!) + await heartbeat(store, CELLS[1]!, { + incarnation: '22222222-2222-4222-8222-222222222222' + }) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + await store.assign(identity) + await store.setCellEnabled('cell-a', false) + const migration = await store.startEvacuation(identity, 'cell-b') + const attempt = { + attemptId: '77777777-7777-4777-8777-777777777777', + cellId: 'cell-a', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + traceValue: '88888888-8888-4888-8888-888888888888', + plannedGraceMs: 120_000 + } + await store.prepareCellDrainAttempt(attempt) + + now = migration.expiresAt + 1 + expect(await store.releaseExpiredActivityLeases()).toBeGreaterThan(0) + await heartbeat(store, CELLS[0]!) + await heartbeat(store, CELLS[1]!, { + incarnation: '22222222-2222-4222-8222-222222222222' + }) + await expect( + store.beginCellDrainSend({ + attemptId: attempt.attemptId, + cellId: attempt.cellId, + cellIncarnation: attempt.cellIncarnation + }) + ).rejects.toThrow('migration_activity_lease_shape_mismatch') + }) + + it('moves a dead assignment only after a completed exact-incarnation fence attempt', async () => { + let now = 100 + const cappedCells = CELLS.map((cell) => ({ + ...cell, + connectionHardCap: 600 as const, + connectionUnobservedBound: 50 + })) + const store = await setupWithHeartbeats(() => now, cappedCells) + await heartbeat(store, cappedCells[0]!) + await heartbeat(store, cappedCells[1]!, { + incarnation: '22222222-2222-4222-8222-222222222222' + }) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + await store.assign(identity) + await store.setCellEnabled('cell-a', false) + now += 45_001 + await heartbeat(store, cappedCells[1]!, { + incarnation: '22222222-2222-4222-8222-222222222222' + }) + + await store.attestCellFence('cell-a', '11111111-1111-4111-8111-111111111111') + expect( + await database!.query( + `SELECT cell_id FROM relay_cell_legacy_fence_adoptions WHERE cell_id = ?`, + ['cell-a'] + ) + ).toEqual([]) + expect(await store.evacuateDeadCells()).toBe(0) + expect( + await database!.query( + `SELECT cell_id FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ cell_id: 'cell-a' }]) + + now += 300_001 + await heartbeat(store, cappedCells[1]!, { + incarnation: '22222222-2222-4222-8222-222222222222' + }) + const evidence = { + attemptId: '22222222-2222-4222-8222-222222222222', + environment: 'production' as const, + cellId: 'cell-a', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + migName: 'orca-relay-c1', + instanceGroup: 'https://compute.example/instanceGroups/orca-relay-c1', + generationIdentity: 'https://compute.example/instanceTemplates/orca-relay-c1-abc', + fenceCommit: 'a'.repeat(40), + planSha256: 'b'.repeat(64), + ...FENCE_PLAN_BINDING + } + await store.prepareCellFenceAttempt(evidence) + await store.bindCellFencePlanGeneration(evidence, evidence.planObjectGeneration) + await store.startCellFenceApply( + evidence, + FENCE_INVOCATION_ID, + FENCE_INVOCATION_REASON + ) + await store.recordCellFenceOperation( + evidence, + FENCE_INVOCATION_ID, + FENCE_INVOCATION_REASON, + 'operation-1' + ) + await store.attestCellFenceAttempt(evidence, 'operation-1') + + expect(await store.evacuateDeadCells()).toBe(1) + expect((await store.resolve(identity))?.cellId).toBe('cell-b') + }) + + it('preserves proven non-delivery and permits one fresh full-grace attempt', async () => { + const store = await setupWithHeartbeats(() => 100, [CELLS[0]!]) + await heartbeat(store) + await store.setCellEnabled('cell-a', false) + const first = { + attemptId: '33333333-3333-4333-8333-333333333333', + cellId: 'cell-a', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + traceValue: '44444444-4444-4444-8444-444444444444', + plannedGraceMs: 120_000 + } + await store.prepareCellDrainAttempt(first) + await store.beginCellDrainSend(first) + await expect(store.proveCellDrainNotDelivered(first)).resolves.toMatchObject({ + state: 'proven-not-delivered', + provenNotDeliveredAt: 100 + }) + await expect( + store.prepareCellDrainAttempt({ + ...first, + attemptId: '55555555-5555-4555-8555-555555555555', + traceValue: '66666666-6666-4666-8666-666666666666' + }) + ).resolves.toMatchObject({ state: 'prepared' }) + }) + + it('binds fence attestation to one durable Terraform attempt', async () => { + let now = 100 + const store = await setupWithHeartbeats(() => now, [CELLS[0]!]) + await heartbeat(store) + await store.setCellEnabled('cell-a', false) + const evidence = { + attemptId: '22222222-2222-4222-8222-222222222222', + environment: 'production' as const, + cellId: 'cell-a', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + migName: 'orca-relay-c1', + instanceGroup: 'https://compute.example/instanceGroups/orca-relay-c1', + generationIdentity: 'https://compute.example/instanceTemplates/orca-relay-c1-abc', + fenceCommit: 'a'.repeat(40), + planSha256: 'b'.repeat(64), + ...FENCE_PLAN_BINDING + } + const prepared = await store.prepareCellFenceAttempt(evidence) + expect(prepared).toMatchObject({ createdAt: 100, expiresAt: 3_600_100 }) + expect(prepared.planObjectGeneration).toBeUndefined() + await expect( + store.bindCellFencePlanGeneration( + evidence, + evidence.planObjectGeneration + ) + ).resolves.toMatchObject({ + planObjectGeneration: evidence.planObjectGeneration + }) + await expect( + store.bindCellFencePlanGeneration( + evidence, + evidence.planObjectGeneration + ) + ).resolves.toMatchObject({ + planObjectGeneration: evidence.planObjectGeneration + }) + await expect( + store.bindCellFencePlanGeneration( + evidence, + '987654321' + ) + ).rejects.toThrow('cell_fence_plan_generation_mismatch') + for (const changed of [ + { attemptId: '33333333-3333-4333-8333-333333333333' }, + { environment: 'staging' as const }, + { cellId: 'cell-b' }, + { cellIncarnation: '33333333-3333-4333-8333-333333333333' }, + { migName: 'orca-relay-other' }, + { instanceGroup: `${evidence.instanceGroup}-other` }, + { generationIdentity: `${evidence.generationIdentity}-other` }, + { fenceCommit: 'c'.repeat(40) }, + { planSha256: 'c'.repeat(64) } + ]) { + await expect( + store.startCellFenceApply( + { ...evidence, ...changed }, + FENCE_INVOCATION_ID, + FENCE_INVOCATION_REASON + ) + ).rejects.toThrow() + } + await store.startCellFenceApply( + evidence, + FENCE_INVOCATION_ID, + FENCE_INVOCATION_REASON + ) + await store.recordCellFenceOperation( + evidence, + FENCE_INVOCATION_ID, + FENCE_INVOCATION_REASON, + 'operation-1' + ) + now += 45_001 + await expect( + store.attestCellFenceAttempt(evidence, 'operation-other') + ).rejects.toThrow('cell_fence_operation_not_attested') + const attested = await store.attestCellFenceAttempt(evidence, 'operation-1') + expect(attested).toMatchObject({ + expiresAt: now + 300_000, + attempt: { ...evidence, gceOperation: 'operation-1', completedAt: now } + }) + await expect(store.attestCellFenceAttempt(evidence, 'operation-1')).resolves.toEqual( + attested + ) + now += 300_001 + await expect( + store.attestCellFenceAttempt(evidence, 'operation-1') + ).resolves.toMatchObject({ + expiresAt: now + 300_000, + attempt: { completedAt: attested.attempt.completedAt } + }) + }) + + it('aborts only a Terraform fence whose apply never started', async () => { + const store = await setupWithHeartbeats(() => 100, [CELLS[0]!]) + await heartbeat(store) + await store.setCellEnabled('cell-a', false) + const evidence = { + attemptId: '22222222-2222-4222-8222-222222222222', + environment: 'production' as const, + cellId: 'cell-a', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + migName: 'orca-relay-c1', + instanceGroup: 'https://compute.example/instanceGroups/orca-relay-c1', + generationIdentity: 'https://compute.example/instanceTemplates/orca-relay-c1-abc', + fenceCommit: 'a'.repeat(40), + planSha256: 'b'.repeat(64), + ...FENCE_PLAN_BINDING + } + await store.prepareCellFenceAttempt(evidence) + await store.bindCellFencePlanGeneration( + evidence, + evidence.planObjectGeneration + ) + await expect(store.abortCellFenceAttempt(evidence)).resolves.toMatchObject({ + abortedAt: 100 + }) + await expect( + store.startCellFenceApply( + evidence, + FENCE_INVOCATION_ID, + FENCE_INVOCATION_REASON + ) + ).rejects.toThrow( + 'cell_fence_attempt_aborted' + ) + }) + + it('adopts a stale disabled legacy fence only when no durable attempt exists', async () => { + let now = 100 + const store = await setupWithHeartbeats(() => now, [CELLS[0]!]) + await heartbeat(store) + await store.setCellEnabled('cell-a', false) + now += 45_001 + + await expect( + store.adoptLegacyCellFence('cell-a', '11111111-1111-4111-8111-111111111111') + ).resolves.toBe(now + 300_000) + await expect( + database!.query( + `SELECT cell_id, cell_incarnation FROM relay_cell_fences WHERE cell_id = ?`, + ['cell-a'] + ) + ).resolves.toEqual([ + { + cell_id: 'cell-a', + cell_incarnation: '11111111-1111-4111-8111-111111111111' + } + ]) + await expect( + database!.query( + `SELECT cell_id, cell_incarnation + FROM relay_cell_legacy_fence_adoptions WHERE cell_id = ?`, + ['cell-a'] + ) + ).resolves.toEqual([]) + await store.commitLegacyCellFenceAdoption( + 'cell-a', + '11111111-1111-4111-8111-111111111111' + ) + await expect( + database!.query( + `SELECT cell_id, cell_incarnation + FROM relay_cell_legacy_fence_adoptions WHERE cell_id = ?`, + ['cell-a'] + ) + ).resolves.toEqual([ + { + cell_id: 'cell-a', + cell_incarnation: '11111111-1111-4111-8111-111111111111' + } + ]) + + now += 1 + await expect( + store.adoptLegacyCellFence('cell-a', '11111111-1111-4111-8111-111111111111') + ).resolves.toBe(now + 300_000) + await store.commitLegacyCellFenceAdoption( + 'cell-a', + '11111111-1111-4111-8111-111111111111' + ) + await expect( + database!.query( + `SELECT attested_at, expires_at + FROM relay_cell_legacy_fence_adoptions WHERE cell_id = ?`, + ['cell-a'] + ) + ).resolves.toEqual([{ attested_at: now, expires_at: now + 300_000 }]) + + await heartbeat(store) + await expect( + database!.query( + `SELECT cell_id FROM relay_cell_legacy_fence_adoptions WHERE cell_id = ?`, + ['cell-a'] + ) + ).resolves.toEqual([]) + }) + + it.each(['active', 'completed', 'aborted'] as const)( + 'rejects legacy adoption after a %s durable attempt', + async (status) => { + let now = 100 + const store = await setupWithHeartbeats(() => now, [CELLS[0]!]) + await heartbeat(store) + await store.setCellEnabled('cell-a', false) + const evidence = cellFenceEvidence() + await store.prepareCellFenceAttempt(evidence) + if (status === 'aborted') await store.abortCellFenceAttempt(evidence) + if (status === 'completed') { + await store.bindCellFencePlanGeneration(evidence, evidence.planObjectGeneration) + await store.startCellFenceApply( + evidence, + FENCE_INVOCATION_ID, + FENCE_INVOCATION_REASON + ) + await store.recordCellFenceOperation( + evidence, + FENCE_INVOCATION_ID, + FENCE_INVOCATION_REASON, + 'operation-1' + ) + } + now += 45_001 + if (status === 'completed') { + await store.attestCellFenceAttempt(evidence, 'operation-1') + } + + await expect( + store.adoptLegacyCellFence('cell-a', '11111111-1111-4111-8111-111111111111') + ).rejects.toThrow('legacy_cell_fence_attempt_exists') + } + ) + + it('serializes legacy adoption against durable attempt preparation', async () => { + let now = 100 + const store = await setupWithHeartbeats(() => now, [CELLS[0]!]) + await heartbeat(store) + await store.setCellEnabled('cell-a', false) + now += 45_001 + + const results = await Promise.allSettled([ + store.adoptLegacyCellFence('cell-a', '11111111-1111-4111-8111-111111111111'), + store.prepareCellFenceAttempt(cellFenceEvidence()) + ]) + expect(results.filter(({ status }) => status === 'fulfilled')).toHaveLength(1) + const attempts = await database!.query( + `SELECT COUNT(*) AS count FROM relay_cell_fence_attempts WHERE cell_id = ?`, + ['cell-a'] + ) + const fences = await database!.query( + `SELECT COUNT(*) AS count FROM relay_cell_fences WHERE cell_id = ?`, + ['cell-a'] + ) + const adoptions = await database!.query( + `SELECT COUNT(*) AS count FROM relay_cell_legacy_fence_adoptions + WHERE cell_id = ?`, + ['cell-a'] + ) + expect(Number(attempts[0]!.count) + Number(fences[0]!.count)).toBe(1) + expect(Number(adoptions[0]!.count)).toBe(0) + }) + + it('rejects an expired durable Terraform fence attempt', async () => { + let now = 100 + const store = await setupWithHeartbeats(() => now, [CELLS[0]!]) + await heartbeat(store) + await store.setCellEnabled('cell-a', false) + const evidence = { + attemptId: '22222222-2222-4222-8222-222222222222', + environment: 'production' as const, + cellId: 'cell-a', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + migName: 'orca-relay-c1', + instanceGroup: 'https://compute.example/instanceGroups/orca-relay-c1', + generationIdentity: 'https://compute.example/instanceTemplates/orca-relay-c1-abc', + fenceCommit: 'a'.repeat(40), + planSha256: 'b'.repeat(64), + ...FENCE_PLAN_BINDING + } + await store.prepareCellFenceAttempt(evidence) + await store.bindCellFencePlanGeneration( + evidence, + evidence.planObjectGeneration + ) + now += 3_600_001 + await expect( + store.startCellFenceApply( + evidence, + FENCE_INVOCATION_ID, + FENCE_INVOCATION_REASON + ) + ).rejects.toThrow( + 'cell_fence_attempt_expired' + ) + }) + + it('keeps a started Terraform fence attempt recoverable after its preparation TTL', async () => { + let now = 100 + const store = await setupWithHeartbeats(() => now, [CELLS[0]!]) + await heartbeat(store) + await store.setCellEnabled('cell-a', false) + const evidence = { + attemptId: '22222222-2222-4222-8222-222222222222', + environment: 'production' as const, + cellId: 'cell-a', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + migName: 'orca-relay-c1', + instanceGroup: 'https://compute.example/instanceGroups/orca-relay-c1', + generationIdentity: 'https://compute.example/instanceTemplates/orca-relay-c1-abc', + fenceCommit: 'a'.repeat(40), + planSha256: 'b'.repeat(64), + ...FENCE_PLAN_BINDING + } + await store.prepareCellFenceAttempt(evidence) + await store.bindCellFencePlanGeneration( + evidence, + evidence.planObjectGeneration + ) + await store.startCellFenceApply( + evidence, + FENCE_INVOCATION_ID, + FENCE_INVOCATION_REASON + ) + now += 3_600_001 + await expect( + store.recordCellFenceOperation( + evidence, + FENCE_INVOCATION_ID, + FENCE_INVOCATION_REASON, + 'operation-1' + ) + ).resolves.toMatchObject({ + attempt: { gceOperation: 'operation-1' }, + invocation: { gceOperation: 'operation-1' } + }) + await expect( + store.prepareCellFenceAttempt({ + ...evidence, + attemptId: '44444444-4444-4444-8444-444444444444', + planObjectName: + 'terraform/state/relay-fence-plans/production/44444444-4444-4444-8444-444444444444.tfplan', + requestReason: + 'orca-relay-fence/44444444-4444-4444-8444-444444444444' + }) + ).rejects.toThrow('cell_fence_attempt_evidence_mismatch') + }) + + it('reports aggregate deployment state and heartbeat freshness without identities', async () => { + let now = 100 + const store = await setupWithHeartbeats(() => now, [CELLS[0]!]) + await heartbeat(store) + const identity = { userId: 'private-user', relayHostId: 'host000000000001' } + const assignment = await store.assign(identity) + expect(await store.cellDeploymentStatus('cell-a')).toMatchObject({ + activityLeases: 1, + activityRequestUnits: 1, + restartBlockingActivityLeases: 0, + restartBlockingActivityRequestUnits: 0, + restartBlockingReservedRequests: 0 + }) + await store.activateControl(identity, { + cellId: 'cell-a', + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + + expect(await store.cellDeploymentStatus('cell-a')).toEqual({ + cellId: 'cell-a', + cellUrl: 'https://relay-a.example.com', + region: 'us-central1', + enabled: true, + admissionState: 'general', + capacityRequests: 2, + reservedRequests: 1, + assignments: 1, + activityLeases: 1, + activityRequestUnits: 1, + restartBlockingActivityLeases: 1, + restartBlockingActivityRequestUnits: 1, + restartBlockingReservedRequests: 1, + outgoingMigrations: 0, + incomingMigrations: 0, + connectionCapacity: null, + runtime: { + cellUrl: 'https://relay-a.example.com', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + lastHeartbeatAt: 100, + heartbeatFresh: true, + regionalRehomeProtocol: 0 + } + }) + now += 45_001 + expect((await store.cellDeploymentStatus('cell-a')).runtime?.heartbeatFresh).toBe(false) + await expect(store.cellDeploymentStatus('missing')).rejects.toThrow('cell_not_found') + }) + + it('keeps concurrent sticky assignment grants restart-safe', async () => { + const store = await setup(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 20 } + ]) + const identity = { userId: 'private-user', relayHostId: 'host000000000001' } + + await Promise.all(Array.from({ length: 20 }, async () => await store.assign(identity))) + + expect(await store.cellDeploymentStatus('cell-a')).toMatchObject({ + activityLeases: 1, + activityRequestUnits: 1, + reservedRequests: 1, + restartBlockingActivityLeases: 0, + restartBlockingActivityRequestUnits: 0, + restartBlockingReservedRequests: 0 + }) + }) + + it('classifies activated and non-control activity as restart-blocking', async () => { + const store = await setup(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 20 } + ]) + const identity = { userId: 'private-user', relayHostId: 'host000000000001' } + const assignment = await store.assign(identity) + await store.activateControl(identity, { + cellId: assignment.cellId, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + for (const kind of ['splice', 'invite', 'install', 'confirmation', 'migration'] as const) { + await store.acquireActivity(identity, { + activityId: `${kind}:test`, + kind, + cellId: assignment.cellId + }) + } + + expect(await store.cellDeploymentStatus('cell-a')).toMatchObject({ + activityLeases: 6, + activityRequestUnits: 7, + reservedRequests: 7, + restartBlockingActivityLeases: 6, + restartBlockingActivityRequestUnits: 7, + restartBlockingReservedRequests: 7 + }) + }) + + it('fails closed for malformed pending controls and unexplained reservations', async () => { + const store = await setup(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 20 } + ]) + await store.assign({ userId: 'private-user', relayHostId: 'host000000000001' }) + await database!.query( + `UPDATE relay_assignment_activity_leases SET request_units = 2 + WHERE cell_id = ?`, + ['cell-a'] + ) + expect(await store.cellDeploymentStatus('cell-a')).toMatchObject({ + restartBlockingActivityLeases: 1, + restartBlockingActivityRequestUnits: 2, + restartBlockingReservedRequests: 1 + }) + + await database!.query( + `UPDATE relay_assignment_activity_leases SET request_units = 1 WHERE cell_id = ?`, + ['cell-a'] + ) + await database!.query( + `UPDATE relay_cells SET reserved_requests = 0 WHERE cell_id = ?`, + ['cell-a'] + ) + expect(await store.cellDeploymentStatus('cell-a')).toMatchObject({ + restartBlockingActivityLeases: 0, + restartBlockingActivityRequestUnits: 0, + restartBlockingReservedRequests: -1 + }) + + await database!.query( + `UPDATE relay_cells SET reserved_requests = 2 WHERE cell_id = ?`, + ['cell-a'] + ) + expect(await store.cellDeploymentStatus('cell-a')).toMatchObject({ + restartBlockingActivityLeases: 0, + restartBlockingActivityRequestUnits: 0, + restartBlockingReservedRequests: 1 + }) + }) + + it('keeps a migration blocking when its target also has a pending control', async () => { + const store = await setup(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 20 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 20 } + ]) + const identity = { userId: 'private-user', relayHostId: 'host000000000001' } + await store.assign(identity) + await store.startEvacuation(identity, 'cell-b') + + expect(await store.cellDeploymentStatus('cell-b')).toMatchObject({ + activityLeases: 2, + restartBlockingActivityLeases: 1, + restartBlockingActivityRequestUnits: 1, + restartBlockingReservedRequests: 1, + incomingMigrations: 1 + }) + }) + + it('starts declared candidates disabled without overwriting later operator state', async () => { + const candidate: RelayCellConfig = { + id: 'candidate', + url: 'https://candidate.example.com', + capacityRequests: 20, + initiallyEnabled: false + } + const store = await setup(() => 100, [candidate]) + expect((await store.cellDeploymentStatus('candidate')).enabled).toBe(false) + await store.setCellEnabled('candidate', true) + await store.reconcileCells([candidate]) + expect((await store.cellDeploymentStatus('candidate')).enabled).toBe(true) + }) + + it('moves a stranded host off an existing-only cell that stopped serving it', async () => { + // Why: C3's decommission created an existing-only cell that rejects + // attaches; the #194 pin then loops returning hosts forever while each + // grant refreshes their own activity (issue #225). C3 has no + // connection-limits row and expired leases are cleaned within a + // maintenance cycle, so the only durable evidence is the assignment row: + // a recent grant with no live real activity behind it. + let now = 100 + const store = await setup(() => now) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + await store.setCellEnabled(first.cellId, false) + + // The pin holds while the last grant is young enough to still attach. + now += 10_000 + const pinned = await store.assign(identity) + expect(pinned.cellId).toBe(first.cellId) + + // The grant had ample time to attach and produced nothing live. + now += 61_000 + const moved = await store.assign(identity) + expect(moved.cellId).not.toBe(first.cellId) + expect(moved.assignmentEpoch).toBe(first.assignmentEpoch + 1) + + // The move is durable: no bounce back to the closed cell. + now += 1_000 + const settled = await store.assign(identity) + expect(settled.cellId).toBe(moved.cellId) + expect(settled.assignmentEpoch).toBe(moved.assignmentEpoch) + }) + + it('keeps a serving existing-only cell pinned for hosts with live activity', async () => { + let now = 100 + const store = await setup(() => now) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + await store.setCellEnabled(first.cellId, false) + // A live claimed control proves the cell still serves this host. + await database!.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, ?, 'control', ?, 1, ?, ?)`, + [identity.userId, identity.relayHostId, 'control:live-1', first.cellId, now + 600_000, now] + ) + now += 61_000 + const pinned = await store.assign(identity) + expect(pinned.cellId).toBe(first.cellId) + }) + + it('keeps migration-only cells pinned inside the stranded window', async () => { + let now = 100 + const store = await setup(() => now) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + await store.setCellAdmissionState(first.cellId, 'migration-only') + now += 61_000 + const pinned = await store.assign(identity) + expect(pinned.cellId).toBe(first.cellId) + }) + + it('leaves quiet existing-only assignments to the normal dormancy rule', async () => { + // Why: outside the 15-minute stranded window there is no active retry + // loop to break; ordinary returns stay pinned until 24h dormancy. + let now = 100 + const store = await setup(() => now) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + await store.setCellEnabled(first.cellId, false) + now += 20 * 60_000 + const pinned = await store.assign(identity) + expect(pinned.cellId).toBe(first.cellId) + }) + + it('reassigns an active assignment from an uncapped stale cell without a fence', async () => { + let now = 100 + const store = await setupWithHeartbeats(() => now) + await heartbeat(store, CELLS[0]!) + await heartbeat(store, CELLS[1]!, { + incarnation: '22222222-2222-4222-8222-222222222222' + }) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + await store.changeActivity(identity, 'invite', 1) + await database!.query( + `INSERT INTO relay_assignment_migrations + (user_id, relay_host_id, source_cell_id, target_cell_id, + previous_epoch, assignment_epoch, source_request_units, + target_reserved_units, expires_at, target_registered_at, + completed_at, aborted_at, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, NULL, NULL, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + 'cell-b', + 'cell-a', + -1, + 0, + 0, + 1, + 90_100, + now, + now + ] + ) + await database!.query( + `INSERT INTO relay_post_drain_migration_pins + (user_id, relay_host_id, assignment_epoch, drain_attempt_id, + source_cell_id, source_cell_incarnation, target_cell_id, + target_cell_incarnation, source_request_units, target_reserved_units, + pinned_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + 0, + '33333333-3333-4333-8333-333333333333', + 'cell-b', + '22222222-2222-4222-8222-222222222222', + 'cell-a', + '11111111-1111-4111-8111-111111111111', + 0, + 1, + now + ] + ) + + now += 45_001 + await heartbeat(store, CELLS[1]!, { + incarnation: '22222222-2222-4222-8222-222222222222' + }) + expect(await store.evacuateDeadCells()).toBe(1) + const replacement = await store.resolve(identity) + expect(replacement).toMatchObject({ + cellId: 'cell-b', + assignmentEpoch: first.assignmentEpoch + 1 + }) + expect( + await database!.query( + `SELECT activity_kind, cell_id FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ activity_kind: 'control', cell_id: 'cell-b' }]) + }) + + it('does not let an unfenced capped cell starve eligible legacy recovery', async () => { + let now = 100 + const cells: RelayCellConfig[] = [ + { + id: 'cell-a', + url: 'https://relay-a.example.com', + capacityRequests: 10, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 }, + { id: 'cell-c', url: 'https://relay-c.example.com', capacityRequests: 10 } + ] + const store = await setupWithHeartbeats(() => now, cells) + for (const cell of cells) await heartbeat(store, cell) + await store.setCellEnabled('cell-b', false) + await store.setCellEnabled('cell-c', false) + const cappedIdentity = { userId: 'a-capped', relayHostId: 'host000000000001' } + expect(await store.assign(cappedIdentity)).toMatchObject({ cellId: 'cell-a' }) + + await store.setCellEnabled('cell-a', false) + await store.setCellEnabled('cell-b', true) + const legacyIdentity = { userId: 'z-legacy', relayHostId: 'host000000000002' } + expect(await store.assign(legacyIdentity)).toMatchObject({ cellId: 'cell-b' }) + await store.setCellEnabled('cell-c', true) + + now += 45_001 + await heartbeat(store, cells[2]!) + expect(await store.evacuateDeadCells(1)).toBe(1) + expect(await store.resolve(legacyIdentity)).toMatchObject({ cellId: 'cell-c' }) + expect( + await database!.query( + `SELECT cell_id FROM relay_assignments WHERE user_id = ?`, + [cappedIdentity.userId] + ) + ).toEqual([{ cell_id: 'cell-a' }]) + }) + + it('refuses evacuation into a cell without a fresh ready heartbeat', async () => { + const store = await setupWithHeartbeats(() => 100) + await heartbeat(store, CELLS[0]!) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + await store.assign(identity) + + await expect(store.startEvacuation(identity, 'cell-b')).rejects.toThrow( + 'target_cell_unavailable' + ) + }) + + it('keeps an active assignment sticky without increasing its epoch or reservation', async () => { + const store = await setup(() => 100) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + const second = await store.assign(identity) + + expect(second).toEqual(first) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 1, 'cell-b': 0 }) + }) + + it('selects the least-loaded cell with a stable cell-id tie break', async () => { + const store = await setup(() => 100) + const first = await store.assign({ userId: 'user-a', relayHostId: 'host000000000001' }) + const second = await store.assign({ userId: 'user-b', relayHostId: 'host000000000002' }) + + expect([first.cellId, second.cellId]).toEqual(['cell-a', 'cell-b']) + }) + + it('admits exactly through the configured capacity boundary', async () => { + const store = await setup(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 2 } + ]) + await store.assign({ userId: 'user-a', relayHostId: 'host000000000001' }) + await store.assign({ userId: 'user-b', relayHostId: 'host000000000002' }) + + await expect( + store.assign({ userId: 'user-c', relayHostId: 'host000000000003' }) + ).rejects.toThrow('relay_capacity_exhausted') + expect(await cellReservations(database!)).toEqual({ 'cell-a': 2 }) + }) + + it('does not oversubscribe when assignments race', async () => { + const store = await setup(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 2 } + ]) + const results = await Promise.allSettled( + Array.from({ length: 8 }, (_, index) => + store.assign({ userId: `user-${index}`, relayHostId: `host${String(index).padStart(12, '0')}` }) + ) + ) + + expect(results.filter(({ status }) => status === 'fulfilled')).toHaveLength(2) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 2 }) + }) + + it('reassigns only after all activity expires and the dormant TTL elapses', async () => { + let now = 100 + const store = await setup(() => now) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + await store.changeActivity(identity, 'invite', 1) + await store.changeActivity(identity, 'control', -1) + + now += ASSIGNMENT_LIMITS.activityLeaseMs + 1 + await store.releaseExpiredActivityLeases() + await store.releaseExpiredActivity() + await database!.query(`UPDATE relay_cells SET observed_requests = 2 WHERE cell_id = ?`, [ + first.cellId + ]) + now += ASSIGNMENT_LIMITS.dormantTtlMs + const reassigned = await store.assign(identity) + + expect(reassigned.cellId).not.toBe(first.cellId) + expect(reassigned.assignmentEpoch).toBe(first.assignmentEpoch + 1) + }) + + it('will not normally move an assignment while any durable activity remains', async () => { + let now = 100 + const store = await setup(() => now) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + await store.changeActivity(identity, 'migration', 1) + await store.changeActivity(identity, 'control', -1) + await database!.query(`UPDATE relay_cells SET observed_requests = 2 WHERE cell_id = ?`, [ + first.cellId + ]) + now += ASSIGNMENT_LIMITS.dormantTtlMs + 1 + + expect(await store.assign(identity)).toMatchObject({ + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch + }) + }) + + it('releases every expired activity reservation without going negative', async () => { + let now = 100 + const store = await setup(() => now, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 } + ]) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + await store.assign(identity) + for (const kind of ['splice', 'invite', 'install', 'confirmation', 'migration'] as const) { + await store.changeActivity(identity, kind, 1) + } + expect(await cellReservations(database!)).toEqual({ 'cell-a': 7 }) + + now += ASSIGNMENT_LIMITS.activityLeaseMs + 1 + expect(await store.releaseExpiredActivityLeases()).toBe(1) + expect(await store.releaseExpiredActivity()).toBe(1) + expect(await store.releaseExpiredActivity()).toBe(0) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 0 }) + }) + + it('isolates assignments by host and verifies both cell and epoch', async () => { + const store = await setup(() => 100) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const assignment = await store.assign(identity) + + await expect( + store.verifyCellAssignment({ ...identity, cellId: assignment.cellId, assignmentEpoch: 1 }) + ).resolves.toBe(true) + await expect( + store.verifyCellAssignment({ ...identity, cellId: 'cell-b', assignmentEpoch: 1 }) + ).resolves.toBe(false) + await expect( + store.verifyCellAssignment({ ...identity, cellId: assignment.cellId, assignmentEpoch: 2 }) + ).resolves.toBe(false) + await expect( + store.resolve({ userId: 'other', relayHostId: identity.relayHostId }) + ).resolves.toBeNull() + }) + + it('converts the director reservation into an idempotent active-control lease', async () => { + const store = await setup(() => 100) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const assignment = await store.assign(identity) + + const activityId = await store.activateControl(identity, { + cellId: assignment.cellId, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + await store.activateControl(identity, { + cellId: assignment.cellId, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + + expect(activityId).toBe(`control:${assignment.cellId}:1`) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 1, 'cell-b': 0 }) + const leases = await database!.query( + `SELECT activity_id FROM relay_assignment_activity_leases` + ) + expect(leases).toEqual([{ activity_id: activityId }]) + }) + + it('transactionally supersedes older controls on only the same cell', async () => { + const cells = [ + { + id: 'cell-a', + url: 'https://relay-a.example.com', + capacityRequests: 10 + }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ] + const store = await setup(() => 100, cells) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + await store.setCellEnabled('cell-b', false) + const assignment = await store.assign(identity) + await store.activateControl(identity, { + cellId: assignment.cellId, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + await store.setCellEnabled('cell-b', true) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await database!.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, ?, 'control', 'cell-b', 1, 90100, 100), + (?, ?, ?, 'control', 'cell-b', 1, 90100, 100)`, + [ + identity.userId, + identity.relayHostId, + 'control:cell-b:2', + identity.userId, + identity.relayHostId, + 'control:cell-b:3' + ] + ) + const latest = await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 4 + }) + + expect( + await database!.query( + `SELECT activity_id, cell_id FROM relay_assignment_activity_leases + WHERE activity_kind = 'control' ORDER BY cell_id` + ) + ).toEqual([ + { activity_id: 'control:cell-a:1', cell_id: 'cell-a' }, + { activity_id: latest, cell_id: 'cell-b' } + ]) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 1, 'cell-b': 2 }) + expect( + await database!.query( + `SELECT reserved_controls FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ reserved_controls: 2 }]) + }) + + it('re-reserves a sticky grant for a host holding no control lease', async () => { + const store = await setup(() => 100) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const assignment = await store.assign(identity) + const control = await store.activateControl(identity, { + cellId: assignment.cellId, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + await store.releaseActivity(identity, control) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 0, 'cell-b': 0 }) + + await expect(store.assign(identity)).resolves.toMatchObject({ + cellId: assignment.cellId, + assignmentEpoch: assignment.assignmentEpoch + }) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 1, 'cell-b': 0 }) + expect( + await database!.query( + `SELECT activity_id FROM relay_assignment_activity_leases` + ) + ).toEqual([{ activity_id: 'control-pending:1' }]) + expect( + await database!.query( + `SELECT reserved_controls FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ reserved_controls: 1 }]) + }) + + it('re-pins a host that took control after the sticky lane read it dormant', async () => { + const now = 100 + const delegate = await openInMemoryRelayDatabase() + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const seeded = new SeedBetweenAssignmentLanesDatabase(delegate, async () => { + await delegate.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, 'control:cell-a:1', 'control', 'cell-a', 1, ?, ?), + (?, ?, 'control-pending:1', 'control', 'cell-a', 1, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + now + ASSIGNMENT_LIMITS.activityLeaseMs, + now, + identity.userId, + identity.relayHostId, + now + ASSIGNMENT_LIMITS.activityLeaseMs, + now + ] + ) + await delegate.query( + `UPDATE relay_assignments SET reserved_controls = 2 + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + await delegate.query( + `UPDATE relay_cells SET reserved_requests = 2 WHERE cell_id = 'cell-a'` + ) + }) + database = seeded + const store = new RelayAssignmentStore(seeded, () => now) + await store.reconcileCells(CELLS) + await delegate.query( + `INSERT INTO relay_assignments + (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, + last_activity_at, reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases) + VALUES (?, ?, 'cell-a', 1, ?, ?, 0, 0, 0, 0, 0, 0)`, + [identity.userId, identity.relayHostId, now, now - ASSIGNMENT_LIMITS.dormantTtlMs] + ) + seeded.arm() + + await expect(store.assign(identity)).resolves.toMatchObject({ + cellId: 'cell-a', + assignmentEpoch: 1 + }) + expect( + await delegate.query( + `SELECT reserved_controls FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ reserved_controls: 2 }]) + expect(await cellReservations(delegate)).toEqual({ 'cell-a': 2, 'cell-b': 0 }) + expect( + await delegate.query( + `SELECT lease_expires_at, last_activity_at FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([ + { lease_expires_at: now + ASSIGNMENT_LIMITS.activityLeaseMs, last_activity_at: now } + ]) + }) + + it('reserves control on the pinned cell when the only lease sits on another', async () => { + const now = 100 + const store = await setup(() => now) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + await database!.query( + `INSERT INTO relay_assignments + (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, + last_activity_at, reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases) + VALUES (?, ?, 'cell-b', 2, ?, ?, 1, 0, 0, 0, 0, 0)`, + [ + identity.userId, + identity.relayHostId, + now + ASSIGNMENT_LIMITS.activityLeaseMs, + now + ] + ) + await database!.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, 'control:cell-a:1', 'control', 'cell-a', 1, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + now + ASSIGNMENT_LIMITS.activityLeaseMs, + now + ] + ) + await database!.query( + `UPDATE relay_cells SET reserved_requests = 1 WHERE cell_id = 'cell-a'` + ) + + await expect(store.assign(identity)).resolves.toMatchObject({ + cellId: 'cell-b', + assignmentEpoch: 2 + }) + expect( + await database!.query( + `SELECT activity_id, cell_id FROM relay_assignment_activity_leases + ORDER BY activity_id ASC` + ) + ).toEqual([ + { activity_id: 'control-pending:2', cell_id: 'cell-b' }, + { activity_id: 'control:cell-a:1', cell_id: 'cell-a' } + ]) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 1, 'cell-b': 1 }) + expect( + await database!.query( + `SELECT reserved_controls FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ reserved_controls: 2 }]) + }) + + it('mints a pending control when the only lease is an older epoch pending', async () => { + const now = 100 + const store = await setup(() => now) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + await database!.query( + `INSERT INTO relay_assignments + (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, + last_activity_at, reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases) + VALUES (?, ?, 'cell-a', 3, ?, ?, 1, 0, 0, 0, 0, 0)`, + [ + identity.userId, + identity.relayHostId, + now + ASSIGNMENT_LIMITS.activityLeaseMs, + now + ] + ) + await database!.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, 'control-pending:1', 'control', 'cell-a', 1, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + now + ASSIGNMENT_LIMITS.activityLeaseMs, + now + ] + ) + await database!.query( + `UPDATE relay_cells SET reserved_requests = 1 WHERE cell_id = 'cell-a'` + ) + + await expect(store.assign(identity)).resolves.toMatchObject({ + cellId: 'cell-a', + assignmentEpoch: 3 + }) + expect( + await database!.query( + `SELECT activity_id FROM relay_assignment_activity_leases + ORDER BY activity_id ASC` + ) + ).toEqual([{ activity_id: 'control-pending:1' }, { activity_id: 'control-pending:3' }]) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 2, 'cell-b': 0 }) + expect( + await database!.query( + `SELECT reserved_controls FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ reserved_controls: 2 }]) + }) + + it('re-reserves a re-pinned grant for a host holding no control lease', async () => { + const now = 100 + const delegate = await openInMemoryRelayDatabase() + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const seeded = new SeedBetweenAssignmentLanesDatabase(delegate, async () => { + await delegate.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, 'splice:connection-1', 'splice', 'cell-a', 1, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + now + ASSIGNMENT_LIMITS.activityLeaseMs, + now + ] + ) + await delegate.query( + `UPDATE relay_assignments SET reserved_splices = 1 + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + await delegate.query( + `UPDATE relay_cells SET reserved_requests = 1 WHERE cell_id = 'cell-a'` + ) + }) + database = seeded + const store = new RelayAssignmentStore(seeded, () => now) + await store.reconcileCells(CELLS) + await delegate.query( + `INSERT INTO relay_assignments + (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, + last_activity_at, reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases) + VALUES (?, ?, 'cell-a', 1, ?, ?, 0, 0, 0, 0, 0, 0)`, + [identity.userId, identity.relayHostId, now, now - ASSIGNMENT_LIMITS.dormantTtlMs] + ) + seeded.arm() + + await expect(store.assign(identity)).resolves.toMatchObject({ + cellId: 'cell-a', + assignmentEpoch: 1 + }) + expect( + await delegate.query( + `SELECT reserved_controls, reserved_splices FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ reserved_controls: 1, reserved_splices: 1 }]) + expect( + await delegate.query( + `SELECT activity_id FROM relay_assignment_activity_leases + ORDER BY activity_id ASC` + ) + ).toEqual([{ activity_id: 'control-pending:1' }, { activity_id: 'splice:connection-1' }]) + expect(await cellReservations(delegate)).toEqual({ 'cell-a': 2, 'cell-b': 0 }) + }) + + it('extends the lease of a control-holding host on a sticky grant', async () => { + let now = 100 + const store = await setup(() => now) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const assignment = await store.assign(identity) + await store.activateControl(identity, { + cellId: assignment.cellId, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + + now = 40_000 + await expect(store.assign(identity)).resolves.toMatchObject({ + cellId: assignment.cellId + }) + expect( + await database!.query( + `SELECT lease_expires_at, last_activity_at FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([ + { lease_expires_at: now + ASSIGNMENT_LIMITS.activityLeaseMs, last_activity_at: now } + ]) + }) + + // The old absolute write shortened a 15-minute migration lease to the 90s + // grant lease; only the touch's monotonic CASE keeps the longer one now. + it('never shortens a migration lease to the grant lease', async () => { + let now = 100 + const store = await setup(() => now) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const assignment = await store.assign(identity) + await store.activateControl(identity, { + cellId: assignment.cellId, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + const before = await database!.query( + `SELECT lease_expires_at FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + expect(before).toEqual([{ lease_expires_at: 100 + ASSIGNMENT_LIMITS.migrationLeaseMs }]) + + now = 1000 + await expect(store.assign(identity)).resolves.toMatchObject({ cellId: 'cell-b' }) + expect( + await database!.query( + `SELECT lease_expires_at, last_activity_at FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([ + { lease_expires_at: 100 + ASSIGNMENT_LIMITS.migrationLeaseMs, last_activity_at: now } + ]) + }) + + it('moves an origin-scoped activity reservation without double counting', async () => { + const store = await setup(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ]) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + await store.assign(identity) + await store.acquireActivity(identity, { + activityId: 'splice:connection-1', + kind: 'splice', + cellId: 'cell-a' + }) + await store.startEvacuation(identity, 'cell-b') + await store.acquireActivity(identity, { + activityId: 'splice:connection-1', + kind: 'splice', + cellId: 'cell-b' + }) + + expect(await cellReservations(database!)).toEqual({ 'cell-a': 1, 'cell-b': 6 }) + await expect(store.releaseActivity(identity, 'splice:connection-1')).resolves.toBe(true) + await expect(store.releaseActivity(identity, 'splice:connection-1')).resolves.toBe(false) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 1, 'cell-b': 4 }) + }) + + it('rejects a durable activity before it can exceed cell capacity', async () => { + const store = await setup(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 2 } + ]) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + await store.assign(identity) + + await expect( + store.acquireActivity(identity, { + activityId: 'splice:connection-1', + kind: 'splice', + cellId: 'cell-a' + }) + ).rejects.toThrow('relay_capacity_exhausted') + expect(await cellReservations(database!)).toEqual({ 'cell-a': 1 }) + }) + + it('moves active assignments target-first and completes only after source drain', async () => { + const store = await setup(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ]) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + const sourceControl = await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + + const migration = await store.startEvacuation(identity, 'cell-b') + expect(await store.startEvacuation(identity, 'cell-b')).toEqual(migration) + expect(migration).toMatchObject({ + sourceCellId: 'cell-a', + targetCellId: 'cell-b', + previousEpoch: 1, + assignmentEpoch: 2 + }) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 1, 'cell-b': 2 }) + + const targetControl = await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: 2, + generation: 1 + }) + await store.markMigrationTargetRegistered(identity, { cellId: 'cell-b', assignmentEpoch: 2 }) + await expect(store.completeEvacuation(identity, 2)).rejects.toThrow( + 'migration_source_still_active' + ) + await store.releaseActivity(identity, sourceControl) + await store.completeEvacuation(identity, 2) + + expect(await cellReservations(database!)).toEqual({ 'cell-a': 0, 'cell-b': 1 }) + expect(await store.resolve(identity)).toMatchObject({ cellId: 'cell-b', assignmentEpoch: 2 }) + expect( + await database!.query( + `SELECT activity_id FROM relay_assignment_activity_leases ORDER BY activity_id` + ) + ).toEqual([{ activity_id: targetControl }]) + await expect( + store.acquireActivity(identity, { + activityId: 'splice:late-source', + kind: 'splice', + cellId: 'cell-a' + }) + ).rejects.toThrow('activity_cell_not_authoritative') + }) + + it('automatically completes only after the source runtime is freshly quiescent', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ] + const store = await setupWithHeartbeats(() => now, cells) + await heartbeat(store, cells[0]!) + await heartbeat(store, cells[1]!) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + const sourceControl = await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + await store.releaseActivity(identity, sourceControl) + await store.setCellEnabled('cell-a', false) + await heartbeat(store, cells[0]!, { observedRequests: 1 }) + + expect(await store.completeReadyEvacuations()).toBe(0) + await heartbeat(store, cells[0]!, { observedRequests: 0 }) + now = 150 + await heartbeat(store, cells[1]!, { + incarnation: '22222222-2222-4222-8222-222222222222', + startedAt: 125 + }) + expect(await store.completeReadyEvacuations()).toBe(0) + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 2 + }) + expect(await store.completeReadyEvacuations()).toBe(1) + }) + + it('completes a fenced migration only after the source heartbeat is stale', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ] + const store = await setupWithHeartbeats(() => now, cells) + await heartbeat(store, cells[0]!, { observedRequests: 1 }) + await heartbeat(store, cells[1]!) + await store.setCellEnabled('cell-b', false) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + const sourceControl = await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + await store.setCellEnabled('cell-b', true) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + await store.releaseActivity(identity, sourceControl) + await store.setCellEnabled('cell-a', false) + + expect(await store.cellEvacuationStatus('cell-a', 'cell-b', true)).toMatchObject({ + inProgress: 1, + blocked: 1 + }) + now += 45_001 + await heartbeat(store, cells[1]!) + await store.attestCellFence( + 'cell-a', + '11111111-1111-4111-8111-111111111111' + ) + expect(await store.cellEvacuationStatus('cell-a', 'cell-b', true)).toMatchObject({ + inProgress: 0, + completed: 1 + }) + }) + + it('blocks aggregate fenced completion on assignment accounting mismatch', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ] + const store = await setupWithHeartbeats(() => now, cells) + await heartbeat(store, cells[0]!) + await heartbeat(store, cells[1]!) + await store.setCellEnabled('cell-b', false) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + const sourceControl = await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + await store.setCellEnabled('cell-b', true) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + await store.releaseActivity(identity, sourceControl) + await store.setCellEnabled('cell-a', false) + await database!.query( + `UPDATE relay_assignments SET reserved_splices = reserved_splices + 1 + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + now += 45_001 + await heartbeat(store, cells[1]!) + await store.attestCellFence( + 'cell-a', + '11111111-1111-4111-8111-111111111111' + ) + + expect(await store.cellEvacuationStatus('cell-a', 'cell-b', true)).toMatchObject({ + inProgress: 1, + blocked: 1 + }) + }) + + it('blocks aggregate fenced completion on an unexpected third-cell lease', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 }, + { id: 'cell-c', url: 'https://relay-c.example.com', capacityRequests: 10 } + ] + const store = await setupWithHeartbeats(() => now, cells) + for (const cell of cells) await heartbeat(store, cell) + await store.setCellEnabled('cell-b', false) + await store.setCellEnabled('cell-c', false) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + const sourceControl = await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + await store.setCellEnabled('cell-b', true) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + await store.releaseActivity(identity, sourceControl) + await store.setCellEnabled('cell-a', false) + await database!.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + 'splice:unexpected-cell', + 'splice', + 'cell-c', + 2, + now + ASSIGNMENT_LIMITS.activityLeaseMs, + now + ] + ) + await database!.query( + `UPDATE relay_assignments SET reserved_splices = reserved_splices + 1 + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + await database!.query( + `UPDATE relay_cells SET reserved_requests = reserved_requests + 2 WHERE cell_id = ?`, + ['cell-c'] + ) + now += 45_001 + await heartbeat(store, cells[1]!) + await store.attestCellFence( + 'cell-a', + '11111111-1111-4111-8111-111111111111' + ) + + expect(await store.cellEvacuationStatus('cell-a', 'cell-b', true)).toMatchObject({ + inProgress: 1, + blocked: 1 + }) + }) + + it('rolls back an unregistered expired evacuation with a strictly newer epoch', async () => { + let now = 100 + const store = await setup(() => now, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ]) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + await store.startEvacuation(identity, 'cell-b') + + now += ASSIGNMENT_LIMITS.migrationLeaseMs + 1 + expect(await store.abortExpiredEvacuations()).toBe(1) + expect(await store.resolve(identity)).toMatchObject({ cellId: 'cell-a', assignmentEpoch: 3 }) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 1, 'cell-b': 0 }) + }) + + // Why: completion needs an enabled target and expiry rollback needs an unregistered one, so a + // target disabled after it registered satisfies neither and the migration never leaves the table. + async function wedgeRegisteredMigration( + now: () => number, + cells: RelayCellConfig[], + disableTarget = true, + disableSource = true, + releaseSourceControl = true + ): Promise<{ + store: RelayAssignmentStore + identity: { userId: string; relayHostId: string } + }> { + const store = await setupWithHeartbeats(now, cells) + await heartbeat(store, cells[0]!) + await heartbeat(store, cells[1]!) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + const sourceControl = await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + const migration = await store.startEvacuation(identity, 'cell-b') + const targetControl = await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + if (releaseSourceControl) await store.releaseActivity(identity, sourceControl) + if (disableSource) await store.setCellEnabled('cell-a', false) + // The desktop leaves the half-migrated target, then an operator disables that cell. + await store.releaseActivity(identity, targetControl) + if (disableTarget) await store.setCellEnabled('cell-b', false) + return { store, identity } + } + + async function insertReservedControlConnection( + identity: { userId: string; relayHostId: string }, + cellId: string, + assignmentEpoch: number, + now: number + ): Promise { + await database!.query( + `INSERT INTO relay_control_connection_reservations + (reservation_id, idempotency_key, user_id, relay_host_id, assignment_epoch, + cell_id, state, created_at, timeout_at, updated_at) + VALUES ('abandoned-reservation', 'abandoned-key', ?, ?, ?, ?, 'reserved', ?, ?, ?)`, + [identity.userId, identity.relayHostId, assignmentEpoch, cellId, now, now + 60_000, now] + ) + } + + it('reaps an expired migration whose registered target cell was disabled', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ] + const { store, identity } = await wedgeRegisteredMigration(() => now, cells) + + now += ASSIGNMENT_LIMITS.migrationLeaseMs + STRANDED_MIGRATION_ABANDON_MS + 1 + // A disabled target can never satisfy completion, so waiting cannot help. + expect(await store.completeReadyEvacuations()).toBe(0) + expect(await store.abortExpiredEvacuations()).toBe(1) + // Rollback lands on the disabled source, so the host resolves to nothing and is placed afresh + // by evacuateDeadCells; the point is that the row is retired rather than reaped every tick. + expect(await store.resolve(identity)).toBeNull() + expect(await store.abortExpiredEvacuations()).toBe(0) + }) + + it('leaves a freshly disabled target for supersede-target to recover', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ] + const { store } = await wedgeRegisteredMigration(() => now, cells) + + // Expired, but a disabled target is what supersede-target itself creates, so the operator + // window must stay open; aborting here would fail their run with migration_already_superseded. + now += ASSIGNMENT_LIMITS.migrationLeaseMs + 1 + expect(await store.abortExpiredEvacuations()).toBe(0) + + now += STRANDED_MIGRATION_ABANDON_MS + expect(await store.abortExpiredEvacuations()).toBe(1) + }) + + it('rolls an abandoned disabled target back while its source remains active', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ] + const { store, identity } = await wedgeRegisteredMigration( + () => now, + cells, + true, + false, + false + ) + + now += ASSIGNMENT_LIMITS.migrationLeaseMs + STRANDED_MIGRATION_ABANDON_MS + 1 + expect(await store.abortExpiredEvacuations()).toBe(1) + expect( + await database!.query( + `SELECT cell_id, assignment_epoch FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ cell_id: 'cell-a', assignment_epoch: 3 }]) + expect( + await database!.query( + `SELECT completed_at, aborted_at FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ completed_at: null, aborted_at: now }]) + }) + + it('starts the abandon window when an old migration target is disabled', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ] + const { store } = await wedgeRegisteredMigration(() => now, cells, false, false) + now += ASSIGNMENT_LIMITS.migrationLeaseMs + STRANDED_MIGRATION_ABANDON_MS + 1 + + await store.setCellEnabled('cell-b', false) + expect(await store.abortExpiredEvacuations()).toBe(0) + + now += STRANDED_MIGRATION_ABANDON_MS + 1 + expect(await store.abortExpiredEvacuations()).toBe(1) + }) + + it('reaps an expired migration from a retired source when its target control is gone', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ] + const { store, identity } = await wedgeRegisteredMigration(() => now, cells, false) + await applyGenerationZeroSelector(store, { + attemptId: 'retired_source_migration_target', + membership: { + existingOnly: ['cell-a'], + migrationOnly: ['cell-b'], + general: [] + } + }) + await insertReservedControlConnection(identity, 'cell-b', 2, now) + + now += ASSIGNMENT_LIMITS.migrationLeaseMs + STRANDED_MIGRATION_ABANDON_MS + 1 + expect(await store.abortExpiredEvacuations()).toBe(1) + expect( + await database!.query( + `SELECT assignment.cell_id, assignment.assignment_epoch, + migration.completed_at, migration.aborted_at + FROM relay_assignments assignment + JOIN relay_assignment_migrations migration + ON migration.user_id = assignment.user_id + AND migration.relay_host_id = assignment.relay_host_id + WHERE assignment.user_id = ? AND assignment.relay_host_id = ? + AND migration.assignment_epoch = ?`, + [identity.userId, identity.relayHostId, 2] + ) + ).toEqual([ + { + cell_id: 'cell-b', + assignment_epoch: 2, + completed_at: now, + aborted_at: null + } + ]) + expect( + await database!.query( + `SELECT reserved_controls, migration_leases FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ reserved_controls: 0, migration_leases: 0 }]) + expect( + await database!.query( + `SELECT activity_id FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([]) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 0, 'cell-b': 0 }) + expect( + await database!.query( + `SELECT state FROM relay_control_connection_reservations + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ state: 'released' }]) + }) + + it('reaps a retired-side migration with only a pending target control', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ] + const store = await setupWithHeartbeats(() => now, cells) + await heartbeat(store, cells[0]!) + await heartbeat(store, cells[1]!) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + const sourceControl = await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + await store.releaseActivity(identity, sourceControl) + await applyGenerationZeroSelector(store, { + attemptId: 'retired_source_pending_target', + membership: { + existingOnly: ['cell-a'], + migrationOnly: ['cell-b'], + general: [] + } + }) + await insertReservedControlConnection(identity, 'cell-b', migration.assignmentEpoch, now) + + now += ASSIGNMENT_LIMITS.migrationLeaseMs + STRANDED_MIGRATION_ABANDON_MS + 1 + await database!.query( + `UPDATE relay_assignment_activity_leases SET expires_at = ? + WHERE user_id = ? AND relay_host_id = ? AND activity_id = ?`, + [now, identity.userId, identity.relayHostId, `control-pending:${migration.assignmentEpoch}`] + ) + expect(await store.abortExpiredEvacuations()).toBe(1) + expect( + await database!.query( + `SELECT cell_id, assignment_epoch FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ cell_id: 'cell-b', assignment_epoch: migration.assignmentEpoch }]) + expect( + await database!.query( + `SELECT reserved_controls, migration_leases FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ reserved_controls: 0, migration_leases: 0 }]) + expect( + await database!.query( + `SELECT activity_id FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([]) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 0, 'cell-b': 0 }) + expect( + await database!.query( + `SELECT state FROM relay_control_connection_reservations + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ state: 'released' }]) + }) + + it('preserves a fresh target grant created after abandoned migration selection', async () => { + let now = 100 + const cells: RelayCellConfig[] = [ + { + id: 'cell-a', + url: 'https://relay-a.example.com', + capacityRequests: 10, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }, + { + id: 'cell-b', + url: 'https://relay-b.example.com', + capacityRequests: 10, + connectionHardCap: 600, + connectionUnobservedBound: 50 + } + ] + const { store, identity } = await wedgeRegisteredMigration(() => now, cells, false) + await applyGenerationZeroSelector(store, { + attemptId: 'retired_source_fresh_target_grant', + membership: { + existingOnly: ['cell-a'], + migrationOnly: ['cell-b'], + general: [] + } + }) + + now += 45_001 + await heartbeat(store, cells[1]!) + await store.adoptLegacyCellFence( + 'cell-a', + '11111111-1111-4111-8111-111111111111' + ) + await store.commitLegacyCellFenceAdoption( + 'cell-a', + '11111111-1111-4111-8111-111111111111' + ) + now += ASSIGNMENT_LIMITS.migrationLeaseMs + 1 + await heartbeat(store, cells[1]!) + const query = database!.query.bind(database) + let grant: Awaited> | undefined + vi.spyOn(database!, 'query').mockImplementationOnce(async (sql, params) => { + const rows = await query(sql, params) + grant = await store.assign(identity) + return rows + }) + + expect(await store.abortExpiredEvacuations()).toBe(0) + expect(grant).toMatchObject({ cellId: 'cell-b', assignmentEpoch: 2 }) + expect( + await database!.query( + `SELECT activity_id, expires_at FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? AND activity_id = ?`, + [identity.userId, identity.relayHostId, 'control-pending:2'] + ) + ).toEqual([{ activity_id: 'control-pending:2', expires_at: grant!.leaseExpiresAt }]) + expect( + await database!.query( + `SELECT state FROM relay_control_connection_reservations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ? + AND cell_id = ? AND state = 'reserved'`, + [identity.userId, identity.relayHostId, 2, 'cell-b'] + ) + ).toEqual([{ state: 'reserved' }]) + + await store.activateControl(identity, { + cellId: grant!.cellId, + assignmentEpoch: grant!.assignmentEpoch, + generation: 2 + }) + expect( + await database!.query( + `SELECT activity_id FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? AND activity_kind = 'control'`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ activity_id: 'control:cell-b:2' }]) + expect( + await database!.query( + `SELECT completed_at FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, 2] + ) + ).toEqual([{ completed_at: null }]) + }) + + it('keeps an abandoned target migration open while its source is still active', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ] + const { store, identity } = await wedgeRegisteredMigration( + () => now, + cells, + false, + true, + false + ) + await applyGenerationZeroSelector(store, { + attemptId: 'retired_source_active', + membership: { + existingOnly: ['cell-a'], + migrationOnly: ['cell-b'], + general: [] + } + }) + + now += ASSIGNMENT_LIMITS.migrationLeaseMs + STRANDED_MIGRATION_ABANDON_MS + 1 + expect(await store.abortExpiredEvacuations()).toBe(0) + expect( + await database!.query( + `SELECT completed_at, aborted_at FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ completed_at: null, aborted_at: null }]) + }) + + it('starts the abandon window when an old migration source is retired', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ] + const { store } = await wedgeRegisteredMigration(() => now, cells, false, false) + now += ASSIGNMENT_LIMITS.migrationLeaseMs + STRANDED_MIGRATION_ABANDON_MS + 1 + + await store.setCellEnabled('cell-a', false) + expect(await store.abortExpiredEvacuations()).toBe(0) + + now += STRANDED_MIGRATION_ABANDON_MS + 1 + expect(await store.abortExpiredEvacuations()).toBe(1) + }) + + it('reaps after an old source migration lease is refreshed and expires', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ] + const { store } = await wedgeRegisteredMigration(() => now, cells, false) + await applyGenerationZeroSelector(store, { + attemptId: 'retired_source_refreshed_lease', + membership: { + existingOnly: ['cell-a'], + migrationOnly: ['cell-b'], + general: [] + } + }) + now += STRANDED_MIGRATION_ABANDON_MS + ASSIGNMENT_LIMITS.migrationLeaseMs + 1 + await database!.query( + `UPDATE relay_assignment_migrations SET expires_at = ?`, + [now - 1] + ) + + expect(await store.abortExpiredEvacuations()).toBe(1) + }) + + it('keeps a freshly retired target open despite an old retired source', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ] + const { store } = await wedgeRegisteredMigration(() => now, cells, false) + now += ASSIGNMENT_LIMITS.migrationLeaseMs + STRANDED_MIGRATION_ABANDON_MS + 1 + + await store.setCellEnabled('cell-b', false) + expect(await store.abortExpiredEvacuations()).toBe(0) + + now += STRANDED_MIGRATION_ABANDON_MS + 1 + expect(await store.abortExpiredEvacuations()).toBe(1) + }) + + it('fails supersession closed when reaping wins after selection', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 }, + { id: 'cell-c', url: 'https://relay-c.example.com', capacityRequests: 10 } + ] + const { store } = await wedgeRegisteredMigration(() => now, cells) + now += ASSIGNMENT_LIMITS.migrationLeaseMs + STRANDED_MIGRATION_ABANDON_MS + 1 + const query = database!.query.bind(database) + vi.spyOn(database!, 'query').mockImplementationOnce(async (sql, params) => { + const rows = await query(sql, params) + expect(await store.abortExpiredEvacuations()).toBe(1) + return rows + }) + + await expect( + store.supersedeRegisteredCellEvacuations('cell-a', 'cell-b', 'cell-c', 100) + ).rejects.toThrow('migration_already_superseded') + }) + + it('fails supersession closed when reaping wins after preparation', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 }, + { id: 'cell-c', url: 'https://relay-c.example.com', capacityRequests: 10 } + ] + const { store } = await wedgeRegisteredMigration(() => now, cells) + now += ASSIGNMENT_LIMITS.migrationLeaseMs + STRANDED_MIGRATION_ABANDON_MS + 1 + await store.attestCellFence('cell-b', '11111111-1111-4111-8111-111111111111') + const supersede = store.supersedeRegisteredEvacuation.bind(store) + vi.spyOn(store, 'supersedeRegisteredEvacuation').mockImplementationOnce(async (...args) => { + expect(await store.abortExpiredEvacuations()).toBe(1) + return await supersede(...args) + }) + + await expect( + store.supersedeRegisteredCellEvacuations('cell-a', 'cell-b', 'cell-c', 100) + ).rejects.toThrow('migration_already_superseded') + }) + + it('fails supersession closed when reaping wins before an accounting retry', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 }, + { id: 'cell-c', url: 'https://relay-c.example.com', capacityRequests: 10 } + ] + const { store } = await wedgeRegisteredMigration(() => now, cells) + now += ASSIGNMENT_LIMITS.migrationLeaseMs + STRANDED_MIGRATION_ABANDON_MS + 1 + await heartbeat(store, cells[0]!) + await heartbeat(store, cells[2]!) + await store.attestCellFence('cell-b', '11111111-1111-4111-8111-111111111111') + await database!.query(`UPDATE relay_cells SET reserved_requests = 1 WHERE cell_id = ?`, [ + 'cell-c' + ]) + const supersede = store.supersedeRegisteredEvacuation.bind(store) + let attempts = 0 + vi.spyOn(store, 'supersedeRegisteredEvacuation').mockImplementation(async (...args) => { + attempts++ + if (attempts === 2) expect(await store.abortExpiredEvacuations()).toBe(1) + return await supersede(...args) + }) + + await expect( + store.supersedeRegisteredCellEvacuations('cell-a', 'cell-b', 'cell-c', 100) + ).rejects.toThrow('migration_already_superseded') + expect(attempts).toBe(2) + }) + + it('reaps a pinned expired migration once its target cell is disabled', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ] + const { store, identity } = await wedgeRegisteredMigration(() => now, cells) + // Drains pin their migrations and nothing ever deletes a pin, so the pin outlives the drain. + await database!.query( + `INSERT INTO relay_post_drain_migration_pins + (user_id, relay_host_id, assignment_epoch, drain_attempt_id, source_cell_id, + source_cell_incarnation, target_cell_id, target_cell_incarnation, + source_request_units, target_reserved_units, pinned_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + 2, + 'attempt-0000', + 'cell-a', + '11111111-1111-4111-8111-111111111111', + 'cell-b', + '11111111-1111-4111-8111-111111111111', + 1, + 1, + now + ] + ) + + now += ASSIGNMENT_LIMITS.migrationLeaseMs + STRANDED_MIGRATION_ABANDON_MS + 1 + expect(await store.completeReadyEvacuations()).toBe(0) + expect(await store.abortExpiredEvacuations()).toBe(1) + expect(await store.abortExpiredEvacuations()).toBe(0) + }) + + it('repairs an expired migration marker from its exact active target control', async () => { + let now = 100 + const store = await setup(() => now, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ]) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + const sourceControl = await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + await store.startEvacuation(identity, 'cell-b') + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: 2, + generation: 1 + }) + + now += ASSIGNMENT_LIMITS.migrationLeaseMs + 1 + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: 2, + generation: 1 + }) + expect(await store.abortExpiredEvacuations()).toBe(0) + expect(await store.cellEvacuationStatus('cell-a', 'cell-b', false)).toEqual({ + inProgress: 1, + oldestExpiresAt: 100 + ASSIGNMENT_LIMITS.migrationLeaseMs, + oldestRemainingMs: -1, + targetRegistered: 1, + registeredSourceActive: 1, + registeredCompletable: 0, + registeredTargetInactive: 0, + completed: 0, + blocked: 0, + ...NO_EXPIRED_EVACUATION_DIAGNOSTICS + }) + await store.releaseActivity(identity, sourceControl) + expect(await store.cellEvacuationStatus('cell-a', 'cell-b', true)).toEqual({ + inProgress: 0, + ...NO_ACTIVE_MIGRATION_LEASE, + targetRegistered: 0, + ...NO_REGISTERED_EVACUATION_DIAGNOSTICS, + completed: 1, + blocked: 0, + ...NO_EXPIRED_EVACUATION_DIAGNOSTICS + }) + expect(await store.resolve(identity)).toMatchObject({ cellId: 'cell-b', assignmentEpoch: 2 }) + }) + + it('keeps a registered migration pending while its proven target control is offline', async () => { + let now = 100 + const store = await setup(() => now, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ]) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + const sourceControl = await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + const migration = await store.startEvacuation(identity, 'cell-b') + const targetControl = await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 2 + }) + expect( + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + ).toBe(true) + expect(await store.completeReadyEvacuations()).toBe(0) + await store.releaseActivity(identity, sourceControl) + now += ASSIGNMENT_LIMITS.activityLeaseMs + 1 + expect(await store.completeReadyEvacuations()).toBe(0) + await expect( + store.completeEvacuation(identity, migration.assignmentEpoch) + ).rejects.toThrow('migration_target_not_active') + await store.releaseActivity(identity, targetControl) + expect(await store.completeReadyEvacuations()).toBe(0) + + expect(await store.cellEvacuationStatus('cell-a', 'cell-b', true)).toEqual({ + inProgress: 1, + oldestExpiresAt: 100 + ASSIGNMENT_LIMITS.migrationLeaseMs, + oldestRemainingMs: + ASSIGNMENT_LIMITS.migrationLeaseMs - ASSIGNMENT_LIMITS.activityLeaseMs - 1, + targetRegistered: 1, + registeredSourceActive: 0, + registeredCompletable: 0, + registeredTargetInactive: 1, + completed: 0, + blocked: 1, + ...NO_EXPIRED_EVACUATION_DIAGNOSTICS + }) + + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 3 + }) + expect(await store.completeReadyEvacuations()).toBe(1) + expect(await store.cellEvacuationStatus('cell-a', 'cell-b', false)).toEqual({ + inProgress: 0, + ...NO_ACTIVE_MIGRATION_LEASE, + targetRegistered: 0, + ...NO_REGISTERED_EVACUATION_DIAGNOSTICS, + completed: 0, + blocked: 0, + ...NO_EXPIRED_EVACUATION_DIAGNOSTICS + }) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 0, 'cell-b': 1 }) + }) + + it('completes a registered migration after its disabled source is proven dead', async () => { + let now = 100 + const store = await setupWithHeartbeats(() => now, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ]) + await heartbeat(store, { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }) + await heartbeat(store, { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 }) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + const sourceControl = await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + await store.setCellEnabled('cell-a', false) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + await store.releaseActivity(identity, sourceControl) + now += 45_001 + await heartbeat( + store, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 }, + { startedAt: 50 } + ) + await store.attestCellFence( + 'cell-a', + '11111111-1111-4111-8111-111111111111' + ) + + const input = { + assignmentEpoch: migration.assignmentEpoch, + sourceCellId: 'cell-a', + targetCellId: 'cell-b' + } + await expect(store.completeEvacuationFromDeadSource(identity, input)).resolves.toEqual({ + changed: true, + assignmentEpoch: 2, + sourceCellId: 'cell-a', + targetCellId: 'cell-b' + }) + await expect(store.completeEvacuationFromDeadSource(identity, input)).resolves.toEqual({ + changed: false, + assignmentEpoch: 2, + sourceCellId: 'cell-a', + targetCellId: 'cell-b' + }) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 0, 'cell-b': 1 }) + expect(await store.cellEvacuationStatus('cell-a', 'cell-b', false)).toMatchObject({ + inProgress: 0 + }) + }) + + it('retires an inactive registered migration after its source is proven dead', async () => { + let now = 100 + const cells = [ + { + id: 'cell-a', + url: 'https://relay-a.example.com', + capacityRequests: 10, + connectionHardCap: 600 as const, + connectionUnobservedBound: 50 + }, + { + id: 'cell-b', + url: 'https://relay-b.example.com', + capacityRequests: 10, + connectionHardCap: 600 as const, + connectionUnobservedBound: 50 + } + ] + const store = await setupWithHeartbeats(() => now, cells) + await heartbeat(store, cells[0]!) + await heartbeat(store, cells[1]!) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + const sourceControl = await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + await store.setCellEnabled('cell-a', false) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + await store.releaseActivity(identity, sourceControl) + now += 45_001 + await heartbeat(store, cells[1]!, { startedAt: 50 }) + await store.attestCellFence('cell-a', '11111111-1111-4111-8111-111111111111') + + await expect( + store.cellEvacuationStatus('cell-a', 'cell-b', true) + ).resolves.toMatchObject({ inProgress: 0, completed: 1, blocked: 0 }) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 0, 'cell-b': 0 }) + expect( + await database!.query( + `SELECT activity_id FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([]) + expect( + await database!.query( + `SELECT state FROM relay_control_connection_reservations + WHERE user_id = ? AND relay_host_id = ? AND cell_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, 'cell-b', migration.assignmentEpoch] + ) + ).toEqual([{ state: 'released' }]) + expect( + await database!.query( + `SELECT cell_id, assignment_epoch, reserved_controls, migration_leases + FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([ + { + cell_id: 'cell-b', + assignment_epoch: 2, + reserved_controls: 0, + migration_leases: 0 + } + ]) + }) + + it.each(['expired', 'previous-incarnation'] as const)( + 'retires a registered migration and its %s target control after the source dies', + async (targetControlState) => { + let now = 100 + const cells = [ + { + id: 'cell-a', + url: 'https://relay-a.example.com', + capacityRequests: 10, + connectionHardCap: 600 as const, + connectionUnobservedBound: 50 + }, + { + id: 'cell-b', + url: 'https://relay-b.example.com', + capacityRequests: 10, + connectionHardCap: 600 as const, + connectionUnobservedBound: 50 + } + ] + const store = await setupWithHeartbeats(() => now, cells) + await heartbeat(store, cells[0]!) + await heartbeat(store, cells[1]!) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + const sourceControl = await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + await store.setCellEnabled('cell-a', false) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + await store.releaseActivity(identity, sourceControl) + now += + targetControlState === 'expired' + ? ASSIGNMENT_LIMITS.activityLeaseMs + 1 + : 45_001 + await heartbeat(store, cells[1]!, { + startedAt: targetControlState === 'previous-incarnation' ? now - 1 : 50, + incarnation: + targetControlState === 'previous-incarnation' + ? '22222222-2222-4222-8222-222222222222' + : '11111111-1111-4111-8111-111111111111' + }) + await store.attestCellFence( + 'cell-a', + '11111111-1111-4111-8111-111111111111' + ) + + await expect( + store.cellEvacuationStatus('cell-a', 'cell-b', false) + ).resolves.toMatchObject({ + inProgress: 1, + registeredCompletable: 0, + registeredTargetInactive: 1 + }) + await expect( + store.cellEvacuationStatus('cell-a', 'cell-b', true) + ).resolves.toMatchObject({ inProgress: 0, completed: 1, blocked: 0 }) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 0, 'cell-b': 0 }) + expect( + await database!.query( + `SELECT activity_id FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([]) + expect( + await database!.query( + `SELECT cell_id, assignment_epoch, reserved_controls, migration_leases + FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([ + { + cell_id: 'cell-b', + assignment_epoch: migration.assignmentEpoch, + reserved_controls: 0, + migration_leases: 0 + } + ]) + } + ) + + it('refuses dead-source completion while the source heartbeat is fresh', async () => { + const store = await setupWithHeartbeats(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ]) + await heartbeat(store, { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }) + await heartbeat(store, { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 }) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + const sourceControl = await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + await store.setCellEnabled('cell-a', false) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + await store.releaseActivity(identity, sourceControl) + + await expect( + store.attestCellFence('cell-a', '11111111-1111-4111-8111-111111111111') + ).rejects.toThrow('cell_fence_runtime_not_stale') + await expect( + store.completeEvacuationFromDeadSource(identity, { + assignmentEpoch: migration.assignmentEpoch, + sourceCellId: 'cell-a', + targetCellId: 'cell-b' + }) + ).rejects.toThrow('cell_fence_attestation_missing') + expect(await cellReservations(database!)).toEqual({ 'cell-a': 0, 'cell-b': 2 }) + }) + + it('supersedes a registered migration after its disabled target is proven unavailable', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 }, + { id: 'cell-c', url: 'https://relay-c.example.com', capacityRequests: 10 } + ] + const store = await setupWithHeartbeats(() => now, cells) + for (const cell of cells) await heartbeat(store, cell) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + await store.setCellEnabled('cell-a', false) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + await database!.query( + `INSERT INTO relay_control_connection_reservations + (reservation_id, idempotency_key, user_id, relay_host_id, + assignment_epoch, cell_id, state, inclusion_watermark, + claim_activity_id, created_at, timeout_at, claimed_at, released_at, + updated_at) + VALUES (?, ?, ?, ?, ?, ?, 'late-arrival-debt', NULL, NULL, ?, ?, NULL, NULL, ?)`, + [ + 'superseded-reservation', + 'superseded-reservation', + identity.userId, + identity.relayHostId, + migration.assignmentEpoch, + 'cell-b', + 100, + 100, + 100 + ] + ) + await store.setCellEnabled('cell-b', false) + now += 45_001 + await heartbeat(store, cells[0]!, { startedAt: 50 }) + await heartbeat(store, cells[2]!, { startedAt: 50 }) + await store.attestCellFence( + 'cell-b', + '11111111-1111-4111-8111-111111111111' + ) + const input = { + assignmentEpoch: migration.assignmentEpoch, + sourceCellId: 'cell-a', + currentTargetCellId: 'cell-b', + replacementTargetCellId: 'cell-c' + } + + const replacement = await store.supersedeRegisteredEvacuation(identity, input) + expect(replacement).toMatchObject({ + sourceCellId: 'cell-a', + targetCellId: 'cell-c', + previousEpoch: 2, + assignmentEpoch: 3 + }) + await expect(store.supersedeRegisteredEvacuation(identity, input)).resolves.toEqual(replacement) + expect(await store.resolve(identity)).toMatchObject({ cellId: 'cell-c', assignmentEpoch: 3 }) + expect( + await database!.query( + `SELECT state FROM relay_control_connection_reservations + WHERE reservation_id = ?`, + ['superseded-reservation'] + ) + ).toEqual([{ state: 'released' }]) + expect(await cellReservations(database!)).toEqual({ + 'cell-a': 1, + 'cell-b': 0, + 'cell-c': 2 + }) + expect( + await database!.query( + `SELECT activity_kind, cell_id FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? ORDER BY cell_id, activity_id`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([ + { activity_kind: 'control', cell_id: 'cell-a' }, + { activity_kind: 'control', cell_id: 'cell-c' }, + { activity_kind: 'migration', cell_id: 'cell-c' } + ]) + expect(await store.cellEvacuationStatus('cell-a', 'cell-b', false)).toMatchObject({ + inProgress: 0 + }) + expect(await store.cellEvacuationStatus('cell-a', 'cell-c', false)).toMatchObject({ + inProgress: 1 + }) + }) + + it('retries registered migration supersession with inventory locked first', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 }, + { id: 'cell-c', url: 'https://relay-c.example.com', capacityRequests: 10 } + ] + const delegate = await openInMemoryRelayDatabase() + const retryDatabase = new OneShotInventoryFailureDatabase(delegate) + database = retryDatabase + const store = new RelayAssignmentStore(retryDatabase, () => now, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + await store.reconcileCells(cells) + for (const cell of cells) await heartbeat(store, cell) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + await store.setCellEnabled('cell-a', false) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + await store.setCellEnabled('cell-b', false) + now += 45_001 + await heartbeat(store, cells[0]!, { startedAt: 50 }) + await heartbeat(store, cells[2]!, { startedAt: 50 }) + await store.attestCellFence( + 'cell-b', + '11111111-1111-4111-8111-111111111111' + ) + retryDatabase.arm() + + await expect( + store.supersedeRegisteredEvacuation(identity, { + assignmentEpoch: migration.assignmentEpoch, + sourceCellId: 'cell-a', + currentTargetCellId: 'cell-b', + replacementTargetCellId: 'cell-c' + }) + ).resolves.toMatchObject({ targetCellId: 'cell-c', assignmentEpoch: 3 }) + expect(retryDatabase.retryLocks).toEqual([ + 'inventory-nowait-failed', + 'inventory-first', + 'assignment-nowait' + ]) + }) + + it('reconciles durable cell accounting once before aggregate supersession', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 }, + { id: 'cell-c', url: 'https://relay-c.example.com', capacityRequests: 10 } + ] + const store = await setupWithHeartbeats(() => now, cells) + for (const cell of cells) await heartbeat(store, cell) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + await store.setCellEnabled('cell-a', false) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + await store.setCellEnabled('cell-b', false) + now += 45_001 + await heartbeat(store, cells[0]!, { startedAt: 50 }) + await heartbeat(store, cells[2]!, { startedAt: 50 }) + await store.attestCellFence( + 'cell-b', + '11111111-1111-4111-8111-111111111111' + ) + await database!.query( + `UPDATE relay_cells + SET reserved_requests = CASE WHEN cell_id = ? THEN 1 ELSE 0 END + WHERE cell_id IN (?, ?)`, + ['cell-c', 'cell-a', 'cell-c'] + ) + + await expect( + store.supersedeRegisteredCellEvacuations('cell-a', 'cell-b', 'cell-c', 100) + ).resolves.toBe(1) + expect(await cellReservations(database!)).toEqual({ + 'cell-a': 1, + 'cell-b': 0, + 'cell-c': 2 + }) + expect( + await database!.query( + `SELECT assignment_epoch, target_cell_id, aborted_at + FROM relay_assignment_migrations + WHERE user_id = ? ORDER BY assignment_epoch`, + [identity.userId] + ) + ).toEqual([ + { assignment_epoch: 2, target_cell_id: 'cell-b', aborted_at: 45_101 }, + { assignment_epoch: 3, target_cell_id: 'cell-c', aborted_at: null } + ]) + }) + + it('preserves healthy third-cell activity during aggregate supersession', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 }, + { id: 'cell-c', url: 'https://relay-c.example.com', capacityRequests: 10 }, + { id: 'cell-d', url: 'https://relay-d.example.com', capacityRequests: 10 } + ] + const store = await setupWithHeartbeats(() => now, cells) + for (const cell of cells) await heartbeat(store, cell) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + const sourceControl = await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + await store.setCellEnabled('cell-a', false) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + await database!.query( + `UPDATE relay_assignment_activity_leases SET cell_id = ? + WHERE user_id = ? AND relay_host_id = ? AND activity_id = ?`, + ['cell-d', identity.userId, identity.relayHostId, sourceControl] + ) + await database!.query( + `UPDATE relay_cells SET reserved_requests = ? WHERE cell_id = ?`, + [5, 'cell-c'] + ) + await store.setCellEnabled('cell-b', false) + now += 45_001 + await heartbeat(store, cells[0]!, { startedAt: 50 }) + await heartbeat(store, cells[2]!, { startedAt: 50 }) + await heartbeat(store, cells[3]!, { startedAt: 50 }) + await store.attestCellFence( + 'cell-b', + '11111111-1111-4111-8111-111111111111' + ) + + await expect( + store.supersedeRegisteredCellEvacuations('cell-a', 'cell-b', 'cell-c', 100) + ).resolves.toBe(1) + expect(await cellReservations(database!)).toEqual({ + 'cell-a': 0, + 'cell-b': 0, + 'cell-c': 1, + 'cell-d': 1 + }) + expect( + await database!.query( + `SELECT activity_kind, cell_id FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? ORDER BY cell_id, activity_id`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([ + { activity_kind: 'control', cell_id: 'cell-c' }, + { activity_kind: 'migration', cell_id: 'cell-c' }, + { activity_kind: 'control', cell_id: 'cell-d' } + ]) + }) + + it('refuses supersession while the registered target remains available', async () => { + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 }, + { id: 'cell-c', url: 'https://relay-c.example.com', capacityRequests: 10 } + ] + const store = await setupWithHeartbeats(() => 100, cells) + for (const cell of cells) await heartbeat(store, cell) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + await store.setCellEnabled('cell-a', false) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + await store.setCellEnabled('cell-b', false) + + await expect( + store.attestCellFence('cell-b', '11111111-1111-4111-8111-111111111111') + ).rejects.toThrow('cell_fence_runtime_not_stale') + await expect( + store.supersedeRegisteredEvacuation(identity, { + assignmentEpoch: migration.assignmentEpoch, + sourceCellId: 'cell-a', + currentTargetCellId: 'cell-b', + replacementTargetCellId: 'cell-c' + }) + ).rejects.toThrow('cell_fence_attestation_missing') + expect( + await database!.query( + `SELECT cell_id, assignment_epoch FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ cell_id: 'cell-b', assignment_epoch: 2 }]) + }) + + it('serializes concurrent retries of registered migration supersession', async () => { + let now = 100 + const cells = [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 }, + { id: 'cell-c', url: 'https://relay-c.example.com', capacityRequests: 10 } + ] + const store = await setupWithHeartbeats(() => now, cells) + for (const cell of cells) await heartbeat(store, cell) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + await store.setCellEnabled('cell-a', false) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + await store.setCellEnabled('cell-b', false) + now += 45_001 + await heartbeat(store, cells[0]!, { startedAt: 50 }) + await heartbeat(store, cells[2]!, { startedAt: 50 }) + await store.attestCellFence( + 'cell-b', + '11111111-1111-4111-8111-111111111111' + ) + const input = { + assignmentEpoch: migration.assignmentEpoch, + sourceCellId: 'cell-a', + currentTargetCellId: 'cell-b', + replacementTargetCellId: 'cell-c' + } + + const results = await Promise.all([ + store.supersedeRegisteredEvacuation(identity, input), + store.supersedeRegisteredEvacuation(identity, input) + ]) + expect(results[0]).toEqual(results[1]) + expect( + await database!.query( + `SELECT assignment_epoch FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ? ORDER BY assignment_epoch`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ assignment_epoch: 2 }, { assignment_epoch: 3 }]) + expect(await cellReservations(database!)).toEqual({ + 'cell-a': 1, + 'cell-b': 0, + 'cell-c': 2 + }) + }) + + it('classifies an expired migration blocked by a newer target assignment', async () => { + let now = 100 + const store = await setup(() => now, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 20 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 20 } + ]) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + await store.startEvacuation(identity, 'cell-b') + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: 2, + generation: 1 + }) + await database!.query( + `UPDATE relay_assignments SET assignment_epoch = ? + WHERE user_id = ? AND relay_host_id = ?`, + [3, identity.userId, identity.relayHostId] + ) + + now += ASSIGNMENT_LIMITS.migrationLeaseMs + 1 + expect(await store.cellEvacuationStatus('cell-a', 'cell-b', false)).toEqual({ + inProgress: 1, + oldestExpiresAt: 100 + ASSIGNMENT_LIMITS.migrationLeaseMs, + oldestRemainingMs: -1, + targetRegistered: 0, + ...NO_REGISTERED_EVACUATION_DIAGNOSTICS, + completed: 0, + blocked: 0, + expiredUnregistered: 1, + repairableExpiredUnregistered: 0, + abortableExpiredUnregistered: 0, + blockedExpiredUnregistered: 1, + blockedExpiredOnNewerTargetAssignment: 1 + }) + }) + + it('does not complete a registered migration after the assignment epoch advances', async () => { + const store = await setup(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 20 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 20 } + ]) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + const sourceControl = await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + const migration = await store.startEvacuation(identity, 'cell-b') + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: migration.assignmentEpoch + }) + await store.releaseActivity(identity, sourceControl) + await database!.query( + `UPDATE relay_assignments SET assignment_epoch = ? + WHERE user_id = ? AND relay_host_id = ?`, + [migration.assignmentEpoch + 1, identity.userId, identity.relayHostId] + ) + + expect(await store.completeReadyEvacuations()).toBe(0) + await expect( + store.completeEvacuation(identity, migration.assignmentEpoch) + ).rejects.toThrow('migration_assignment_mismatch') + expect(await store.cellEvacuationStatus('cell-a', 'cell-b', false)).toMatchObject({ + inProgress: 1, + targetRegistered: 1 + }) + }) + + it('retires an expired migration without rewriting its newer target assignment', async () => { + let now = 100 + const store = await setup(() => now, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 20 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 20 } + ]) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + await store.startEvacuation(identity, 'cell-b') + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: 2, + generation: 1 + }) + await database!.query( + `UPDATE relay_assignments SET assignment_epoch = ? + WHERE user_id = ? AND relay_host_id = ?`, + [3, identity.userId, identity.relayHostId] + ) + + now += ASSIGNMENT_LIMITS.migrationLeaseMs + 1 + expect(await store.abortExpiredEvacuations()).toBe(1) + expect(await store.resolve(identity)).toMatchObject({ + cellId: 'cell-b', + assignmentEpoch: 3 + }) + expect( + await database!.query( + `SELECT activity_id FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? ORDER BY activity_id`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([ + { activity_id: 'control:cell-a:1' }, + { activity_id: 'control:cell-b:1' } + ]) + expect(await store.cellEvacuationStatus('cell-a', 'cell-b', false)).toMatchObject({ + inProgress: 0 + }) + }) + + it.each([ + { assignmentEpoch: 1, removeAssignment: false, expected: 'migration_assignment_mismatch' }, + { assignmentEpoch: 2, removeAssignment: true, expected: 'migration_assignment_missing' } + ])( + 'fails closed when expired migration assignment state is not superseding: $expected', + async ({ assignmentEpoch, removeAssignment, expected }) => { + let now = 100 + const store = await setup(() => now, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 20 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 20 } + ]) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + await store.assign(identity) + await store.startEvacuation(identity, 'cell-b') + if (removeAssignment) { + await database!.query( + `DELETE FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + } else { + await database!.query( + `UPDATE relay_assignments SET assignment_epoch = ? + WHERE user_id = ? AND relay_host_id = ?`, + [assignmentEpoch, identity.userId, identity.relayHostId] + ) + } + + now += ASSIGNMENT_LIMITS.migrationLeaseMs + 1 + await expect(store.abortExpiredEvacuations()).rejects.toThrow(expected) + expect(await store.cellEvacuationStatus('cell-a', 'cell-b', false)).toMatchObject({ + inProgress: 1, + targetRegistered: 0 + }) + } + ) + + it('reconciles duplicated assignment counters and cell reservations after evacuation', async () => { + const store = await setup(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 20 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 20 } + ]) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + const sourceControl = await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + await store.startEvacuation(identity, 'cell-b') + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: 2, + generation: 1 + }) + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: 2 + }) + await store.releaseActivity(identity, sourceControl) + await database!.query( + `UPDATE relay_assignments SET reserved_controls = 9, reserved_splices = 4, + reserved_invites = 3, pending_installs = 2, pending_confirmations = 2, + migration_leases = 7 WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + await database!.query( + `UPDATE relay_cells SET reserved_requests = + CASE WHEN cell_id = ? THEN 11 ELSE 3 END`, + ['cell-a'] + ) + + expect(await store.cellEvacuationStatus('cell-a', 'cell-b', true)).toEqual({ + inProgress: 0, + ...NO_ACTIVE_MIGRATION_LEASE, + targetRegistered: 0, + ...NO_REGISTERED_EVACUATION_DIAGNOSTICS, + completed: 1, + blocked: 0, + ...NO_EXPIRED_EVACUATION_DIAGNOSTICS + }) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 0, 'cell-b': 1 }) + expect( + await database!.query( + `SELECT reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases + FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([ + { + reserved_controls: 1, + reserved_splices: 0, + reserved_invites: 0, + pending_installs: 0, + pending_confirmations: 0, + migration_leases: 0 + } + ]) + }) + + it('refuses reservation reconciliation when an activity lease has no assignment', async () => { + const store = await setup(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 20 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 20 } + ]) + await database!.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, + ['orphan-user', 'orphanhost000001', 'control:orphan', 'control', 'cell-a', 1, 200, 100] + ) + + await expect(store.cellEvacuationStatus('cell-a', 'cell-b', true)).rejects.toThrow( + 'activity_lease_assignment_missing' + ) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 0, 'cell-b': 0 }) + }) + + it('refuses reservation reconciliation when an activity lease names no cell', async () => { + const store = await setup(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 20 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 20 } + ]) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + await store.assign(identity) + await database!.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + 'splice:missing-cell', + 'splice', + 'missing-cell', + 2, + 200, + 100 + ] + ) + + await expect(store.cellEvacuationStatus('cell-a', 'cell-b', true)).rejects.toThrow( + 'activity_lease_cell_missing' + ) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 1, 'cell-b': 0 }) + }) + + it('leaves accounting for cells outside the selected evacuation pair unchanged', async () => { + const store = await setup(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 20 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 20 }, + { id: 'cell-c', url: 'https://relay-c.example.com', capacityRequests: 20 } + ]) + await database!.query( + `INSERT INTO relay_assignments + (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, + last_activity_at, reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + ['unrelated-user', 'unrelatedhost001', 'cell-c', 1, 200, 100, 9, 0, 0, 0, 0, 0] + ) + await database!.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, + [ + 'unrelated-user', + 'unrelatedhost001', + 'control:unrelated', + 'control', + 'cell-c', + 1, + 200, + 100 + ] + ) + await database!.query( + `UPDATE relay_cells SET reserved_requests = 9 WHERE cell_id = ?`, + ['cell-c'] + ) + + expect(await store.cellEvacuationStatus('cell-a', 'cell-b', true)).toMatchObject({ + inProgress: 0 + }) + expect(await cellReservations(database!)).toEqual({ + 'cell-a': 0, + 'cell-b': 0, + 'cell-c': 9 + }) + expect( + await database!.query( + `SELECT reserved_controls FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + ['unrelated-user', 'unrelatedhost001'] + ) + ).toEqual([{ reserved_controls: 9 }]) + }) + + it('rebalances only a fully inactive assignment after the dormant TTL', async () => { + let now = 100 + const store = await setup(() => now, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ]) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const first = await store.assign(identity) + const control = await store.activateControl(identity, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + await expect(store.rebalanceDormant(identity, 'cell-b')).rejects.toThrow('assignment_active') + await store.releaseActivity(identity, control) + now += ASSIGNMENT_LIMITS.dormantTtlMs + 1 + + await expect(store.rebalanceDormant(identity, 'cell-b')).resolves.toMatchObject({ + cellId: 'cell-b', + assignmentEpoch: 2 + }) + expect(await cellReservations(database!)).toEqual({ 'cell-a': 0, 'cell-b': 1 }) + }) + + it('durably removes an unhealthy cell from new assignment without moving existing hosts', async () => { + const store = await setup(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ]) + const existing = { userId: 'user-a', relayHostId: 'host000000000001' } + expect(await store.assign(existing)).toMatchObject({ cellId: 'cell-a' }) + await store.setCellEnabled('cell-a', false) + await store.reconcileCells([ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 } + ]) + expect(await store.resolve(existing)).toMatchObject({ cellId: 'cell-a' }) + await expect( + store.assign({ userId: 'user-b', relayHostId: 'host000000000002' }) + ).resolves.toMatchObject({ cellId: 'cell-b' }) + await store.setCellEnabled('cell-a', true) + await expect(store.setCellEnabled('missing', false)).rejects.toThrow('cell_not_found') + }) + + it('bulk-migrates active controls without exposing assignment identities', async () => { + const store = await setup(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 20 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 20 } + ]) + await store.setCellEnabled('cell-b', false) + const identities = [1, 2, 3].map((index) => ({ + userId: `user-${index}`, + relayHostId: `host${String(index).padStart(12, '0')}` + })) + const sourceControls: string[] = [] + for (const identity of identities) { + const assignment = await store.assign(identity) + sourceControls.push( + await store.activateControl(identity, { + cellId: 'cell-a', + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + ) + } + await store.setCellEnabled('cell-b', true) + expect(await store.cellEvacuationCapacity('cell-a', 'cell-b')).toEqual({ + sourceAssignments: 3, + requiredTargetUnits: 6, + availableTargetUnits: 20 + }) + expect(await store.startActiveCellEvacuations('cell-a', 'cell-b', 2)).toBe(2) + expect(await store.startActiveCellEvacuations('cell-a', 'cell-b', 2)).toBe(1) + expect(await store.startActiveCellEvacuations('cell-a', 'cell-b', 2)).toBe(0) + expect(await store.cellEvacuationStatus('cell-a', 'cell-b', false)).toEqual({ + inProgress: 3, + oldestExpiresAt: 100 + ASSIGNMENT_LIMITS.migrationLeaseMs, + oldestRemainingMs: ASSIGNMENT_LIMITS.migrationLeaseMs, + targetRegistered: 0, + ...NO_REGISTERED_EVACUATION_DIAGNOSTICS, + completed: 0, + blocked: 0, + ...NO_EXPIRED_EVACUATION_DIAGNOSTICS + }) + for (const identity of identities) { + const assignment = await store.resolve(identity) + await store.activateControl(identity, { + cellId: 'cell-b', + assignmentEpoch: assignment!.assignmentEpoch, + generation: 2 + }) + await store.markMigrationTargetRegistered(identity, { + cellId: 'cell-b', + assignmentEpoch: assignment!.assignmentEpoch + }) + } + expect(await store.cellEvacuationStatus('cell-a', 'cell-b', true)).toEqual({ + inProgress: 3, + oldestExpiresAt: 100 + ASSIGNMENT_LIMITS.migrationLeaseMs, + oldestRemainingMs: ASSIGNMENT_LIMITS.migrationLeaseMs, + targetRegistered: 3, + registeredSourceActive: 3, + registeredCompletable: 0, + registeredTargetInactive: 0, + completed: 0, + blocked: 3, + ...NO_EXPIRED_EVACUATION_DIAGNOSTICS + }) + for (let index = 0; index < identities.length; index++) { + await store.releaseActivity(identities[index]!, sourceControls[index]!) + } + expect(await store.cellEvacuationStatus('cell-a', 'cell-b', true)).toEqual({ + inProgress: 0, + ...NO_ACTIVE_MIGRATION_LEASE, + targetRegistered: 0, + ...NO_REGISTERED_EVACUATION_DIAGNOSTICS, + completed: 3, + blocked: 0, + ...NO_EXPIRED_EVACUATION_DIAGNOSTICS + }) + }) + + it.each(['cell-b', 'cell-c'])( + 'skips a batch row that concurrently moved from the selected source to %s', + async (movedCellId) => { + const store = await setup(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 20 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 20 }, + { id: 'cell-c', url: 'https://relay-c.example.com', capacityRequests: 20 } + ]) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const assignment = await store.assign(identity) + await store.activateControl(identity, { + cellId: 'cell-a', + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + const startEvacuation = store.startEvacuation.bind(store) + vi.spyOn(store, 'startEvacuation').mockImplementationOnce(async (...args) => { + await database!.query( + `UPDATE relay_assignments SET cell_id = ? WHERE user_id = ? AND relay_host_id = ?`, + [movedCellId, identity.userId, identity.relayHostId] + ) + return await startEvacuation(...args) + }) + + expect(await store.startActiveCellEvacuations('cell-a', 'cell-b', 10)).toBe(0) + expect(await store.resolve(identity)).toMatchObject({ cellId: movedCellId }) + expect(await database!.query(`SELECT * FROM relay_assignment_migrations`)).toEqual([]) + } + ) + + it('does not count an existing migration after a stale batch row moves to its target', async () => { + const store = await setup(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 20 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 20 } + ]) + const identity = { userId: 'user-a', relayHostId: 'host000000000001' } + const assignment = await store.assign(identity) + await store.activateControl(identity, { + cellId: 'cell-a', + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + const startEvacuation = store.startEvacuation.bind(store) + vi.spyOn(store, 'startEvacuation').mockImplementationOnce(async (...args) => { + await startEvacuation(identity, 'cell-b') + return await startEvacuation(...args) + }) + + expect(await store.startActiveCellEvacuations('cell-a', 'cell-b', 10)).toBe(0) + expect(await store.resolve(identity)).toMatchObject({ cellId: 'cell-b' }) + expect(await store.cellEvacuationStatus('cell-a', 'cell-b', false)).toMatchObject({ + inProgress: 1 + }) + }) + + it('preserves operator-owned tagged URLs and admission state across reconciliation', async () => { + const store = await setup(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 20 } + ]) + await store.configureCell( + { id: 'cell-a', url: 'https://old---relay-a.example.com', capacityRequests: 20 }, + false + ) + await store.reconcileCells([ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 20 } + ]) + const rows = await database!.query( + `SELECT cell_url, enabled, capacity_requests FROM relay_cells WHERE cell_id = ?`, + ['cell-a'] + ) + expect(rows).toEqual([ + { + cell_url: 'https://old---relay-a.example.com', + enabled: 0, + capacity_requests: 20 + } + ]) + }) + + it('bulk-migrates activity even when its source control is temporarily absent', async () => { + const store = await setup(() => 100, [ + { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 20 }, + { id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 20 } + ]) + const identity = { userId: 'user-1', relayHostId: 'host000000000001' } + const assignment = await store.assign(identity) + const control = await store.activateControl(identity, { + cellId: 'cell-a', + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + await store.acquireActivity(identity, { + activityId: 'splice-without-control', + kind: 'splice', + cellId: 'cell-a' + }) + await store.releaseActivity(identity, control) + expect(await store.cellEvacuationCapacity('cell-a', 'cell-b')).toEqual({ + sourceAssignments: 1, + requiredTargetUnits: 3, + availableTargetUnits: 20 + }) + expect(await store.startActiveCellEvacuations('cell-a', 'cell-b', 10)).toBe(1) + expect(await store.resolve(identity)).toMatchObject({ cellId: 'cell-b', assignmentEpoch: 2 }) + }) +}) + +async function cellReservations(database: RelayDatabase): Promise> { + const rows = await database.query( + `SELECT cell_id, reserved_requests FROM relay_cells ORDER BY cell_id ASC` + ) + return Object.fromEntries(rows.map((row) => [String(row.cell_id), Number(row.reserved_requests)])) +} diff --git a/cloud/apps/relay/src/assignment-store.ts b/cloud/apps/relay/src/assignment-store.ts new file mode 100644 index 00000000000..0b2b1ef72a9 --- /dev/null +++ b/cloud/apps/relay/src/assignment-store.ts @@ -0,0 +1,8347 @@ +import { randomUUID } from 'node:crypto' +import { performance } from 'node:perf_hooks' +import { + ASSIGNMENT_LIMITS, + mayNormallyReassign, + RELAY_ADMISSION_BUDGETS, + RELAY_DEFAULT_REGION, + RELAY_REGIONS, + RELAY_PROTOCOL_LIMITS, + type RelayRegion +} from '@orca-cloud/relay-contract' +import { + cellAdmissionState, + cellAdmissionStates, + ensureCellAdmission, + parseCellAdmissionState, + RelayCellAdmissionSelector, + setCellAdmissionBeforeBoundary, + stateFromEnabled, + synchronizeCellAdmissionBoundary, + type CellAdmissionMembership, + type CellAdmissionSelectorInspection, + type CellAdmissionState +} from './cell-admission-selector.js' +import { + RelayMigrationCellRegistrar, + type MigrationCellRegistration +} from './cell-admission-migration-registration.js' +import { + ASSIGNMENT_CONNECTION_HEADROOM_QUERY +} from './assignment-connection-headroom-query.js' +import { AssignmentIdentityQueue } from './assignment-identity-queue.js' +import type { RelayCellConfig } from './config.js' +import type { RelayDatabase, RelayTransactionOptions, SqlRow } from './database.js' +import type { RegionalRehomeSafetySnapshot } from './relay-observability.js' +import { + combineRegionalRehomeSafety, + REGIONAL_REHOME_RECONNECTS_PER_CELL_LIMIT, + REGIONAL_REHOME_SQL_FAILURES_PER_CELL_LIMIT, + regionalRehomePoolPressure, + regionalRehomeSafetyFailure +} from './regional-rehome-safety.js' +import { + ABANDONED_REGISTERED_MIGRATION, + DURABLY_FENCED_MIGRATION_SOURCE, + REGISTERED_MIGRATION_ABANDON_MS +} from './registered-migration-abandonment.js' + +type AssignmentIdentity = { userId: string; relayHostId: string } +type CellHeartbeat = { + cellId: string + cellUrl: string + cellIncarnation: string + startedAt: number + ready: boolean + observedRequests: number + region?: RelayRegion + totalConnections?: number + inFlightConnections?: number + reservedConnectionUnits?: number + enforcedConnectionUnits?: number + connectionInclusionWatermark?: number + connectionHardCap?: number + connectionUnobservedBound?: number +} + +type CellRegionalRehomeStatus = { + cellId: string + cellIncarnation: string + regionalRehomeProtocol: number + safety: RegionalRehomeSafetySnapshot +} + +type RelayAssignmentStoreOptions = { + requireLiveCells?: boolean + heartbeatTtlMs?: number + recordControlRenewal?: (durationMs: number, outcome: ControlRenewalOutcome) => void +} + +export type ControlRenewalOutcome = + | 'renewed' + | 'assignment_not_found' + | 'activity_cell_not_authoritative' + | 'control_activity_not_found' + | 'control_activity_moved' + | 'database_error' + +const CONTROL_RENEWAL_OUTCOMES = new Set([ + 'renewed', + 'assignment_not_found', + 'activity_cell_not_authoritative', + 'control_activity_not_found', + 'control_activity_moved' +]) +export type RelayAssignment = AssignmentIdentity & { + cellId: string + cellUrl: string + assignmentEpoch: number + leaseExpiresAt: number + region?: RelayRegion +} + +export type RelayRegionCatalogEntry = { + region: RelayRegion + probeOrigins: string[] +} + +export type RelayAssignmentMigration = AssignmentIdentity & { + sourceCellId: string + targetCellId: string + previousEpoch: number + assignmentEpoch: number + expiresAt: number + targetRegisteredAt?: number +} + +export type RegionalRehomeAttempt = AssignmentIdentity & { + attemptId: string + preferredRegion: 'asia-east2' + sourceCellId: string + sourceCellUrl: string + sourceCellIncarnation: string + targetCellId: string + targetCellIncarnation: string + previousEpoch: number + assignmentEpoch: number + drainGraceMs: number + sendAttempts: number +} + +export type RegionalHostDrainOutcome = + | 'accepted' + | 'already-accepted' + | 'host-not-connected' + +export type RegionalRehomeFleetSafety = RegionalRehomeSafetySnapshot & { + requiredCells: number + missingCells: number + maxReconnects: number +} + +export type RegionalRehomeControl = { + generation: number + enabled: boolean + observationStartedAt: number + notBefore: number + ratePerMinute: number + preferenceMaxAgeMs: number + drainGraceMs: number +} + +type RegisteredEvacuationSupersessionInput = { + assignmentEpoch: number + sourceCellId: string + currentTargetCellId: string + replacementTargetCellId: string +} + +export type DeadSourceCompletionResult = { + changed: boolean + assignmentEpoch: number + sourceCellId: string + targetCellId: string +} + +export type CellEvacuationStatus = { + inProgress: number + oldestExpiresAt: number | null + oldestRemainingMs: number | null + targetRegistered: number + registeredSourceActive: number + registeredCompletable: number + registeredTargetInactive: number + completed: number + blocked: number + expiredUnregistered: number + repairableExpiredUnregistered: number + abortableExpiredUnregistered: number + blockedExpiredUnregistered: number + blockedExpiredOnNewerTargetAssignment: number +} + +export type CellEvacuationCapacity = { + sourceAssignments: number + requiredTargetUnits: number + availableTargetUnits: number +} + +export type CellDeploymentStatus = { + cellId: string + cellUrl: string + region: RelayRegion + enabled: boolean + admissionState: CellAdmissionState + capacityRequests: number + reservedRequests: number + assignments: number + activityLeases: number + activityRequestUnits: number + restartBlockingActivityLeases: number + restartBlockingActivityRequestUnits: number + restartBlockingReservedRequests: number + outgoingMigrations: number + incomingMigrations: number + connectionCapacity: null | { + hardCap: number + controlRebindReserve: number + ordinaryConnectionLimit: number + unobservedBound: number + normalAdmissionPause: number + observedConnections: number + inFlightConnections: number + reservedConnectionUnits: number + enforcedConnectionUnits: number + pendingControlReservations: number + heartbeatFresh: boolean + } + runtime: null | { + cellUrl: string + cellIncarnation: string + startedAt: number + ready: boolean + observedRequests: number + lastHeartbeatAt: number + heartbeatFresh: boolean + regionalRehomeProtocol: number + } +} + +export type CellFenceAttemptEvidence = { + attemptId: string + environment: 'staging' | 'production' + cellId: string + cellIncarnation: string + migName: string + instanceGroup: string + generationIdentity: string + fenceCommit: string + planSha256: string + planObjectName: string + planObjectGeneration?: string + varFileSha256: string + terraformStateLineage: string + terraformStateSerial: number + terraformStateObjectGeneration: string + terraformStateObjectSha256: string + requestReason: string +} + +export type CellFenceApplyInvocation = { + invocationId: string + requestReason: string + startedAt: number + gceOperation?: string +} + +export type CellFenceAttempt = CellFenceAttemptEvidence & { + applyInvocations?: CellFenceApplyInvocation[] + gceOperation?: string + createdAt: number + expiresAt: number + applyStartedAt?: number + completedAt?: number + abortedAt?: number +} + +export type CellDrainAttemptState = + | 'prepared' + | 'send-may-have-started' + | 'application-receipt' + | 'proven-not-delivered' + +export type CellDrainAttempt = { + attemptId: string + cellId: string + cellIncarnation: string + traceValue: string + plannedGraceMs: number + state: CellDrainAttemptState + preparedAt: number + sendMayHaveStartedAt?: number + sendPermitExpiresAt?: number + applicationReceiptAt?: number + backendSuccessStatus?: number + backendInstance?: string + receiptCellIncarnation?: string + retryAfter?: number + recoverForwardAttemptedAt?: number + provenNotDeliveredAt?: number +} + +export type AssignmentActivityKind = + | 'control' + | 'splice' + | 'invite' + | 'install' + | 'confirmation' + | 'migration' + +const ACTIVITY_COLUMN: Record = { + control: 'reserved_controls', + splice: 'reserved_splices', + invite: 'reserved_invites', + install: 'pending_installs', + confirmation: 'pending_confirmations', + migration: 'migration_leases' +} + +const ACTIVITY_REQUEST_UNITS: Record = { + control: 1, + splice: 2, + invite: 1, + install: 1, + confirmation: 1, + migration: 1 +} + +const ASSIGNMENT_LOCK_RETRY_DEADLINE_MS = 15_000 +// Why: stranded detection (issue #225) needs a grant old enough that a real +// attach would have registered (the 90s activity lease covers dial + +// activation), yet recent enough to prove an active retry loop rather than +// ordinary dormancy — which stays governed by the 24h rule. +const STRANDED_MIN_GRANT_AGE_MS = 60_000 +const STRANDED_RECENT_ACTIVITY_MS = 15 * 60_000 +const REGION_PREFERENCE_RETENTION_MS = 30 * 24 * 60 * 60_000 +const REGIONAL_REHOME_UNREGISTERED_REFRESH_MS = 5 * 60_000 +const REGIONAL_REHOME_MAX_REFRESH_MS = 24 * 60 * 60_000 +// Drain grace is enforced by session-scoped cell state that any control +// reconnect sheds, so receipted attempts can stall dual-homed past grace; +// redrains re-dispatch them with the elapsed (zero) grace until they detach. +export const REGIONAL_REHOME_REDRAIN_INTERVAL_MS = 60_000 +export const REGIONAL_REHOME_REDRAIN_SEND_LIMIT = 20 +export const REGIONAL_REHOME_QUARANTINE_FAILURES = 3 +export const REGIONAL_REHOME_QUARANTINE_MS = 15 * 60_000 +const REGIONAL_REHOME_QUARANTINE_EXCLUSION_LIMIT = 50 +const REGIONAL_REHOME_QUARANTINE_MEMORY_LIMIT = 1_000 +const REGIONAL_REHOME_OBSERVATION_MS = 24 * 60 * 60_000 +const ASSIGNMENT_LOCK_RETRY_MAX_DELAY_MS = 50 +type AssignmentInventoryScope = 'none' | 'general' | 'all' +type RetriedAssignmentInventoryScope = Exclude + +class AssignmentInventoryLockUnavailable extends Error { + constructor(readonly inventoryScope: RetriedAssignmentInventoryScope) { + super('database_lock_unavailable') + } +} + +class AssignmentInventoryScopeChanged extends Error { + constructor() { + super('assignment_inventory_scope_changed') + } +} + +// Debt holds connection headroom for a control that may still arrive shortly +// after its director-side timeout. Nothing legitimately arrives minutes late +// (attach deadline 10s, orphan grace 30s); unretired debt from hosts that +// never return otherwise starves connection headroom fleet-wide and turns +// every placement into relay_capacity_exhausted. +const LATE_ARRIVAL_DEBT_RETENTION_MS = 10 * 60 * 1_000 +const CELL_FENCE_TTL_MS = 5 * 60 * 1_000 +const CELL_FENCE_ATTEMPT_TTL_MS = 60 * 60 * 1_000 +const CELL_DRAIN_SEND_PERMIT_MS = 30_000 +const MAX_DRAIN_ACCOUNTING_REPAIR_ATTEMPTS = 3 +const CELL_FENCE_ATTEMPT_SELECT = ` + SELECT attempts.*, bindings.plan_object_name, bindings.plan_object_generation, + bindings.var_file_sha256, bindings.terraform_state_lineage, + bindings.terraform_state_serial, bindings.terraform_state_object_generation, + bindings.terraform_state_object_sha256, bindings.request_reason + FROM relay_cell_fence_attempts attempts + JOIN relay_cell_fence_plan_bindings bindings + ON bindings.attempt_id = attempts.attempt_id` +export const STRANDED_MIGRATION_ABANDON_MS = REGISTERED_MIGRATION_ABANDON_MS +const ABORTABLE_EXPIRED_MIGRATION = `( + ( + migration.target_registered_at IS NULL + AND NOT EXISTS ( + SELECT 1 FROM relay_post_drain_migration_pins pin + WHERE pin.user_id = migration.user_id + AND pin.relay_host_id = migration.relay_host_id + AND pin.assignment_epoch = migration.assignment_epoch + ) + ) + OR ${ABANDONED_REGISTERED_MIGRATION} +)` + +export class RelayAssignmentStore { + private readonly requireLiveCells: boolean + private readonly heartbeatTtlMs: number + // Poisoned attempts never complete or abort and stay the oldest rows, so + // unquarantined they eventually fill the sweeps' LIMIT page and starve + // every healthy candidate. Process-local on purpose: it resets on deploy + // and each director relearns within a few ticks. + private readonly regionalRehomeCandidateQuarantine = new Map< + string, + { failures: number; until: number } + >() + private readonly recordControlRenewal?: RelayAssignmentStoreOptions['recordControlRenewal'] + private readonly admissionSelector: RelayCellAdmissionSelector + private readonly migrationCellRegistrar: RelayMigrationCellRegistrar + private readonly activityQueue = new AssignmentIdentityQueue() + private assignmentTail: Promise = Promise.resolve() + private pendingRegionalRehomeDisableLog: Record | null = null + + constructor( + private readonly database: RelayDatabase, + private readonly now: () => number = Date.now, + options: RelayAssignmentStoreOptions = {} + ) { + this.requireLiveCells = options.requireLiveCells ?? false + this.heartbeatTtlMs = options.heartbeatTtlMs ?? 45_000 + this.recordControlRenewal = options.recordControlRenewal + this.admissionSelector = new RelayCellAdmissionSelector(database, now) + this.migrationCellRegistrar = new RelayMigrationCellRegistrar(database, now) + } + + async reconcileCells(cells: RelayCellConfig[], disableMissing = true): Promise { + await this.reconcileCellsWithOptions(cells, disableMissing) + } + + async reconcileCellsAtStartup(cells: RelayCellConfig[]): Promise { + await this.reconcileCellsWithOptions(cells, false, { reportRetries: false }) + } + + private async reconcileCellsWithOptions( + cells: RelayCellConfig[], + disableMissing: boolean, + transactionOptions?: RelayTransactionOptions + ): Promise { + const now = this.now() + await this.database.transaction(async (transaction) => { + await transaction.queryLocked(`SELECT cell_id FROM relay_cells ORDER BY cell_id ASC`) + // Capacity rows share one global lock order with assignment transactions. + for (const cell of [...cells].sort((left, right) => left.id.localeCompare(right.id))) { + // Operator-owned enabled state and tagged URLs must survive revision restarts. + await transaction.query( + `INSERT INTO relay_cells + (cell_id, cell_url, enabled, capacity_requests, reserved_requests, + observed_requests, last_heartbeat_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT (cell_id) DO UPDATE SET + capacity_requests = excluded.capacity_requests, + last_heartbeat_at = excluded.last_heartbeat_at, + updated_at = excluded.updated_at`, + [ + cell.id, + cell.url, + cell.initiallyEnabled === false ? 0 : 1, + cell.capacityRequests, + 0, + 0, + now, + now + ] + ) + await transaction.query( + `INSERT INTO relay_cell_regions (cell_id, region) VALUES (?, ?) + ON CONFLICT (cell_id) DO UPDATE SET region = excluded.region`, + [cell.id, cell.region ?? RELAY_DEFAULT_REGION] + ) + await ensureCellAdmission( + transaction, + cell.id, + 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] + ) + } + } + } + const current = await transaction.query( + `SELECT cell_id, enabled FROM relay_cells ORDER BY cell_id ASC` + ) + for (const row of current) { + await ensureCellAdmission( + transaction, + text(row, 'cell_id'), + stateFromEnabled(integer(row, 'enabled') === 1), + now + ) + } + if (disableMissing) { + const ids = new Set(cells.map(({ id }) => id)) + for (const row of current) { + const id = text(row, 'cell_id') + if ( + !ids.has(id) && + (await cellAdmissionState(transaction, id)) !== 'existing-only' + ) { + await setCellAdmissionBeforeBoundary(transaction, id, 'existing-only', now) + } + } + } + await synchronizeCellAdmissionBoundary(transaction, now) + }, transactionOptions) + } + + async assign( + identity: AssignmentIdentity, + preferredRegion?: RelayRegion, + placementRegion: RelayRegion = preferredRegion ?? RELAY_DEFAULT_REGION + ): Promise { + const sticky = await this.assignStickyWithLockRetry(identity, preferredRegion) + if (sticky) return sticky + // Only placement needs the global inventory critical section; queueing those + // attempts locally avoids turning true placement bursts into NOWAIT storms. + return await this.serializeAssignment( + async () => await this.assignWithLockRetry(identity, preferredRegion, placementRegion) + ) + } + + private async assignStickyWithLockRetry( + identity: AssignmentIdentity, + preferredRegion?: RelayRegion + ): Promise { + return await this.withAssignmentLockRetry( + async (inventoryFirst) => + await this.assignStickyOnce(identity, inventoryFirst, preferredRegion) + ) + } + + private async assignWithLockRetry( + identity: AssignmentIdentity, + preferredRegion?: RelayRegion, + placementRegion: RelayRegion = preferredRegion ?? RELAY_DEFAULT_REGION + ): Promise { + const deadline = Date.now() + ASSIGNMENT_LOCK_RETRY_DEADLINE_MS + let inventoryScope: AssignmentInventoryScope = 'none' + while (true) { + try { + return await this.assignOnce(identity, inventoryScope, preferredRegion, placementRegion) + } catch (error) { + if (error instanceof AssignmentInventoryScopeChanged) { + inventoryScope = 'all' + continue + } + if ( + !(error instanceof AssignmentInventoryLockUnavailable) || + Date.now() >= deadline + ) { + throw error + } + inventoryScope = error.inventoryScope + } + await waitForAssignmentLockRetry(deadline) + } + } + + private async withAssignmentLockRetry( + operation: (inventoryFirst: boolean) => Promise + ): Promise { + const deadline = Date.now() + ASSIGNMENT_LOCK_RETRY_DEADLINE_MS + let inventoryFirst = false + while (true) { + try { + return await operation(inventoryFirst) + } catch (error) { + if (!isDatabaseLockUnavailable(error) || Date.now() >= deadline) throw error + inventoryFirst = true + } + // The retry joins the cell queue without holding later legacy-cycle locks. + await waitForAssignmentLockRetry(deadline) + } + } + + private async assignStickyOnce( + identity: AssignmentIdentity, + inventoryFirst: boolean, + preferredRegion?: RelayRegion + ): Promise { + const now = this.now() + return await this.database.transaction(async (transaction) => { + const lockedCells = inventoryFirst + ? await this.lockCellInventory(transaction) + : undefined + const existing = await this.assignmentRow(transaction, identity, inventoryFirst) + if (!existing) return null + const activityLeases = await this.lockAssignmentActivities(transaction, identity, true) + await this.recordRegionPreference(transaction, identity, preferredRegion, now) + if (mayNormallyReassign(activity(existing), now)) return null + // Why: a stranded host must fall through to placement — re-granting the + // pinned cell here is what refreshes its own activity and sustains the + // loop (issue #225). + if (await this.assignmentStrandedOnUnservedCell(transaction, identity, existing, now)) { + return null + } + + const currentCellId = text(existing, 'cell_id') + const hadControl = holdsControlLease( + activityLeases, + currentCellId, + integer(existing, 'assignment_epoch') + ) + const currentRow = lockedCells + ? lockedCells.find((row) => text(row, 'cell_id') === currentCellId) + : hadControl + ? ( + await transaction.query(`SELECT * FROM relay_cells WHERE cell_id = ?`, [ + currentCellId + ]) + )[0] + : ( + await transaction.queryLocked( + `SELECT * FROM relay_cells WHERE cell_id = ?`, + [currentCellId], + { failIfUnavailable: true } + ) + )[0] + if (!currentRow) throw new Error('assigned_cell_missing') + if ( + this.requireLiveCells && + !(await this.cellIsLive(transaction, currentCellId, now)) + ) { + return null + } + + if ( + !hadControl && + !(await this.cellHasConnectionHeadroom(transaction, currentCellId)) + ) { + if (requestUnits(existing) === 0) return null + throw new Error('relay_connection_headroom_exhausted') + } + + const leaseExpiresAt = now + ASSIGNMENT_LIMITS.activityLeaseMs + if (hadControl) { + await this.touchAssignment(transaction, identity, leaseExpiresAt, now) + } else { + const nextReservation = integer(currentRow, 'reserved_requests') + 1 + if (nextReservation > integer(currentRow, 'capacity_requests')) { + throw new Error('relay_capacity_exhausted') + } + await transaction.query( + `UPDATE relay_cells SET reserved_requests = ?, updated_at = ? WHERE cell_id = ?`, + [nextReservation, now, currentCellId] + ) + await this.adjustActivityCount(transaction, identity, 'control', 1, leaseExpiresAt, now) + await this.insertPendingControlLease( + transaction, + identity, + currentCellId, + integer(existing, 'assignment_epoch'), + now + ) + } + return this.result( + identity, + existing, + cell(currentRow, await this.cellRegion(transaction, currentCellId)), + leaseExpiresAt + ) + }) + } + + // Why: PR #194 pins assignments on existing-only cells so capacity pressure + // cannot scatter existing hosts — assuming those cells still serve them. A + // decommissioning cell (C3) broke that assumption: existing-only AND + // rejecting attaches, so pinned hosts loop forever while each grant + // refreshes their own last_activity_at, keeping dormancy-based + // reassignment permanently out of reach (issue #225). "Stranded" therefore + // requires proof the cell is not serving this host: a recent grant + // (last_activity_at inside the window) that had ample time to attach and + // still produced no live real activity. The evidence must come from the + // assignment row itself — reservation rows do not exist for cells without + // connection limits (C3), and expired pending leases are cleaned within a + // maintenance cycle, so neither reliably survives until the next grant. + private async assignmentStrandedOnUnservedCell( + transaction: RelayDatabase, + identity: AssignmentIdentity, + existing: SqlRow, + now: number + ): Promise { + const lastActivityAt = integer(existing, 'last_activity_at') + if ( + now < lastActivityAt + STRANDED_MIN_GRANT_AGE_MS || + now >= lastActivityAt + STRANDED_RECENT_ACTIVITY_MS + ) { + return false + } + const admissionRow = ( + await transaction.query( + `SELECT admission_state FROM relay_cell_admission WHERE cell_id = ?`, + [text(existing, 'cell_id')] + ) + )[0] + // Unknown or missing admission fails safe: the pin stays. + if (admissionRow?.['admission_state'] !== 'existing-only') { + return false + } + const liveLeases = ( + await transaction.query( + `SELECT COUNT(*) AS live FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? AND expires_at > ? + AND activity_id NOT LIKE 'control-pending:%'`, + [identity.userId, identity.relayHostId, now] + ) + )[0] + return integer(liveLeases!, 'live') === 0 + } + + private async serializeAssignment(operation: () => Promise): Promise { + const previous = this.assignmentTail + let release!: () => void + this.assignmentTail = new Promise((resolve) => (release = resolve)) + await previous + try { + return await operation() + } finally { + release() + } + } + + private async assignOnce( + identity: AssignmentIdentity, + inventoryScope: AssignmentInventoryScope, + preferredRegion?: RelayRegion, + placementRegion: RelayRegion = preferredRegion ?? RELAY_DEFAULT_REGION + ): Promise { + const now = this.now() + let retryScope: RetriedAssignmentInventoryScope = + inventoryScope === 'all' ? 'all' : 'general' + return await this.database.transaction(async (transaction) => { + let lockedCells = + inventoryScope === 'all' + ? await this.lockCellInventory(transaction) + : inventoryScope === 'general' + ? await this.lockGeneralCellInventory(transaction) + : undefined + const existing = await this.assignmentRow( + transaction, + identity, + inventoryScope !== 'none' + ) + retryScope = existing ? 'all' : 'general' + if (inventoryScope === 'general' && existing) { + throw new AssignmentInventoryScopeChanged() + } + const activityLeases = await this.lockAssignmentActivities(transaction, identity, true) + await this.recordRegionPreference(transaction, identity, preferredRegion, now) + let forcedDeadReassignment = false + let connectionHeadroomReassignment = false + let strandedReassignment = false + if (existing && !mayNormallyReassign(activity(existing), now)) { + lockedCells ??= await this.lockCellInventory(transaction, true) + const admission = await cellAdmissionStates(transaction) + const currentRow = lockedCells.find( + (row) => text(row, 'cell_id') === text(existing, 'cell_id') + ) + if (!currentRow) throw new Error('assigned_cell_missing') + const current = cell( + currentRow, + await this.cellRegion(transaction, text(existing, 'cell_id')) + ) + // Why: an existing-only cell that rejects attaches (C3's decommission + // posture) must not re-pin the hosts it refuses; the dead-cell fence + // path below does not apply either — the cell is live, just unwilling. + strandedReassignment = await this.assignmentStrandedOnUnservedCell( + transaction, + identity, + existing, + now + ) + if ( + !strandedReassignment && + (!this.requireLiveCells || (await this.cellIsLive(transaction, current.cellId, now))) + ) { + const hadControl = holdsControlLease( + activityLeases, + current.cellId, + integer(existing, 'assignment_epoch') + ) + const hasConnectionHeadroom = + hadControl || + (await this.cellHasConnectionHeadroom(transaction, current.cellId)) + if (hasConnectionHeadroom) { + const leaseExpiresAt = now + ASSIGNMENT_LIMITS.activityLeaseMs + if (hadControl) { + await this.touchAssignment(transaction, identity, leaseExpiresAt, now) + } else { + await this.adjustCellReservation(transaction, current.cellId, 1) + await this.adjustActivityCount( + transaction, + identity, + 'control', + 1, + leaseExpiresAt, + now + ) + await this.insertPendingControlLease( + transaction, + identity, + current.cellId, + integer(existing, 'assignment_epoch'), + now + ) + } + return this.result(identity, existing, current, leaseExpiresAt) + } + if (requestUnits(existing) > 0) { + throw new Error('relay_connection_headroom_exhausted') + } + if (admission.get(current.cellId) !== 'general') { + throw new Error('relay_connection_headroom_exhausted') + } + connectionHeadroomReassignment = true + } + if (!strandedReassignment && !connectionHeadroomReassignment) { + if ( + (await this.deadCellRequiresCommittedFence( + transaction, + identity, + current.cellId, + integer(existing, 'assignment_epoch') + )) && + !(await this.cellHasCommittedFence(transaction, current.cellId, now)) + ) { + throw new Error('relay_capacity_exhausted') + } + forcedDeadReassignment = true + } + } + + lockedCells ??= existing + ? await this.lockCellInventory(transaction, true) + : await this.lockGeneralCellInventory(transaction, true) + const target = await this.leastLoadedCell( + transaction, + lockedCells, + placementRegion + ) + if (!target) throw new Error('relay_capacity_exhausted') + const previousUnits = existing ? requestUnits(existing) : 0 + if (existing) { + await this.adjustCellReservation(transaction, text(existing, 'cell_id'), -previousUnits) + if (forcedDeadReassignment || strandedReassignment) { + // A fenced/dead incarnation cannot own drainable work, and a + // stranded host's only leases are the unclaimed grant artifacts of + // its own loop. Removing them prevents late expiry from + // decrementing the replacement. + await transaction.query( + `DELETE FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + } + } + await this.adjustCellReservation(transaction, target.cellId, 1) + const assignmentEpoch = existing ? integer(existing, 'assignment_epoch') + 1 : 1 + const leaseExpiresAt = now + ASSIGNMENT_LIMITS.activityLeaseMs + await transaction.query( + `INSERT INTO relay_assignments + (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, + last_activity_at, reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT (user_id, relay_host_id) DO UPDATE SET + cell_id = excluded.cell_id, + assignment_epoch = excluded.assignment_epoch, + lease_expires_at = excluded.lease_expires_at, + last_activity_at = excluded.last_activity_at, + reserved_controls = excluded.reserved_controls, + reserved_splices = excluded.reserved_splices, + reserved_invites = excluded.reserved_invites, + pending_installs = excluded.pending_installs, + pending_confirmations = excluded.pending_confirmations, + migration_leases = excluded.migration_leases`, + [ + identity.userId, + identity.relayHostId, + target.cellId, + assignmentEpoch, + leaseExpiresAt, + now, + 1, + 0, + 0, + 0, + 0, + 0 + ] + ) + if (existing) { + await this.releaseSupersededControlConnectionReservations( + transaction, + identity, + text(existing, 'cell_id'), + integer(existing, 'assignment_epoch'), + now + ) + } + await this.insertPendingControlLease( + transaction, + identity, + target.cellId, + assignmentEpoch, + now + ) + return { ...identity, ...target, assignmentEpoch, leaseExpiresAt } + }).catch((error: unknown) => { + if (isDatabaseLockUnavailable(error)) { + throw new AssignmentInventoryLockUnavailable(retryScope) + } + throw error + }) + } + + async resolve(identity: AssignmentIdentity): Promise { + const rows = await this.database.query( + `SELECT assignment.*, cell.cell_url, region.region + FROM relay_assignments assignment + JOIN relay_cells cell ON cell.cell_id = assignment.cell_id + LEFT JOIN relay_cell_regions region ON region.cell_id = cell.cell_id + WHERE assignment.user_id = ? AND assignment.relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + const row = rows[0] + if ( + row && + this.requireLiveCells && + !(await this.cellIsLive(this.database, text(row, 'cell_id'), this.now())) + ) { + return null + } + return row + ? { + ...identity, + cellId: text(row, 'cell_id'), + cellUrl: text(row, 'cell_url'), + region: optionalRelayRegion(row, 'region') ?? RELAY_DEFAULT_REGION, + assignmentEpoch: integer(row, 'assignment_epoch'), + leaseExpiresAt: integer(row, 'lease_expires_at') + } + : null + } + + async setCellEnabled(cellId: string, enabled: boolean): Promise { + await this.setCellAdmissionState(cellId, stateFromEnabled(enabled)) + } + + async setCellAdmissionState(cellId: string, state: CellAdmissionState): Promise { + await this.database.transaction( + async (transaction) => + await setCellAdmissionBeforeBoundary(transaction, cellId, state, this.now()) + ) + } + + async applyCellAdmissionSelector(input: { + attemptId: string + expectedGeneration: number + expectedMembershipSha256?: string + membership: CellAdmissionMembership + }): Promise<{ + changed: boolean + selector: { + generation: number + attemptId: string | null + membership: CellAdmissionMembership + } + }> { + return await this.admissionSelector.apply(input) + } + + async inspectCellAdmissionSelector( + attemptId?: string + ): Promise { + return await this.admissionSelector.inspect(attemptId) + } + + async addMigrationCells(input: { + attemptId: string + expectedGeneration: number + cells: MigrationCellRegistration[] + }): Promise<{ + changed: boolean + selector: { + generation: number + attemptId: string | null + membership: CellAdmissionMembership + } + }> { + return await this.migrationCellRegistrar.add(input) + } + + async recordCellHeartbeat(input: CellHeartbeat): Promise { + const now = this.now() + let inclusionWatermark: number | undefined + await this.database.transaction(async (transaction) => { + const configured = ( + await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [input.cellId]) + )[0] + if (!configured) throw new Error('cell_not_found') + if (text(configured, 'cell_url') !== input.cellUrl) throw new Error('cell_origin_mismatch') + if ( + (await this.cellRegion(transaction, input.cellId)) !== + (input.region ?? RELAY_DEFAULT_REGION) + ) { + throw new Error('cell_region_mismatch') + } + const current = ( + await transaction.queryLocked(`SELECT * FROM relay_cell_runtime WHERE cell_id = ?`, [ + input.cellId + ]) + )[0] + if ( + current && + ((text(current, 'cell_incarnation') === input.cellIncarnation && + input.startedAt !== integer(current, 'started_at')) || + (text(current, 'cell_incarnation') !== input.cellIncarnation && + input.startedAt <= integer(current, 'started_at'))) + ) { + throw new Error('stale_cell_incarnation') + } + const connectionLimit = ( + await transaction.queryLocked( + `SELECT * FROM relay_cell_connection_limits WHERE cell_id = ?`, + [input.cellId] + ) + )[0] + if (connectionLimit) { + if ( + input.totalConnections === undefined || + input.inFlightConnections === undefined || + input.reservedConnectionUnits === undefined || + input.enforcedConnectionUnits === undefined || + input.enforcedConnectionUnits !== + input.totalConnections + + input.inFlightConnections + + input.reservedConnectionUnits || + input.connectionHardCap !== integer(connectionLimit, 'hard_cap') || + input.connectionUnobservedBound !== + integer(connectionLimit, 'unobserved_bound') + ) { + throw new Error('cell_connection_telemetry_mismatch') + } + } else if ( + input.totalConnections !== undefined || + input.inFlightConnections !== undefined || + input.reservedConnectionUnits !== undefined || + input.enforcedConnectionUnits !== undefined || + input.connectionInclusionWatermark !== undefined || + input.connectionHardCap !== undefined || + input.connectionUnobservedBound !== undefined + ) { + throw new Error('cell_connection_limit_not_configured') + } + await transaction.query( + `INSERT INTO relay_cell_runtime + (cell_id, cell_url, cell_incarnation, started_at, ready, + observed_requests, last_heartbeat_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT (cell_id) DO UPDATE SET + cell_url = excluded.cell_url, + cell_incarnation = excluded.cell_incarnation, + started_at = excluded.started_at, + ready = excluded.ready, + observed_requests = excluded.observed_requests, + last_heartbeat_at = excluded.last_heartbeat_at, + updated_at = excluded.updated_at`, + [ + input.cellId, + input.cellUrl, + input.cellIncarnation, + input.startedAt, + input.ready ? 1 : 0, + input.observedRequests, + now, + now + ] + ) + // Live directors read relay_cell_runtime; keep heartbeats off the placement capacity lock. + if (!this.requireLiveCells) { + await transaction.query( + `UPDATE relay_cells SET observed_requests = ?, last_heartbeat_at = ?, updated_at = ? + WHERE cell_id = ?`, + [input.observedRequests, now, now, input.cellId] + ) + } + if (connectionLimit) { + const currentSnapshot = ( + await transaction.queryLocked( + `SELECT * FROM relay_cell_connection_snapshots WHERE cell_id = ?`, + [input.cellId] + ) + )[0] + const previousWatermark = + currentSnapshot && + text(currentSnapshot, 'cell_incarnation') === input.cellIncarnation + ? integer(currentSnapshot, 'inclusion_watermark') + : -1 + inclusionWatermark = input.connectionInclusionWatermark ?? previousWatermark + 1 + if ( + input.connectionInclusionWatermark !== undefined && + inclusionWatermark <= previousWatermark + ) { + throw new Error('stale_connection_snapshot') + } + await transaction.query( + `INSERT INTO relay_cell_connection_runtime + (cell_id, cell_incarnation, total_connections, in_flight_connections, + reserved_connection_units, enforced_connection_units, last_heartbeat_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT (cell_id) DO UPDATE SET + cell_incarnation = excluded.cell_incarnation, + total_connections = excluded.total_connections, + in_flight_connections = excluded.in_flight_connections, + reserved_connection_units = excluded.reserved_connection_units, + enforced_connection_units = excluded.enforced_connection_units, + last_heartbeat_at = excluded.last_heartbeat_at, + updated_at = excluded.updated_at`, + [ + input.cellId, + input.cellIncarnation, + input.totalConnections, + input.inFlightConnections, + input.reservedConnectionUnits, + input.enforcedConnectionUnits, + now, + now + ] + ) + await transaction.query( + `INSERT INTO relay_cell_connection_snapshots + (cell_id, cell_incarnation, inclusion_watermark, total_connections, + in_flight_connections, reserved_connection_units, + enforced_connection_units, snapshot_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT (cell_id) DO UPDATE SET + cell_incarnation = excluded.cell_incarnation, + inclusion_watermark = excluded.inclusion_watermark, + total_connections = excluded.total_connections, + in_flight_connections = excluded.in_flight_connections, + reserved_connection_units = excluded.reserved_connection_units, + enforced_connection_units = excluded.enforced_connection_units, + snapshot_at = excluded.snapshot_at`, + [ + input.cellId, + input.cellIncarnation, + inclusionWatermark, + input.totalConnections, + input.inFlightConnections, + input.reservedConnectionUnits, + input.enforcedConnectionUnits, + now + ] + ) + } + await transaction.query(`DELETE FROM relay_cell_fences WHERE cell_id = ?`, [input.cellId]) + await transaction.query(`DELETE FROM relay_cell_committed_fences WHERE cell_id = ?`, [ + input.cellId + ]) + await transaction.query( + `DELETE FROM relay_cell_legacy_fence_adoptions WHERE cell_id = ?`, + [input.cellId] + ) + }) + if (inclusionWatermark !== undefined) { + await this.releaseHeartbeatConnectionReservationDebt( + input.cellId, + input.cellIncarnation, + inclusionWatermark, + now + ) + } + } + + async recordCellRegionalRehomeStatus(input: CellRegionalRehomeStatus): Promise { + const now = this.now() + await this.database.transaction(async (transaction) => { + const runtime = ( + await transaction.queryLocked( + `SELECT cell_incarnation FROM relay_cell_runtime WHERE cell_id = ?`, + [input.cellId] + ) + )[0] + if (!runtime || text(runtime, 'cell_incarnation') !== input.cellIncarnation) { + throw new Error('stale_cell_incarnation') + } + await transaction.query( + `INSERT INTO relay_cell_capabilities + (cell_id, cell_incarnation, regional_rehome_protocol, last_heartbeat_at) + VALUES (?, ?, ?, ?) + ON CONFLICT (cell_id) DO UPDATE SET + cell_incarnation = excluded.cell_incarnation, + regional_rehome_protocol = excluded.regional_rehome_protocol, + last_heartbeat_at = excluded.last_heartbeat_at`, + [input.cellId, input.cellIncarnation, input.regionalRehomeProtocol, now] + ) + const safety = input.safety + await transaction.query( + `INSERT INTO relay_cell_rehome_safety + (cell_id, cell_incarnation, observed_at, sql_failures, reconnects, + control_activity_recovery_failures, database_pool_waiting, + database_pool_waiters_max, database_pool_wait_ms_max) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT (cell_id) DO UPDATE SET + cell_incarnation = excluded.cell_incarnation, + observed_at = excluded.observed_at, + sql_failures = excluded.sql_failures, + reconnects = excluded.reconnects, + control_activity_recovery_failures = + excluded.control_activity_recovery_failures, + database_pool_waiting = excluded.database_pool_waiting, + database_pool_waiters_max = excluded.database_pool_waiters_max, + database_pool_wait_ms_max = excluded.database_pool_wait_ms_max`, + [ + input.cellId, + input.cellIncarnation, + safety.observedAt, + safety.sqlFailures, + safety.reconnects, + safety.controlActivityRecoveryFailures, + safety.databasePoolWaiting, + safety.databasePoolWaitersMax, + safety.databasePoolWaitMsMax + ] + ) + }) + } + + private async releaseHeartbeatConnectionReservationDebt( + cellId: string, + cellIncarnation: string, + inclusionWatermark: number, + now: number + ): Promise { + await this.database.transaction(async (transaction) => { + const runtime = ( + await transaction.queryLocked( + `SELECT * FROM relay_cell_runtime WHERE cell_id = ?`, + [cellId] + ) + )[0] + const snapshot = ( + await transaction.queryLocked( + `SELECT * FROM relay_cell_connection_snapshots WHERE cell_id = ?`, + [cellId] + ) + )[0] + if ( + !runtime || + !snapshot || + text(runtime, 'cell_incarnation') !== cellIncarnation || + text(snapshot, 'cell_incarnation') !== cellIncarnation || + integer(snapshot, 'inclusion_watermark') !== inclusionWatermark + ) { + return + } + await transaction.query( + `UPDATE relay_control_connection_reservations + SET state = 'released', released_at = ?, updated_at = ? + WHERE cell_id = ? AND state = 'claimed' + AND inclusion_watermark IS NOT NULL + AND inclusion_watermark <= ?`, + [now, now, cellId, inclusionWatermark] + ) + await transaction.query( + `UPDATE relay_control_connection_reservations + SET state = 'released', released_at = ?, updated_at = ? + WHERE reservation_id IN ( + SELECT duplicate.reservation_id + FROM relay_control_connection_reservations duplicate + WHERE duplicate.cell_id = ? + AND duplicate.state = 'late-arrival-debt' + AND duplicate.claim_activity_id IS NULL + AND duplicate.timeout_at <= ? + AND EXISTS ( + SELECT 1 + FROM relay_control_connection_reservations retained + WHERE retained.user_id = duplicate.user_id + AND retained.relay_host_id = duplicate.relay_host_id + AND retained.assignment_epoch = duplicate.assignment_epoch + AND retained.cell_id = duplicate.cell_id + AND retained.state = 'late-arrival-debt' + AND retained.claim_activity_id IS NULL + AND retained.timeout_at <= ? + AND ( + retained.created_at < duplicate.created_at OR + ( + retained.created_at = duplicate.created_at AND + retained.reservation_id < duplicate.reservation_id + ) + ) + ) + )`, + [now, now, cellId, now, now] + ) + await transaction.query( + `UPDATE relay_control_connection_reservations + SET state = 'released', released_at = ?, updated_at = ? + WHERE cell_id = ? AND state = 'late-arrival-debt' + AND claim_activity_id IS NULL AND timeout_at <= ? + AND EXISTS ( + SELECT 1 FROM relay_assignments assignment + WHERE assignment.user_id = + relay_control_connection_reservations.user_id + AND assignment.relay_host_id = + relay_control_connection_reservations.relay_host_id + AND assignment.assignment_epoch > + relay_control_connection_reservations.assignment_epoch + ) + AND EXISTS ( + SELECT 1 FROM relay_assignment_migrations migration + WHERE migration.user_id = + relay_control_connection_reservations.user_id + AND migration.relay_host_id = + relay_control_connection_reservations.relay_host_id + AND migration.assignment_epoch = + relay_control_connection_reservations.assignment_epoch + AND migration.target_cell_id = + relay_control_connection_reservations.cell_id + AND migration.aborted_at IS NOT NULL + )`, + [now, now, cellId, now] + ) + }) + } + + async attestCellFence(cellId: string, cellIncarnation: string): Promise { + const now = this.now() + return await this.database.transaction(async (transaction) => { + const cell = ( + await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [cellId]) + )[0] + if (!cell) throw new Error('cell_not_found') + if (integer(cell, 'enabled') !== 0) throw new Error('cell_fence_admission_enabled') + const runtime = ( + await transaction.queryLocked(`SELECT * FROM relay_cell_runtime WHERE cell_id = ?`, [ + cellId + ]) + )[0] + if ( + !runtime || + text(runtime, 'cell_incarnation') !== cellIncarnation || + integer(runtime, 'last_heartbeat_at') > now - this.heartbeatTtlMs + ) { + throw new Error('cell_fence_runtime_not_stale') + } + const expiresAt = now + CELL_FENCE_TTL_MS + await transaction.query( + `INSERT INTO relay_cell_fences + (cell_id, cell_incarnation, attested_at, expires_at) + VALUES (?, ?, ?, ?) + ON CONFLICT (cell_id) DO UPDATE SET + cell_incarnation = excluded.cell_incarnation, + attested_at = excluded.attested_at, + expires_at = excluded.expires_at`, + [cellId, cellIncarnation, now, expiresAt] + ) + return expiresAt + }) + } + + async adoptLegacyCellFence(cellId: string, cellIncarnation: string): Promise { + const now = this.now() + return await this.database.transaction(async (transaction) => { + const cell = ( + await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [cellId]) + )[0] + if (!cell) throw new Error('cell_not_found') + if (integer(cell, 'enabled') !== 0) throw new Error('cell_fence_admission_enabled') + const runtime = ( + await transaction.queryLocked(`SELECT * FROM relay_cell_runtime WHERE cell_id = ?`, [ + cellId + ]) + )[0] + if ( + !runtime || + text(runtime, 'cell_incarnation') !== cellIncarnation || + integer(runtime, 'last_heartbeat_at') > now - this.heartbeatTtlMs + ) { + throw new Error('cell_fence_runtime_not_stale') + } + const attempt = ( + await transaction.queryLocked( + `SELECT attempt_id FROM relay_cell_fence_attempts + WHERE cell_id = ? ORDER BY created_at DESC LIMIT 1`, + [cellId] + ) + )[0] + if (attempt) throw new Error('legacy_cell_fence_attempt_exists') + const expiresAt = now + CELL_FENCE_TTL_MS + await transaction.query( + `INSERT INTO relay_cell_fences + (cell_id, cell_incarnation, attested_at, expires_at) + VALUES (?, ?, ?, ?) + ON CONFLICT (cell_id) DO UPDATE SET + cell_incarnation = excluded.cell_incarnation, + attested_at = excluded.attested_at, + expires_at = excluded.expires_at`, + [cellId, cellIncarnation, now, expiresAt] + ) + return expiresAt + }) + } + + async commitLegacyCellFenceAdoption( + cellId: string, + cellIncarnation: string + ): Promise { + const now = this.now() + await this.database.transaction(async (transaction) => { + const cell = ( + await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [cellId]) + )[0] + if (!cell) throw new Error('cell_not_found') + if (integer(cell, 'enabled') !== 0) throw new Error('cell_fence_admission_enabled') + const runtime = ( + await transaction.queryLocked(`SELECT * FROM relay_cell_runtime WHERE cell_id = ?`, [ + cellId + ]) + )[0] + const fence = ( + await transaction.queryLocked(`SELECT * FROM relay_cell_fences WHERE cell_id = ?`, [ + cellId + ]) + )[0] + const attempt = ( + await transaction.queryLocked( + `SELECT attempt_id FROM relay_cell_fence_attempts + WHERE cell_id = ? ORDER BY created_at DESC LIMIT 1`, + [cellId] + ) + )[0] + if (attempt) throw new Error('legacy_cell_fence_attempt_exists') + if ( + !runtime || + !fence || + text(runtime, 'cell_incarnation') !== cellIncarnation || + text(fence, 'cell_incarnation') !== cellIncarnation || + integer(runtime, 'last_heartbeat_at') > now - this.heartbeatTtlMs || + integer(fence, 'attested_at') < integer(runtime, 'last_heartbeat_at') || + integer(fence, 'expires_at') <= now + ) { + throw new Error('legacy_cell_fence_not_active') + } + await transaction.query( + `INSERT INTO relay_cell_legacy_fence_adoptions + (cell_id, cell_incarnation, attested_at, expires_at) + VALUES (?, ?, ?, ?) + ON CONFLICT (cell_id) DO UPDATE SET + cell_incarnation = excluded.cell_incarnation, + attested_at = excluded.attested_at, + expires_at = excluded.expires_at`, + [cellId, cellIncarnation, now, integer(fence, 'expires_at')] + ) + }) + } + + async prepareCellFenceAttempt(input: CellFenceAttemptEvidence): Promise { + const now = this.now() + return await this.database.transaction(async (transaction) => { + let createdAt = now + const cell = ( + await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [ + input.cellId + ]) + )[0] + if (!cell) throw new Error('cell_not_found') + if (integer(cell, 'enabled') !== 0) throw new Error('cell_fence_admission_enabled') + const runtime = ( + await transaction.queryLocked(`SELECT * FROM relay_cell_runtime WHERE cell_id = ?`, [ + input.cellId + ]) + )[0] + if (!runtime || text(runtime, 'cell_incarnation') !== input.cellIncarnation) { + throw new Error('cell_fence_runtime_mismatch') + } + const existing = ( + await transaction.queryLocked( + `${CELL_FENCE_ATTEMPT_SELECT} + WHERE attempts.cell_id = ? ORDER BY attempts.created_at DESC LIMIT 1`, + [input.cellId] + ) + )[0] + if (existing) { + const attempt = cellFenceAttempt(existing) + if (!attempt.completedAt && !attempt.abortedAt) { + assertCellFenceAttemptBase(existing, input) + return attempt + } + createdAt = Math.max(now, attempt.createdAt + 1) + } else { + const activeFence = ( + await transaction.queryLocked( + `SELECT cell_id FROM relay_cell_fences + WHERE cell_id = ? AND cell_incarnation = ? AND expires_at > ?`, + [input.cellId, input.cellIncarnation, now] + ) + )[0] + if (activeFence) throw new Error('cell_fence_already_attested') + } + const expiresAt = createdAt + CELL_FENCE_ATTEMPT_TTL_MS + await transaction.query( + `INSERT INTO relay_cell_fence_attempts + (attempt_id, environment, cell_id, cell_incarnation, mig_name, instance_group, + generation_identity, fence_commit, plan_sha256, gce_operation, created_at, + expires_at, apply_started_at, completed_at, aborted_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, ?, ?, NULL, NULL, NULL)`, + [ + input.attemptId, + input.environment, + input.cellId, + input.cellIncarnation, + input.migName, + input.instanceGroup, + input.generationIdentity, + input.fenceCommit, + input.planSha256, + createdAt, + expiresAt + ] + ) + await transaction.query( + `INSERT INTO relay_cell_fence_plan_bindings + (attempt_id, plan_object_name, plan_object_generation, var_file_sha256, + terraform_state_lineage, terraform_state_serial, + terraform_state_object_generation, terraform_state_object_sha256, + request_reason) + VALUES (?, ?, NULL, ?, ?, ?, ?, ?, ?)`, + [ + input.attemptId, + input.planObjectName, + input.varFileSha256, + input.terraformStateLineage, + input.terraformStateSerial, + input.terraformStateObjectGeneration, + input.terraformStateObjectSha256, + input.requestReason + ] + ) + const { planObjectGeneration: _ignored, ...prepared } = input + return { ...prepared, createdAt, expiresAt } + }) + } + + async bindCellFencePlanGeneration( + input: CellFenceAttemptEvidence, + planObjectGeneration: string + ): Promise { + const now = this.now() + return await this.database.transaction(async (transaction) => { + const row = await lockedCellFenceAttempt(transaction, input.attemptId) + assertCellFenceAttemptBase(row, input) + if (integer(row, 'expires_at') <= now) throw new Error('cell_fence_attempt_expired') + if (row.apply_started_at !== null) throw new Error('cell_fence_apply_already_started') + if (row.completed_at !== null || row.aborted_at !== null) { + throw new Error('cell_fence_attempt_terminal') + } + const existing = optionalText(row, 'plan_object_generation') + if (existing && existing !== planObjectGeneration) { + throw new Error('cell_fence_plan_generation_mismatch') + } + if (!existing) { + await transaction.query( + `UPDATE relay_cell_fence_plan_bindings + SET plan_object_generation = ? WHERE attempt_id = ?`, + [planObjectGeneration, input.attemptId] + ) + row.plan_object_generation = planObjectGeneration + } + return cellFenceAttempt(row) + }) + } + + async startCellFenceApply( + input: CellFenceAttemptEvidence, + invocationId: string, + requestReason: string + ): Promise<{ attempt: CellFenceAttempt; invocation: CellFenceApplyInvocation }> { + const now = this.now() + return await this.database.transaction(async (transaction) => { + if (requestReason !== `${input.requestReason}/${invocationId}`) { + throw new Error('cell_fence_invocation_mismatch') + } + const row = await lockedCellFenceAttempt(transaction, input.attemptId) + assertActiveCellFenceAttempt(row, input, now) + const existing = ( + await transaction.queryLocked( + `SELECT * FROM relay_cell_fence_apply_invocations WHERE invocation_id = ?`, + [invocationId] + ) + )[0] + if (existing) { + if ( + text(existing, 'attempt_id') !== input.attemptId || + text(existing, 'request_reason') !== requestReason + ) { + throw new Error('cell_fence_invocation_mismatch') + } + return { + attempt: cellFenceAttempt(row), + invocation: cellFenceApplyInvocation(existing) + } + } + if (row.apply_started_at === null) { + await transaction.query( + `UPDATE relay_cell_fence_attempts SET apply_started_at = ? WHERE attempt_id = ?`, + [now, input.attemptId] + ) + row.apply_started_at = now + } + await transaction.query( + `INSERT INTO relay_cell_fence_apply_invocations + (invocation_id, attempt_id, request_reason, started_at, gce_operation) + VALUES (?, ?, ?, ?, NULL)`, + [invocationId, input.attemptId, requestReason, now] + ) + return { + attempt: cellFenceAttempt(row), + invocation: { invocationId, requestReason, startedAt: now } + } + }) + } + + async recordCellFenceOperation( + input: CellFenceAttemptEvidence, + invocationId: string, + requestReason: string, + gceOperation: string + ): Promise<{ attempt: CellFenceAttempt; invocation: CellFenceApplyInvocation }> { + const now = this.now() + return await this.database.transaction(async (transaction) => { + if (requestReason !== `${input.requestReason}/${invocationId}`) { + throw new Error('cell_fence_invocation_mismatch') + } + const row = await lockedCellFenceAttempt(transaction, input.attemptId) + assertActiveCellFenceAttempt(row, input, now) + if (row.apply_started_at === null) throw new Error('cell_fence_apply_not_started') + const invocation = ( + await transaction.queryLocked( + `SELECT * FROM relay_cell_fence_apply_invocations WHERE invocation_id = ?`, + [invocationId] + ) + )[0] + if ( + !invocation || + text(invocation, 'attempt_id') !== input.attemptId || + text(invocation, 'request_reason') !== requestReason + ) { + throw new Error('cell_fence_invocation_mismatch') + } + const existing = invocation.gce_operation + if (existing !== null && existing !== gceOperation) { + throw new Error('cell_fence_operation_mismatch') + } + await transaction.query( + `UPDATE relay_cell_fence_apply_invocations + SET gce_operation = ? WHERE invocation_id = ?`, + [gceOperation, invocationId] + ) + invocation.gce_operation = gceOperation + await transaction.query( + `UPDATE relay_cell_fence_attempts SET gce_operation = ? WHERE attempt_id = ?`, + [gceOperation, input.attemptId] + ) + row.gce_operation = gceOperation + return { + attempt: cellFenceAttempt(row), + invocation: cellFenceApplyInvocation(invocation) + } + }) + } + + async cellFenceAttempt(cellId: string): Promise { + const rows = await this.database.query( + `${CELL_FENCE_ATTEMPT_SELECT} + WHERE attempts.cell_id = ? ORDER BY attempts.created_at DESC LIMIT 1`, + [cellId] + ) + if (rows.length === 0) return null + const attempt = cellFenceAttempt(rows[0]!) + const invocations = await this.database.query( + `SELECT * FROM relay_cell_fence_apply_invocations + WHERE attempt_id = ? ORDER BY started_at, invocation_id`, + [attempt.attemptId] + ) + return { + ...attempt, + applyInvocations: invocations.map(cellFenceApplyInvocation) + } + } + + async abortCellFenceAttempt(input: CellFenceAttemptEvidence): Promise { + return await this.database.transaction(async (transaction) => { + const row = await lockedCellFenceAttempt(transaction, input.attemptId) + assertCellFenceAttemptBase(row, input) + if (row.completed_at !== null) throw new Error('cell_fence_attempt_completed') + if (row.aborted_at !== null) throw new Error('cell_fence_attempt_aborted') + if (row.apply_started_at !== null || row.gce_operation !== null) { + throw new Error('cell_fence_apply_may_have_started') + } + const abortedAt = this.now() + await transaction.query( + `UPDATE relay_cell_fence_attempts SET aborted_at = ? WHERE attempt_id = ?`, + [abortedAt, input.attemptId] + ) + row.aborted_at = abortedAt + return cellFenceAttempt(row) + }) + } + + async attestCellFenceAttempt( + input: CellFenceAttemptEvidence, + gceOperation: string + ): Promise<{ expiresAt: number; attempt: CellFenceAttempt }> { + const now = this.now() + return await this.database.transaction(async (transaction) => { + const attemptRow = await lockedCellFenceAttempt(transaction, input.attemptId) + assertCellFenceAttemptEvidence(attemptRow, input) + if (attemptRow.completed_at !== null) { + if (attemptRow.gce_operation !== gceOperation) { + throw new Error('cell_fence_operation_not_attested') + } + const cell = ( + await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [ + input.cellId + ]) + )[0] + if (!cell) throw new Error('cell_not_found') + if (integer(cell, 'enabled') !== 0) throw new Error('cell_fence_admission_enabled') + const runtime = ( + await transaction.queryLocked(`SELECT * FROM relay_cell_runtime WHERE cell_id = ?`, [ + input.cellId + ]) + )[0] + if ( + !runtime || + text(runtime, 'cell_incarnation') !== input.cellIncarnation || + integer(runtime, 'last_heartbeat_at') > now - this.heartbeatTtlMs + ) { + throw new Error('cell_fence_runtime_not_stale') + } + const expiresAt = now + CELL_FENCE_TTL_MS + await transaction.query( + `INSERT INTO relay_cell_fences + (cell_id, cell_incarnation, attested_at, expires_at) + VALUES (?, ?, ?, ?) + ON CONFLICT (cell_id) DO UPDATE SET + cell_incarnation = excluded.cell_incarnation, + attested_at = excluded.attested_at, + expires_at = excluded.expires_at`, + [input.cellId, input.cellIncarnation, now, expiresAt] + ) + await this.recordCommittedCellFence( + transaction, + input.cellId, + input.cellIncarnation, + input.attemptId, + now, + expiresAt + ) + return { + expiresAt, + attempt: cellFenceAttempt(attemptRow) + } + } + assertActiveCellFenceAttempt(attemptRow, input, now) + if ( + attemptRow.apply_started_at === null || + attemptRow.gce_operation !== gceOperation + ) { + throw new Error('cell_fence_operation_not_attested') + } + const cell = ( + await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [ + input.cellId + ]) + )[0] + if (!cell) throw new Error('cell_not_found') + if (integer(cell, 'enabled') !== 0) throw new Error('cell_fence_admission_enabled') + const runtime = ( + await transaction.queryLocked(`SELECT * FROM relay_cell_runtime WHERE cell_id = ?`, [ + input.cellId + ]) + )[0] + if ( + !runtime || + text(runtime, 'cell_incarnation') !== input.cellIncarnation || + integer(runtime, 'last_heartbeat_at') > now - this.heartbeatTtlMs + ) { + throw new Error('cell_fence_runtime_not_stale') + } + const expiresAt = now + CELL_FENCE_TTL_MS + await transaction.query( + `INSERT INTO relay_cell_fences + (cell_id, cell_incarnation, attested_at, expires_at) + VALUES (?, ?, ?, ?) + ON CONFLICT (cell_id) DO UPDATE SET + cell_incarnation = excluded.cell_incarnation, + attested_at = excluded.attested_at, + expires_at = excluded.expires_at`, + [input.cellId, input.cellIncarnation, now, expiresAt] + ) + await this.recordCommittedCellFence( + transaction, + input.cellId, + input.cellIncarnation, + input.attemptId, + now, + expiresAt + ) + await transaction.query( + `UPDATE relay_cell_fence_attempts SET completed_at = ? WHERE attempt_id = ?`, + [now, input.attemptId] + ) + attemptRow.completed_at = now + return { expiresAt, attempt: cellFenceAttempt(attemptRow) } + }) + } + + async prepareCellDrainAttempt(input: { + attemptId: string + cellId: string + cellIncarnation: string + traceValue: string + plannedGraceMs: number + }): Promise { + const now = this.now() + return await this.database.transaction(async (transaction) => { + await this.assertDrainCellGeneration( + transaction, + input.cellId, + input.cellIncarnation + ) + const existing = ( + await transaction.queryLocked( + `SELECT * FROM relay_cell_drain_attempt_states + WHERE cell_id = ? ORDER BY prepared_at DESC LIMIT 1`, + [input.cellId] + ) + )[0] + if (existing) { + if (text(existing, 'attempt_id') === input.attemptId) { + if ( + text(existing, 'cell_incarnation') !== input.cellIncarnation || + text(existing, 'trace_value') !== input.traceValue || + integer(existing, 'planned_grace_ms') !== input.plannedGraceMs + ) { + throw new Error('drain_attempt_generation_mismatch') + } + return { ...cellDrainAttempt(existing), shouldSend: false } + } + if ( + text(existing, 'state') !== 'proven-not-delivered' || + text(existing, 'cell_incarnation') !== input.cellIncarnation + ) { + throw new Error('drain_attempt_generation_mismatch') + } + } + await transaction.query( + `INSERT INTO relay_cell_drain_attempt_states + (attempt_id, cell_id, cell_incarnation, trace_value, planned_grace_ms, + state, prepared_at, send_may_have_started_at, send_permit_expires_at, + application_receipt_at, backend_success_status, backend_instance, + receipt_cell_incarnation, retry_after, recover_forward_attempted_at, + proven_not_delivered_at) + VALUES (?, ?, ?, ?, ?, 'prepared', ?, NULL, NULL, NULL, NULL, NULL, + NULL, NULL, NULL, NULL)`, + [ + input.attemptId, + input.cellId, + input.cellIncarnation, + input.traceValue, + input.plannedGraceMs, + now + ] + ) + return { + attemptId: input.attemptId, + cellId: input.cellId, + cellIncarnation: input.cellIncarnation, + traceValue: input.traceValue, + plannedGraceMs: input.plannedGraceMs, + state: 'prepared', + preparedAt: now, + shouldSend: false + } + }) + } + + async beginCellDrainSend(input: { + attemptId: string + cellId: string + cellIncarnation: string + }): Promise { + try { + return await this.beginCellDrainSendOnce(input) + } catch (error) { + if ( + !(error instanceof Error) || + error.message !== 'migration_activity_accounting_mismatch' + ) { + throw error + } + const activityCells = await this.database.query( + `SELECT DISTINCT lease.cell_id + FROM relay_assignment_activity_leases lease + WHERE EXISTS ( + SELECT 1 FROM relay_assignment_migrations migration + WHERE migration.user_id = lease.user_id + AND migration.relay_host_id = lease.relay_host_id + AND migration.source_cell_id = ? + AND migration.completed_at IS NULL + AND migration.aborted_at IS NULL + ) + ORDER BY lease.cell_id`, + [input.cellId] + ) + const cellIds = activityCells + .map((row) => text(row, 'cell_id')) + .filter((cellId) => cellId !== input.cellId) + for (const cellId of cellIds) { + await this.reconcileReservationAccounting(input.cellId, cellId) + } + await this.refreshDrainMigrationLeases(input) + return await this.beginCellDrainSendOnce(input) + } + } + + private async refreshDrainMigrationLeases(input: { + cellId: string + cellIncarnation: string + }): Promise { + await this.retireObsoleteDrainMigrations(input.cellId) + for (let repairAttempt = 0; ; repairAttempt++) { + try { + await this.refreshDrainMigrationLeasesOnce(input) + return + } catch (error) { + if ( + !(error instanceof Error) || + error.message !== 'migration_activity_accounting_mismatch' || + repairAttempt >= MAX_DRAIN_ACCOUNTING_REPAIR_ATTEMPTS + ) { + throw error + } + await this.reconcileDrainMigrationAccounting(input.cellId) + } + } + } + + private async reconcileDrainMigrationAccounting(sourceCellId: string): Promise { + const activityCells = await this.database.query( + `SELECT DISTINCT lease.cell_id + FROM relay_assignment_activity_leases lease + WHERE EXISTS ( + SELECT 1 FROM relay_assignment_migrations migration + WHERE migration.user_id = lease.user_id + AND migration.relay_host_id = lease.relay_host_id + AND migration.source_cell_id = ? + AND migration.completed_at IS NULL + AND migration.aborted_at IS NULL + ) + ORDER BY lease.cell_id`, + [sourceCellId] + ) + for (const cellId of activityCells + .map((row) => text(row, 'cell_id')) + .filter((cellId) => cellId !== sourceCellId)) { + await this.reconcileReservationAccounting(sourceCellId, cellId) + } + } + + private async retireObsoleteDrainMigrations(sourceCellId: string): Promise { + const candidates = await this.database.query( + `SELECT migration.user_id, migration.relay_host_id, + migration.assignment_epoch + FROM relay_assignment_migrations migration + JOIN relay_assignments assignment + ON assignment.user_id = migration.user_id + AND assignment.relay_host_id = migration.relay_host_id + WHERE migration.source_cell_id = ? + AND migration.completed_at IS NULL AND migration.aborted_at IS NULL + AND assignment.assignment_epoch > migration.assignment_epoch + ORDER BY migration.user_id, migration.relay_host_id`, + [sourceCellId] + ) + for (const candidate of candidates) { + await this.retireObsoleteDrainMigration(sourceCellId, candidate) + } + } + + private async retireObsoleteDrainMigration( + sourceCellId: string, + candidate: SqlRow + ): Promise { + const now = this.now() + await this.database.transaction(async (transaction) => { + const identity = { + userId: text(candidate, 'user_id'), + relayHostId: text(candidate, 'relay_host_id') + } + const assignment = await this.assignmentRow(transaction, identity) + const assignmentEpoch = integer(candidate, 'assignment_epoch') + const migrationRow = ( + await transaction.queryLocked( + `SELECT * FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ? + AND source_cell_id = ? + AND completed_at IS NULL AND aborted_at IS NULL`, + [identity.userId, identity.relayHostId, assignmentEpoch, sourceCellId] + ) + )[0] + if (!migrationRow) return + if (!assignment || integer(assignment, 'assignment_epoch') <= assignmentEpoch) { + throw new Error('migration_assignment_mismatch') + } + const leases = await this.lockAssignmentActivities(transaction, identity) + const obsoleteActivityIds = new Set([ + pendingControlActivityId(assignmentEpoch), + migrationActivityId(assignmentEpoch) + ]) + if (leases.some((lease) => obsoleteActivityIds.has(text(lease, 'activity_id')))) { + throw new Error('migration_activity_topology_mismatch') + } + await this.releaseSupersededControlConnectionReservations( + transaction, + identity, + text(migrationRow, 'target_cell_id'), + assignmentEpoch, + now + ) + await transaction.query( + `UPDATE relay_assignment_migrations SET aborted_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [now, now, identity.userId, identity.relayHostId, assignmentEpoch] + ) + }) + } + + private async refreshDrainMigrationLeasesOnce(input: { + cellId: string + cellIncarnation: string + }): Promise { + const now = this.now() + await this.database.transaction(async (transaction) => { + await this.assertDrainCellGeneration( + transaction, + input.cellId, + input.cellIncarnation + ) + const assignments = await transaction.queryLocked( + `SELECT assignment.* + FROM relay_assignments assignment + WHERE EXISTS ( + SELECT 1 FROM relay_assignment_migrations migration + WHERE migration.user_id = assignment.user_id + AND migration.relay_host_id = assignment.relay_host_id + AND migration.source_cell_id = ? + AND migration.completed_at IS NULL + AND migration.aborted_at IS NULL + ) + ORDER BY assignment.user_id, assignment.relay_host_id`, + [input.cellId] + ) + const activityLeases = await transaction.queryLocked( + `SELECT lease.* + FROM relay_assignment_activity_leases lease + WHERE EXISTS ( + SELECT 1 FROM relay_assignment_migrations migration + WHERE migration.user_id = lease.user_id + AND migration.relay_host_id = lease.relay_host_id + AND migration.source_cell_id = ? + AND migration.completed_at IS NULL + AND migration.aborted_at IS NULL + ) + ORDER BY lease.user_id, lease.relay_host_id, lease.activity_id`, + [input.cellId] + ) + const migrations = await transaction.queryLocked( + `SELECT migration.* + FROM relay_assignment_migrations migration + WHERE migration.source_cell_id = ? + AND migration.completed_at IS NULL + AND migration.aborted_at IS NULL + ORDER BY migration.user_id, migration.relay_host_id`, + [input.cellId] + ) + const cells = await this.lockCellInventory(transaction) + for (const migrationRow of migrations) { + const identity = { + userId: text(migrationRow, 'user_id'), + relayHostId: text(migrationRow, 'relay_host_id') + } + const assignment = assignments.find( + (candidate) => + text(candidate, 'user_id') === identity.userId && + text(candidate, 'relay_host_id') === identity.relayHostId + ) + assertCurrentMigrationAssignment(assignment, migrationRow) + const leases = activityLeases.filter( + (lease) => + text(lease, 'user_id') === identity.userId && + text(lease, 'relay_host_id') === identity.relayHostId + ) + const migrationLeases = leases.filter( + (lease) => text(lease, 'activity_kind') === 'migration' + ) + const assignmentEpoch = integer(migrationRow, 'assignment_epoch') + const sourceRequestUnits = integer(migrationRow, 'source_request_units') + const targetCellId = text(migrationRow, 'target_cell_id') + const expiresAt = now + ASSIGNMENT_LIMITS.migrationLeaseMs + if (migrationLeases.length > 0) { + assertAssignmentActivityAccounting(assignment, leases, migrationRow) + await transaction.query( + `UPDATE relay_assignment_activity_leases SET expires_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? + AND activity_id IN (?, ?)`, + [ + expiresAt, + now, + identity.userId, + identity.relayHostId, + pendingControlActivityId(assignmentEpoch), + migrationActivityId(assignmentEpoch) + ] + ) + await transaction.query( + `UPDATE relay_assignments SET + lease_expires_at = CASE WHEN lease_expires_at > ? THEN lease_expires_at ELSE ? END, + last_activity_at = ? + WHERE user_id = ? AND relay_host_id = ?`, + [expiresAt, expiresAt, now, identity.userId, identity.relayHostId] + ) + await transaction.query( + `UPDATE relay_assignment_migrations SET expires_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [expiresAt, now, identity.userId, identity.relayHostId, assignmentEpoch] + ) + await this.refreshPendingControlReservation( + transaction, + identity, + targetCellId, + assignmentEpoch, + expiresAt, + now + ) + continue + } + if ( + optionalInteger(migrationRow, 'target_registered_at') === undefined || + integer(migrationRow, 'expires_at') > now || + sourceRequestUnits < 0 || + integer(migrationRow, 'target_reserved_units') !== sourceRequestUnits + 1 || + activityLeaseById(leases, migrationActivityId(assignmentEpoch)) || + !cells.some((cell) => text(cell, 'cell_id') === targetCellId) + ) { + throw new Error('migration_activity_lease_shape_mismatch') + } + await this.adjustCellReservation(transaction, targetCellId, sourceRequestUnits) + await transaction.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, ?, 'migration', ?, ?, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + migrationActivityId(assignmentEpoch), + targetCellId, + sourceRequestUnits, + expiresAt, + now + ] + ) + await transaction.query( + `UPDATE relay_assignments SET migration_leases = migration_leases + 1, + lease_expires_at = CASE WHEN lease_expires_at > ? THEN lease_expires_at ELSE ? END, + last_activity_at = ? + WHERE user_id = ? AND relay_host_id = ?`, + [expiresAt, expiresAt, now, identity.userId, identity.relayHostId] + ) + await transaction.query( + `UPDATE relay_assignment_migrations SET expires_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [expiresAt, now, identity.userId, identity.relayHostId, assignmentEpoch] + ) + await this.refreshPendingControlReservation( + transaction, + identity, + targetCellId, + assignmentEpoch, + expiresAt, + now + ) + } + }) + } + + private async beginCellDrainSendOnce(input: { + attemptId: string + cellId: string + cellIncarnation: string + }): Promise { + const now = this.now() + return await this.database.transaction(async (transaction) => { + await this.assertDrainCellGeneration( + transaction, + input.cellId, + input.cellIncarnation + ) + const row = ( + await transaction.queryLocked( + `SELECT * FROM relay_cell_drain_attempt_states + WHERE attempt_id = ? AND cell_id = ?`, + [input.attemptId, input.cellId] + ) + )[0] + if (!row || text(row, 'cell_incarnation') !== input.cellIncarnation) { + throw new Error('drain_attempt_not_found') + } + if (text(row, 'state') !== 'prepared') { + return { ...cellDrainAttempt(row), shouldSend: false } + } + const assignments = await transaction.queryLocked( + `SELECT assignment.* + FROM relay_assignments assignment + JOIN relay_assignment_migrations migration + ON migration.user_id = assignment.user_id + AND migration.relay_host_id = assignment.relay_host_id + WHERE migration.source_cell_id = ? + AND migration.completed_at IS NULL AND migration.aborted_at IS NULL + ORDER BY assignment.user_id, assignment.relay_host_id`, + [input.cellId] + ) + const activityLeases = await transaction.queryLocked( + `SELECT lease.* + FROM relay_assignment_activity_leases lease + WHERE EXISTS ( + SELECT 1 FROM relay_assignment_migrations migration + WHERE migration.user_id = lease.user_id + AND migration.relay_host_id = lease.relay_host_id + AND migration.source_cell_id = ? + AND migration.completed_at IS NULL + AND migration.aborted_at IS NULL + ) + ORDER BY lease.user_id, lease.relay_host_id, lease.activity_id`, + [input.cellId] + ) + const migrationIncarnations = await transaction.queryLocked( + `SELECT migration.*, + ( + SELECT incarnation.source_cell_incarnation + FROM relay_assignment_migration_incarnations incarnation + WHERE incarnation.user_id = migration.user_id + AND incarnation.relay_host_id = migration.relay_host_id + AND incarnation.assignment_epoch = migration.assignment_epoch + ) AS source_cell_incarnation, + ( + SELECT incarnation.target_cell_incarnation + FROM relay_assignment_migration_incarnations incarnation + WHERE incarnation.user_id = migration.user_id + AND incarnation.relay_host_id = migration.relay_host_id + AND incarnation.assignment_epoch = migration.assignment_epoch + ) AS target_cell_incarnation + FROM relay_assignment_migrations migration + WHERE migration.source_cell_id = ? + AND migration.completed_at IS NULL AND migration.aborted_at IS NULL + ORDER BY migration.user_id, migration.relay_host_id`, + [input.cellId] + ) + if ( + migrationIncarnations.some( + (migration) => + optionalText(migration, 'source_cell_incarnation') !== input.cellIncarnation + ) + ) { + throw new Error('drain_migration_source_incarnation_mismatch') + } + for (const migrationRow of migrationIncarnations) { + const assignment = assignments.find( + (candidate) => + text(candidate, 'user_id') === text(migrationRow, 'user_id') && + text(candidate, 'relay_host_id') === text(migrationRow, 'relay_host_id') + ) + assertCurrentMigrationAssignment(assignment, migrationRow) + assertAssignmentActivityAccounting( + assignment, + activityLeases.filter( + (lease) => + text(lease, 'user_id') === text(migrationRow, 'user_id') && + text(lease, 'relay_host_id') === text(migrationRow, 'relay_host_id') + ), + migrationRow + ) + } + const sendPermitExpiresAt = now + CELL_DRAIN_SEND_PERMIT_MS + await transaction.query( + `UPDATE relay_cell_drain_attempt_states + SET state = 'send-may-have-started', send_may_have_started_at = ?, + send_permit_expires_at = ? + WHERE attempt_id = ?`, + [now, sendPermitExpiresAt, input.attemptId] + ) + await transaction.query( + `INSERT INTO relay_post_drain_migration_pins + (user_id, relay_host_id, assignment_epoch, drain_attempt_id, + source_cell_id, source_cell_incarnation, target_cell_id, + target_cell_incarnation, source_request_units, target_reserved_units, + pinned_at) + SELECT migration.user_id, migration.relay_host_id, + migration.assignment_epoch, ?, migration.source_cell_id, + incarnation.source_cell_incarnation, migration.target_cell_id, + incarnation.target_cell_incarnation, migration.source_request_units, + migration.target_reserved_units, ? + FROM relay_assignment_migrations migration + JOIN relay_assignment_migration_incarnations incarnation + ON incarnation.user_id = migration.user_id + AND incarnation.relay_host_id = migration.relay_host_id + AND incarnation.assignment_epoch = migration.assignment_epoch + WHERE migration.source_cell_id = ? + AND migration.completed_at IS NULL AND migration.aborted_at IS NULL + ON CONFLICT (user_id, relay_host_id, assignment_epoch) DO NOTHING`, + [input.attemptId, now, input.cellId] + ) + row.state = 'send-may-have-started' + row.send_may_have_started_at = now + row.send_permit_expires_at = sendPermitExpiresAt + return { ...cellDrainAttempt(row), shouldSend: true } + }) + } + + async recordCellDrainApplicationReceipt(input: { + attemptId: string + cellId: string + cellIncarnation: string + traceValue: string + backendStatus: number + backendInstance?: string + }): Promise { + const now = this.now() + return await this.database.transaction(async (transaction) => { + await this.assertDrainCellGeneration( + transaction, + input.cellId, + input.cellIncarnation + ) + const row = ( + await transaction.queryLocked( + `SELECT * FROM relay_cell_drain_attempt_states + WHERE attempt_id = ? AND cell_id = ?`, + [input.attemptId, input.cellId] + ) + )[0] + if ( + !row || + text(row, 'cell_incarnation') !== input.cellIncarnation || + text(row, 'trace_value') !== input.traceValue + ) { + throw new Error('drain_attempt_not_found') + } + if (text(row, 'state') === 'application-receipt') return cellDrainAttempt(row) + if ( + text(row, 'state') !== 'send-may-have-started' || + input.backendStatus < 200 || + input.backendStatus >= 300 + ) { + throw new Error('drain_application_receipt_invalid') + } + const retryAfter = now + integer(row, 'planned_grace_ms') + 30_000 + await transaction.query( + `UPDATE relay_cell_drain_attempt_states + SET state = 'application-receipt', application_receipt_at = ?, + backend_success_status = ?, backend_instance = ?, + receipt_cell_incarnation = ?, retry_after = ? + WHERE attempt_id = ?`, + [ + now, + input.backendStatus, + input.backendInstance, + input.cellIncarnation, + retryAfter, + input.attemptId + ] + ) + row.state = 'application-receipt' + row.application_receipt_at = now + row.backend_success_status = input.backendStatus + row.backend_instance = input.backendInstance ?? null + row.receipt_cell_incarnation = input.cellIncarnation + row.retry_after = retryAfter + return cellDrainAttempt(row) + }) + } + + async proveCellDrainNotDelivered(input: { + attemptId: string + cellId: string + cellIncarnation: string + }): Promise { + const now = this.now() + return await this.database.transaction(async (transaction) => { + const row = ( + await transaction.queryLocked( + `SELECT * FROM relay_cell_drain_attempt_states + WHERE attempt_id = ? AND cell_id = ?`, + [input.attemptId, input.cellId] + ) + )[0] + if (!row || text(row, 'cell_incarnation') !== input.cellIncarnation) { + throw new Error('drain_attempt_not_found') + } + if (text(row, 'state') === 'proven-not-delivered') return cellDrainAttempt(row) + if (text(row, 'state') !== 'send-may-have-started') { + throw new Error('drain_delivery_proof_invalid') + } + await transaction.query( + `UPDATE relay_cell_drain_attempt_states + SET state = 'proven-not-delivered', proven_not_delivered_at = ? + WHERE attempt_id = ?`, + [now, input.attemptId] + ) + row.state = 'proven-not-delivered' + row.proven_not_delivered_at = now + return cellDrainAttempt(row) + }) + } + + async prepareCellDrainRecovery(input: { + attemptId?: string + cellId: string + cellIncarnation: string + }): Promise<{ + shouldSend: boolean + retryAfter: number + preparedAttempt?: CellDrainAttempt + }> { + await this.refreshDrainMigrationLeases(input) + const now = this.now() + return await this.database.transaction(async (transaction) => { + await this.assertDrainCellGeneration( + transaction, + input.cellId, + input.cellIncarnation + ) + const attempt = ( + await transaction.queryLocked( + `SELECT * FROM relay_cell_drain_attempt_states + WHERE cell_id = ? + ORDER BY prepared_at DESC LIMIT 1`, + [input.cellId] + ) + )[0] + if (!attempt || (input.attemptId && text(attempt, 'attempt_id') !== input.attemptId)) { + throw new Error('drain_application_receipt_missing') + } + if (text(attempt, 'state') === 'prepared') { + if (text(attempt, 'cell_incarnation') !== input.cellIncarnation) { + throw new Error('drain_application_receipt_missing') + } + return { + shouldSend: false, + retryAfter: now, + preparedAttempt: cellDrainAttempt(attempt) + } + } + const applicationReceiptAt = optionalInteger(attempt, 'application_receipt_at') + const backendSuccessStatus = optionalInteger(attempt, 'backend_success_status') + const retryAfter = optionalInteger(attempt, 'retry_after') + if ( + text(attempt, 'state') !== 'application-receipt' || + optionalText(attempt, 'receipt_cell_incarnation') !== + text(attempt, 'cell_incarnation') || + applicationReceiptAt === undefined || + backendSuccessStatus === undefined || + backendSuccessStatus < 200 || + backendSuccessStatus >= 300 || + retryAfter === undefined + ) { + throw new Error('drain_application_receipt_missing') + } + if (now < retryAfter) throw new Error('drain_recovery_too_early') + const attemptId = text(attempt, 'attempt_id') + if (text(attempt, 'cell_incarnation') !== input.cellIncarnation) { + const priorRecovery = ( + await transaction.queryLocked( + `SELECT attempted_at FROM relay_cell_drain_recovery_attempts + WHERE drain_attempt_id = ? AND cell_incarnation = ?`, + [attemptId, input.cellIncarnation] + ) + )[0] + if (priorRecovery) return { shouldSend: false, retryAfter } + await transaction.query( + `INSERT INTO relay_cell_drain_recovery_attempts + (drain_attempt_id, cell_incarnation, attempted_at) VALUES (?, ?, ?)`, + [attemptId, input.cellIncarnation, now] + ) + return { shouldSend: true, retryAfter } + } + if (optionalInteger(attempt, 'recover_forward_attempted_at') !== undefined) { + return { shouldSend: false, retryAfter } + } + await transaction.query( + `UPDATE relay_cell_drain_attempt_states SET recover_forward_attempted_at = ? + WHERE attempt_id = ? AND recover_forward_attempted_at IS NULL`, + [now, attemptId] + ) + return { shouldSend: true, retryAfter } + }) + } + + async evacuateDeadCells(limit = 100): Promise { + if (!this.requireLiveCells) return 0 + const cutoff = this.now() - this.heartbeatTtlMs + const rows = await this.database.query( + `SELECT assignment.user_id, assignment.relay_host_id, assignment.cell_id + FROM relay_assignments assignment + JOIN relay_cells cell ON cell.cell_id = assignment.cell_id + LEFT JOIN relay_cell_committed_fences committed + ON committed.cell_id = assignment.cell_id + LEFT JOIN relay_cell_fence_attempts attempt + ON attempt.attempt_id = committed.attempt_id + LEFT JOIN relay_cell_fences fence ON fence.cell_id = assignment.cell_id + LEFT JOIN relay_cell_runtime runtime ON runtime.cell_id = cell.cell_id + WHERE (runtime.cell_id IS NULL OR runtime.ready != ? OR runtime.last_heartbeat_at <= ?) + AND ( + ( + cell.enabled = 1 + AND NOT EXISTS ( + SELECT 1 FROM relay_cell_connection_limits limits + WHERE limits.cell_id = assignment.cell_id + ) + AND NOT EXISTS ( + SELECT 1 FROM relay_post_drain_migration_pins pin + JOIN relay_assignment_migrations migration + ON migration.user_id = pin.user_id + AND migration.relay_host_id = pin.relay_host_id + AND migration.assignment_epoch = pin.assignment_epoch + WHERE pin.user_id = assignment.user_id + AND pin.relay_host_id = assignment.relay_host_id + AND pin.assignment_epoch = assignment.assignment_epoch + AND pin.target_cell_id = assignment.cell_id + AND migration.completed_at IS NULL + AND migration.aborted_at IS NULL + ) + ) + OR ( + cell.enabled = 0 + AND ( + EXISTS ( + SELECT 1 FROM relay_cell_connection_limits limits + WHERE limits.cell_id = assignment.cell_id + ) + OR EXISTS ( + SELECT 1 FROM relay_post_drain_migration_pins pin + JOIN relay_assignment_migrations migration + ON migration.user_id = pin.user_id + AND migration.relay_host_id = pin.relay_host_id + AND migration.assignment_epoch = pin.assignment_epoch + WHERE pin.user_id = assignment.user_id + AND pin.relay_host_id = assignment.relay_host_id + AND pin.assignment_epoch = assignment.assignment_epoch + AND pin.target_cell_id = assignment.cell_id + AND migration.completed_at IS NULL + AND migration.aborted_at IS NULL + ) + ) + AND attempt.completed_at IS NOT NULL + AND attempt.aborted_at IS NULL + AND committed.cell_incarnation = runtime.cell_incarnation + AND fence.cell_incarnation = committed.cell_incarnation + AND committed.attested_at >= runtime.last_heartbeat_at + AND committed.expires_at > ? + AND fence.expires_at > ? + ) + ) + ORDER BY assignment.user_id, assignment.relay_host_id LIMIT ?`, + [1, cutoff, this.now(), this.now(), limit] + ) + let moved = 0 + for (const row of rows) { + try { + const assignment = await this.assign({ + userId: text(row, 'user_id'), + relayHostId: text(row, 'relay_host_id') + }) + if (assignment.cellId !== text(row, 'cell_id')) moved++ + } catch (error) { + if (!(error instanceof Error && error.message === 'relay_capacity_exhausted')) throw error + } + } + return moved + } + + async configureCell( + cell: RelayCellConfig, + admission: boolean | CellAdmissionState + ): Promise { + const now = this.now() + const state = typeof admission === 'boolean' ? stateFromEnabled(admission) : admission + await this.database.transaction(async (transaction) => { + const cells = await transaction.queryLocked( + `SELECT * FROM relay_cells ORDER BY cell_id ASC` + ) + const current = cells.find((row) => text(row, 'cell_id') === cell.id) + if (current && integer(current, 'reserved_requests') > cell.capacityRequests) { + throw new Error('cell_capacity_below_reserved') + } + // Deploy automation owns tagged URLs; a restarted revision must not overwrite them. + await transaction.query( + `INSERT INTO relay_cells + (cell_id, cell_url, enabled, capacity_requests, reserved_requests, + observed_requests, last_heartbeat_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT (cell_id) DO UPDATE SET + cell_url = excluded.cell_url, + enabled = excluded.enabled, + capacity_requests = excluded.capacity_requests, + updated_at = excluded.updated_at`, + [ + cell.id, + cell.url, + state === 'existing-only' ? 0 : 1, + cell.capacityRequests, + 0, + 0, + now, + now + ] + ) + await transaction.query( + `INSERT INTO relay_cell_regions (cell_id, region) VALUES (?, ?) + ON CONFLICT (cell_id) DO UPDATE SET region = excluded.region`, + [cell.id, cell.region ?? RELAY_DEFAULT_REGION] + ) + 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] + ) + } + } + }) + } + + async startActiveCellEvacuations( + sourceCellId: string, + targetCellId: string, + limit: number + ): Promise { + const rows = await this.database.query( + `SELECT user_id, relay_host_id FROM relay_assignments + WHERE cell_id = ? AND + (reserved_controls > 0 OR reserved_splices > 0 OR reserved_invites > 0 OR + pending_installs > 0 OR pending_confirmations > 0 OR migration_leases > 0) + ORDER BY user_id, relay_host_id LIMIT ?`, + [sourceCellId, limit] + ) + let started = 0 + for (const row of rows) { + const migration = await this.startEvacuation( + { userId: text(row, 'user_id'), relayHostId: text(row, 'relay_host_id') }, + targetCellId, + sourceCellId + ) + if (migration) started++ + } + return started + } + + async cellEvacuationCapacity( + sourceCellId: string, + targetCellId: string + ): Promise { + const cells = await this.database.query( + `SELECT * FROM relay_cells WHERE cell_id IN (?, ?) ORDER BY cell_id`, + [sourceCellId, targetCellId] + ) + if (!cells.some((row) => text(row, 'cell_id') === sourceCellId)) { + throw new Error('source_cell_not_found') + } + const target = cells.find((row) => text(row, 'cell_id') === targetCellId) + if (!target) throw new Error('target_cell_not_found') + if (!(await this.cellIsLive(this.database, targetCellId, this.now()))) { + throw new Error('target_cell_unavailable') + } + const summary = ( + await this.database.query( + `SELECT COUNT(*) AS assignments, + COALESCE(SUM((SELECT COALESCE(SUM(lease.request_units), 0) + FROM relay_assignment_activity_leases lease + WHERE lease.user_id = assignment.user_id + AND lease.relay_host_id = assignment.relay_host_id)), 0) AS source_units + FROM relay_assignments assignment + WHERE assignment.cell_id = ? AND + (assignment.reserved_controls > 0 OR assignment.reserved_splices > 0 OR + assignment.reserved_invites > 0 OR assignment.pending_installs > 0 OR + assignment.pending_confirmations > 0 OR assignment.migration_leases > 0)`, + [sourceCellId] + ) + )[0]! + const sourceAssignments = integer(summary, 'assignments') + return { + sourceAssignments, + // Each migration reserves its future target control plus all source-owned units. + requiredTargetUnits: integer(summary, 'source_units') + sourceAssignments, + availableTargetUnits: + integer(target, 'capacity_requests') - integer(target, 'reserved_requests') + } + } + + async cellEvacuationStatus( + sourceCellId: string, + targetCellId: string, + completeReady: boolean + ): Promise { + let completed = 0 + let blocked = 0 + const sourceFenced = await this.cellHasActiveFence(sourceCellId) + if (completeReady) { + const candidates = await this.activeCellMigrations(sourceCellId, targetCellId) + for (const [index, row] of candidates.entries()) { + try { + const identity = { + userId: text(row, 'user_id'), + relayHostId: text(row, 'relay_host_id') + } + const assignmentEpoch = integer(row, 'assignment_epoch') + if (sourceFenced) { + await this.completeEvacuationFromDeadSource(identity, { + assignmentEpoch, + sourceCellId, + targetCellId + }) + } else { + await this.completeEvacuation(identity, assignmentEpoch) + } + completed++ + } catch (error) { + if (error instanceof Error && error.message === 'migration_cell_inventory_busy') { + blocked += candidates.length - index + break + } + if (isIncompleteMigration(error)) blocked++ + else if (!(error instanceof Error && error.message === 'migration_not_found')) throw error + } + } + } + const active = await this.activeCellMigrations(sourceCellId, targetCellId) + const oldestExpiresAt = + active.length === 0 ? null : Math.min(...active.map((row) => integer(row, 'expires_at'))) + if (completeReady && active.length === 0) { + await this.reconcileReservationAccounting(sourceCellId, targetCellId) + } + const registered = active.filter( + (row) => optionalInteger(row, 'target_registered_at') !== undefined + ) + const registeredSourceActive = registered.filter( + (row) => integer(row, 'source_activity_units') > 0 + ).length + const registeredCompletable = registered.filter( + (row) => + integer(row, 'source_activity_units') === 0 && + integer(row, 'target_control_current') === 1 + ).length + const registeredTargetInactive = registered.filter( + (row) => + integer(row, 'source_activity_units') === 0 && + integer(row, 'target_control_current') === 0 + ).length + const expiredUnregistered = active.filter( + (row) => + optionalInteger(row, 'target_registered_at') === undefined && + integer(row, 'expires_at') <= this.now() + ) + const repairableExpiredUnregistered = expiredUnregistered.filter( + (row) => migrationHasExactActiveTarget(row, targetCellId) + ).length + const abortableExpiredUnregistered = expiredUnregistered.filter( + (row) => integer(row, 'target_control_active') === 0 + ).length + const blockedExpiredUnregistered = + expiredUnregistered.length - + repairableExpiredUnregistered - + abortableExpiredUnregistered + return { + inProgress: active.length, + oldestExpiresAt, + oldestRemainingMs: oldestExpiresAt === null ? null : oldestExpiresAt - this.now(), + targetRegistered: registered.length, + registeredSourceActive, + registeredCompletable, + registeredTargetInactive, + completed, + blocked, + expiredUnregistered: expiredUnregistered.length, + repairableExpiredUnregistered, + abortableExpiredUnregistered, + blockedExpiredUnregistered, + blockedExpiredOnNewerTargetAssignment: expiredUnregistered.filter( + (row) => + integer(row, 'target_control_active') === 1 && + optionalText(row, 'current_cell_id') === targetCellId && + (optionalInteger(row, 'current_assignment_epoch') ?? 0) > + integer(row, 'assignment_epoch') + ).length + } + } + + async completeReadyEvacuations(limit = 100): Promise { + if (!Number.isSafeInteger(limit) || limit < 1 || limit > 100) { + throw new Error('invalid_evacuation_completion_limit') + } + const now = this.now() + const runtimeSafety = this.requireLiveCells + ? `AND EXISTS ( + SELECT 1 FROM relay_cells source_cell + WHERE source_cell.cell_id = migration.source_cell_id + AND source_cell.enabled = 0 + ) + AND EXISTS ( + SELECT 1 FROM relay_cells target_cell + WHERE target_cell.cell_id = migration.target_cell_id + AND target_cell.enabled = 1 + ) + AND EXISTS ( + SELECT 1 FROM relay_cell_runtime source_runtime + WHERE source_runtime.cell_id = migration.source_cell_id + AND source_runtime.ready = 1 + AND source_runtime.observed_requests = 0 + AND source_runtime.last_heartbeat_at > ? + ) + AND EXISTS ( + SELECT 1 FROM relay_cell_runtime target_runtime + WHERE target_runtime.cell_id = migration.target_cell_id + AND target_runtime.ready = 1 + AND target_runtime.last_heartbeat_at > ? + )` + : '' + const targetIncarnation = this.requireLiveCells + ? `AND target_control.updated_at >= ( + SELECT target_runtime.started_at FROM relay_cell_runtime target_runtime + WHERE target_runtime.cell_id = migration.target_cell_id + )` + : '' + // Selection avoids polling offline desktops; completeEvacuation rechecks + // current target activity and zero source ownership under row locks. + const candidates = await this.database.query( + `SELECT migration.user_id, migration.relay_host_id, migration.source_cell_id, + migration.target_cell_id, migration.assignment_epoch + FROM relay_assignment_migrations migration + WHERE migration.target_registered_at IS NOT NULL + AND migration.completed_at IS NULL AND migration.aborted_at IS NULL + ${runtimeSafety} + AND NOT EXISTS ( + SELECT 1 FROM relay_assignment_activity_leases source_lease + WHERE source_lease.user_id = migration.user_id + AND source_lease.relay_host_id = migration.relay_host_id + AND source_lease.cell_id = migration.source_cell_id + ) + AND EXISTS ( + SELECT 1 FROM relay_assignment_activity_leases target_control + WHERE target_control.user_id = migration.user_id + AND target_control.relay_host_id = migration.relay_host_id + AND target_control.cell_id = migration.target_cell_id + AND target_control.activity_kind = 'control' + AND target_control.activity_id NOT LIKE 'control-pending:%' + AND target_control.expires_at > ? + ${targetIncarnation} + ) + ORDER BY migration.user_id, migration.relay_host_id + LIMIT ?`, + [ + ...(this.requireLiveCells + ? [now - this.heartbeatTtlMs, now - this.heartbeatTtlMs] + : []), + now, + limit + ] + ) + const pairs = new Map() + let completed = 0 + for (const row of candidates) { + const sourceCellId = text(row, 'source_cell_id') + const targetCellId = text(row, 'target_cell_id') + pairs.set(JSON.stringify([sourceCellId, targetCellId]), { sourceCellId, targetCellId }) + try { + await this.completeEvacuation( + { userId: text(row, 'user_id'), relayHostId: text(row, 'relay_host_id') }, + integer(row, 'assignment_epoch') + ) + completed++ + } catch (error) { + if (!isIncompleteMigration(error) && !isMissingMigration(error)) throw error + } + } + for (const { sourceCellId, targetCellId } of pairs.values()) { + if ((await this.activeCellMigrations(sourceCellId, targetCellId)).length === 0) { + await this.reconcileReservationAccounting(sourceCellId, targetCellId) + } + } + return completed + } + + async cellDeploymentStatus(cellId: string): Promise { + const cellRow = ( + await this.database.query( + `WITH activity AS ( + SELECT COUNT(*) AS activity_lease_count, + COALESCE(SUM(request_units), 0) AS activity_request_units, + COALESCE(SUM(CASE WHEN activity_kind = 'control' + AND activity_id LIKE 'control-pending:%' + AND request_units = 1 + THEN 1 ELSE 0 END), 0) AS pending_control_units + FROM relay_assignment_activity_leases WHERE cell_id = ? + ) + SELECT cell.*, runtime.cell_url AS runtime_cell_url, + admission.admission_state, + runtime.cell_incarnation AS runtime_cell_incarnation, + runtime.started_at AS runtime_started_at, runtime.ready AS runtime_ready, + runtime.observed_requests AS runtime_observed_requests, + runtime.last_heartbeat_at AS runtime_last_heartbeat_at, + COALESCE(capabilities.regional_rehome_protocol, 0) + AS runtime_regional_rehome_protocol, + activity.activity_lease_count, activity.activity_request_units, + activity.activity_lease_count - activity.pending_control_units + AS restart_blocking_activity_leases, + activity.activity_request_units - activity.pending_control_units + AS restart_blocking_activity_request_units, + cell.reserved_requests - activity.pending_control_units + AS restart_blocking_reserved_requests + FROM relay_cells cell + CROSS JOIN activity + LEFT JOIN relay_cell_admission admission ON admission.cell_id = cell.cell_id + LEFT JOIN relay_cell_runtime runtime ON runtime.cell_id = cell.cell_id + LEFT JOIN relay_cell_capabilities capabilities + ON capabilities.cell_id = runtime.cell_id + AND capabilities.cell_incarnation = runtime.cell_incarnation + WHERE cell.cell_id = ?`, + [cellId, cellId] + ) + )[0] + if (!cellRow) throw new Error('cell_not_found') + const assignmentRow = ( + await this.database.query( + `SELECT COUNT(*) AS assignment_count FROM relay_assignments WHERE cell_id = ?`, + [cellId] + ) + )[0]! + const migrationRow = ( + await this.database.query( + `SELECT + COALESCE(SUM(CASE WHEN source_cell_id = ? THEN 1 ELSE 0 END), 0) + AS outgoing_migrations, + COALESCE(SUM(CASE WHEN target_cell_id = ? THEN 1 ELSE 0 END), 0) + AS incoming_migrations + FROM relay_assignment_migrations + WHERE completed_at IS NULL AND aborted_at IS NULL + AND (source_cell_id = ? OR target_cell_id = ?)`, + [cellId, cellId, cellId, cellId] + ) + )[0]! + const connectionRow = ( + await this.database.query( + `SELECT limits.hard_cap, limits.unobserved_bound, + connection_runtime.total_connections, connection_runtime.in_flight_connections, + connection_runtime.reserved_connection_units, + connection_runtime.enforced_connection_units, + connection_runtime.last_heartbeat_at, + current_runtime.cell_incarnation AS current_incarnation, + connection_runtime.cell_incarnation AS connection_incarnation, + (SELECT COUNT(*) FROM relay_control_connection_reservations reservation + WHERE reservation.cell_id = limits.cell_id + AND reservation.state IN + ('reserved', 'late-arrival-debt', 'claimed')) AS outstanding_reservations + FROM relay_cell_connection_limits limits + LEFT JOIN relay_cell_connection_runtime connection_runtime + ON connection_runtime.cell_id = limits.cell_id + LEFT JOIN relay_cell_runtime current_runtime + ON current_runtime.cell_id = limits.cell_id + WHERE limits.cell_id = ?`, + [cellId] + ) + )[0] + const runtimeHeartbeat = optionalInteger(cellRow, 'runtime_last_heartbeat_at') + const connectionHeartbeat = connectionRow + ? optionalInteger(connectionRow, 'last_heartbeat_at') + : undefined + return { + cellId, + cellUrl: text(cellRow, 'cell_url'), + region: await this.cellRegion(this.database, cellId), + enabled: integer(cellRow, 'enabled') === 1, + admissionState: parseCellAdmissionState(text(cellRow, 'admission_state')), + capacityRequests: integer(cellRow, 'capacity_requests'), + reservedRequests: integer(cellRow, 'reserved_requests'), + assignments: integer(assignmentRow, 'assignment_count'), + activityLeases: integer(cellRow, 'activity_lease_count'), + activityRequestUnits: integer(cellRow, 'activity_request_units'), + restartBlockingActivityLeases: integer( + cellRow, + 'restart_blocking_activity_leases' + ), + restartBlockingActivityRequestUnits: integer( + cellRow, + 'restart_blocking_activity_request_units' + ), + restartBlockingReservedRequests: integer( + cellRow, + 'restart_blocking_reserved_requests' + ), + outgoingMigrations: integer(migrationRow, 'outgoing_migrations'), + incomingMigrations: integer(migrationRow, 'incoming_migrations'), + 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'), + observedConnections: + optionalInteger(connectionRow, 'total_connections') ?? 0, + inFlightConnections: + optionalInteger(connectionRow, 'in_flight_connections') ?? 0, + reservedConnectionUnits: + optionalInteger(connectionRow, 'reserved_connection_units') ?? 0, + enforcedConnectionUnits: + optionalInteger(connectionRow, 'enforced_connection_units') ?? 0, + pendingControlReservations: integer( + connectionRow, + 'outstanding_reservations' + ), + heartbeatFresh: + connectionHeartbeat !== undefined && + connectionHeartbeat > this.now() - this.heartbeatTtlMs && + optionalText(connectionRow, 'current_incarnation') === + optionalText(connectionRow, 'connection_incarnation') + } + : null, + runtime: + runtimeHeartbeat === undefined + ? null + : { + cellUrl: text(cellRow, 'runtime_cell_url'), + cellIncarnation: text(cellRow, 'runtime_cell_incarnation'), + startedAt: integer(cellRow, 'runtime_started_at'), + ready: integer(cellRow, 'runtime_ready') === 1, + observedRequests: integer(cellRow, 'runtime_observed_requests'), + lastHeartbeatAt: runtimeHeartbeat, + heartbeatFresh: runtimeHeartbeat > this.now() - this.heartbeatTtlMs, + regionalRehomeProtocol: integer( + cellRow, + 'runtime_regional_rehome_protocol' + ) + } + } + } + + async verifyCellAssignment(input: AssignmentIdentity & { + cellId: string + assignmentEpoch: number + }): Promise { + const rows = await this.database.query( + `SELECT cell_id, assignment_epoch FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [input.userId, input.relayHostId] + ) + return Boolean( + rows[0] && + text(rows[0], 'cell_id') === input.cellId && + integer(rows[0], 'assignment_epoch') === input.assignmentEpoch + ) + } + + async changeActivity( + identity: AssignmentIdentity, + kind: AssignmentActivityKind, + delta: 1 | -1 + ): Promise { + await this.activityQueue.run(identity, async () => { + const column = ACTIVITY_COLUMN[kind] + const now = this.now() + await this.database.transaction(async (transaction) => { + const row = ( + await transaction.queryLocked( + `SELECT * FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + )[0] + if (!row) return + const before = integer(row, column) + const after = Math.max(0, before + delta) + await transaction.query( + `UPDATE relay_assignments SET ${column} = ?, lease_expires_at = ?, last_activity_at = ? + WHERE user_id = ? AND relay_host_id = ?`, + [after, now + ASSIGNMENT_LIMITS.activityLeaseMs, now, identity.userId, identity.relayHostId] + ) + const requestDelta = ACTIVITY_REQUEST_UNITS[kind] * (after - before) + if (requestDelta !== 0) { + await this.lockCellInventory(transaction) + await this.adjustCellReservation(transaction, text(row, 'cell_id'), requestDelta) + } + }) + }) + } + + async acquireActivity( + identity: AssignmentIdentity, + input: { + activityId: string + kind: AssignmentActivityKind + cellId: string + expiresAt?: number + } + ): Promise { + validateActivityId(input.activityId) + await this.activityQueue.run(identity, async () => { + const now = this.now() + await this.database.transaction(async (transaction) => { + const assignment = await this.assignmentRow(transaction, identity) + if (!assignment) throw new Error('assignment_not_found') + const assignmentCellId = text(assignment, 'cell_id') + if (input.cellId !== assignmentCellId) { + // Origin controls may renew work only while the exact forward + // migration is active; completion must fence late source activity. + const migration = ( + await transaction.queryLocked( + `SELECT assignment_epoch FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ? + AND source_cell_id = ? AND target_cell_id = ? + AND assignment_epoch = ? + AND completed_at IS NULL AND aborted_at IS NULL`, + [ + identity.userId, + identity.relayHostId, + input.cellId, + assignmentCellId, + integer(assignment, 'assignment_epoch') + ] + ) + )[0] + if (!migration) throw new Error('activity_cell_not_authoritative') + } + const activityLeases = await this.lockAssignmentActivities(transaction, identity) + const existing = activityLeaseById(activityLeases, input.activityId) + const expiresAt = input.expiresAt ?? now + ASSIGNMENT_LIMITS.activityLeaseMs + if (!Number.isSafeInteger(expiresAt) || expiresAt <= now) { + throw new Error('invalid_activity_expiry') + } + if (existing && text(existing, 'activity_kind') === input.kind && text(existing, 'cell_id') === input.cellId) { + await transaction.query( + `UPDATE relay_assignment_activity_leases SET expires_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND activity_id = ?`, + [expiresAt, now, identity.userId, identity.relayHostId, input.activityId] + ) + await this.touchAssignment(transaction, identity, expiresAt, now) + return + } + const units = ACTIVITY_REQUEST_UNITS[input.kind] + if (existing) { + await this.lockCellInventory(transaction) + await this.removeActivityLease(transaction, identity, existing, now) + await this.adjustCellReservation(transaction, input.cellId, units) + } + await this.adjustActivityCount(transaction, identity, input.kind, 1, expiresAt, now) + await transaction.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + input.activityId, + input.kind, + input.cellId, + units, + expiresAt, + now + ] + ) + // Keep the contended cell row locked for only the final write and commit. + if (!existing) { + await this.adjustCellReservationAtomically(transaction, input.cellId, units) + } + }) + }) + } + + async renewControlActivity( + identity: AssignmentIdentity, + input: { activityId: string; cellId: string; expiresAt: number } + ): Promise { + validateActivityId(input.activityId) + const now = this.now() + const maximumExpiresAt = + now + + ASSIGNMENT_LIMITS.activityLeaseMs + + RELAY_PROTOCOL_LIMITS.controlPingIntervalMs * 2 + if ( + !Number.isSafeInteger(input.expiresAt) || + input.expiresAt <= now || + input.expiresAt > maximumExpiresAt + ) { + throw new Error('invalid_activity_expiry') + } + const startedAt = performance.now() + let outcome: ControlRenewalOutcome = 'database_error' + try { + outcome = + this.database.dialect === 'postgres' + ? await this.renewPostgresControlActivity(identity, input, now) + : await this.renewTransactionalControlActivity(identity, input, now) + if (outcome !== 'renewed') throw new Error(outcome) + } catch (error) { + const message = String((error as { message?: unknown }).message) + if (CONTROL_RENEWAL_OUTCOMES.has(message as ControlRenewalOutcome)) { + outcome = message as ControlRenewalOutcome + } + throw error + } finally { + this.recordControlRenewal?.(performance.now() - startedAt, outcome) + } + } + + private async renewPostgresControlActivity( + identity: AssignmentIdentity, + input: { activityId: string; cellId: string; expiresAt: number }, + now: number + ): Promise { + const row = ( + await this.database.query( + `WITH assignment_state AS MATERIALIZED ( + SELECT cell_id, assignment_epoch + FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ? + FOR UPDATE + ), migration_state AS MATERIALIZED ( + SELECT migration.assignment_epoch + FROM relay_assignment_migrations migration + JOIN assignment_state assignment + ON migration.target_cell_id = assignment.cell_id + AND migration.assignment_epoch = assignment.assignment_epoch + WHERE migration.user_id = ? AND migration.relay_host_id = ? + AND migration.source_cell_id = ? + AND migration.completed_at IS NULL AND migration.aborted_at IS NULL + FOR UPDATE OF migration + ), authorization_state AS MATERIALIZED ( + SELECT 1 AS authorized + FROM assignment_state assignment + WHERE assignment.cell_id = ? OR EXISTS (SELECT 1 FROM migration_state) + ), lease_state AS MATERIALIZED ( + SELECT lease.activity_kind, lease.cell_id + FROM relay_assignment_activity_leases lease + CROSS JOIN authorization_state + WHERE lease.user_id = ? AND lease.relay_host_id = ? AND lease.activity_id = ? + FOR UPDATE OF lease + ), renewed_lease AS ( + UPDATE relay_assignment_activity_leases lease + SET expires_at = GREATEST(lease.expires_at, ?), + updated_at = GREATEST(lease.updated_at, ?) + FROM lease_state state + WHERE lease.user_id = ? AND lease.relay_host_id = ? AND lease.activity_id = ? + AND state.activity_kind = 'control' AND state.cell_id = ? + RETURNING 1 + ), renewed_assignment AS ( + UPDATE relay_assignments assignment + SET lease_expires_at = GREATEST(assignment.lease_expires_at, ?), + last_activity_at = GREATEST(assignment.last_activity_at, ?) + WHERE assignment.user_id = ? AND assignment.relay_host_id = ? + AND EXISTS (SELECT 1 FROM renewed_lease) + RETURNING 1 + ) + SELECT CASE + WHEN NOT EXISTS (SELECT 1 FROM assignment_state) + THEN 'assignment_not_found' + WHEN NOT EXISTS (SELECT 1 FROM authorization_state) + THEN 'activity_cell_not_authoritative' + WHEN NOT EXISTS (SELECT 1 FROM lease_state) + THEN 'control_activity_not_found' + WHEN EXISTS ( + SELECT 1 FROM lease_state + WHERE activity_kind <> 'control' OR cell_id <> ? + ) THEN 'control_activity_moved' + WHEN EXISTS (SELECT 1 FROM renewed_assignment) THEN 'renewed' + ELSE 'control_activity_not_found' + END AS outcome`, + [ + identity.userId, + identity.relayHostId, + identity.userId, + identity.relayHostId, + input.cellId, + input.cellId, + identity.userId, + identity.relayHostId, + input.activityId, + input.expiresAt, + now, + identity.userId, + identity.relayHostId, + input.activityId, + input.cellId, + input.expiresAt, + now, + identity.userId, + identity.relayHostId, + input.cellId + ] + ) + )[0] + if (!row) throw new Error('missing_control_renewal_outcome') + const outcome = text(row, 'outcome') as ControlRenewalOutcome + if (!CONTROL_RENEWAL_OUTCOMES.has(outcome)) throw new Error('invalid_control_renewal_outcome') + return outcome + } + + private async renewTransactionalControlActivity( + identity: AssignmentIdentity, + input: { activityId: string; cellId: string; expiresAt: number }, + now: number + ): Promise { + // Bypass the process queue so a network-stalled activity call cannot suppress renewal. + await this.database.transaction(async (transaction) => { + const assignment = await this.assignmentRow(transaction, identity) + if (!assignment) throw new Error('assignment_not_found') + const assignmentCellId = text(assignment, 'cell_id') + if (input.cellId !== assignmentCellId) { + const migration = ( + await transaction.queryLocked( + `SELECT assignment_epoch FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ? + AND source_cell_id = ? AND target_cell_id = ? + AND assignment_epoch = ? + AND completed_at IS NULL AND aborted_at IS NULL`, + [ + identity.userId, + identity.relayHostId, + input.cellId, + assignmentCellId, + integer(assignment, 'assignment_epoch') + ] + ) + )[0] + if (!migration) throw new Error('activity_cell_not_authoritative') + } + const lease = ( + await transaction.queryLocked( + `SELECT * FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? AND activity_id = ?`, + [identity.userId, identity.relayHostId, input.activityId] + ) + )[0] + if (!lease) throw new Error('control_activity_not_found') + if ( + text(lease, 'activity_kind') !== 'control' || + text(lease, 'cell_id') !== input.cellId + ) { + throw new Error('control_activity_moved') + } + await transaction.query( + `UPDATE relay_assignment_activity_leases SET + expires_at = CASE WHEN expires_at > ? THEN expires_at ELSE ? END, + updated_at = CASE WHEN updated_at > ? THEN updated_at ELSE ? END + WHERE user_id = ? AND relay_host_id = ? AND activity_id = ?`, + [ + input.expiresAt, + input.expiresAt, + now, + now, + identity.userId, + identity.relayHostId, + input.activityId + ] + ) + await transaction.query( + `UPDATE relay_assignments SET + lease_expires_at = CASE WHEN lease_expires_at > ? THEN lease_expires_at ELSE ? END, + last_activity_at = CASE WHEN last_activity_at > ? THEN last_activity_at ELSE ? END + WHERE user_id = ? AND relay_host_id = ?`, + [input.expiresAt, input.expiresAt, now, now, identity.userId, identity.relayHostId] + ) + }) + return 'renewed' + } + + async releaseActivity(identity: AssignmentIdentity, activityId: string): Promise { + validateActivityId(activityId) + return await this.activityQueue.run(identity, async () => { + const now = this.now() + return await this.database.transaction(async (transaction) => { + await this.assignmentRow(transaction, identity) + const activityLeases = await this.lockAssignmentActivities(transaction, identity) + const existing = activityLeaseById(activityLeases, activityId) + if (!existing) return false + await transaction.query( + `DELETE FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? AND activity_id = ?`, + [identity.userId, identity.relayHostId, activityId] + ) + await this.adjustActivityCount( + transaction, + identity, + activityKind(existing), + -1, + now, + now + ) + // Release paths can safely defer the cell-row lock until their final write. + await this.adjustCellReservationAtomically( + transaction, + text(existing, 'cell_id'), + -integer(existing, 'request_units') + ) + return true + }) + }) + } + + async activateControl( + identity: AssignmentIdentity, + input: { + cellId: string + assignmentEpoch: number + generation: number + connectionInclusionWatermark?: number + } + ): Promise { + const activityId = `control:${input.cellId}:${input.generation}` + validateActivityId(activityId) + return await this.activityQueue.run(identity, async () => { + const now = this.now() + return await this.database.transaction(async (transaction) => { + const assignment = await this.assignmentRow(transaction, identity) + if ( + !assignment || + text(assignment, 'cell_id') !== input.cellId || + integer(assignment, 'assignment_epoch') !== input.assignmentEpoch + ) { + throw new Error('wrong_assignment') + } + const expiresAt = now + ASSIGNMENT_LIMITS.activityLeaseMs + const activityLeases = await this.lockAssignmentActivities(transaction, identity) + const existing = activityLeaseById(activityLeases, activityId) + const pendingId = pendingControlActivityId(input.assignmentEpoch) + const pending = activityLeaseById(activityLeases, pendingId) + const retainedActivityId = + existing || (pending && text(pending, 'cell_id') === input.cellId) + ? existing + ? activityId + : pendingId + : activityId + await this.removeSupersededSameCellControls( + transaction, + identity, + activityLeases, + input.cellId, + retainedActivityId, + now + ) + if (existing) { + await transaction.query( + `UPDATE relay_assignment_activity_leases SET expires_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND activity_id = ?`, + [expiresAt, now, identity.userId, identity.relayHostId, activityId] + ) + await this.touchAssignment(transaction, identity, expiresAt, now) + } else { + if (pending && text(pending, 'cell_id') === input.cellId) { + await transaction.query( + `UPDATE relay_assignment_activity_leases + SET activity_id = ?, expires_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND activity_id = ?`, + [activityId, expiresAt, now, identity.userId, identity.relayHostId, pendingId] + ) + await this.touchAssignment(transaction, identity, expiresAt, now) + } else { + await this.lockCellInventory(transaction) + await this.adjustCellReservation(transaction, input.cellId, 1) + await this.adjustActivityCount(transaction, identity, 'control', 1, expiresAt, now) + await transaction.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + activityId, + 'control', + input.cellId, + 1, + expiresAt, + now + ] + ) + } + } + await this.claimControlConnectionReservation( + transaction, + identity, + input.cellId, + input.assignmentEpoch, + activityId, + input.connectionInclusionWatermark, + now + ) + return activityId + }) + }) + } + + async startEvacuation( + identity: AssignmentIdentity, + targetCellId: string + ): Promise + async startEvacuation( + identity: AssignmentIdentity, + targetCellId: string, + expectedSourceCellId: string + ): Promise + async startEvacuation( + identity: AssignmentIdentity, + targetCellId: string, + expectedSourceCellId?: string + ): Promise { + const now = this.now() + return await this.database.transaction(async (transaction) => { + const assignment = await this.assignmentRow(transaction, identity) + if (!assignment) throw new Error('assignment_not_found') + const sourceCellId = text(assignment, 'cell_id') + // Bulk selection is intentionally staleable; fence it before considering + // an existing migration that already moved the assignment to its target. + if (expectedSourceCellId && sourceCellId !== expectedSourceCellId) return null + const existing = ( + await transaction.queryLocked( + `SELECT * FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ? AND completed_at IS NULL + AND aborted_at IS NULL AND expires_at > ? + ORDER BY assignment_epoch DESC LIMIT 1`, + [identity.userId, identity.relayHostId, now] + ) + )[0] + if (existing) { + if (text(existing, 'target_cell_id') !== targetCellId) { + throw new Error('migration_in_progress') + } + return migration(identity, existing) + } + if (sourceCellId === targetCellId) throw new Error('target_matches_source') + await this.lockAssignmentActivities(transaction, identity) + const cells = await this.lockCellInventory(transaction) + const target = cells.find((row) => text(row, 'cell_id') === targetCellId) + if (!target || integer(target, 'enabled') !== 1) throw new Error('target_cell_unavailable') + if (!(await this.cellIsLive(transaction, targetCellId, now))) { + throw new Error('target_cell_unavailable') + } + let sourceCellIncarnation: string | undefined + let targetCellIncarnation: string | undefined + if (this.requireLiveCells) { + const runtimes = await transaction.queryLocked( + `SELECT * FROM relay_cell_runtime WHERE cell_id IN (?, ?) ORDER BY cell_id`, + [sourceCellId, targetCellId] + ) + const sourceRuntime = runtimes.find( + (runtime) => text(runtime, 'cell_id') === sourceCellId + ) + const targetRuntime = runtimes.find( + (runtime) => text(runtime, 'cell_id') === targetCellId + ) + if ( + !sourceRuntime || + integer(sourceRuntime, 'ready') !== 1 || + integer(sourceRuntime, 'last_heartbeat_at') <= now - this.heartbeatTtlMs + ) { + throw new Error('source_cell_unavailable') + } + if ( + !targetRuntime || + integer(targetRuntime, 'ready') !== 1 || + integer(targetRuntime, 'last_heartbeat_at') <= now - this.heartbeatTtlMs + ) { + throw new Error('target_cell_unavailable') + } + sourceCellIncarnation = text(sourceRuntime, 'cell_incarnation') + targetCellIncarnation = text(targetRuntime, 'cell_incarnation') + } + await this.assertCellConnectionHeadroom(transaction, targetCellId) + const sourceUnitsRow = ( + await transaction.query( + `SELECT COALESCE(SUM(request_units), 0) AS units + FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? AND cell_id = ?`, + [identity.userId, identity.relayHostId, sourceCellId] + ) + )[0] + const sourceRequestUnits = integer(sourceUnitsRow!, 'units') + const targetReservedUnits = sourceRequestUnits + 1 + await this.adjustCellReservation(transaction, targetCellId, targetReservedUnits) + const previousEpoch = integer(assignment, 'assignment_epoch') + const assignmentEpoch = previousEpoch + 1 + const expiresAt = now + ASSIGNMENT_LIMITS.migrationLeaseMs + await transaction.query( + `UPDATE relay_assignments SET cell_id = ?, assignment_epoch = ?, + reserved_controls = reserved_controls + 1, + migration_leases = migration_leases + 1, + lease_expires_at = ?, last_activity_at = ? + WHERE user_id = ? AND relay_host_id = ?`, + [ + targetCellId, + assignmentEpoch, + expiresAt, + now, + identity.userId, + identity.relayHostId + ] + ) + await transaction.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?), (?, ?, ?, ?, ?, ?, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + pendingControlActivityId(assignmentEpoch), + 'control', + targetCellId, + 1, + expiresAt, + now, + identity.userId, + identity.relayHostId, + migrationActivityId(assignmentEpoch), + 'migration', + targetCellId, + sourceRequestUnits, + expiresAt, + now + ] + ) + await this.insertControlConnectionReservation( + transaction, + identity, + targetCellId, + assignmentEpoch, + expiresAt, + now + ) + await transaction.query( + `INSERT INTO relay_assignment_migrations + (user_id, relay_host_id, source_cell_id, target_cell_id, + previous_epoch, assignment_epoch, source_request_units, + target_reserved_units, expires_at, target_registered_at, + completed_at, aborted_at, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, NULL, NULL, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + sourceCellId, + targetCellId, + previousEpoch, + assignmentEpoch, + sourceRequestUnits, + targetReservedUnits, + expiresAt, + now, + now + ] + ) + if (sourceCellIncarnation && targetCellIncarnation) { + await transaction.query( + `INSERT INTO relay_assignment_migration_incarnations + (user_id, relay_host_id, assignment_epoch, source_cell_incarnation, + target_cell_incarnation) + VALUES (?, ?, ?, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + assignmentEpoch, + sourceCellIncarnation, + targetCellIncarnation + ] + ) + } + return { + ...identity, + sourceCellId, + targetCellId, + previousEpoch, + assignmentEpoch, + expiresAt + } + }) + } + + async markMigrationTargetRegistered( + identity: AssignmentIdentity, + input: { cellId: string; assignmentEpoch: number } + ): Promise { + const now = this.now() + return await this.database.transaction(async (transaction) => { + const rows = await transaction.queryLocked( + `SELECT * FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ? + AND target_cell_id = ? AND completed_at IS NULL AND aborted_at IS NULL`, + [ + identity.userId, + identity.relayHostId, + input.assignmentEpoch, + input.cellId + ] + ) + if (!rows[0]) return false + await transaction.query( + `UPDATE relay_assignment_migrations + SET target_registered_at = COALESCE(target_registered_at, ?), updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [now, now, identity.userId, identity.relayHostId, input.assignmentEpoch] + ) + return true + }) + } + + async completeEvacuationFromDeadSource( + identity: AssignmentIdentity, + input: { + assignmentEpoch: number + sourceCellId: string + targetCellId: string + } + ): Promise { + try { + return await this.withAssignmentLockRetry( + async (inventoryFirst) => + await this.completeEvacuationFromDeadSourceOnce(identity, input, inventoryFirst) + ) + } catch (error) { + if (isDatabaseLockUnavailable(error)) { + throw new Error('migration_cell_inventory_busy') + } + throw error + } + } + + private async completeEvacuationFromDeadSourceOnce( + identity: AssignmentIdentity, + input: { + assignmentEpoch: number + sourceCellId: string + targetCellId: string + }, + inventoryFirst: boolean + ): Promise { + const now = this.now() + return await this.database.transaction(async (transaction) => { + let lockedCells: SqlRow[] | undefined + if (inventoryFirst) { + try { + lockedCells = await this.lockCellInventory(transaction) + } catch (error) { + if (isDatabaseLockTimeout(error)) { + throw new Error('database_lock_unavailable') + } + throw error + } + } + const assignment = await this.assignmentRow(transaction, identity, inventoryFirst) + const row = ( + await transaction.queryLocked( + `SELECT * FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, input.assignmentEpoch] + ) + )[0] + if (!row) throw new Error('migration_not_found') + assertMigrationPair(row, input.sourceCellId, input.targetCellId) + if (optionalInteger(row, 'completed_at') !== undefined) { + return deadSourceCompletionResult(row, false) + } + if (optionalInteger(row, 'aborted_at') !== undefined) { + throw new Error('migration_already_aborted') + } + if (optionalInteger(row, 'target_registered_at') === undefined) { + throw new Error('migration_target_not_registered') + } + assertCurrentMigrationAssignment(assignment, row) + const activityLeases = await this.lockAssignmentActivities(transaction, identity) + assertAssignmentActivityAccounting(assignment, activityLeases, row) + if ( + activityLeases.some( + (lease) => + ![input.sourceCellId, input.targetCellId].includes(text(lease, 'cell_id')) + ) + ) { + throw new Error('migration_activity_topology_mismatch') + } + if (activityUnitsForCell(activityLeases, input.sourceCellId) > 0) { + throw new Error('migration_source_still_active') + } + const cells = lockedCells ?? (await this.lockCellInventory(transaction, true)) + const source = cells.find((cell) => text(cell, 'cell_id') === input.sourceCellId) + const target = cells.find((cell) => text(cell, 'cell_id') === input.targetCellId) + if (!source || integer(source, 'enabled') !== 0) { + throw new Error('migration_source_admission_changed') + } + if (!target || integer(target, 'enabled') !== 1) { + throw new Error('migration_target_admission_changed') + } + await assertCellReservationAccounting(transaction, cells, [ + input.sourceCellId, + input.targetCellId + ]) + const runtimes = await transaction.queryLocked( + `SELECT * FROM relay_cell_runtime WHERE cell_id IN (?, ?) ORDER BY cell_id`, + [input.sourceCellId, input.targetCellId] + ) + const sourceRuntime = runtimes.find( + (runtime) => text(runtime, 'cell_id') === input.sourceCellId + ) + const targetRuntime = runtimes.find( + (runtime) => text(runtime, 'cell_id') === input.targetCellId + ) + const freshAfter = now - this.heartbeatTtlMs + await this.requireCellFence(transaction, input.sourceCellId, sourceRuntime, now) + if (!sourceRuntime || integer(sourceRuntime, 'last_heartbeat_at') > freshAfter) { + throw new Error('migration_source_runtime_not_dead') + } + if ( + !targetRuntime || + integer(targetRuntime, 'ready') !== 1 || + integer(targetRuntime, 'last_heartbeat_at') <= freshAfter + ) { + throw new Error('migration_target_runtime_not_ready') + } + const targetIsActive = activityLeases.some( + (lease) => + text(lease, 'cell_id') === input.targetCellId && + text(lease, 'activity_kind') === 'control' && + !text(lease, 'activity_id').startsWith('control-pending:') && + integer(lease, 'expires_at') > now && + integer(lease, 'updated_at') >= integer(targetRuntime, 'started_at') + ) + if (!targetIsActive) { + const inactiveTargetControls = activityLeases.filter( + (lease) => + text(lease, 'cell_id') === input.targetCellId && + text(lease, 'activity_kind') === 'control' + ) + for (const lease of inactiveTargetControls) { + await this.removeActivityLease(transaction, identity, lease, now) + } + await this.releaseSupersededControlConnectionReservations( + transaction, + identity, + input.targetCellId, + input.assignmentEpoch, + now + ) + } + const migrationLease = activityLeaseById( + activityLeases, + migrationActivityId(input.assignmentEpoch) + ) + if (migrationLease) { + await this.removeActivityLease(transaction, identity, migrationLease, now) + } + await transaction.query( + `UPDATE relay_assignment_migrations SET completed_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [now, now, identity.userId, identity.relayHostId, input.assignmentEpoch] + ) + return deadSourceCompletionResult(row, true) + }) + } + + async supersedeRegisteredEvacuation( + identity: AssignmentIdentity, + input: RegisteredEvacuationSupersessionInput, + preservedActivityCellIds: readonly string[] = [] + ): Promise { + if (input.assignmentEpoch >= Number.MAX_SAFE_INTEGER) { + throw new Error('assignment_epoch_exhausted') + } + return await this.withAssignmentLockRetry( + async (inventoryFirst) => + await this.supersedeRegisteredEvacuationOnce( + identity, + input, + inventoryFirst, + preservedActivityCellIds + ) + ) + } + + private async supersedeRegisteredEvacuationOnce( + identity: AssignmentIdentity, + input: RegisteredEvacuationSupersessionInput, + inventoryFirst: boolean, + preservedActivityCellIds: readonly string[] + ): Promise { + const now = this.now() + return await this.database.transaction(async (transaction) => { + const lockedCells = inventoryFirst + ? await this.lockCellInventory(transaction) + : undefined + const assignment = await this.assignmentRow(transaction, identity, inventoryFirst) + const existing = ( + await transaction.queryLocked( + `SELECT * FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, input.assignmentEpoch] + ) + )[0] + if (!existing) throw new Error('migration_not_found') + assertMigrationPair(existing, input.sourceCellId, input.currentTargetCellId) + if (optionalInteger(existing, 'completed_at') !== undefined) { + throw new Error('migration_already_completed') + } + if (optionalInteger(existing, 'aborted_at') !== undefined) { + const successor = ( + await transaction.queryLocked( + `SELECT * FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ? AND previous_epoch = ? + AND source_cell_id = ? AND target_cell_id = ? + ORDER BY assignment_epoch DESC LIMIT 1`, + [ + identity.userId, + identity.relayHostId, + input.assignmentEpoch, + input.sourceCellId, + input.replacementTargetCellId + ] + ) + )[0] + if (!successor) throw new Error('migration_already_superseded') + return migration(identity, successor) + } + if (optionalInteger(existing, 'target_registered_at') === undefined) { + throw new Error('migration_target_not_registered') + } + const existingIncarnation = ( + await transaction.queryLocked( + `SELECT * FROM relay_assignment_migration_incarnations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, input.assignmentEpoch] + ) + )[0] + const existingPin = ( + await transaction.queryLocked( + `SELECT * FROM relay_post_drain_migration_pins + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, input.assignmentEpoch] + ) + )[0] + if (existingPin && !existingIncarnation) { + throw new Error('drain_migration_source_incarnation_mismatch') + } + assertCurrentMigrationAssignment(assignment, existing) + if (input.currentTargetCellId === input.replacementTargetCellId) { + throw new Error('target_matches_source') + } + const activityLeases = await this.lockAssignmentActivities(transaction, identity) + assertAssignmentActivityAccounting(assignment, activityLeases, existing) + const additionalActivityCellIds = new Set( + activityLeases + .map((lease) => text(lease, 'cell_id')) + .filter( + (cellId) => + ![input.sourceCellId, input.currentTargetCellId].includes(cellId) + ) + ) + if ( + preservedActivityCellIds.includes(input.replacementTargetCellId) || + !matchesExactCellSet(additionalActivityCellIds, preservedActivityCellIds) + ) { + throw new Error('migration_activity_topology_mismatch') + } + const cells = lockedCells ?? (await this.lockCellInventory(transaction, true)) + const source = cells.find((cell) => text(cell, 'cell_id') === input.sourceCellId) + const currentTarget = cells.find( + (cell) => text(cell, 'cell_id') === input.currentTargetCellId + ) + const replacement = cells.find( + (cell) => text(cell, 'cell_id') === input.replacementTargetCellId + ) + if (!source || integer(source, 'enabled') !== 0) { + throw new Error('migration_source_admission_changed') + } + if (!currentTarget || integer(currentTarget, 'enabled') !== 0) { + throw new Error('migration_target_still_enabled') + } + if (!replacement || integer(replacement, 'enabled') !== 1) { + throw new Error('replacement_target_unavailable') + } + await assertCellReservationAccounting(transaction, cells, [ + input.sourceCellId, + input.currentTargetCellId, + input.replacementTargetCellId, + ...preservedActivityCellIds + ]) + const runtimeCellIds = [ + input.currentTargetCellId, + input.replacementTargetCellId, + ...preservedActivityCellIds + ] + const runtimes = await transaction.queryLocked( + `SELECT * FROM relay_cell_runtime + WHERE cell_id IN (${runtimeCellIds.map(() => '?').join(', ')}) + ORDER BY cell_id`, + runtimeCellIds + ) + const currentRuntime = runtimes.find( + (runtime) => text(runtime, 'cell_id') === input.currentTargetCellId + ) + const replacementRuntime = runtimes.find( + (runtime) => text(runtime, 'cell_id') === input.replacementTargetCellId + ) + const freshAfter = now - this.heartbeatTtlMs + await this.requireCellFence( + transaction, + input.currentTargetCellId, + currentRuntime, + now + ) + if ( + currentRuntime && + integer(currentRuntime, 'ready') === 1 && + integer(currentRuntime, 'last_heartbeat_at') > freshAfter + ) { + throw new Error('migration_target_still_available') + } + if ( + !replacementRuntime || + integer(replacementRuntime, 'ready') !== 1 || + integer(replacementRuntime, 'last_heartbeat_at') <= freshAfter + ) { + throw new Error('replacement_target_unavailable') + } + if ( + preservedActivityCellIds.some((cellId) => { + const runtime = runtimes.find((row) => text(row, 'cell_id') === cellId) + return ( + !runtime || + integer(runtime, 'ready') !== 1 || + integer(runtime, 'last_heartbeat_at') <= freshAfter + ) + }) + ) { + throw new Error('migration_preserved_activity_cell_unavailable') + } + await this.assertCellConnectionHeadroom( + transaction, + input.replacementTargetCellId + ) + for (const lease of activityLeases.filter( + (candidate) => text(candidate, 'cell_id') === input.currentTargetCellId + )) { + await this.removeActivityLease(transaction, identity, lease, now) + } + const sourceRequestUnits = activityUnitsForCell(activityLeases, input.sourceCellId) + const targetReservedUnits = sourceRequestUnits + 1 + await this.adjustCellReservation( + transaction, + input.replacementTargetCellId, + targetReservedUnits + ) + const assignmentEpoch = input.assignmentEpoch + 1 + const expiresAt = now + ASSIGNMENT_LIMITS.migrationLeaseMs + await transaction.query( + `UPDATE relay_assignments SET cell_id = ?, assignment_epoch = ?, + reserved_controls = reserved_controls + 1, + migration_leases = migration_leases + 1, + lease_expires_at = ?, last_activity_at = ? + WHERE user_id = ? AND relay_host_id = ?`, + [ + input.replacementTargetCellId, + assignmentEpoch, + expiresAt, + now, + identity.userId, + identity.relayHostId + ] + ) + await transaction.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?), (?, ?, ?, ?, ?, ?, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + pendingControlActivityId(assignmentEpoch), + 'control', + input.replacementTargetCellId, + 1, + expiresAt, + now, + identity.userId, + identity.relayHostId, + migrationActivityId(assignmentEpoch), + 'migration', + input.replacementTargetCellId, + sourceRequestUnits, + expiresAt, + now + ] + ) + await this.insertControlConnectionReservation( + transaction, + identity, + input.replacementTargetCellId, + assignmentEpoch, + expiresAt, + now + ) + await this.releaseSupersededControlConnectionReservations( + transaction, + identity, + input.currentTargetCellId, + input.assignmentEpoch, + now + ) + await transaction.query( + `UPDATE relay_assignment_migrations SET aborted_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [now, now, identity.userId, identity.relayHostId, input.assignmentEpoch] + ) + await transaction.query( + `INSERT INTO relay_assignment_migrations + (user_id, relay_host_id, source_cell_id, target_cell_id, + previous_epoch, assignment_epoch, source_request_units, + target_reserved_units, expires_at, target_registered_at, + completed_at, aborted_at, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, NULL, NULL, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + input.sourceCellId, + input.replacementTargetCellId, + input.assignmentEpoch, + assignmentEpoch, + sourceRequestUnits, + targetReservedUnits, + expiresAt, + now, + now + ] + ) + if (existingIncarnation) { + await transaction.query( + `INSERT INTO relay_assignment_migration_incarnations + (user_id, relay_host_id, assignment_epoch, source_cell_incarnation, + target_cell_incarnation) + VALUES (?, ?, ?, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + assignmentEpoch, + text(existingIncarnation, 'source_cell_incarnation'), + text(replacementRuntime, 'cell_incarnation') + ] + ) + } + if (existingPin) { + await transaction.query( + `INSERT INTO relay_post_drain_migration_pins + (user_id, relay_host_id, assignment_epoch, drain_attempt_id, + source_cell_id, source_cell_incarnation, target_cell_id, + target_cell_incarnation, source_request_units, target_reserved_units, + pinned_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + assignmentEpoch, + text(existingPin, 'drain_attempt_id'), + input.sourceCellId, + text(existingPin, 'source_cell_incarnation'), + input.replacementTargetCellId, + text(replacementRuntime, 'cell_incarnation'), + sourceRequestUnits, + targetReservedUnits, + now + ] + ) + } + return { + ...identity, + sourceCellId: input.sourceCellId, + targetCellId: input.replacementTargetCellId, + previousEpoch: input.assignmentEpoch, + assignmentEpoch, + expiresAt + } + }) + } + + async supersedeRegisteredCellEvacuations( + sourceCellId: string, + currentTargetCellId: string, + replacementTargetCellId: string, + limit: number + ): Promise { + if (!Number.isSafeInteger(limit) || limit < 1 || limit > 100) { + throw new Error('invalid_evacuation_limit') + } + const rows = await this.database.query( + `SELECT user_id, relay_host_id, assignment_epoch + FROM relay_assignment_migrations + WHERE source_cell_id = ? AND target_cell_id = ? + AND target_registered_at IS NOT NULL + AND completed_at IS NULL AND aborted_at IS NULL + AND NOT EXISTS ( + SELECT 1 FROM relay_assignment_migrations earlier + WHERE earlier.user_id = relay_assignment_migrations.user_id + AND earlier.relay_host_id = relay_assignment_migrations.relay_host_id + AND earlier.source_cell_id = relay_assignment_migrations.source_cell_id + AND earlier.target_cell_id = relay_assignment_migrations.target_cell_id + AND earlier.target_registered_at IS NOT NULL + AND earlier.completed_at IS NULL AND earlier.aborted_at IS NULL + AND earlier.assignment_epoch < relay_assignment_migrations.assignment_epoch + ) + ORDER BY user_id, relay_host_id, assignment_epoch + LIMIT ?`, + [sourceCellId, currentTargetCellId, limit] + ) + let superseded = 0 + let reconciledAccounting = false + for (const row of rows) { + const identity = { + userId: text(row, 'user_id'), + relayHostId: text(row, 'relay_host_id') + } + const input = { + assignmentEpoch: integer(row, 'assignment_epoch'), + sourceCellId, + currentTargetCellId, + replacementTargetCellId + } + const prepared = await this.prepareRegisteredCellSupersession(identity, input) + if (prepared.retired) { + input.assignmentEpoch = integer(row, 'assignment_epoch') + await this.supersedeRegisteredEvacuation(identity, input) + superseded++ + continue + } + input.assignmentEpoch = prepared.assignmentEpoch + try { + await this.supersedeRegisteredEvacuation(identity, input) + } catch (error) { + if ( + error instanceof Error && + error.message === 'migration_activity_topology_mismatch' + ) { + const preservedActivityCellIds = await this.additionalActivityCellIds( + identity, + sourceCellId, + currentTargetCellId + ) + if ( + preservedActivityCellIds.length === 0 || + preservedActivityCellIds.includes(replacementTargetCellId) + ) { + throw error + } + await this.reconcileReservationAccounting(sourceCellId, currentTargetCellId) + await this.reconcileReservationAccounting(sourceCellId, replacementTargetCellId) + for (const cellId of preservedActivityCellIds) { + await this.reconcileReservationAccounting(sourceCellId, cellId) + } + await this.supersedeRegisteredEvacuation( + identity, + input, + preservedActivityCellIds + ) + superseded++ + continue + } + if ( + reconciledAccounting || + !(error instanceof Error) || + error.message !== 'migration_cell_reservation_accounting_mismatch' + ) { + throw error + } + await this.reconcileReservationAccounting(sourceCellId, currentTargetCellId) + await this.reconcileReservationAccounting(sourceCellId, replacementTargetCellId) + reconciledAccounting = true + await this.supersedeRegisteredEvacuation(identity, input) + } + superseded++ + } + return superseded + } + + private async prepareRegisteredCellSupersession( + identity: AssignmentIdentity, + input: RegisteredEvacuationSupersessionInput + ): Promise<{ assignmentEpoch: number; retired: boolean }> { + const now = this.now() + return await this.database.transaction(async (transaction) => { + const assignment = await this.assignmentRow(transaction, identity) + const staleMigration = ( + await transaction.queryLocked( + `SELECT * FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, input.assignmentEpoch] + ) + )[0] + if (!assignment || !staleMigration) throw new Error('migration_assignment_mismatch') + assertMigrationPair( + staleMigration, + input.sourceCellId, + input.currentTargetCellId + ) + if ( + optionalInteger(staleMigration, 'completed_at') !== undefined || + optionalInteger(staleMigration, 'aborted_at') !== undefined + ) { + return { assignmentEpoch: integer(assignment, 'assignment_epoch'), retired: true } + } + if (optionalInteger(staleMigration, 'target_registered_at') === undefined) { + throw new Error('migration_not_registered_for_supersession') + } + const assignmentEpoch = integer(assignment, 'assignment_epoch') + if (assignmentEpoch < input.assignmentEpoch) { + throw new Error('migration_assignment_mismatch') + } + const assignmentCellId = text(assignment, 'cell_id') + if ( + ![input.currentTargetCellId, input.replacementTargetCellId].includes( + assignmentCellId + ) + ) { + throw new Error('migration_assignment_mismatch') + } + const leases = await this.lockAssignmentActivities(transaction, identity) + const staleIncarnation = ( + await transaction.queryLocked( + `SELECT * FROM relay_assignment_migration_incarnations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, input.assignmentEpoch] + ) + )[0] + const stalePin = ( + await transaction.queryLocked( + `SELECT * FROM relay_post_drain_migration_pins + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, input.assignmentEpoch] + ) + )[0] + assertMigrationRecoveryMetadata(staleMigration, staleIncarnation, stalePin) + const failedRuntime = ( + await transaction.queryLocked( + `SELECT * FROM relay_cell_runtime WHERE cell_id = ?`, + [input.currentTargetCellId] + ) + )[0] + await this.requireCellFence( + transaction, + input.currentTargetCellId, + failedRuntime, + now + ) + if ( + assignmentEpoch > input.assignmentEpoch && + assignmentCellId === input.replacementTargetCellId + ) { + const obsoleteActivityIds = new Set([ + pendingControlActivityId(input.assignmentEpoch), + migrationActivityId(input.assignmentEpoch) + ]) + const obsoleteLeases = leases.filter((lease) => + obsoleteActivityIds.has(text(lease, 'activity_id')) + ) + assertAssignmentActivityCounts( + assignment, + leases, + leases.filter((lease) => activityKind(lease) === 'migration').length + ) + for (const lease of obsoleteLeases) { + const kind = activityKind(lease) + const expectedUnits = + kind === 'migration' + ? integer(staleMigration, 'source_request_units') + : ACTIVITY_REQUEST_UNITS.control + if ( + text(lease, 'cell_id') !== input.currentTargetCellId || + !['control', 'migration'].includes(kind) || + integer(lease, 'request_units') !== expectedUnits + ) { + throw new Error('migration_activity_topology_mismatch') + } + } + if (obsoleteLeases.length > 0) await this.lockCellInventory(transaction) + for (const lease of obsoleteLeases) { + await this.removeActivityLease(transaction, identity, lease, now) + } + await this.releaseSupersededControlConnectionReservations( + transaction, + identity, + input.currentTargetCellId, + input.assignmentEpoch, + now + ) + await transaction.query( + `UPDATE relay_assignment_migrations SET aborted_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [now, now, identity.userId, identity.relayHostId, input.assignmentEpoch] + ) + return { assignmentEpoch, retired: true } + } + let migrationRow = staleMigration + if (assignmentEpoch > input.assignmentEpoch) { + const existingCurrent = ( + await transaction.queryLocked( + `SELECT * FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, assignmentEpoch] + ) + )[0] + if (existingCurrent) { + assertMigrationPair( + existingCurrent, + input.sourceCellId, + input.currentTargetCellId + ) + if ( + optionalInteger(existingCurrent, 'completed_at') !== undefined || + optionalInteger(existingCurrent, 'aborted_at') !== undefined || + optionalInteger(existingCurrent, 'target_registered_at') === undefined + ) { + throw new Error('migration_activity_topology_mismatch') + } + assertCurrentMigrationAssignment(assignment, existingCurrent) + const currentIncarnation = ( + await transaction.queryLocked( + `SELECT * FROM relay_assignment_migration_incarnations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, assignmentEpoch] + ) + )[0] + const currentPin = ( + await transaction.queryLocked( + `SELECT * FROM relay_post_drain_migration_pins + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, assignmentEpoch] + ) + )[0] + assertMigrationRecoveryMetadata( + existingCurrent, + currentIncarnation, + currentPin + ) + migrationRow = existingCurrent + } else { + if (leases.length !== 0) { + throw new Error('migration_activity_topology_mismatch') + } + assertAssignmentActivityCounts(assignment, leases, 0) + const currentIncarnations = await transaction.queryLocked( + `SELECT assignment_epoch FROM relay_assignment_migration_incarnations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, assignmentEpoch] + ) + const currentPins = await transaction.queryLocked( + `SELECT assignment_epoch FROM relay_post_drain_migration_pins + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, assignmentEpoch] + ) + if (currentIncarnations.length > 0 || currentPins.length > 0) { + throw new Error('migration_activity_topology_mismatch') + } + await transaction.query( + `INSERT INTO relay_assignment_migrations + (user_id, relay_host_id, source_cell_id, target_cell_id, + previous_epoch, assignment_epoch, source_request_units, + target_reserved_units, expires_at, target_registered_at, + completed_at, aborted_at, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, 0, 1, ?, ?, NULL, NULL, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + input.sourceCellId, + input.currentTargetCellId, + input.assignmentEpoch, + assignmentEpoch, + now - 1, + now, + now, + now + ] + ) + migrationRow = { + ...staleMigration, + assignment_epoch: assignmentEpoch, + previous_epoch: input.assignmentEpoch, + source_request_units: 0, + target_reserved_units: 1, + expires_at: now - 1, + target_registered_at: now, + completed_at: null, + aborted_at: null + } + if (staleIncarnation) { + await transaction.query( + `INSERT INTO relay_assignment_migration_incarnations + (user_id, relay_host_id, assignment_epoch, + source_cell_incarnation, target_cell_incarnation) + VALUES (?, ?, ?, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + assignmentEpoch, + text(staleIncarnation, 'source_cell_incarnation'), + text(staleIncarnation, 'target_cell_incarnation') + ] + ) + } + if (stalePin) { + await transaction.query( + `INSERT INTO relay_post_drain_migration_pins + (user_id, relay_host_id, assignment_epoch, drain_attempt_id, + source_cell_id, source_cell_incarnation, target_cell_id, + target_cell_incarnation, source_request_units, + target_reserved_units, pinned_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, 0, 1, ?)`, + [ + identity.userId, + identity.relayHostId, + assignmentEpoch, + text(stalePin, 'drain_attempt_id'), + input.sourceCellId, + text(stalePin, 'source_cell_incarnation'), + input.currentTargetCellId, + text(stalePin, 'target_cell_incarnation'), + now + ] + ) + } + } + await this.releaseSupersededControlConnectionReservations( + transaction, + identity, + input.currentTargetCellId, + input.assignmentEpoch, + now + ) + await transaction.query( + `UPDATE relay_assignment_migrations SET aborted_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [now, now, identity.userId, identity.relayHostId, input.assignmentEpoch] + ) + } + const migrationLeases = leases.filter( + (lease) => activityKind(lease) === 'migration' + ) + if (migrationLeases.length > 0) { + assertAssignmentActivityAccounting(assignment, leases, migrationRow) + return { assignmentEpoch, retired: false } + } + assertAssignmentActivityCounts(assignment, leases, 0) + const sourceRequestUnits = integer(migrationRow, 'source_request_units') + if ( + integer(migrationRow, 'expires_at') > now || + sourceRequestUnits < 0 || + integer(migrationRow, 'target_reserved_units') !== sourceRequestUnits + 1 + ) { + throw new Error('migration_activity_lease_shape_mismatch') + } + await this.lockCellInventory(transaction) + await this.adjustCellReservation( + transaction, + input.currentTargetCellId, + sourceRequestUnits + ) + const expiresAt = now + ASSIGNMENT_LIMITS.migrationLeaseMs + await transaction.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, ?, 'migration', ?, ?, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + migrationActivityId(assignmentEpoch), + input.currentTargetCellId, + sourceRequestUnits, + expiresAt, + now + ] + ) + await transaction.query( + `UPDATE relay_assignments SET migration_leases = migration_leases + 1, + lease_expires_at = CASE WHEN lease_expires_at > ? THEN lease_expires_at ELSE ? END, + last_activity_at = ? + WHERE user_id = ? AND relay_host_id = ?`, + [expiresAt, expiresAt, now, identity.userId, identity.relayHostId] + ) + await transaction.query( + `UPDATE relay_assignment_migrations SET expires_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [expiresAt, now, identity.userId, identity.relayHostId, assignmentEpoch] + ) + return { assignmentEpoch, retired: false } + }) + } + + private async additionalActivityCellIds( + identity: AssignmentIdentity, + sourceCellId: string, + currentTargetCellId: string + ): Promise { + const rows = await this.database.query( + `SELECT DISTINCT cell_id FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? + AND cell_id NOT IN (?, ?) + ORDER BY cell_id`, + [identity.userId, identity.relayHostId, sourceCellId, currentTargetCellId] + ) + return rows.map((row) => text(row, 'cell_id')) + } + + async completeEvacuation( + identity: AssignmentIdentity, + assignmentEpoch: number + ): Promise { + const now = this.now() + await this.database.transaction(async (transaction) => { + const assignment = await this.assignmentRow(transaction, identity) + const row = ( + await transaction.queryLocked( + `SELECT * FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ? + AND completed_at IS NULL AND aborted_at IS NULL`, + [identity.userId, identity.relayHostId, assignmentEpoch] + ) + )[0] + if (!row) throw new Error('migration_not_found') + if (optionalInteger(row, 'target_registered_at') === undefined) { + throw new Error('migration_target_not_registered') + } + const sourceCellId = text(row, 'source_cell_id') + const targetCellId = text(row, 'target_cell_id') + if ( + !assignment || + text(assignment, 'cell_id') !== targetCellId || + integer(assignment, 'assignment_epoch') !== assignmentEpoch + ) { + throw new Error('migration_assignment_mismatch') + } + const activityLeases = await this.lockAssignmentActivities(transaction, identity) + const sourceUnits = activityUnitsForCell(activityLeases, sourceCellId) + if (sourceUnits > 0) throw new Error('migration_source_still_active') + let cellsLocked = false + let targetStartedAt = 0 + if (this.requireLiveCells) { + let cells: SqlRow[] + try { + cells = await this.lockCellInventory(transaction, true) + } catch (error) { + if (isDatabaseLockUnavailable(error)) { + // Mixed-version workers may still hold a cell-first lock; defer + // instead of waiting long enough to form their legacy lock cycle. + throw new Error('migration_cell_inventory_busy') + } + throw error + } + cellsLocked = true + const source = cells.find((cell) => text(cell, 'cell_id') === sourceCellId) + const target = cells.find((cell) => text(cell, 'cell_id') === targetCellId) + if (!source || integer(source, 'enabled') !== 0) { + throw new Error('migration_source_admission_changed') + } + if (!target || integer(target, 'enabled') !== 1) { + throw new Error('migration_target_admission_changed') + } + const runtimes = await transaction.queryLocked( + `SELECT * FROM relay_cell_runtime WHERE cell_id IN (?, ?) ORDER BY cell_id`, + [sourceCellId, targetCellId] + ) + const sourceRuntime = runtimes.find( + (runtime) => text(runtime, 'cell_id') === sourceCellId + ) + const targetRuntime = runtimes.find( + (runtime) => text(runtime, 'cell_id') === targetCellId + ) + const freshAfter = now - this.heartbeatTtlMs + if ( + !sourceRuntime || + integer(sourceRuntime, 'ready') !== 1 || + integer(sourceRuntime, 'observed_requests') !== 0 || + integer(sourceRuntime, 'last_heartbeat_at') <= freshAfter + ) { + throw new Error('migration_source_runtime_not_quiescent') + } + if ( + !targetRuntime || + integer(targetRuntime, 'ready') !== 1 || + integer(targetRuntime, 'last_heartbeat_at') <= freshAfter + ) { + throw new Error('migration_target_runtime_not_ready') + } + targetStartedAt = integer(targetRuntime, 'started_at') + } + const targetIsActive = activityLeases.some( + (lease) => + text(lease, 'cell_id') === targetCellId && + text(lease, 'activity_kind') === 'control' && + !text(lease, 'activity_id').startsWith('control-pending:') && + integer(lease, 'expires_at') > now && + integer(lease, 'updated_at') >= targetStartedAt + ) + if (!targetIsActive) throw new Error('migration_target_not_active') + const lease = activityLeaseById(activityLeases, migrationActivityId(assignmentEpoch)) + if (lease && !cellsLocked) await this.lockCellInventory(transaction) + if (lease) await this.removeActivityLease(transaction, identity, lease, now) + await transaction.query( + `UPDATE relay_assignment_migrations SET completed_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [now, now, identity.userId, identity.relayHostId, assignmentEpoch] + ) + }) + } + + async rebalanceDormant( + identity: AssignmentIdentity, + targetCellId: string + ): Promise { + const now = this.now() + return await this.database.transaction(async (transaction) => { + const assignment = await this.assignmentRow(transaction, identity) + if (!assignment) throw new Error('assignment_not_found') + if (!mayNormallyReassign(activity(assignment), now)) throw new Error('assignment_active') + const sourceCellId = text(assignment, 'cell_id') + if (sourceCellId === targetCellId) throw new Error('target_matches_source') + await this.lockAssignmentActivities(transaction, identity) + const cells = await this.lockCellInventory(transaction) + const admission = await cellAdmissionStates(transaction) + const targetRow = cells.find( + (row) => + text(row, 'cell_id') === targetCellId && + admission.get(targetCellId) === 'general' + ) + if (!targetRow) throw new Error('target_cell_unavailable') + if (!(await this.cellIsLive(transaction, targetCellId, now))) { + throw new Error('target_cell_unavailable') + } + await this.assertCellConnectionHeadroom(transaction, targetCellId) + const target = cell(targetRow, await this.cellRegion(transaction, targetCellId)) + await this.adjustCellReservation(transaction, targetCellId, 1) + const assignmentEpoch = integer(assignment, 'assignment_epoch') + 1 + const leaseExpiresAt = now + ASSIGNMENT_LIMITS.activityLeaseMs + await transaction.query( + `UPDATE relay_assignments SET cell_id = ?, assignment_epoch = ?, + reserved_controls = 1, lease_expires_at = ?, last_activity_at = ? + WHERE user_id = ? AND relay_host_id = ?`, + [ + targetCellId, + assignmentEpoch, + leaseExpiresAt, + now, + identity.userId, + identity.relayHostId + ] + ) + await this.releaseSupersededControlConnectionReservations( + transaction, + identity, + sourceCellId, + assignmentEpoch - 1, + now + ) + await this.insertPendingControlLease( + transaction, + identity, + targetCellId, + assignmentEpoch, + now + ) + return { + ...identity, + ...target, + assignmentEpoch, + leaseExpiresAt + } + }) + } + + async inspectRegionalRehomeControl(): Promise { + const now = this.now() + return await this.database.transaction(async (transaction) => { + await this.initializeRegionalRehomeControl(transaction, now) + const row = ( + await transaction.query( + `SELECT * FROM relay_region_rehome_control WHERE control_id = 'global'` + ) + )[0]! + return regionalRehomeControl(row) + }) + } + + async applyRegionalRehomeControl(input: { + expectedGeneration: number + enabled: boolean + notBefore: number + ratePerMinute: number + preferenceMaxAgeMs: number + drainGraceMs: number + }): Promise { + if (!Number.isSafeInteger(input.expectedGeneration) || input.expectedGeneration < 0) { + throw new Error('invalid_regional_rehome_generation') + } + if (!Number.isSafeInteger(input.notBefore) || input.notBefore < 0) { + throw new Error('invalid_regional_rehome_not_before') + } + if (!Number.isSafeInteger(input.ratePerMinute) || input.ratePerMinute < 1 || input.ratePerMinute > 120) { + throw new Error('invalid_regional_rehome_rate') + } + if ( + !Number.isSafeInteger(input.preferenceMaxAgeMs) || + input.preferenceMaxAgeMs < 60_000 || + input.preferenceMaxAgeMs > 30 * 24 * 60 * 60_000 + ) { + throw new Error('invalid_regional_rehome_preference_age') + } + if ( + !Number.isSafeInteger(input.drainGraceMs) || + input.drainGraceMs < 60_000 || + input.drainGraceMs > 60 * 60_000 + ) { + throw new Error('invalid_regional_rehome_drain_grace') + } + const now = this.now() + return await this.database.transaction(async (transaction) => { + await this.initializeRegionalRehomeControl(transaction, now) + const current = ( + await transaction.queryLocked( + `SELECT * FROM relay_region_rehome_control WHERE control_id = 'global'` + ) + )[0]! + if (integer(current, 'generation') !== input.expectedGeneration) { + throw new Error('regional_rehome_generation_mismatch') + } + if ( + input.enabled && + input.notBefore < + integer(current, 'observation_started_at') + REGIONAL_REHOME_OBSERVATION_MS + ) { + throw new Error('regional_rehome_observation_window_incomplete') + } + await transaction.query( + `UPDATE relay_region_rehome_control + SET generation = generation + 1, enabled = ?, not_before = ?, + rate_per_minute = ?, preference_max_age_ms = ?, drain_grace_ms = ?, + updated_at = ? + WHERE control_id = 'global'`, + [ + input.enabled ? 1 : 0, + input.notBefore, + input.ratePerMinute, + input.preferenceMaxAgeMs, + input.drainGraceMs, + now + ] + ) + const updated = ( + await transaction.query( + `SELECT * FROM relay_region_rehome_control WHERE control_id = 'global'` + ) + )[0]! + return regionalRehomeControl(updated) + }) + } + + async disableRegionalRehomeControl(): Promise { + const now = this.now() + const result = await this.database.query( + `UPDATE relay_region_rehome_control + SET generation = generation + 1, enabled = 0, updated_at = ? + WHERE control_id = 'global' AND enabled = 1`, + [now] + ) + return integer(result[0]!, 'changes') === 1 + } + + private async initializeRegionalRehomeControl( + database: RelayDatabase, + now: number + ): Promise { + await database.query( + `INSERT INTO relay_region_rehome_control + (control_id, generation, enabled, observation_started_at, not_before, + rate_per_minute, preference_max_age_ms, drain_grace_ms, updated_at) + VALUES ('global', 0, 0, ?, 0, 10, ?, ?, ?) + ON CONFLICT (control_id) DO NOTHING`, + [now, 24 * 60 * 60_000, 60 * 60_000, now] + ) + } + + async regionalRehomeFleetSafety(): Promise { + return await this.readRegionalRehomeFleetSafety(this.database, this.now()) + } + + private async readRegionalRehomeFleetSafety( + database: RelayDatabase, + now: number + ): Promise { + const rows = await database.query( + `SELECT runtime.ready, runtime.last_heartbeat_at, + safety.cell_id AS safety_cell_id, safety.observed_at, + safety.sql_failures, safety.reconnects, + safety.control_activity_recovery_failures, + safety.database_pool_waiting, safety.database_pool_waiters_max, + safety.database_pool_wait_ms_max + FROM relay_cells cell + JOIN relay_cell_admission admission ON admission.cell_id = cell.cell_id + JOIN relay_cell_regions region ON region.cell_id = cell.cell_id + LEFT JOIN relay_cell_runtime runtime ON runtime.cell_id = cell.cell_id + LEFT JOIN relay_cell_capabilities capability ON capability.cell_id = cell.cell_id + LEFT JOIN relay_cell_rehome_safety safety + ON safety.cell_id = runtime.cell_id + AND safety.cell_incarnation = runtime.cell_incarnation + WHERE cell.enabled = 1 AND admission.admission_state = 'general' + AND ( + region.region = 'asia-east2' OR + (region.region = 'us-central1' AND capability.regional_rehome_protocol >= 1) + )` + ) + const valid = rows.filter( + (row) => + integer(row, 'ready') === 1 && + integer(row, 'last_heartbeat_at') > now - this.heartbeatTtlMs && + optionalText(row, 'safety_cell_id') !== undefined && + integer(row, 'observed_at') > now - 60_000 + ) + const missingCells = rows.length === 0 ? 1 : rows.length - valid.length + return { + requiredCells: rows.length, + missingCells, + observedAt: + missingCells > 0 + ? 0 + : Math.min(...valid.map((row) => integer(row, 'observed_at'))), + sqlFailures: valid.reduce((total, row) => total + integer(row, 'sql_failures'), 0), + reconnects: valid.reduce((total, row) => total + integer(row, 'reconnects'), 0), + maxReconnects: Math.max(0, ...valid.map((row) => integer(row, 'reconnects'))), + controlActivityRecoveryFailures: valid.reduce( + (total, row) => total + integer(row, 'control_activity_recovery_failures'), + 0 + ), + databasePoolWaiting: Math.max( + 0, + ...valid.map((row) => integer(row, 'database_pool_waiting')) + ), + databasePoolWaitersMax: Math.max( + 0, + ...valid.map((row) => integer(row, 'database_pool_waiters_max')) + ), + databasePoolWaitMsMax: Math.max( + 0, + ...valid.map((row) => integer(row, 'database_pool_wait_ms_max')) + ) + } + } + + async claimRegionalRehome( + processSafety?: RegionalRehomeSafetySnapshot + ): Promise { + const now = this.now() + // Directors poll every second; avoid taking the global worker-row lock while disabled. + const control = ( + await this.database.query( + `SELECT enabled, not_before + FROM relay_region_rehome_control + WHERE control_id = 'global'` + ) + )[0] + if (!control) { + await this.initializeRegionalRehomeControl(this.database, now) + return null + } + if (integer(control, 'enabled') !== 1 || integer(control, 'not_before') > now) { + return null + } + this.pendingRegionalRehomeDisableLog = null + const candidateSkips: RegionalRehomeCandidateSkip[] = [] + const claimResult = await this.database.transaction(async (transaction) => { + candidateSkips.length = 0 + await this.initializeRegionalRehomeControl(transaction, now) + const control = ( + await transaction.queryLocked( + `SELECT * FROM relay_region_rehome_control WHERE control_id = 'global'` + ) + )[0]! + if (integer(control, 'enabled') !== 1 || integer(control, 'not_before') > now) { + return null + } + const intervalMs = Math.ceil(60_000 / integer(control, 'rate_per_minute')) + const preferenceCutoff = now - integer(control, 'preference_max_age_ms') + await transaction.query( + `INSERT INTO relay_region_rehome_worker_state + (worker_id, next_dispatch_at, paused_until, consecutive_failures, updated_at) + VALUES ('global', 0, 0, 0, ?) + ON CONFLICT (worker_id) DO NOTHING`, + [now] + ) + const worker = ( + await transaction.queryLocked( + `SELECT * FROM relay_region_rehome_worker_state WHERE worker_id = 'global'` + ) + )[0]! + if ( + integer(worker, 'paused_until') > now || + integer(worker, 'next_dispatch_at') > now + ) { + return null + } + const effectiveProcessSafety = processSafety ?? cleanRegionalRehomeSafety(now) + const fleetSafety = await this.readRegionalRehomeFleetSafety(transaction, now) + if ( + !(await this.regionalRehomeSafetyAllowsClaim( + transaction, + worker, + effectiveProcessSafety, + fleetSafety, + now + )) + ) { + return null + } + const retry = ( + await transaction.queryLocked( + `SELECT attempt.*, source.cell_url AS source_cell_url + FROM relay_region_rehome_attempts attempt + JOIN relay_cells source ON source.cell_id = attempt.source_cell_id + JOIN relay_cell_runtime runtime ON runtime.cell_id = attempt.source_cell_id + JOIN relay_cell_capabilities capability + ON capability.cell_id = runtime.cell_id + AND capability.cell_incarnation = runtime.cell_incarnation + JOIN relay_assignment_migrations migration + ON migration.user_id = attempt.user_id + AND migration.relay_host_id = attempt.relay_host_id + AND migration.assignment_epoch = attempt.assignment_epoch + WHERE attempt.drain_receipt_at IS NULL + AND attempt.completed_at IS NULL AND attempt.aborted_at IS NULL + AND attempt.send_attempts < 10 + AND (attempt.last_send_attempt_at IS NULL OR attempt.last_send_attempt_at <= ?) + AND runtime.cell_incarnation = attempt.source_cell_incarnation + AND runtime.ready = 1 AND runtime.last_heartbeat_at > ? + AND capability.regional_rehome_protocol >= 1 + AND migration.completed_at IS NULL AND migration.aborted_at IS NULL + ORDER BY attempt.created_at, attempt.attempt_id + LIMIT 1`, + [now - 30_000, now - this.heartbeatTtlMs] + ) + )[0] + if (retry) { + const fleetSafety = await this.lockedRegionalRehomeFleetSafety(transaction, now) + if ( + !(await this.regionalRehomeSafetyAllowsClaim( + transaction, + worker, + effectiveProcessSafety, + fleetSafety, + now + )) + ) { + return null + } + await this.markRegionalRehomeDispatchClaimed( + transaction, + text(retry, 'attempt_id'), + now, + intervalMs + ) + retry.send_attempts = integer(retry, 'send_attempts') + 1 + return regionalRehomeAttempt(retry) + } + + // A drain receipt is not convergence: grace enforcement lives only in + // source-cell session state, and attempts have been observed stalled + // dual-homed well past grace with source leases still renewing. Such + // attempts are re-dispatched with the remaining (zero) grace so the + // source force-closes and the host re-resolves onto its registered + // target. + const redrain = ( + await transaction.queryLocked( + `SELECT attempt.*, source.cell_url AS source_cell_url + FROM relay_region_rehome_attempts attempt + JOIN relay_cells source ON source.cell_id = attempt.source_cell_id + JOIN relay_cell_runtime runtime ON runtime.cell_id = attempt.source_cell_id + JOIN relay_cell_capabilities capability + ON capability.cell_id = runtime.cell_id + AND capability.cell_incarnation = runtime.cell_incarnation + JOIN relay_assignment_migrations migration + ON migration.user_id = attempt.user_id + AND migration.relay_host_id = attempt.relay_host_id + AND migration.assignment_epoch = attempt.assignment_epoch + WHERE attempt.drain_receipt_at IS NOT NULL + AND attempt.completed_at IS NULL AND attempt.aborted_at IS NULL + AND attempt.created_at + attempt.drain_grace_ms <= ? + AND attempt.send_attempts < ? + AND (attempt.last_send_attempt_at IS NULL OR attempt.last_send_attempt_at <= ?) + AND runtime.cell_incarnation = attempt.source_cell_incarnation + AND runtime.ready = 1 AND runtime.last_heartbeat_at > ? + AND capability.regional_rehome_protocol >= 1 + AND migration.completed_at IS NULL AND migration.aborted_at IS NULL + AND migration.target_registered_at IS NOT NULL + AND EXISTS ( + SELECT 1 FROM relay_assignment_activity_leases source_lease + WHERE source_lease.user_id = attempt.user_id + AND source_lease.relay_host_id = attempt.relay_host_id + AND source_lease.cell_id = attempt.source_cell_id + ) + ORDER BY attempt.created_at, attempt.attempt_id + LIMIT 1`, + [ + now, + REGIONAL_REHOME_REDRAIN_SEND_LIMIT, + now - REGIONAL_REHOME_REDRAIN_INTERVAL_MS, + now - this.heartbeatTtlMs + ] + ) + )[0] + if (redrain) { + const fleetSafety = await this.lockedRegionalRehomeFleetSafety(transaction, now) + if ( + !(await this.regionalRehomeSafetyAllowsClaim( + transaction, + worker, + effectiveProcessSafety, + fleetSafety, + now + )) + ) { + return null + } + await this.markRegionalRehomeDispatchClaimed( + transaction, + text(redrain, 'attempt_id'), + now, + intervalMs + ) + redrain.send_attempts = integer(redrain, 'send_attempts') + 1 + redrain.drain_grace_ms = 0 + return regionalRehomeAttempt(redrain) + } + + const candidates = await transaction.query( + `SELECT preference.user_id, preference.relay_host_id, + preference.observed_at, assignment.cell_id AS source_cell_id, + assignment.assignment_epoch + FROM relay_assignment_region_preferences preference + JOIN relay_assignments assignment + ON assignment.user_id = preference.user_id + AND assignment.relay_host_id = preference.relay_host_id + JOIN relay_cell_regions region ON region.cell_id = assignment.cell_id + JOIN relay_cell_admission admission ON admission.cell_id = assignment.cell_id + JOIN relay_cell_runtime runtime ON runtime.cell_id = assignment.cell_id + JOIN relay_cell_capabilities capability + ON capability.cell_id = runtime.cell_id + AND capability.cell_incarnation = runtime.cell_incarnation + WHERE preference.preferred_region = 'asia-east2' + AND preference.observed_at >= ? + AND region.region = 'us-central1' + AND admission.admission_state = 'general' + AND runtime.ready = 1 AND runtime.last_heartbeat_at > ? + AND capability.regional_rehome_protocol >= 1 + AND EXISTS ( + SELECT 1 FROM relay_assignment_activity_leases control + WHERE control.user_id = assignment.user_id + AND control.relay_host_id = assignment.relay_host_id + AND control.cell_id = assignment.cell_id + AND control.activity_kind = 'control' + AND control.activity_id NOT LIKE 'control-pending:%' + AND control.expires_at > ? + AND control.updated_at >= runtime.started_at + ) + AND NOT EXISTS ( + SELECT 1 FROM relay_assignment_migrations migration + WHERE migration.user_id = assignment.user_id + AND migration.relay_host_id = assignment.relay_host_id + AND migration.completed_at IS NULL AND migration.aborted_at IS NULL + ) + ORDER BY preference.observed_at, preference.user_id, preference.relay_host_id + LIMIT 10`, + [preferenceCutoff, now - this.heartbeatTtlMs, now] + ) + for (const candidate of candidates) { + const claimed = await this.startRegionalRehomeCandidate(transaction, { + identity: { + userId: text(candidate, 'user_id'), + relayHostId: text(candidate, 'relay_host_id') + }, + sourceCellId: text(candidate, 'source_cell_id'), + assignmentEpoch: integer(candidate, 'assignment_epoch'), + preferenceCutoff, + drainGraceMs: integer(control, 'drain_grace_ms'), + processSafety: effectiveProcessSafety, + worker, + now, + skips: candidateSkips + }) + if (!claimed) continue + await this.markRegionalRehomeDispatchClaimed( + transaction, + claimed.attemptId, + now, + intervalMs + ) + return { ...claimed, sendAttempts: 1 } + } + if (candidates.length > 0) { + // Skipped candidates still cost all-rows FOR UPDATE inventory scans; + // charge the dispatch interval so skips are rate-limited like claims. + await this.markRegionalRehomeTickSkipped(transaction, now, intervalMs) + } + return null + }) + const pendingDisableLog = this.pendingRegionalRehomeDisableLog + this.pendingRegionalRehomeDisableLog = null + if (pendingDisableLog) console.warn(JSON.stringify(pendingDisableLog)) + if (claimResult === null && candidateSkips.length > 0) { + console.warn(JSON.stringify(aggregateRegionalRehomeCandidateSkips(candidateSkips))) + } + return claimResult + } + + private async startRegionalRehomeCandidate( + transaction: RelayDatabase, + input: { + identity: AssignmentIdentity + sourceCellId: string + assignmentEpoch: number + preferenceCutoff: number + drainGraceMs: number + processSafety: RegionalRehomeSafetySnapshot + worker: SqlRow + now: number + skips: RegionalRehomeCandidateSkip[] + } + ): Promise | null> { + const assignment = await this.assignmentRow(transaction, input.identity) + if ( + !assignment || + text(assignment, 'cell_id') !== input.sourceCellId || + integer(assignment, 'assignment_epoch') !== input.assignmentEpoch + ) { + input.skips.push({ reason: 'candidate_stale' }) + return null + } + const preference = ( + await transaction.queryLocked( + `SELECT * FROM relay_assignment_region_preferences + WHERE user_id = ? AND relay_host_id = ?`, + [input.identity.userId, input.identity.relayHostId] + ) + )[0] + if ( + !preference || + text(preference, 'preferred_region') !== 'asia-east2' || + integer(preference, 'observed_at') < input.preferenceCutoff + ) { + input.skips.push({ reason: 'candidate_stale' }) + return null + } + const activeMigration = await transaction.queryLocked( + `SELECT assignment_epoch FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ? + AND completed_at IS NULL AND aborted_at IS NULL`, + [input.identity.userId, input.identity.relayHostId] + ) + if (activeMigration.length > 0) { + input.skips.push({ reason: 'candidate_stale' }) + return null + } + const activityLeases = await this.lockAssignmentActivities(transaction, input.identity) + assertAssignmentActivityCounts(assignment, activityLeases, 0) + const cells = await this.lockCellInventory(transaction) + const admission = await cellAdmissionStates(transaction) + const regions = new Map( + (await transaction.query(`SELECT cell_id, region FROM relay_cell_regions`)).map((row) => [ + text(row, 'cell_id'), + relayRegion(row, 'region') + ]) + ) + const runtimes = await transaction.queryLocked( + `SELECT * FROM relay_cell_runtime ORDER BY cell_id` + ) + const capabilities = await transaction.queryLocked( + `SELECT * FROM relay_cell_capabilities ORDER BY cell_id` + ) + const safetyRows = await transaction.queryLocked( + `SELECT * FROM relay_cell_rehome_safety ORDER BY cell_id` + ) + const source = cells.find((row) => text(row, 'cell_id') === input.sourceCellId) + const sourceRuntime = runtimes.find( + (row) => text(row, 'cell_id') === input.sourceCellId + ) + const sourceCapability = capabilities.find( + (row) => text(row, 'cell_id') === input.sourceCellId + ) + const sourceSafety = safetyRows.find( + (row) => text(row, 'cell_id') === input.sourceCellId + ) + const fleetSafety = regionalRehomeFleetSafetyFromInventory({ + cells, + admission, + regions, + runtimes, + capabilities, + safetyRows, + now: input.now, + heartbeatTtlMs: this.heartbeatTtlMs + }) + const safetyFailure = regionalRehomeFleetSafetyFailure( + input.processSafety, + fleetSafety, + input.now + ) + if (safetyFailure) { + await this.pauseRegionalRehomeForSafety( + transaction, + input.worker, + input.now, + safetyFailure, + fleetSafety + ) + return null + } + if ( + !source || + integer(source, 'enabled') !== 1 || + admission.get(input.sourceCellId) !== 'general' || + regions.get(input.sourceCellId) !== RELAY_DEFAULT_REGION || + !sourceRuntime || + integer(sourceRuntime, 'ready') !== 1 || + integer(sourceRuntime, 'last_heartbeat_at') <= input.now - this.heartbeatTtlMs || + !sourceCapability || + text(sourceCapability, 'cell_incarnation') !== + text(sourceRuntime, 'cell_incarnation') || + integer(sourceCapability, 'regional_rehome_protocol') < 1 + ) { + input.skips.push({ reason: 'source_ineligible', cellId: input.sourceCellId }) + return null + } + if (!regionalRehomeCellSafetyIsClean(sourceSafety, sourceRuntime, input.now)) { + input.skips.push(cellUncleanSkip('source_unclean', input.sourceCellId, sourceSafety)) + return null + } + const sourceControlActive = activityLeases.some( + (lease) => + text(lease, 'cell_id') === input.sourceCellId && + text(lease, 'activity_kind') === 'control' && + !text(lease, 'activity_id').startsWith('control-pending:') && + integer(lease, 'expires_at') > input.now && + integer(lease, 'updated_at') >= integer(sourceRuntime, 'started_at') + ) + if (!sourceControlActive) { + input.skips.push({ reason: 'source_control_inactive', cellId: input.sourceCellId }) + return null + } + const connectionHeadroom = await this.connectionHeadroomByCell(transaction) + const eligibleTargets = cells.filter((row) => { + const cellId = text(row, 'cell_id') + const runtime = runtimes.find((candidate) => text(candidate, 'cell_id') === cellId) + return ( + cellId !== input.sourceCellId && + integer(row, 'enabled') === 1 && + admission.get(cellId) === 'general' && + regions.get(cellId) === 'asia-east2' && + runtime !== undefined && + integer(runtime, 'ready') === 1 && + integer(runtime, 'last_heartbeat_at') > input.now - this.heartbeatTtlMs + ) + }) + const targetIsClean = (row: SqlRow): boolean => { + const cellId = text(row, 'cell_id') + return regionalRehomeCellSafetyIsClean( + safetyRows.find((safety) => text(safety, 'cell_id') === cellId), + runtimes.find((candidate) => text(candidate, 'cell_id') === cellId)!, + input.now + ) + } + const targetCandidates = eligibleTargets.filter( + (row) => targetIsClean(row) && connectionHeadroom.get(text(row, 'cell_id')) !== false + ) + if (targetCandidates.length === 0) { + const unclean = eligibleTargets.filter((row) => !targetIsClean(row)) + for (const row of unclean) { + const cellId = text(row, 'cell_id') + input.skips.push( + cellUncleanSkip( + 'target_unclean', + cellId, + safetyRows.find((safety) => text(safety, 'cell_id') === cellId) + ) + ) + } + if (unclean.length === 0) { + input.skips.push({ + reason: eligibleTargets.length === 0 ? 'no_eligible_target' : 'no_target_headroom' + }) + } + return null + } + targetCandidates.sort((left, right) => { + const leftRuntime = runtimes.find( + (runtime) => text(runtime, 'cell_id') === text(left, 'cell_id') + )! + const rightRuntime = runtimes.find( + (runtime) => text(runtime, 'cell_id') === text(right, 'cell_id') + )! + const leftLoad = + (integer(left, 'reserved_requests') + integer(leftRuntime, 'observed_requests')) / + integer(left, 'capacity_requests') + const rightLoad = + (integer(right, 'reserved_requests') + integer(rightRuntime, 'observed_requests')) / + integer(right, 'capacity_requests') + return leftLoad - rightLoad || + text(left, 'cell_id').localeCompare(text(right, 'cell_id')) + }) + const sourceRequestUnits = activityUnitsForCell(activityLeases, input.sourceCellId) + const targetReservedUnits = sourceRequestUnits + 1 + const target = targetCandidates.find( + (row) => + integer(row, 'reserved_requests') + targetReservedUnits <= + integer(row, 'capacity_requests') + ) + if (!target) { + input.skips.push({ reason: 'no_target_headroom' }) + return null + } + const targetCellId = text(target, 'cell_id') + const targetRuntime = runtimes.find( + (runtime) => text(runtime, 'cell_id') === targetCellId + )! + await this.adjustCellReservation(transaction, targetCellId, targetReservedUnits) + const previousEpoch = integer(assignment, 'assignment_epoch') + const assignmentEpoch = previousEpoch + 1 + const expiresAt = input.now + ASSIGNMENT_LIMITS.migrationLeaseMs + await transaction.query( + `UPDATE relay_assignments SET cell_id = ?, assignment_epoch = ?, + reserved_controls = reserved_controls + 1, + migration_leases = migration_leases + 1, + lease_expires_at = ?, last_activity_at = ? + WHERE user_id = ? AND relay_host_id = ?`, + [ + targetCellId, + assignmentEpoch, + expiresAt, + input.now, + input.identity.userId, + input.identity.relayHostId + ] + ) + await transaction.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, ?, 'control', ?, 1, ?, ?), + (?, ?, ?, 'migration', ?, ?, ?, ?)`, + [ + input.identity.userId, + input.identity.relayHostId, + pendingControlActivityId(assignmentEpoch), + targetCellId, + expiresAt, + input.now, + input.identity.userId, + input.identity.relayHostId, + migrationActivityId(assignmentEpoch), + targetCellId, + sourceRequestUnits, + expiresAt, + input.now + ] + ) + await this.insertControlConnectionReservation( + transaction, + input.identity, + targetCellId, + assignmentEpoch, + expiresAt, + input.now + ) + await transaction.query( + `INSERT INTO relay_assignment_migrations + (user_id, relay_host_id, source_cell_id, target_cell_id, + previous_epoch, assignment_epoch, source_request_units, + target_reserved_units, expires_at, target_registered_at, + completed_at, aborted_at, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, NULL, NULL, ?, ?)`, + [ + input.identity.userId, + input.identity.relayHostId, + input.sourceCellId, + targetCellId, + previousEpoch, + assignmentEpoch, + sourceRequestUnits, + targetReservedUnits, + expiresAt, + input.now, + input.now + ] + ) + await transaction.query( + `INSERT INTO relay_assignment_migration_incarnations + (user_id, relay_host_id, assignment_epoch, source_cell_incarnation, + target_cell_incarnation) + VALUES (?, ?, ?, ?, ?)`, + [ + input.identity.userId, + input.identity.relayHostId, + assignmentEpoch, + text(sourceRuntime, 'cell_incarnation'), + text(targetRuntime, 'cell_incarnation') + ] + ) + const attemptId = randomUUID() + await transaction.query( + `INSERT INTO relay_region_rehome_attempts + (attempt_id, user_id, relay_host_id, preferred_region, + source_cell_id, source_cell_incarnation, target_cell_id, + target_cell_incarnation, previous_epoch, assignment_epoch, + drain_grace_ms, send_attempts, last_send_attempt_at, + drain_receipt_at, drain_outcome, completed_at, aborted_at, + created_at, updated_at) + VALUES (?, ?, ?, 'asia-east2', ?, ?, ?, ?, ?, ?, ?, 0, NULL, + NULL, NULL, NULL, NULL, ?, ?)`, + [ + attemptId, + input.identity.userId, + input.identity.relayHostId, + input.sourceCellId, + text(sourceRuntime, 'cell_incarnation'), + targetCellId, + text(targetRuntime, 'cell_incarnation'), + previousEpoch, + assignmentEpoch, + input.drainGraceMs, + input.now, + input.now + ] + ) + return { + ...input.identity, + attemptId, + preferredRegion: 'asia-east2', + sourceCellId: input.sourceCellId, + sourceCellUrl: text(source, 'cell_url'), + sourceCellIncarnation: text(sourceRuntime, 'cell_incarnation'), + targetCellId, + targetCellIncarnation: text(targetRuntime, 'cell_incarnation'), + previousEpoch, + assignmentEpoch, + drainGraceMs: input.drainGraceMs + } + } + + private async lockedRegionalRehomeFleetSafety( + transaction: RelayDatabase, + now: number + ): Promise { + const cells = await this.lockCellInventory(transaction) + const admission = await cellAdmissionStates(transaction) + const regions = new Map( + (await transaction.query(`SELECT cell_id, region FROM relay_cell_regions`)).map((row) => [ + text(row, 'cell_id'), + relayRegion(row, 'region') + ]) + ) + const runtimes = await transaction.queryLocked( + `SELECT * FROM relay_cell_runtime ORDER BY cell_id` + ) + const capabilities = await transaction.queryLocked( + `SELECT * FROM relay_cell_capabilities ORDER BY cell_id` + ) + const safetyRows = await transaction.queryLocked( + `SELECT * FROM relay_cell_rehome_safety ORDER BY cell_id` + ) + return regionalRehomeFleetSafetyFromInventory({ + cells, + admission, + regions, + runtimes, + capabilities, + safetyRows, + now, + heartbeatTtlMs: this.heartbeatTtlMs + }) + } + + private async regionalRehomeSafetyAllowsClaim( + transaction: RelayDatabase, + worker: SqlRow, + processSafety: RegionalRehomeSafetySnapshot, + fleetSafety: RegionalRehomeFleetSafety, + now: number + ): Promise { + const failure = regionalRehomeFleetSafetyFailure(processSafety, fleetSafety, now) + if (!failure) { + return true + } + await this.pauseRegionalRehomeForSafety(transaction, worker, now, failure, fleetSafety) + return false + } + + private async pauseRegionalRehomeForSafety( + transaction: RelayDatabase, + worker: SqlRow, + now: number, + reason: string, + fleetSafety: RegionalRehomeFleetSafety + ): Promise { + const disabled = await transaction.query( + `UPDATE relay_region_rehome_control + SET generation = generation + 1, enabled = 0, updated_at = ? + WHERE control_id = 'global' AND enabled = 1 + RETURNING generation`, + [now] + ) + // The durable disable is otherwise invisible: nothing else records why + // claims stopped and inspection only shows enabled=false. Logged after + // the transaction commits so a rollback cannot fabricate the record. + if (disabled.length > 0) { + this.pendingRegionalRehomeDisableLog = { + event: 'orca_relay_regional_rehome_safety_disabled', + reason, + controlGeneration: integer(disabled[0]!, 'generation'), + now, + requiredCells: fleetSafety.requiredCells, + missingCells: fleetSafety.missingCells, + observedAt: fleetSafety.observedAt, + sqlFailures: fleetSafety.sqlFailures, + reconnects: fleetSafety.reconnects, + maxReconnects: fleetSafety.maxReconnects, + controlActivityRecoveryFailures: fleetSafety.controlActivityRecoveryFailures, + databasePoolWaiting: fleetSafety.databasePoolWaiting, + databasePoolWaitersMax: fleetSafety.databasePoolWaitersMax, + databasePoolWaitMsMax: fleetSafety.databasePoolWaitMsMax + } + } + await this.incrementRegionalRehomeWorkerFailure(transaction, worker, now) + } + + private async markRegionalRehomeTickSkipped( + transaction: RelayDatabase, + now: number, + intervalMs: number + ): Promise { + await transaction.query( + `UPDATE relay_region_rehome_worker_state + SET next_dispatch_at = ?, updated_at = ? WHERE worker_id = 'global'`, + [now + intervalMs, now] + ) + } + + private async markRegionalRehomeDispatchClaimed( + transaction: RelayDatabase, + attemptId: string, + now: number, + intervalMs: number + ): Promise { + await transaction.query( + `UPDATE relay_region_rehome_attempts + SET send_attempts = send_attempts + 1, last_send_attempt_at = ?, updated_at = ? + WHERE attempt_id = ?`, + [now, now, attemptId] + ) + await transaction.query( + `UPDATE relay_region_rehome_worker_state + SET next_dispatch_at = ?, updated_at = ? WHERE worker_id = 'global'`, + [now + intervalMs, now] + ) + } + + async recordRegionalRehomeDrainReceipt( + attemptId: string, + outcome: RegionalHostDrainOutcome + ): Promise { + const now = this.now() + return await this.database.transaction(async (transaction) => { + const worker = ( + await transaction.queryLocked( + `SELECT * FROM relay_region_rehome_worker_state WHERE worker_id = 'global'` + ) + )[0] + const attempt = ( + await transaction.queryLocked( + `SELECT * FROM relay_region_rehome_attempts WHERE attempt_id = ?`, + [attemptId] + ) + )[0] + if (!attempt) throw new Error('regional_rehome_attempt_not_found') + // Any receipt proves the source cell answered: reset the failure budget + // even when a redrain repeats the stored outcome; otherwise a + // redrain-dominated stream lets scattered transient failures reach the + // durable three-failure disable. + if (worker) { + await transaction.query( + `UPDATE relay_region_rehome_worker_state + SET consecutive_failures = 0, paused_until = 0, updated_at = ? + WHERE worker_id = 'global'`, + [now] + ) + } + const existingOutcome = optionalText(attempt, 'drain_outcome') + if (existingOutcome === outcome) return false + // Redrains produce one receipt per dispatch; the latest outcome wins. + await transaction.query( + `UPDATE relay_region_rehome_attempts + SET drain_receipt_at = ?, drain_outcome = ?, updated_at = ? + WHERE attempt_id = ?`, + [now, outcome, now, attemptId] + ) + return true + }) + } + + async recordRegionalRehomeDispatchFailure(attemptId: string): Promise { + const now = this.now() + await this.database.transaction(async (transaction) => { + const worker = ( + await transaction.queryLocked( + `SELECT * FROM relay_region_rehome_worker_state WHERE worker_id = 'global'` + ) + )[0] + const attempt = ( + await transaction.queryLocked( + `SELECT attempt_id FROM relay_region_rehome_attempts WHERE attempt_id = ?`, + [attemptId] + ) + )[0] + if (!worker || !attempt) return + await this.incrementRegionalRehomeWorkerFailure(transaction, worker, now) + }) + } + + async recordRegionalRehomeWorkerFailure(): Promise { + const now = this.now() + await this.database.transaction(async (transaction) => { + await transaction.query( + `INSERT INTO relay_region_rehome_worker_state + (worker_id, next_dispatch_at, paused_until, consecutive_failures, updated_at) + VALUES ('global', 0, 0, 0, ?) + ON CONFLICT (worker_id) DO NOTHING`, + [now] + ) + const worker = ( + await transaction.queryLocked( + `SELECT * FROM relay_region_rehome_worker_state WHERE worker_id = 'global'` + ) + )[0]! + await this.incrementRegionalRehomeWorkerFailure(transaction, worker, now) + }) + } + + private async incrementRegionalRehomeWorkerFailure( + transaction: RelayDatabase, + worker: SqlRow, + now: number + ): Promise { + const failures = integer(worker, 'consecutive_failures') + 1 + await transaction.query( + `UPDATE relay_region_rehome_worker_state + SET consecutive_failures = ?, paused_until = ?, updated_at = ? + WHERE worker_id = 'global'`, + [failures, failures >= 3 ? now + 5 * 60_000 : 0, now] + ) + if (failures >= 3) { + await transaction.query( + `UPDATE relay_region_rehome_control + SET generation = generation + 1, enabled = 0, updated_at = ? + WHERE control_id = 'global' AND enabled = 1`, + [now] + ) + } + } + + async completeReadyRegionalRehomes(limit = 10): Promise { + const now = this.now() + const quarantined = this.quarantinedRegionalRehomeAttemptIds(now) + const exclusion = quarantined.length + ? ` AND attempt.attempt_id NOT IN (${quarantined.map(() => '?').join(', ')})` + : '' + const candidates = await this.database.query( + `SELECT attempt.attempt_id, attempt.user_id, attempt.relay_host_id, + attempt.assignment_epoch + FROM relay_region_rehome_attempts attempt + JOIN relay_assignment_migrations migration + ON migration.user_id = attempt.user_id + AND migration.relay_host_id = attempt.relay_host_id + AND migration.assignment_epoch = attempt.assignment_epoch + WHERE attempt.completed_at IS NULL AND attempt.aborted_at IS NULL + AND migration.target_registered_at IS NOT NULL + AND migration.completed_at IS NULL AND migration.aborted_at IS NULL + AND NOT EXISTS ( + SELECT 1 FROM relay_assignment_activity_leases source_lease + WHERE source_lease.user_id = attempt.user_id + AND source_lease.relay_host_id = attempt.relay_host_id + AND source_lease.cell_id = attempt.source_cell_id + )${exclusion} + ORDER BY attempt.created_at, attempt.attempt_id + LIMIT ?`, + [...quarantined, limit] + ) + let completed = 0 + for (const candidate of candidates) { + // One poisoned row must not stall every later candidate: an invariant + // throw here blocked fleet completions head-of-line in production. + const attemptId = text(candidate, 'attempt_id') + try { + const changed = await this.completeRegionalRehomeCandidate( + { + userId: text(candidate, 'user_id'), + relayHostId: text(candidate, 'relay_host_id') + }, + integer(candidate, 'assignment_epoch'), + now + ) + if (changed) completed++ + this.regionalRehomeCandidateQuarantine.delete(attemptId) + } catch (error) { + this.recordRegionalRehomeCandidateFailure('complete', attemptId, now, error) + } + } + return completed + } + + // Repeated invariant failures quarantine the attempt out of the sweeps' + // LIMIT pages: poisoned rows are permanent and always the oldest, so + // without exclusion they eventually starve every healthy candidate. + private recordRegionalRehomeCandidateFailure( + operation: 'complete' | 'abort', + attemptId: string, + now: number, + error: unknown + ): void { + const entry = this.regionalRehomeCandidateQuarantine.get(attemptId) ?? { + failures: 0, + until: 0 + } + entry.failures++ + if (entry.failures >= REGIONAL_REHOME_QUARANTINE_FAILURES) { + entry.until = now + REGIONAL_REHOME_QUARANTINE_MS + } + this.regionalRehomeCandidateQuarantine.delete(attemptId) + this.regionalRehomeCandidateQuarantine.set(attemptId, entry) + if (this.regionalRehomeCandidateQuarantine.size > REGIONAL_REHOME_QUARANTINE_MEMORY_LIMIT) { + // Evict a non-quarantined entry first: mid-quarantine rows are excluded + // from the candidate pages and losing one returns it to the page with a + // reset counter. + let evict = this.regionalRehomeCandidateQuarantine.keys().next().value! + for (const [key, candidate] of this.regionalRehomeCandidateQuarantine) { + if (candidate.until <= now) { + evict = key + break + } + } + this.regionalRehomeCandidateQuarantine.delete(evict) + } + warnRegionalRehomeCandidateFailure(operation, attemptId, error) + } + + private quarantinedRegionalRehomeAttemptIds(now: number): string[] { + const excluded: string[] = [] + for (const [attemptId, entry] of this.regionalRehomeCandidateQuarantine) { + if (entry.until > now) excluded.push(attemptId) + if (excluded.length >= REGIONAL_REHOME_QUARANTINE_EXCLUSION_LIMIT) break + } + return excluded + } + + async refreshRegionalRehomeLeases(limit = 100): Promise { + const now = this.now() + const candidates = await this.database.query( + `SELECT user_id, relay_host_id, assignment_epoch + FROM relay_region_rehome_attempts + WHERE completed_at IS NULL AND aborted_at IS NULL + ORDER BY created_at, attempt_id LIMIT ?`, + [limit] + ) + let refreshed = 0 + for (const candidate of candidates) { + const identity = { + userId: text(candidate, 'user_id'), + relayHostId: text(candidate, 'relay_host_id') + } + const assignmentEpoch = integer(candidate, 'assignment_epoch') + const changed = await this.database.transaction(async (transaction) => { + const assignment = await this.assignmentRow(transaction, identity) + const attempt = ( + await transaction.queryLocked( + `SELECT * FROM relay_region_rehome_attempts + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, assignmentEpoch] + ) + )[0] + const migration = ( + await transaction.queryLocked( + `SELECT * FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, assignmentEpoch] + ) + )[0] + if ( + !assignment || + !attempt || + !migration || + optionalInteger(attempt, 'completed_at') !== undefined || + optionalInteger(attempt, 'aborted_at') !== undefined || + optionalInteger(migration, 'completed_at') !== undefined || + optionalInteger(migration, 'aborted_at') !== undefined + ) { + return false + } + const attemptAgeMs = now - integer(attempt, 'created_at') + if (attemptAgeMs >= REGIONAL_REHOME_MAX_REFRESH_MS) { + return false + } + if ( + optionalInteger(migration, 'target_registered_at') === undefined && + attemptAgeMs >= REGIONAL_REHOME_UNREGISTERED_REFRESH_MS + ) { + await transaction.query( + `UPDATE relay_assignment_activity_leases + SET expires_at = CASE WHEN expires_at < ? THEN expires_at ELSE ? END, + updated_at = ? + WHERE user_id = ? AND relay_host_id = ? + AND activity_id IN (?, ?)`, + [ + now, + now, + now, + identity.userId, + identity.relayHostId, + pendingControlActivityId(assignmentEpoch), + migrationActivityId(assignmentEpoch) + ] + ) + await transaction.query( + `UPDATE relay_assignment_migrations + SET expires_at = CASE WHEN expires_at < ? THEN expires_at ELSE ? END, + updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [now, now, now, identity.userId, identity.relayHostId, assignmentEpoch] + ) + return false + } + const leases = await this.lockAssignmentActivities(transaction, identity) + const protectedIds = new Set([ + pendingControlActivityId(assignmentEpoch), + migrationActivityId(assignmentEpoch) + ]) + const protectedLeases = leases.filter((lease) => + protectedIds.has(text(lease, 'activity_id')) + ) + if (protectedLeases.length === 0) return false + const expiresAt = now + ASSIGNMENT_LIMITS.migrationLeaseMs + await transaction.query( + `UPDATE relay_assignment_activity_leases + SET expires_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? + AND activity_id IN (?, ?)`, + [ + expiresAt, + now, + identity.userId, + identity.relayHostId, + pendingControlActivityId(assignmentEpoch), + migrationActivityId(assignmentEpoch) + ] + ) + await transaction.query( + `UPDATE relay_assignment_migrations SET expires_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [expiresAt, now, identity.userId, identity.relayHostId, assignmentEpoch] + ) + await transaction.query( + `UPDATE relay_assignments SET lease_expires_at = + CASE WHEN lease_expires_at > ? THEN lease_expires_at ELSE ? END, + last_activity_at = ? + WHERE user_id = ? AND relay_host_id = ?`, + [expiresAt, expiresAt, now, identity.userId, identity.relayHostId] + ) + return true + }) + if (changed) refreshed++ + } + return refreshed + } + + private async completeRegionalRehomeCandidate( + identity: AssignmentIdentity, + assignmentEpoch: number, + now: number + ): Promise { + return await this.database.transaction(async (transaction) => { + const assignment = await this.assignmentRow(transaction, identity) + const attempt = ( + await transaction.queryLocked( + `SELECT * FROM relay_region_rehome_attempts + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, assignmentEpoch] + ) + )[0] + const migration = ( + await transaction.queryLocked( + `SELECT * FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, assignmentEpoch] + ) + )[0] + if ( + !attempt || + !migration || + optionalInteger(attempt, 'completed_at') !== undefined || + optionalInteger(attempt, 'aborted_at') !== undefined || + optionalInteger(migration, 'target_registered_at') === undefined || + optionalInteger(migration, 'completed_at') !== undefined || + optionalInteger(migration, 'aborted_at') !== undefined + ) { + return false + } + const sourceCellId = text(attempt, 'source_cell_id') + const targetCellId = text(attempt, 'target_cell_id') + if ( + !assignment || + text(assignment, 'cell_id') !== targetCellId || + integer(assignment, 'assignment_epoch') !== assignmentEpoch || + text(migration, 'source_cell_id') !== sourceCellId || + text(migration, 'target_cell_id') !== targetCellId + ) { + throw new Error('regional_rehome_assignment_mismatch') + } + const leases = await this.lockAssignmentActivities(transaction, identity) + if (activityUnitsForCell(leases, sourceCellId) !== 0) return false + await this.repairAssignmentActivityCounts( + transaction, + identity, + text(attempt, 'attempt_id'), + assignment, + leases, + migration + ) + const cells = await this.lockCellInventory(transaction) + const target = cells.find((cell) => text(cell, 'cell_id') === targetCellId) + const admission = await cellAdmissionStates(transaction) + if ( + !target || + integer(target, 'enabled') !== 1 || + !['general', 'migration-only'].includes(admission.get(targetCellId) ?? '') + ) { + return false + } + const targetRuntime = ( + await transaction.queryLocked( + `SELECT * FROM relay_cell_runtime WHERE cell_id = ?`, + [targetCellId] + ) + )[0] + if ( + !targetRuntime || + integer(targetRuntime, 'ready') !== 1 || + integer(targetRuntime, 'last_heartbeat_at') <= now - this.heartbeatTtlMs + ) { + return false + } + const targetActive = leases.some( + (lease) => + text(lease, 'cell_id') === targetCellId && + text(lease, 'activity_kind') === 'control' && + !text(lease, 'activity_id').startsWith('control-pending:') && + integer(lease, 'expires_at') > now && + integer(lease, 'updated_at') >= integer(targetRuntime, 'started_at') + ) + if (!targetActive) return false + const migrationLease = activityLeaseById( + leases, + migrationActivityId(assignmentEpoch) + ) + if (migrationLease) { + await this.removeActivityLease(transaction, identity, migrationLease, now) + } + await transaction.query( + `UPDATE relay_assignment_migrations SET completed_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [now, now, identity.userId, identity.relayHostId, assignmentEpoch] + ) + await transaction.query( + `UPDATE relay_region_rehome_attempts SET completed_at = ?, updated_at = ? + WHERE attempt_id = ?`, + [now, now, text(attempt, 'attempt_id')] + ) + return true + }) + } + + // Why: counter skew left by a pre-fix sticky grant is reconstructible from the + // locked lease rows; lease shape and migration topology are not, so only a + // count mismatch is repaired and the re-assert still throws on anything else. + // reconcileReservationAccounting cannot stand in: it opens its own + // transaction, while this has to run inside the sweep's on rows it holds. + private async repairAssignmentActivityCounts( + transaction: RelayDatabase, + identity: AssignmentIdentity, + attemptId: string, + assignment: SqlRow, + leases: SqlRow[], + migrationRow: SqlRow + ): Promise { + try { + assertAssignmentActivityAccounting(assignment, leases, migrationRow) + return + } catch (error) { + if ( + !(error instanceof Error) || + error.message !== 'migration_activity_accounting_mismatch' + ) { + throw error + } + } + const counts = activityCounts(leases) + await transaction.query( + `UPDATE relay_assignments SET reserved_controls = ?, reserved_splices = ?, + reserved_invites = ?, pending_installs = ?, pending_confirmations = ?, + migration_leases = ? + WHERE user_id = ? AND relay_host_id = ?`, + [ + counts.control, + counts.splice, + counts.invite, + counts.install, + counts.confirmation, + counts.migration, + identity.userId, + identity.relayHostId + ] + ) + const repaired = await this.assignmentRow(transaction, identity) + if (!repaired) throw new Error('regional_rehome_assignment_mismatch') + assertAssignmentActivityAccounting(repaired, leases, migrationRow) + noteRegionalRehomeActivityCountsRepaired(attemptId) + } + + async reapRegionalRehomeAttempts(): Promise { + const now = this.now() + const completed = await this.database.query( + `UPDATE relay_region_rehome_attempts + SET completed_at = COALESCE(completed_at, ?), updated_at = ? + WHERE completed_at IS NULL AND aborted_at IS NULL + AND EXISTS ( + SELECT 1 FROM relay_assignment_migrations migration + WHERE migration.user_id = relay_region_rehome_attempts.user_id + AND migration.relay_host_id = relay_region_rehome_attempts.relay_host_id + AND migration.assignment_epoch = relay_region_rehome_attempts.assignment_epoch + AND migration.completed_at IS NOT NULL + )`, + [now, now] + ) + const aborted = await this.database.query( + `UPDATE relay_region_rehome_attempts + SET aborted_at = COALESCE(aborted_at, ?), updated_at = ? + WHERE completed_at IS NULL AND aborted_at IS NULL + AND EXISTS ( + SELECT 1 FROM relay_assignment_migrations migration + WHERE migration.user_id = relay_region_rehome_attempts.user_id + AND migration.relay_host_id = relay_region_rehome_attempts.relay_host_id + AND migration.assignment_epoch = relay_region_rehome_attempts.assignment_epoch + AND migration.aborted_at IS NOT NULL + )`, + [now, now] + ) + if (integer(aborted[0]!, 'changes') > 0) { + await this.disableRegionalRehomeControl() + } + return integer(completed[0]!, 'changes') + integer(aborted[0]!, 'changes') + } + + async abortExpiredRegionalRehomes(limit = 100): Promise { + const now = this.now() + const quarantined = this.quarantinedRegionalRehomeAttemptIds(now) + const exclusion = quarantined.length + ? ` AND attempt_id NOT IN (${quarantined.map(() => '?').join(', ')})` + : '' + const candidates = await this.database.query( + `SELECT attempt_id, user_id, relay_host_id, assignment_epoch + FROM relay_region_rehome_attempts + WHERE completed_at IS NULL AND aborted_at IS NULL + AND created_at <= ?${exclusion} + ORDER BY created_at, attempt_id LIMIT ?`, + [now - REGIONAL_REHOME_MAX_REFRESH_MS, ...quarantined, limit] + ) + let aborted = 0 + for (const candidate of candidates) { + const identity = { + userId: text(candidate, 'user_id'), + relayHostId: text(candidate, 'relay_host_id') + } + const assignmentEpoch = integer(candidate, 'assignment_epoch') + const attemptId = text(candidate, 'attempt_id') + // Isolated like the completion sweep: one poisoned row must not stall + // every later candidate. + let changed = false + try { + changed = await this.database.transaction(async (transaction) => { + const assignment = await this.assignmentRow(transaction, identity) + const attempt = ( + await transaction.queryLocked( + `SELECT * FROM relay_region_rehome_attempts + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, assignmentEpoch] + ) + )[0] + const migration = ( + await transaction.queryLocked( + `SELECT * FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [identity.userId, identity.relayHostId, assignmentEpoch] + ) + )[0] + if ( + !assignment || + !attempt || + !migration || + optionalInteger(attempt, 'completed_at') !== undefined || + optionalInteger(attempt, 'aborted_at') !== undefined || + integer(attempt, 'created_at') > now - REGIONAL_REHOME_MAX_REFRESH_MS || + optionalInteger(migration, 'completed_at') !== undefined || + optionalInteger(migration, 'aborted_at') !== undefined + ) { + return false + } + const sourceCellId = text(attempt, 'source_cell_id') + const targetCellId = text(attempt, 'target_cell_id') + if ( + text(assignment, 'cell_id') !== targetCellId || + integer(assignment, 'assignment_epoch') !== assignmentEpoch + ) { + throw new Error('regional_rehome_assignment_mismatch') + } + const leases = await this.lockAssignmentActivities(transaction, identity) + if (activityUnitsForCell(leases, sourceCellId) > 0) return false + const targetActive = leases.some( + (lease) => + text(lease, 'cell_id') === targetCellId && + text(lease, 'activity_kind') === 'control' && + !text(lease, 'activity_id').startsWith('control-pending:') && + integer(lease, 'expires_at') > now + ) + if (targetActive) return false + const cells = await this.lockCellInventory(transaction) + const source = cells.find((cell) => text(cell, 'cell_id') === sourceCellId) + const admission = await cellAdmissionStates(transaction) + if ( + !source || + integer(source, 'enabled') !== 1 || + admission.get(sourceCellId) !== 'general' || + !(await this.cellIsLive(transaction, sourceCellId, now)) + ) { + return false + } + for (const activityId of [ + pendingControlActivityId(assignmentEpoch), + migrationActivityId(assignmentEpoch) + ]) { + const lease = activityLeaseById(leases, activityId) + if (lease) await this.removeActivityLease(transaction, identity, lease, now) + } + await this.releaseSupersededControlConnectionReservations( + transaction, + identity, + targetCellId, + assignmentEpoch, + now + ) + await transaction.query( + `UPDATE relay_assignments SET cell_id = ?, assignment_epoch = ?, + lease_expires_at = ?, last_activity_at = ? + WHERE user_id = ? AND relay_host_id = ?`, + [ + sourceCellId, + assignmentEpoch + 1, + now + ASSIGNMENT_LIMITS.activityLeaseMs, + now, + identity.userId, + identity.relayHostId + ] + ) + await transaction.query( + `UPDATE relay_assignment_migrations SET aborted_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [now, now, identity.userId, identity.relayHostId, assignmentEpoch] + ) + await transaction.query( + `UPDATE relay_region_rehome_attempts SET aborted_at = ?, updated_at = ? + WHERE attempt_id = ?`, + [now, now, text(attempt, 'attempt_id')] + ) + await transaction.query( + `UPDATE relay_region_rehome_control + SET generation = generation + 1, enabled = 0, updated_at = ? + WHERE control_id = 'global' AND enabled = 1`, + [now] + ) + return true + }) + this.regionalRehomeCandidateQuarantine.delete(attemptId) + } catch (error) { + this.recordRegionalRehomeCandidateFailure('abort', attemptId, now, error) + } + if (changed) aborted++ + } + return aborted + } + + async abortExpiredEvacuations(): Promise { + const now = this.now() + const abandonedBefore = now - STRANDED_MIGRATION_ABANDON_MS + const candidates = await this.database.query( + `SELECT migration.user_id, migration.relay_host_id, migration.assignment_epoch + FROM relay_assignment_migrations migration + WHERE migration.expires_at <= ? + AND migration.completed_at IS NULL AND migration.aborted_at IS NULL + AND ${ABORTABLE_EXPIRED_MIGRATION} + ORDER BY user_id, relay_host_id, assignment_epoch`, + [now, now, abandonedBefore, abandonedBefore] + ) + let aborted = 0 + for (const candidate of candidates) { + const didAbort = await this.database.transaction(async (transaction) => { + const identity = { + userId: text(candidate, 'user_id'), + relayHostId: text(candidate, 'relay_host_id') + } + // Migration cleanup follows the same assignment-first order as evacuation. + const assignment = await this.assignmentRow(transaction, identity) + const assignmentEpoch = integer(candidate, 'assignment_epoch') + const regionalAttempt = ( + await transaction.queryLocked( + `SELECT attempt_id FROM relay_region_rehome_attempts + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ? + AND completed_at IS NULL AND aborted_at IS NULL`, + [identity.userId, identity.relayHostId, assignmentEpoch] + ) + )[0] + const row = ( + await transaction.queryLocked( + `SELECT migration.* FROM relay_assignment_migrations migration + WHERE migration.user_id = ? AND migration.relay_host_id = ? + AND migration.assignment_epoch = ? + AND migration.expires_at <= ? + AND migration.completed_at IS NULL AND migration.aborted_at IS NULL + AND ${ABORTABLE_EXPIRED_MIGRATION}`, + [ + identity.userId, + identity.relayHostId, + assignmentEpoch, + now, + now, + abandonedBefore, + abandonedBefore + ] + ) + )[0] + if (!row) return false + const targetCellId = text(row, 'target_cell_id') + const activityLeases = await this.lockAssignmentActivities(transaction, identity) + if (!assignment) throw new Error('migration_assignment_missing') + const currentAssignmentEpoch = integer(assignment, 'assignment_epoch') + const assignmentEpochMatches = + text(assignment, 'cell_id') === targetCellId && + currentAssignmentEpoch === assignmentEpoch + const pendingTargetControl = activityLeaseById( + activityLeases, + pendingControlActivityId(assignmentEpoch) + ) + const targetGrantIsFresh = + assignmentEpochMatches && + pendingTargetControl !== undefined && + text(pendingTargetControl, 'cell_id') === targetCellId && + text(pendingTargetControl, 'activity_kind') === 'control' && + integer(pendingTargetControl, 'expires_at') > now + const targetIsActive = activityLeases.some( + (lease) => + text(lease, 'cell_id') === targetCellId && + text(lease, 'activity_kind') === 'control' && + text(lease, 'activity_id') !== pendingControlActivityId(assignmentEpoch) + ) + if (targetGrantIsFresh) return false + if (targetIsActive && assignmentEpochMatches) { + // A committed target control is stronger evidence than a failed follow-up + // write; repair the marker instead of rolling a live desktop backward. + await transaction.query( + `UPDATE relay_assignment_migrations + SET target_registered_at = COALESCE(target_registered_at, ?), updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [now, now, identity.userId, identity.relayHostId, assignmentEpoch] + ) + return false + } + if (!assignmentEpochMatches) { + if (currentAssignmentEpoch <= assignmentEpoch) { + throw new Error('migration_assignment_mismatch') + } + // A newer assignment is authoritative regardless of where it landed. + // Retire only this obsolete migration; never rewrite the newer epoch. + const obsoleteLeases = [ + pendingControlActivityId(assignmentEpoch), + migrationActivityId(assignmentEpoch) + ] + .map((activityId) => activityLeaseById(activityLeases, activityId)) + .filter((lease): lease is SqlRow => lease !== undefined) + if (obsoleteLeases.length > 0) await this.lockCellInventory(transaction) + for (const lease of obsoleteLeases) { + await this.removeActivityLease(transaction, identity, lease, now) + } + await this.releaseSupersededControlConnectionReservations( + transaction, + identity, + targetCellId, + assignmentEpoch, + now + ) + await transaction.query( + `UPDATE relay_assignment_migrations SET aborted_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [now, now, identity.userId, identity.relayHostId, assignmentEpoch] + ) + return true + } + const cells = await this.lockCellInventory(transaction) + const sourceCellId = text(row, 'source_cell_id') + const admissionRows = await transaction.query( + `SELECT cell_id, admission_state, updated_at FROM relay_cell_admission + WHERE cell_id IN (?, ?)`, + [sourceCellId, targetCellId] + ) + const sourceCell = cells.find((cell) => text(cell, 'cell_id') === sourceCellId) + const targetCell = cells.find((cell) => text(cell, 'cell_id') === targetCellId) + const sourceAdmission = admissionRows.find( + (admission) => text(admission, 'cell_id') === sourceCellId + ) + const targetAdmission = admissionRows.find( + (admission) => text(admission, 'cell_id') === targetCellId + ) + const registered = optionalInteger(row, 'target_registered_at') !== undefined + const sourceIsDurablyFenced = + registered && + ( + await transaction.query( + `SELECT 1 FROM relay_assignment_migrations migration + WHERE migration.user_id = ? AND migration.relay_host_id = ? + AND migration.assignment_epoch = ? + AND ${DURABLY_FENCED_MIGRATION_SOURCE}`, + [identity.userId, identity.relayHostId, assignmentEpoch] + ) + ).length === 1 + const retireOnTarget = + registered && + activityUnitsForCell(activityLeases, sourceCellId) === 0 && + sourceCell !== undefined && + integer(sourceCell, 'enabled') === 0 && + sourceAdmission !== undefined && + text(sourceAdmission, 'admission_state') === 'existing-only' && + (integer(sourceAdmission, 'updated_at') <= abandonedBefore || + sourceIsDurablyFenced) && + targetCell !== undefined && + integer(targetCell, 'enabled') === 1 && + targetAdmission !== undefined && + ['migration-only', 'general'].includes(text(targetAdmission, 'admission_state')) + const rollbackReason = + !registered || + (targetCell !== undefined && + integer(targetCell, 'enabled') === 0 && + targetAdmission !== undefined && + text(targetAdmission, 'admission_state') === 'existing-only' && + integer(targetAdmission, 'updated_at') <= abandonedBefore) + const regionalRollbackSourceAvailable = + !regionalAttempt || + (sourceCell !== undefined && + integer(sourceCell, 'enabled') === 1 && + sourceAdmission !== undefined && + text(sourceAdmission, 'admission_state') === 'general' && + (await this.cellIsLive(transaction, sourceCellId, now))) + const rollbackToSource = rollbackReason && regionalRollbackSourceAvailable + if (!retireOnTarget && !rollbackToSource) return false + for (const activityId of [ + pendingControlActivityId(assignmentEpoch), + migrationActivityId(assignmentEpoch) + ]) { + const lease = activityLeaseById(activityLeases, activityId) + if (lease) await this.removeActivityLease(transaction, identity, lease, now) + } + await this.releaseSupersededControlConnectionReservations( + transaction, + identity, + targetCellId, + assignmentEpoch, + now + ) + if (retireOnTarget) { + await transaction.query( + `UPDATE relay_assignment_migrations SET completed_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [now, now, identity.userId, identity.relayHostId, assignmentEpoch] + ) + return true + } + await transaction.query( + `UPDATE relay_assignments SET cell_id = ?, assignment_epoch = ?, + lease_expires_at = ?, last_activity_at = ? + WHERE user_id = ? AND relay_host_id = ?`, + [ + sourceCellId, + assignmentEpoch + 1, + now + ASSIGNMENT_LIMITS.activityLeaseMs, + now, + identity.userId, + identity.relayHostId + ] + ) + await transaction.query( + `UPDATE relay_assignment_migrations SET aborted_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [now, now, identity.userId, identity.relayHostId, assignmentEpoch] + ) + return true + }) + if (didAbort) aborted++ + } + return aborted + } + + async releaseExpiredActivityLeases(): Promise { + const now = this.now() + await this.database.query( + `UPDATE relay_control_connection_reservations + SET state = 'late-arrival-debt', updated_at = ? + WHERE state = 'reserved' AND timeout_at <= ?`, + [now, now] + ) + await this.database.query( + `UPDATE relay_control_connection_reservations + SET state = 'released', released_at = ?, updated_at = ? + WHERE state = 'late-arrival-debt' + AND claim_activity_id IS NULL + AND timeout_at <= ?`, + [now, now, now - LATE_ARRIVAL_DEBT_RETENTION_MS] + ) + const candidates = await this.database.query( + `SELECT user_id, relay_host_id, activity_id + FROM relay_assignment_activity_leases lease + WHERE expires_at <= ? + AND NOT EXISTS ( + SELECT 1 FROM relay_region_rehome_attempts rehome + WHERE rehome.user_id = lease.user_id + AND rehome.relay_host_id = lease.relay_host_id + AND rehome.completed_at IS NULL AND rehome.aborted_at IS NULL + AND lease.activity_id IN ( + 'control-pending:' || CAST(rehome.assignment_epoch AS TEXT), + 'migration:' || CAST(rehome.assignment_epoch AS TEXT) + ) + ) + AND NOT EXISTS ( + SELECT 1 FROM relay_post_drain_migration_pins pin + JOIN relay_assignment_migrations migration + ON migration.user_id = pin.user_id + AND migration.relay_host_id = pin.relay_host_id + AND migration.assignment_epoch = pin.assignment_epoch + WHERE pin.user_id = lease.user_id + AND pin.relay_host_id = lease.relay_host_id + AND migration.completed_at IS NULL + AND migration.aborted_at IS NULL + AND lease.activity_id IN ( + 'control-pending:' || CAST(pin.assignment_epoch AS TEXT), + 'migration:' || CAST(pin.assignment_epoch AS TEXT) + ) + )`, + [now] + ) + let released = 0 + for (const candidate of candidates) { + let didRelease: boolean + try { + didRelease = await this.database.transaction(async (transaction) => { + const identity = { + userId: text(candidate, 'user_id'), + relayHostId: text(candidate, 'relay_host_id') + } + // Why: re-check under the canonical assignment-first lock so a concurrent + // renewal wins without cleanup reaping its refreshed lease. + // Several directors sweep the same expired rows. Skip a row another + // director is settling instead of waiting and retrying the transaction. + await this.assignmentRow(transaction, identity, true) + const activityLeases = await this.lockAssignmentActivities(transaction, identity, true) + const lease = activityLeaseById(activityLeases, text(candidate, 'activity_id')) + if (!lease || integer(lease, 'expires_at') > now) return false + await this.lockCellInventory(transaction, true) + await this.removeActivityLease(transaction, identity, lease, now) + return true + }) + } catch (error) { + // Released cell images can still hold their legacy cell-first lock; + // expiry is durable, so a later maintenance sweep can safely retry it. + if (isDatabaseLockUnavailable(error)) continue + throw error + } + if (didRelease) released++ + } + return released + } + + async releaseExpiredActivity(): Promise { + const now = this.now() + try { + return await this.database.transaction(async (transaction) => { + const expired = await transaction.queryLocked( + `SELECT * FROM relay_assignments WHERE lease_expires_at <= ? AND + (reserved_controls > 0 OR reserved_splices > 0 OR reserved_invites > 0 OR + pending_installs > 0 OR pending_confirmations > 0 OR migration_leases > 0) + AND NOT EXISTS ( + SELECT 1 FROM relay_region_rehome_attempts rehome + WHERE rehome.user_id = relay_assignments.user_id + AND rehome.relay_host_id = relay_assignments.relay_host_id + AND rehome.completed_at IS NULL AND rehome.aborted_at IS NULL + ) + AND NOT EXISTS ( + SELECT 1 FROM relay_post_drain_migration_pins pin + JOIN relay_assignment_migrations migration + ON migration.user_id = pin.user_id + AND migration.relay_host_id = pin.relay_host_id + AND migration.assignment_epoch = pin.assignment_epoch + WHERE pin.user_id = relay_assignments.user_id + AND pin.relay_host_id = relay_assignments.relay_host_id + AND migration.completed_at IS NULL + AND migration.aborted_at IS NULL + ) + ORDER BY user_id, relay_host_id`, + [now], + { failIfUnavailable: true } + ) + if (expired.length > 0) await this.lockCellInventory(transaction, true) + for (const row of expired) { + await this.adjustCellReservation(transaction, text(row, 'cell_id'), -requestUnits(row)) + await transaction.query( + `UPDATE relay_assignments SET reserved_controls = 0, reserved_splices = 0, + reserved_invites = 0, pending_installs = 0, pending_confirmations = 0, + migration_leases = 0 + WHERE user_id = ? AND relay_host_id = ?`, + [text(row, 'user_id'), text(row, 'relay_host_id')] + ) + } + return expired.length + }) + } catch (error) { + // Aggregate expiry is reconstructible from durable leases; never wait + // long enough to form a mixed-version cell/assignment lock cycle. + if (isDatabaseLockUnavailable(error)) return 0 + throw error + } + } + + async releaseExpiredRegionPreferences(): Promise { + const result = await this.database.query( + `DELETE FROM relay_assignment_region_preferences WHERE observed_at < ?`, + [this.now() - REGION_PREFERENCE_RETENTION_MS] + ) + return integer(result[0]!, 'changes') + } + + private async reconcileReservationAccounting( + sourceCellId: string, + targetCellId: string + ): Promise { + const now = this.now() + await this.database.transaction(async (transaction) => { + // Reconciliation takes the same assignment→activity→cell order as live + // mutations so correcting drift never races a credential or socket lease. + const assignments = await transaction.queryLocked( + `SELECT assignment.* FROM relay_assignments assignment + WHERE assignment.cell_id IN (?, ?) OR EXISTS ( + SELECT 1 FROM relay_assignment_activity_leases scoped_lease + WHERE scoped_lease.user_id = assignment.user_id + AND scoped_lease.relay_host_id = assignment.relay_host_id + AND scoped_lease.cell_id IN (?, ?) + ) + ORDER BY assignment.user_id, assignment.relay_host_id`, + [sourceCellId, targetCellId, sourceCellId, targetCellId] + ) + const leases = await transaction.queryLocked( + `SELECT lease.* FROM relay_assignment_activity_leases lease + WHERE lease.cell_id IN (?, ?) OR EXISTS ( + SELECT 1 FROM relay_assignments assignment + WHERE assignment.user_id = lease.user_id + AND assignment.relay_host_id = lease.relay_host_id + AND ( + assignment.cell_id IN (?, ?) OR EXISTS ( + SELECT 1 FROM relay_assignment_activity_leases scoped_lease + WHERE scoped_lease.user_id = lease.user_id + AND scoped_lease.relay_host_id = lease.relay_host_id + AND scoped_lease.cell_id IN (?, ?) + ) + ) + ) + ORDER BY lease.user_id, lease.relay_host_id, lease.activity_id`, + [ + sourceCellId, + targetCellId, + sourceCellId, + targetCellId, + sourceCellId, + targetCellId + ] + ) + const cells = await this.lockCellInventory(transaction) + const assignmentKeys = new Set( + assignments.map((row) => + assignmentKey(text(row, 'user_id'), text(row, 'relay_host_id')) + ) + ) + const cellIds = new Set(cells.map((row) => text(row, 'cell_id'))) + const assignmentCounts = new Map< + string, + { counts: Record; leaseExpiresAt: number } + >() + const cellUnits = new Map() + + for (const lease of leases) { + const key = assignmentKey(text(lease, 'user_id'), text(lease, 'relay_host_id')) + if (!assignmentKeys.has(key)) throw new Error('activity_lease_assignment_missing') + const cellId = text(lease, 'cell_id') + if (!cellIds.has(cellId)) throw new Error('activity_lease_cell_missing') + const current = + assignmentCounts.get(key) ?? { + counts: emptyActivityCounts(), + leaseExpiresAt: 0 + } + const kind = activityKind(lease) + current.counts[kind]++ + current.leaseExpiresAt = Math.max(current.leaseExpiresAt, integer(lease, 'expires_at')) + assignmentCounts.set(key, current) + cellUnits.set(cellId, (cellUnits.get(cellId) ?? 0) + integer(lease, 'request_units')) + } + + for (const row of assignments) { + const current = assignmentCounts.get( + assignmentKey(text(row, 'user_id'), text(row, 'relay_host_id')) + ) + const counts = current?.counts ?? emptyActivityCounts() + const differs = (Object.keys(ACTIVITY_COLUMN) as AssignmentActivityKind[]).some( + (kind) => integer(row, ACTIVITY_COLUMN[kind]) !== counts[kind] + ) + const expiryDiffers = + current !== undefined && integer(row, 'lease_expires_at') !== current.leaseExpiresAt + if (!differs && !expiryDiffers) continue + await transaction.query( + `UPDATE relay_assignments SET reserved_controls = ?, reserved_splices = ?, + reserved_invites = ?, pending_installs = ?, pending_confirmations = ?, + migration_leases = ?, lease_expires_at = ? + WHERE user_id = ? AND relay_host_id = ?`, + [ + counts.control, + counts.splice, + counts.invite, + counts.install, + counts.confirmation, + counts.migration, + current?.leaseExpiresAt ?? integer(row, 'lease_expires_at'), + text(row, 'user_id'), + text(row, 'relay_host_id') + ] + ) + } + + for (const row of cells.filter((cell) => + [sourceCellId, targetCellId].includes(text(cell, 'cell_id')) + )) { + const cellId = text(row, 'cell_id') + const expected = cellUnits.get(cellId) ?? 0 + if (expected > integer(row, 'capacity_requests')) { + throw new Error('relay_capacity_exhausted') + } + if (integer(row, 'reserved_requests') === expected) continue + await transaction.query( + `UPDATE relay_cells SET reserved_requests = ?, updated_at = ? WHERE cell_id = ?`, + [expected, now, cellId] + ) + } + }) + } + + private async lockCellInventory( + database: RelayDatabase, + failIfUnavailable = false + ): Promise { + // Every capacity-changing assignment takes the tiny cell inventory in one + // order; dynamically locking only the selected target allowed cross-cell cycles. + return await database.queryLocked( + `SELECT * FROM relay_cells ORDER BY cell_id ASC`, + [], + { failIfUnavailable } + ) + } + + private async lockGeneralCellInventory( + database: RelayDatabase, + failIfUnavailable = false + ): Promise { + return await database.queryLocked( + `SELECT * FROM relay_cells + WHERE cell_id IN ( + SELECT cell_id FROM relay_cell_admission WHERE admission_state = 'general' + ) + ORDER BY cell_id ASC`, + [], + { failIfUnavailable } + ) + } + + private async leastLoadedCell( + database: RelayDatabase, + lockedCells: SqlRow[] | undefined, + preferredRegion: RelayRegion + ): Promise { + const rows = lockedCells ?? (await this.lockCellInventory(database)) + const regions = new Map( + (await database.query(`SELECT cell_id, region FROM relay_cell_regions`)).map((row) => [ + text(row, 'cell_id'), + relayRegion(row, 'region') + ]) + ) + const admission = await cellAdmissionStates(database) + const runtimeLoad = this.requireLiveCells + ? new Map( + ( + await database.query( + `SELECT cell_id, observed_requests FROM relay_cell_runtime + WHERE ready = ? AND last_heartbeat_at > ?`, + [1, this.now() - this.heartbeatTtlMs] + ) + ).map((row) => [text(row, 'cell_id'), integer(row, 'observed_requests')]) + ) + : null + const connectionHeadroom = await this.connectionHeadroomByCell(database) + const candidates = rows.filter((row) => { + const cellId = text(row, 'cell_id') + return ( + admission.get(cellId) === 'general' && + integer(row, 'reserved_requests') < integer(row, 'capacity_requests') && + (!runtimeLoad || runtimeLoad.has(cellId)) && + connectionHeadroom.get(cellId) !== false + ) + }) + candidates.sort((left, right) => { + const leftLoad = + integer(left, 'reserved_requests') + + (runtimeLoad?.get(text(left, 'cell_id')) ?? integer(left, 'observed_requests')) + const rightLoad = + integer(right, 'reserved_requests') + + (runtimeLoad?.get(text(right, 'cell_id')) ?? integer(right, 'observed_requests')) + const loadDifference = + leftLoad / integer(left, 'capacity_requests') - + rightLoad / integer(right, 'capacity_requests') + return loadDifference || text(left, 'cell_id').localeCompare(text(right, 'cell_id')) + }) + const preferred = candidates.filter( + (candidate) => + (regions.get(text(candidate, 'cell_id')) ?? RELAY_DEFAULT_REGION) === preferredRegion + ) + const selected = preferred[0] ?? candidates[0] + return selected + ? cell(selected, regions.get(text(selected, 'cell_id')) ?? RELAY_DEFAULT_REGION) + : null + } + + async regionCatalog(): Promise { + const rows = await this.database.query( + `SELECT cell.cell_id, cell.cell_url, region.region + FROM relay_cells cell + LEFT JOIN relay_cell_regions region ON region.cell_id = cell.cell_id + JOIN relay_cell_admission admission ON admission.cell_id = cell.cell_id + JOIN relay_cell_runtime runtime ON runtime.cell_id = cell.cell_id + WHERE cell.enabled = 1 + AND admission.admission_state = 'general' + AND runtime.ready = ? AND runtime.last_heartbeat_at > ? + ORDER BY cell.cell_id ASC`, + [1, this.now() - this.heartbeatTtlMs] + ) + const origins = new Map() + for (const row of rows) { + const region = optionalRelayRegion(row, 'region') ?? RELAY_DEFAULT_REGION + const current = origins.get(region) ?? [] + if (current.length < 2) current.push(text(row, 'cell_url')) + origins.set(region, current) + } + return RELAY_REGIONS.flatMap((region) => { + const probeOrigins = origins.get(region) + return probeOrigins?.length ? [{ region, probeOrigins }] : [] + }) + } + + private async recordRegionPreference( + database: RelayDatabase, + identity: AssignmentIdentity, + preferredRegion: RelayRegion | undefined, + observedAt: number + ): Promise { + if (!preferredRegion) return + await database.query( + `INSERT INTO relay_assignment_region_preferences + (user_id, relay_host_id, preferred_region, observed_at) + VALUES (?, ?, ?, ?) + ON CONFLICT (user_id, relay_host_id) DO UPDATE SET + preferred_region = CASE + WHEN excluded.observed_at >= relay_assignment_region_preferences.observed_at + THEN excluded.preferred_region + ELSE relay_assignment_region_preferences.preferred_region + END, + observed_at = CASE + WHEN excluded.observed_at >= relay_assignment_region_preferences.observed_at + THEN excluded.observed_at + ELSE relay_assignment_region_preferences.observed_at + END`, + [identity.userId, identity.relayHostId, preferredRegion, observedAt] + ) + } + + private async cellRegion(database: RelayDatabase, cellId: string): Promise { + const row = ( + await database.query(`SELECT region FROM relay_cell_regions WHERE cell_id = ?`, [cellId]) + )[0] + return row ? relayRegion(row, 'region') : RELAY_DEFAULT_REGION + } + + private async connectionHeadroomByCell( + database: RelayDatabase + ): Promise> { + const now = this.now() + const rows = await database.query(ASSIGNMENT_CONNECTION_HEADROOM_QUERY) + return new Map( + rows.map((row) => { + const heartbeat = optionalInteger(row, 'last_heartbeat_at') + const hasFreshTelemetry = + heartbeat !== undefined && + heartbeat > now - this.heartbeatTtlMs && + optionalText(row, 'connection_incarnation') === + optionalText(row, 'current_incarnation') + const hasHeadroom = + hasFreshTelemetry && + integer(row, 'enforced_connection_units') + + integer(row, 'outstanding_reservations') + + integer(row, 'unobserved_bound') < + integer(row, 'hard_cap') - + RELAY_ADMISSION_BUDGETS.reservedHostControls + return [text(row, 'cell_id'), hasHeadroom] as const + }) + ) + } + + private async assertCellConnectionHeadroom( + database: RelayDatabase, + cellId: string + ): Promise { + if (!(await this.cellHasConnectionHeadroom(database, cellId))) { + throw new Error('relay_connection_headroom_exhausted') + } + } + + private async cellHasConnectionHeadroom( + database: RelayDatabase, + cellId: string + ): Promise { + return (await this.connectionHeadroomByCell(database)).get(cellId) !== false + } + + private async cellIsLive( + database: RelayDatabase, + cellId: string, + now: number + ): Promise { + if (!this.requireLiveCells) return true + const rows = await database.query( + `SELECT cell.cell_id FROM relay_cells cell + JOIN relay_cell_runtime runtime ON runtime.cell_id = cell.cell_id + WHERE cell.cell_id = ? AND runtime.ready = ? AND runtime.last_heartbeat_at > ?`, + [cellId, 1, now - this.heartbeatTtlMs] + ) + return rows.length === 1 + } + + private async cellHasActiveFence(cellId: string): Promise { + const rows = await this.database.query( + `SELECT fence.cell_id FROM relay_cell_fences fence + JOIN relay_cells cell ON cell.cell_id = fence.cell_id + JOIN relay_cell_runtime runtime ON runtime.cell_id = fence.cell_id + WHERE fence.cell_id = ? AND cell.enabled = 0 + AND fence.cell_incarnation = runtime.cell_incarnation + AND fence.attested_at >= runtime.last_heartbeat_at + AND fence.expires_at > ?`, + [cellId, this.now()] + ) + return rows.length === 1 + } + + private async recordCommittedCellFence( + transaction: RelayDatabase, + cellId: string, + cellIncarnation: string, + attemptId: string, + attestedAt: number, + expiresAt: number + ): Promise { + await transaction.query( + `INSERT INTO relay_cell_committed_fences + (cell_id, attempt_id, cell_incarnation, attested_at, expires_at) + VALUES (?, ?, ?, ?, ?) + ON CONFLICT (cell_id) DO UPDATE SET + attempt_id = excluded.attempt_id, + cell_incarnation = excluded.cell_incarnation, + attested_at = excluded.attested_at, + expires_at = excluded.expires_at`, + [cellId, attemptId, cellIncarnation, attestedAt, expiresAt] + ) + } + + private async cellHasCommittedFence( + database: RelayDatabase, + cellId: string, + now: number + ): Promise { + const rows = await database.query( + `SELECT committed.cell_id + FROM relay_cell_committed_fences committed + JOIN relay_cell_fence_attempts attempt + ON attempt.attempt_id = committed.attempt_id + JOIN relay_cell_fences fence ON fence.cell_id = committed.cell_id + JOIN relay_cell_runtime runtime ON runtime.cell_id = committed.cell_id + JOIN relay_cells cell ON cell.cell_id = committed.cell_id + WHERE committed.cell_id = ? + AND attempt.completed_at IS NOT NULL + AND attempt.aborted_at IS NULL + AND cell.enabled = 0 + AND committed.cell_incarnation = runtime.cell_incarnation + AND fence.cell_incarnation = committed.cell_incarnation + AND committed.attested_at >= runtime.last_heartbeat_at + AND committed.expires_at > ? + AND fence.expires_at > ?`, + [cellId, now, now] + ) + return rows.length === 1 + } + + private async deadCellRequiresCommittedFence( + database: RelayDatabase, + identity: AssignmentIdentity, + cellId: string, + assignmentEpoch: number + ): Promise { + const row = ( + await database.query( + `SELECT + CASE WHEN EXISTS ( + SELECT 1 FROM relay_cell_connection_limits limits + WHERE limits.cell_id = ? + ) THEN 1 ELSE 0 END AS capped, + CASE WHEN EXISTS ( + SELECT 1 FROM relay_post_drain_migration_pins pin + JOIN relay_assignment_migrations migration + ON migration.user_id = pin.user_id + AND migration.relay_host_id = pin.relay_host_id + AND migration.assignment_epoch = pin.assignment_epoch + WHERE pin.user_id = ? AND pin.relay_host_id = ? + AND pin.assignment_epoch = ? + AND pin.target_cell_id = ? + AND migration.completed_at IS NULL + AND migration.aborted_at IS NULL + ) THEN 1 ELSE 0 END AS pinned`, + [cellId, identity.userId, identity.relayHostId, assignmentEpoch, cellId] + ) + )[0] + if (!row) return false + return integer(row, 'capped') === 1 || integer(row, 'pinned') === 1 + } + + private async assertDrainCellGeneration( + transaction: RelayDatabase, + cellId: string, + cellIncarnation: string + ): Promise { + const cell = ( + await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [cellId]) + )[0] + const runtime = ( + await transaction.queryLocked(`SELECT * FROM relay_cell_runtime WHERE cell_id = ?`, [ + cellId + ]) + )[0] + if (!cell || integer(cell, 'enabled') !== 0) { + throw new Error('drain_attempt_admission_enabled') + } + if (!runtime || text(runtime, 'cell_incarnation') !== cellIncarnation) { + throw new Error('drain_attempt_generation_mismatch') + } + } + + private async requireCellFence( + transaction: RelayDatabase, + cellId: string, + runtime: SqlRow | undefined, + now: number + ): Promise { + const fence = ( + await transaction.queryLocked(`SELECT * FROM relay_cell_fences WHERE cell_id = ?`, [ + cellId + ]) + )[0] + if ( + !fence || + !runtime || + text(fence, 'cell_incarnation') !== text(runtime, 'cell_incarnation') || + integer(fence, 'attested_at') < integer(runtime, 'last_heartbeat_at') || + integer(fence, 'expires_at') <= now + ) { + throw new Error('cell_fence_attestation_missing') + } + } + + private async activeCellMigrations(sourceCellId: string, targetCellId: string): Promise { + const targetRuntimeSafety = this.requireLiveCells + ? `AND EXISTS ( + SELECT 1 FROM relay_cell_runtime target_runtime + WHERE target_runtime.cell_id = migration.target_cell_id + AND lease.updated_at >= target_runtime.started_at + )` + : '' + return await this.database.query( + `SELECT migration.*, assignment.cell_id AS current_cell_id, + assignment.assignment_epoch AS current_assignment_epoch, + COALESCE(( + SELECT SUM(source_lease.request_units) + FROM relay_assignment_activity_leases source_lease + WHERE source_lease.user_id = migration.user_id + AND source_lease.relay_host_id = migration.relay_host_id + AND source_lease.cell_id = migration.source_cell_id + ), 0) AS source_activity_units, + CASE WHEN EXISTS ( + SELECT 1 FROM relay_assignment_activity_leases lease + WHERE lease.user_id = migration.user_id + AND lease.relay_host_id = migration.relay_host_id + AND lease.cell_id = migration.target_cell_id + AND lease.activity_kind = 'control' + AND lease.activity_id NOT LIKE 'control-pending:%' + ) THEN 1 ELSE 0 END AS target_control_active, + CASE WHEN EXISTS ( + SELECT 1 FROM relay_assignment_activity_leases lease + WHERE lease.user_id = migration.user_id + AND lease.relay_host_id = migration.relay_host_id + AND lease.cell_id = migration.target_cell_id + AND lease.activity_kind = 'control' + AND lease.activity_id NOT LIKE 'control-pending:%' + AND lease.expires_at > ? + ${targetRuntimeSafety} + ) THEN 1 ELSE 0 END AS target_control_current + FROM relay_assignment_migrations migration + LEFT JOIN relay_assignments assignment + ON assignment.user_id = migration.user_id + AND assignment.relay_host_id = migration.relay_host_id + WHERE migration.source_cell_id = ? AND migration.target_cell_id = ? + AND migration.completed_at IS NULL AND migration.aborted_at IS NULL + ORDER BY migration.user_id, migration.relay_host_id`, + [this.now(), sourceCellId, targetCellId] + ) + } + + private async insertPendingControlLease( + database: RelayDatabase, + identity: AssignmentIdentity, + cellId: string, + assignmentEpoch: number, + now: number + ): Promise { + const timeoutAt = now + ASSIGNMENT_LIMITS.activityLeaseMs + await database.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT (user_id, relay_host_id, activity_id) DO UPDATE SET + expires_at = excluded.expires_at, updated_at = excluded.updated_at`, + [ + identity.userId, + identity.relayHostId, + pendingControlActivityId(assignmentEpoch), + 'control', + cellId, + 1, + timeoutAt, + now + ] + ) + await this.insertControlConnectionReservation( + database, + identity, + cellId, + assignmentEpoch, + timeoutAt, + now + ) + } + + private async insertControlConnectionReservation( + database: RelayDatabase, + identity: AssignmentIdentity, + cellId: string, + assignmentEpoch: number, + timeoutAt: number, + now: number + ): Promise { + const existing = await database.queryLocked( + `SELECT reservation_id + FROM relay_control_connection_reservations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ? + AND cell_id = ? AND state IN ('reserved', 'late-arrival-debt') + AND claim_activity_id IS NULL + ORDER BY created_at ASC, reservation_id ASC + LIMIT 1`, + [identity.userId, identity.relayHostId, assignmentEpoch, cellId] + ) + if (existing.length > 0) return + const reservationId = randomUUID() + await database.query( + `INSERT INTO relay_control_connection_reservations + (reservation_id, idempotency_key, user_id, relay_host_id, + assignment_epoch, cell_id, state, inclusion_watermark, + claim_activity_id, created_at, timeout_at, claimed_at, released_at, + updated_at) + SELECT ?, ?, ?, ?, ?, ?, 'reserved', NULL, NULL, ?, ?, NULL, NULL, ? + FROM relay_cell_connection_limits WHERE cell_id = ?`, + [ + reservationId, + reservationId, + identity.userId, + identity.relayHostId, + assignmentEpoch, + cellId, + now, + timeoutAt, + now, + cellId + ] + ) + } + + private async refreshPendingControlReservation( + database: RelayDatabase, + identity: AssignmentIdentity, + cellId: string, + assignmentEpoch: number, + timeoutAt: number, + now: number + ): Promise { + await database.query( + `UPDATE relay_control_connection_reservations + SET state = 'reserved', timeout_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ? + AND cell_id = ? AND state IN ('reserved', 'late-arrival-debt') + AND claim_activity_id IS NULL`, + [ + timeoutAt, + now, + identity.userId, + identity.relayHostId, + assignmentEpoch, + cellId + ] + ) + } + + private async claimControlConnectionReservation( + database: RelayDatabase, + identity: AssignmentIdentity, + cellId: string, + assignmentEpoch: number, + activityId: string, + inclusionWatermark: number | undefined, + now: number + ): Promise { + await database.query( + `UPDATE relay_control_connection_reservations + SET claim_activity_id = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ? + AND cell_id = ? AND state = 'claimed' + AND claim_activity_id IS NOT NULL AND claim_activity_id <> ?`, + [ + activityId, + now, + identity.userId, + identity.relayHostId, + assignmentEpoch, + cellId, + activityId + ] + ) + // A released row's claim_activity_id is history, not a live claim: a control that + // reconnects under the same generation would otherwise leave its fresh reservation + // unclaimable, decaying through debt while the connection is already enforced. + const alreadyClaimed = await database.queryLocked( + `SELECT reservation_id FROM relay_control_connection_reservations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ? + AND cell_id = ? AND claim_activity_id = ? AND state <> 'released' + ORDER BY created_at ASC, reservation_id ASC`, + [ + identity.userId, + identity.relayHostId, + assignmentEpoch, + cellId, + activityId + ] + ) + if (alreadyClaimed.length > 0) return + const reservation = ( + await database.queryLocked( + `SELECT * FROM relay_control_connection_reservations + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ? + AND cell_id = ? AND state IN ('reserved', 'late-arrival-debt') + AND claim_activity_id IS NULL + ORDER BY created_at ASC, reservation_id ASC`, + [identity.userId, identity.relayHostId, assignmentEpoch, cellId] + ) + )[0] + if (!reservation) return + await database.query( + `UPDATE relay_control_connection_reservations + SET state = 'claimed', inclusion_watermark = ?, claim_activity_id = ?, + claimed_at = ?, updated_at = ? + WHERE reservation_id = ?`, + [ + inclusionWatermark, + activityId, + now, + now, + text(reservation, 'reservation_id') + ] + ) + } + + private async releaseSupersededControlConnectionReservations( + database: RelayDatabase, + identity: AssignmentIdentity, + cellId: string, + assignmentEpoch: number, + now: number + ): Promise { + await database.query( + `UPDATE relay_control_connection_reservations + SET state = 'released', released_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND cell_id = ? + AND assignment_epoch = ? AND state <> 'released'`, + [ + now, + now, + identity.userId, + identity.relayHostId, + cellId, + assignmentEpoch + ] + ) + } + + private async assignmentRow( + database: RelayDatabase, + identity: AssignmentIdentity, + failIfUnavailable = false + ): Promise { + return ( + await database.queryLocked( + `SELECT * FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId], + { failIfUnavailable } + ) + )[0] + } + + private async lockAssignmentActivities( + database: RelayDatabase, + identity: AssignmentIdentity, + failIfUnavailable = false + ): Promise { + // Lease rows always precede the globally ordered capacity rows. + return await database.queryLocked( + `SELECT * FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? ORDER BY activity_id ASC`, + [identity.userId, identity.relayHostId], + { failIfUnavailable } + ) + } + + private async removeActivityLease( + database: RelayDatabase, + identity: AssignmentIdentity, + lease: SqlRow, + now: number + ): Promise { + const kind = activityKind(lease) + await this.adjustCellReservation( + database, + text(lease, 'cell_id'), + -integer(lease, 'request_units') + ) + await database.query( + `DELETE FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? AND activity_id = ?`, + [identity.userId, identity.relayHostId, text(lease, 'activity_id')] + ) + await this.adjustActivityCount(database, identity, kind, -1, now, now) + } + + private async removeSupersededSameCellControls( + database: RelayDatabase, + identity: AssignmentIdentity, + leases: SqlRow[], + cellId: string, + retainedActivityId: string, + now: number + ): Promise { + const superseded = leases.filter( + (lease) => + activityKind(lease) === 'control' && + text(lease, 'cell_id') === cellId && + text(lease, 'activity_id') !== retainedActivityId + ) + if (superseded.length === 0) return + if ( + superseded.some( + (lease) => integer(lease, 'request_units') !== ACTIVITY_REQUEST_UNITS.control + ) + ) { + throw new Error('activity_lease_shape_mismatch') + } + const cells = await this.lockCellInventory(database) + await database.query( + `DELETE FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? AND activity_kind = 'control' + AND cell_id = ? AND activity_id <> ?`, + [identity.userId, identity.relayHostId, cellId, retainedActivityId] + ) + const remainingControls = + leases.filter((lease) => activityKind(lease) === 'control').length - superseded.length + await database.query( + `UPDATE relay_assignments SET reserved_controls = ?, last_activity_at = ? + WHERE user_id = ? AND relay_host_id = ?`, + [remainingControls, now, identity.userId, identity.relayHostId] + ) + const cellUnitsRow = ( + await database.query( + `SELECT COALESCE(SUM(request_units), 0) AS request_units + FROM relay_assignment_activity_leases WHERE cell_id = ?`, + [cellId] + ) + )[0]! + const cellRow = cells.find((cell) => text(cell, 'cell_id') === cellId) + const cellUnits = integer(cellUnitsRow, 'request_units') + if (!cellRow) throw new Error('assigned_cell_missing') + if (cellUnits > integer(cellRow, 'capacity_requests')) { + throw new Error('relay_capacity_exhausted') + } + await database.query( + `UPDATE relay_cells SET reserved_requests = ?, updated_at = ? WHERE cell_id = ?`, + [cellUnits, now, cellId] + ) + } + + private async adjustActivityCount( + database: RelayDatabase, + identity: AssignmentIdentity, + kind: AssignmentActivityKind, + delta: 1 | -1, + leaseExpiresAt: number, + now: number + ): Promise { + const column = ACTIVITY_COLUMN[kind] + await database.query( + `UPDATE relay_assignments SET ${column} = + CASE WHEN ${column} + ? < 0 THEN 0 ELSE ${column} + ? END, + lease_expires_at = + CASE WHEN lease_expires_at > ? THEN lease_expires_at ELSE ? END, + last_activity_at = ? + WHERE user_id = ? AND relay_host_id = ?`, + [delta, delta, leaseExpiresAt, leaseExpiresAt, now, identity.userId, identity.relayHostId] + ) + } + + private async touchAssignment( + database: RelayDatabase, + identity: AssignmentIdentity, + leaseExpiresAt: number, + now: number + ): Promise { + await database.query( + `UPDATE relay_assignments SET lease_expires_at = + CASE WHEN lease_expires_at > ? THEN lease_expires_at ELSE ? END, + last_activity_at = ? + WHERE user_id = ? AND relay_host_id = ?`, + [leaseExpiresAt, leaseExpiresAt, now, identity.userId, identity.relayHostId] + ) + } + + private async adjustCellReservation( + database: RelayDatabase, + cellId: string, + delta: number + ): Promise { + const row = (await database.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [cellId]))[0] + if (!row) throw new Error('assigned_cell_missing') + const next = integer(row, 'reserved_requests') + delta + if (next > integer(row, 'capacity_requests')) throw new Error('relay_capacity_exhausted') + await database.query( + `UPDATE relay_cells SET reserved_requests = + CASE WHEN reserved_requests + ? < 0 THEN 0 ELSE reserved_requests + ? END, + updated_at = ? WHERE cell_id = ?`, + [delta, delta, this.now(), cellId] + ) + } + + private async adjustCellReservationAtomically( + database: RelayDatabase, + cellId: string, + delta: number + ): Promise { + const rows = await database.query( + `UPDATE relay_cells SET reserved_requests = + CASE WHEN reserved_requests + ? < 0 THEN 0 ELSE reserved_requests + ? END, + updated_at = ? + WHERE cell_id = ? + AND (? <= 0 OR reserved_requests + ? <= capacity_requests) + RETURNING cell_id`, + [delta, delta, this.now(), cellId, delta, delta] + ) + if (rows.length > 0) return + const cell = ( + await database.query(`SELECT cell_id FROM relay_cells WHERE cell_id = ?`, [cellId]) + )[0] + if (!cell) throw new Error('assigned_cell_missing') + throw new Error('relay_capacity_exhausted') + } + + private result( + identity: AssignmentIdentity, + row: SqlRow, + cellRow: CellRow, + leaseExpiresAt: number + ): RelayAssignment { + return { + ...identity, + ...cellRow, + assignmentEpoch: integer(row, 'assignment_epoch'), + leaseExpiresAt + } + } +} + +type CellRow = { cellId: string; cellUrl: string; region: RelayRegion } + +function cell(row: SqlRow, region: RelayRegion): CellRow { + return { cellId: text(row, 'cell_id'), cellUrl: text(row, 'cell_url'), region } +} + +function relayRegion(row: SqlRow, field: string): RelayRegion { + const value = text(row, field) + if (!RELAY_REGIONS.includes(value as RelayRegion)) throw new Error(`invalid region field ${field}`) + return value as RelayRegion +} + +function optionalRelayRegion(row: SqlRow, field: string): RelayRegion | undefined { + const value = optionalText(row, field) + if (value === undefined) return undefined + if (!RELAY_REGIONS.includes(value as RelayRegion)) throw new Error(`invalid region field ${field}`) + return value as RelayRegion +} + +function activityLeaseById(rows: SqlRow[], activityId: string): SqlRow | undefined { + return rows.find((row) => text(row, 'activity_id') === activityId) +} + +// Why: mid-rehome a host legitimately holds the source control plus the +// target's pending control, so only the lease rows — never the counter a grant +// would overwrite — can say whether this cell is already reserved for it. +function holdsControlLease( + rows: SqlRow[], + cellId: string, + assignmentEpoch: number +): boolean { + const pendingId = pendingControlActivityId(assignmentEpoch) + return rows.some((row) => { + if (activityKind(row) !== 'control' || text(row, 'cell_id') !== cellId) return false + const activityId = text(row, 'activity_id') + return activityId === pendingId || !activityId.startsWith('control-pending:') + }) +} + +function activityCounts(rows: SqlRow[]): Record { + const counts = emptyActivityCounts() + for (const row of rows) counts[activityKind(row)]++ + return counts +} + +function activityUnitsForCell(rows: SqlRow[], cellId: string): number { + return rows.reduce( + (total, row) => + text(row, 'cell_id') === cellId ? total + integer(row, 'request_units') : total, + 0 + ) +} + +function activity(row: SqlRow) { + return { + relayHostId: text(row, 'relay_host_id'), + cellId: text(row, 'cell_id'), + assignmentEpoch: integer(row, 'assignment_epoch'), + leaseExpiresAt: integer(row, 'lease_expires_at'), + lastActivityAt: integer(row, 'last_activity_at'), + reservedControls: integer(row, 'reserved_controls'), + reservedSplices: integer(row, 'reserved_splices'), + reservedInvites: integer(row, 'reserved_invites'), + pendingInstalls: integer(row, 'pending_installs'), + pendingConfirmations: integer(row, 'pending_confirmations'), + migrationLeases: integer(row, 'migration_leases') + } +} + +function requestUnits(row: SqlRow): number { + return ( + integer(row, 'reserved_controls') + + 2 * integer(row, 'reserved_splices') + + integer(row, 'reserved_invites') + + integer(row, 'pending_installs') + + integer(row, 'pending_confirmations') + + integer(row, 'migration_leases') + ) +} + +function integer(row: SqlRow, field: string): number { + const value = Number(row[field]) + if (!Number.isSafeInteger(value)) throw new Error(`invalid_${field}`) + return value +} + +function text(row: SqlRow, field: string): string { + const value = row[field] + if (typeof value !== 'string') throw new Error(`invalid_${field}`) + return value +} + +function cellFenceAttempt(row: SqlRow): CellFenceAttempt { + const environment = text(row, 'environment') + if (!['staging', 'production'].includes(environment)) { + throw new Error('invalid_environment') + } + return { + attemptId: text(row, 'attempt_id'), + environment: environment as CellFenceAttempt['environment'], + cellId: text(row, 'cell_id'), + cellIncarnation: text(row, 'cell_incarnation'), + migName: text(row, 'mig_name'), + instanceGroup: text(row, 'instance_group'), + generationIdentity: text(row, 'generation_identity'), + fenceCommit: text(row, 'fence_commit'), + planSha256: text(row, 'plan_sha256'), + planObjectName: text(row, 'plan_object_name'), + planObjectGeneration: optionalText(row, 'plan_object_generation'), + varFileSha256: text(row, 'var_file_sha256'), + terraformStateLineage: text(row, 'terraform_state_lineage'), + terraformStateSerial: integer(row, 'terraform_state_serial'), + terraformStateObjectGeneration: text( + row, + 'terraform_state_object_generation' + ), + terraformStateObjectSha256: text(row, 'terraform_state_object_sha256'), + requestReason: text(row, 'request_reason'), + gceOperation: optionalText(row, 'gce_operation'), + createdAt: integer(row, 'created_at'), + expiresAt: integer(row, 'expires_at'), + applyStartedAt: optionalInteger(row, 'apply_started_at'), + completedAt: optionalInteger(row, 'completed_at'), + abortedAt: optionalInteger(row, 'aborted_at') + } +} + +function cellFenceApplyInvocation(row: SqlRow): CellFenceApplyInvocation { + return { + invocationId: text(row, 'invocation_id'), + requestReason: text(row, 'request_reason'), + startedAt: integer(row, 'started_at'), + gceOperation: optionalText(row, 'gce_operation') + } +} + +function assertCellFenceAttemptBase( + row: SqlRow, + input: CellFenceAttemptEvidence +): void { + const fields: [keyof CellFenceAttemptEvidence, string][] = [ + ['attemptId', 'attempt_id'], + ['environment', 'environment'], + ['cellId', 'cell_id'], + ['cellIncarnation', 'cell_incarnation'], + ['migName', 'mig_name'], + ['instanceGroup', 'instance_group'], + ['generationIdentity', 'generation_identity'], + ['fenceCommit', 'fence_commit'], + ['planSha256', 'plan_sha256'], + ['planObjectName', 'plan_object_name'], + ['varFileSha256', 'var_file_sha256'], + ['terraformStateLineage', 'terraform_state_lineage'], + ['terraformStateObjectGeneration', 'terraform_state_object_generation'], + ['terraformStateObjectSha256', 'terraform_state_object_sha256'], + ['requestReason', 'request_reason'] + ] + if ( + fields.some(([inputField, rowField]) => input[inputField] !== text(row, rowField)) || + input.terraformStateSerial !== integer(row, 'terraform_state_serial') + ) { + throw new Error('cell_fence_attempt_evidence_mismatch') + } +} + +function assertCellFenceAttemptEvidence( + row: SqlRow, + input: CellFenceAttemptEvidence +): void { + assertCellFenceAttemptBase(row, input) + if ( + !input.planObjectGeneration || + input.planObjectGeneration !== optionalText(row, 'plan_object_generation') + ) { + throw new Error('cell_fence_attempt_evidence_mismatch') + } +} + +function assertActiveCellFenceAttempt( + row: SqlRow, + input: CellFenceAttemptEvidence, + now: number +): void { + assertCellFenceAttemptEvidence(row, input) + if ( + integer(row, 'expires_at') <= now && + optionalInteger(row, 'apply_started_at') === undefined + ) { + throw new Error('cell_fence_attempt_expired') + } + if (row.completed_at !== null) throw new Error('cell_fence_attempt_completed') + if (row.aborted_at !== null) throw new Error('cell_fence_attempt_aborted') +} + +async function lockedCellFenceAttempt( + transaction: RelayDatabase, + attemptId: string +): Promise { + const row = ( + await transaction.queryLocked( + `${CELL_FENCE_ATTEMPT_SELECT} WHERE attempts.attempt_id = ?`, + [attemptId] + ) + )[0] + if (!row) throw new Error('cell_fence_attempt_not_found') + return row +} + +function pendingControlActivityId(assignmentEpoch: number): string { + return `control-pending:${assignmentEpoch}` +} + +function validateActivityId(activityId: string): void { + if (!activityId || activityId.length > 256) throw new Error('invalid_activity_id') +} + +function activityKind(row: SqlRow): AssignmentActivityKind { + const value = text(row, 'activity_kind') + if (!(value in ACTIVITY_REQUEST_UNITS)) throw new Error('invalid_activity_kind') + return value as AssignmentActivityKind +} + +function migrationActivityId(assignmentEpoch: number): string { + return `migration:${assignmentEpoch}` +} + +function isIncompleteMigration(error: unknown): boolean { + return ( + error instanceof Error && + [ + 'migration_target_not_registered', + 'migration_source_still_active', + 'migration_target_not_active', + 'migration_assignment_mismatch', + 'migration_source_admission_changed', + 'migration_target_admission_changed', + 'migration_source_runtime_not_quiescent', + 'migration_source_runtime_not_dead', + 'migration_target_runtime_not_ready', + 'migration_cell_inventory_busy', + 'migration_activity_accounting_mismatch', + 'migration_activity_topology_mismatch', + 'migration_activity_lease_shape_mismatch', + 'migration_cell_reservation_accounting_mismatch', + 'cell_fence_attestation_missing' + ].includes(error.message) + ) +} + +function isMissingMigration(error: unknown): boolean { + return error instanceof Error && error.message === 'migration_not_found' +} + +function isDatabaseLockUnavailable(error: unknown): boolean { + return error instanceof Error && error.message === 'database_lock_unavailable' +} + +function isDatabaseLockTimeout(error: unknown): boolean { + return String((error as { code?: unknown }).code) === '55P03' +} + +async function waitForAssignmentLockRetry(deadline: number): Promise { + const remainingMs = Math.max(0, deadline - Date.now()) + const maxDelayMs = Math.min(ASSIGNMENT_LOCK_RETRY_MAX_DELAY_MS, remainingMs) + const delayMs = Math.floor(Math.random() * (maxDelayMs + 1)) + await new Promise((resolve) => setTimeout(resolve, delayMs)) +} + +function migration(identity: AssignmentIdentity, row: SqlRow): RelayAssignmentMigration { + const targetRegisteredAt = optionalInteger(row, 'target_registered_at') + return { + ...identity, + sourceCellId: text(row, 'source_cell_id'), + targetCellId: text(row, 'target_cell_id'), + previousEpoch: integer(row, 'previous_epoch'), + assignmentEpoch: integer(row, 'assignment_epoch'), + expiresAt: integer(row, 'expires_at'), + ...(targetRegisteredAt === undefined ? {} : { targetRegisteredAt }) + } +} + +// Attempt ids are server-minted UUIDs and this codebase's invariant messages +// are snake_case slugs; anything else could carry secrets and logs redacted. +function warnRegionalRehomeCandidateFailure( + operation: 'complete' | 'abort', + attemptId: string, + error: unknown +): void { + const message = error instanceof Error ? error.message : '' + console.warn( + JSON.stringify({ + event: 'orca_relay_regional_rehome_candidate_failed', + operation, + attemptId, + reason: /^[a-z0-9_]{1,64}$/.test(message) ? message : 'redacted' + }) + ) +} + +// The attempt id is the only field: counts would say which host holds what. +function noteRegionalRehomeActivityCountsRepaired(attemptId: string): void { + console.warn( + JSON.stringify({ + event: 'orca_relay_regional_rehome_activity_counts_repaired', + attemptId + }) + ) +} + +function regionalRehomeAttempt(row: SqlRow): RegionalRehomeAttempt { + return { + attemptId: text(row, 'attempt_id'), + userId: text(row, 'user_id'), + relayHostId: text(row, 'relay_host_id'), + preferredRegion: 'asia-east2', + sourceCellId: text(row, 'source_cell_id'), + sourceCellUrl: text(row, 'source_cell_url'), + sourceCellIncarnation: text(row, 'source_cell_incarnation'), + targetCellId: text(row, 'target_cell_id'), + targetCellIncarnation: text(row, 'target_cell_incarnation'), + previousEpoch: integer(row, 'previous_epoch'), + assignmentEpoch: integer(row, 'assignment_epoch'), + drainGraceMs: integer(row, 'drain_grace_ms'), + sendAttempts: integer(row, 'send_attempts') + } +} + +function regionalRehomeControl(row: SqlRow): RegionalRehomeControl { + return { + generation: integer(row, 'generation'), + enabled: integer(row, 'enabled') === 1, + observationStartedAt: integer(row, 'observation_started_at'), + notBefore: integer(row, 'not_before'), + ratePerMinute: integer(row, 'rate_per_minute'), + preferenceMaxAgeMs: integer(row, 'preference_max_age_ms'), + drainGraceMs: integer(row, 'drain_grace_ms') + } +} + +function cleanRegionalRehomeSafety(now: number): RegionalRehomeSafetySnapshot { + return { + observedAt: now, + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + } +} + +function regionalRehomeFleetSafetyFromInventory(input: { + cells: SqlRow[] + admission: ReadonlyMap + regions: ReadonlyMap + runtimes: SqlRow[] + capabilities: SqlRow[] + safetyRows: SqlRow[] + now: number + heartbeatTtlMs: number +}): RegionalRehomeFleetSafety { + const capabilities = new Map( + input.capabilities.map((row) => [text(row, 'cell_id'), row]) + ) + const runtimes = new Map(input.runtimes.map((row) => [text(row, 'cell_id'), row])) + const safetyRows = new Map( + input.safetyRows.map((row) => [text(row, 'cell_id'), row]) + ) + const required = input.cells.filter((row) => { + const cellId = text(row, 'cell_id') + const capability = capabilities.get(cellId) + return ( + integer(row, 'enabled') === 1 && + input.admission.get(cellId) === 'general' && + (input.regions.get(cellId) === 'asia-east2' || + (input.regions.get(cellId) === RELAY_DEFAULT_REGION && + capability !== undefined && + integer(capability, 'regional_rehome_protocol') >= 1)) + ) + }) + const valid = required.flatMap((row) => { + const cellId = text(row, 'cell_id') + const runtime = runtimes.get(cellId) + const safety = safetyRows.get(cellId) + if ( + !runtime || + integer(runtime, 'ready') !== 1 || + integer(runtime, 'last_heartbeat_at') <= input.now - input.heartbeatTtlMs || + !safety || + text(safety, 'cell_incarnation') !== text(runtime, 'cell_incarnation') || + integer(safety, 'observed_at') <= input.now - 60_000 + ) { + return [] + } + return [safety] + }) + const missingCells = required.length === 0 ? 1 : required.length - valid.length + return { + requiredCells: required.length, + missingCells, + observedAt: + missingCells > 0 + ? 0 + : Math.min(...valid.map((row) => integer(row, 'observed_at'))), + sqlFailures: valid.reduce((total, row) => total + integer(row, 'sql_failures'), 0), + reconnects: valid.reduce((total, row) => total + integer(row, 'reconnects'), 0), + maxReconnects: Math.max(0, ...valid.map((row) => integer(row, 'reconnects'))), + controlActivityRecoveryFailures: valid.reduce( + (total, row) => total + integer(row, 'control_activity_recovery_failures'), + 0 + ), + databasePoolWaiting: Math.max( + 0, + ...valid.map((row) => integer(row, 'database_pool_waiting')) + ), + databasePoolWaitersMax: Math.max( + 0, + ...valid.map((row) => integer(row, 'database_pool_waiters_max')) + ), + databasePoolWaitMsMax: Math.max( + 0, + ...valid.map((row) => integer(row, 'database_pool_wait_ms_max')) + ) + } +} + +function regionalRehomeFleetSafetyFailure( + processSafety: RegionalRehomeSafetySnapshot, + fleetSafety: RegionalRehomeFleetSafety, + now: number +): string | null { + if (fleetSafety.maxReconnects > REGIONAL_REHOME_RECONNECTS_PER_CELL_LIMIT) { + return 'elevated_reconnects' + } + return regionalRehomeSafetyFailure( + combineRegionalRehomeSafety(processSafety, fleetSafety), + now, + fleetSafety.requiredCells + ) +} + +type RegionalRehomeCandidateSkip = { + reason: + | 'candidate_stale' + | 'source_ineligible' + | 'source_unclean' + | 'source_control_inactive' + | 'target_unclean' + | 'no_eligible_target' + | 'no_target_headroom' + cellId?: string + sqlFailures?: number + reconnects?: number +} + +function cellUncleanSkip( + reason: 'source_unclean' | 'target_unclean', + cellId: string, + safety: SqlRow | undefined +): RegionalRehomeCandidateSkip { + return { + reason, + cellId, + sqlFailures: safety === undefined ? undefined : integer(safety, 'sql_failures'), + reconnects: safety === undefined ? undefined : integer(safety, 'reconnects') + } +} + +// Candidate skips are otherwise invisible: they neither latch the control off +// nor produce attempts, so an operator cannot tell "skipping" from "idle". +// Cell ids and counters only — never free-form error text. +function aggregateRegionalRehomeCandidateSkips( + skips: readonly RegionalRehomeCandidateSkip[] +): Record { + // `candidates` counts skipped candidate iterations, not distinct cells: one + // unclean cell blocking six candidates reports candidates=6 on one cellId. + const aggregated = new Map() + for (const skip of skips) { + const key = `${skip.reason}:${skip.cellId ?? ''}` + const entry = aggregated.get(key) + if (entry) entry.candidates += 1 + else aggregated.set(key, { ...skip, candidates: 1 }) + } + return { + event: 'orca_relay_regional_rehome_candidates_skipped', + skips: [...aggregated.values()] + } +} + +function regionalRehomeCellSafetyIsClean( + safety: SqlRow | undefined, + runtime: SqlRow, + now: number +): boolean { + return ( + safety !== undefined && + text(safety, 'cell_incarnation') === text(runtime, 'cell_incarnation') && + integer(safety, 'observed_at') > now - 60_000 && + integer(safety, 'sql_failures') <= REGIONAL_REHOME_SQL_FAILURES_PER_CELL_LIMIT && + integer(safety, 'reconnects') <= REGIONAL_REHOME_RECONNECTS_PER_CELL_LIMIT && + integer(safety, 'control_activity_recovery_failures') === 0 && + !regionalRehomePoolPressure({ + databasePoolWaitersMax: integer(safety, 'database_pool_waiters_max'), + databasePoolWaitMsMax: integer(safety, 'database_pool_wait_ms_max') + }) + ) +} + +function assertMigrationPair(row: SqlRow, sourceCellId: string, targetCellId: string): void { + if ( + text(row, 'source_cell_id') !== sourceCellId || + text(row, 'target_cell_id') !== targetCellId + ) { + throw new Error('migration_pair_mismatch') + } +} + +function matchesExactCellSet( + actual: ReadonlySet, + expected: readonly string[] +): boolean { + const expectedSet = new Set(expected) + return ( + expectedSet.size === expected.length && + actual.size === expectedSet.size && + [...actual].every((cellId) => expectedSet.has(cellId)) + ) +} + +function assertCurrentMigrationAssignment( + assignment: SqlRow | undefined, + migrationRow: SqlRow +): asserts assignment is SqlRow { + if ( + !assignment || + text(assignment, 'cell_id') !== text(migrationRow, 'target_cell_id') || + integer(assignment, 'assignment_epoch') !== integer(migrationRow, 'assignment_epoch') + ) { + throw new Error('migration_assignment_mismatch') + } +} + +function assertMigrationRecoveryMetadata( + migrationRow: SqlRow, + incarnation: SqlRow | undefined, + pin: SqlRow | undefined +): void { + if (pin && !incarnation) { + throw new Error('drain_migration_source_incarnation_mismatch') + } + if ( + pin && + incarnation && + (text(pin, 'source_cell_id') !== text(migrationRow, 'source_cell_id') || + text(pin, 'target_cell_id') !== text(migrationRow, 'target_cell_id') || + text(pin, 'source_cell_incarnation') !== + text(incarnation, 'source_cell_incarnation') || + text(pin, 'target_cell_incarnation') !== + text(incarnation, 'target_cell_incarnation') || + integer(pin, 'source_request_units') !== + integer(migrationRow, 'source_request_units') || + integer(pin, 'target_reserved_units') !== + integer(migrationRow, 'target_reserved_units')) + ) { + throw new Error('migration_activity_topology_mismatch') + } +} + +function assertAssignmentActivityAccounting( + assignment: SqlRow, + leases: SqlRow[], + migrationRow: SqlRow +): void { + if ( + integer(migrationRow, 'target_reserved_units') !== + integer(migrationRow, 'source_request_units') + 1 + ) { + throw new Error('migration_activity_lease_shape_mismatch') + } + const counts = emptyActivityCounts() + for (const lease of leases) { + const kind = activityKind(lease) + const requestUnits = integer(lease, 'request_units') + if (kind === 'migration') { + if ( + text(lease, 'activity_id') !== + migrationActivityId(integer(migrationRow, 'assignment_epoch')) || + text(lease, 'cell_id') !== text(migrationRow, 'target_cell_id') || + requestUnits !== integer(migrationRow, 'source_request_units') + ) { + throw new Error('migration_activity_lease_shape_mismatch') + } + } else if (requestUnits !== ACTIVITY_REQUEST_UNITS[kind]) { + throw new Error('migration_activity_lease_shape_mismatch') + } + counts[kind]++ + } + assertAssignmentActivityCounts(assignment, leases, 1) +} + +function assertAssignmentActivityCounts( + assignment: SqlRow, + leases: SqlRow[], + expectedMigrations: number +): void { + const counts = activityCounts(leases) + if ( + (Object.keys(ACTIVITY_COLUMN) as AssignmentActivityKind[]).some( + (kind) => integer(assignment, ACTIVITY_COLUMN[kind]) !== counts[kind] + ) || + counts.migration !== expectedMigrations + ) { + throw new Error('migration_activity_accounting_mismatch') + } +} + +async function assertCellReservationAccounting( + transaction: RelayDatabase, + cells: SqlRow[], + cellIds: string[] +): Promise { + const uniqueCellIds = [...new Set(cellIds)] + const rows = await transaction.query( + `SELECT cell_id, COALESCE(SUM(request_units), 0) AS request_units + FROM relay_assignment_activity_leases + WHERE cell_id IN (${uniqueCellIds.map(() => '?').join(', ')}) + GROUP BY cell_id`, + uniqueCellIds + ) + const unitsByCell = new Map( + rows.map((row) => [text(row, 'cell_id'), integer(row, 'request_units')]) + ) + for (const cellId of uniqueCellIds) { + const cell = cells.find((row) => text(row, 'cell_id') === cellId) + if (!cell || integer(cell, 'reserved_requests') !== (unitsByCell.get(cellId) ?? 0)) { + throw new Error('migration_cell_reservation_accounting_mismatch') + } + } +} + +function deadSourceCompletionResult( + row: SqlRow, + changed: boolean +): DeadSourceCompletionResult { + return { + changed, + assignmentEpoch: integer(row, 'assignment_epoch'), + sourceCellId: text(row, 'source_cell_id'), + targetCellId: text(row, 'target_cell_id') + } +} + +function cellDrainAttempt(row: SqlRow): CellDrainAttempt { + const state = text(row, 'state') + if ( + ![ + 'prepared', + 'send-may-have-started', + 'application-receipt', + 'proven-not-delivered' + ].includes(state) + ) { + throw new Error('invalid_drain_attempt_state') + } + return { + attemptId: text(row, 'attempt_id'), + cellId: text(row, 'cell_id'), + cellIncarnation: text(row, 'cell_incarnation'), + traceValue: text(row, 'trace_value'), + plannedGraceMs: integer(row, 'planned_grace_ms'), + state: state as CellDrainAttemptState, + preparedAt: integer(row, 'prepared_at'), + sendMayHaveStartedAt: optionalInteger(row, 'send_may_have_started_at'), + sendPermitExpiresAt: optionalInteger(row, 'send_permit_expires_at'), + applicationReceiptAt: optionalInteger(row, 'application_receipt_at'), + backendSuccessStatus: optionalInteger(row, 'backend_success_status'), + backendInstance: optionalText(row, 'backend_instance'), + receiptCellIncarnation: optionalText(row, 'receipt_cell_incarnation'), + retryAfter: optionalInteger(row, 'retry_after'), + recoverForwardAttemptedAt: optionalInteger(row, 'recover_forward_attempted_at'), + provenNotDeliveredAt: optionalInteger(row, 'proven_not_delivered_at') + } +} + +function optionalInteger(row: SqlRow, field: string): number | undefined { + if (row[field] === null || row[field] === undefined) return undefined + return integer(row, field) +} + +function optionalText(row: SqlRow, field: string): string | undefined { + if (row[field] === null || row[field] === undefined) return undefined + return text(row, field) +} + +function migrationHasExactActiveTarget(row: SqlRow, targetCellId: string): boolean { + return ( + integer(row, 'target_control_active') === 1 && + optionalText(row, 'current_cell_id') === targetCellId && + optionalInteger(row, 'current_assignment_epoch') === integer(row, 'assignment_epoch') + ) +} + +function assignmentKey(userId: string, relayHostId: string): string { + return `${userId}\u0000${relayHostId}` +} + +function emptyActivityCounts(): Record { + return { + control: 0, + splice: 0, + invite: 0, + install: 0, + confirmation: 0, + migration: 0 + } +} diff --git a/cloud/apps/relay/src/cell-admission-migration-registration.ts b/cloud/apps/relay/src/cell-admission-migration-registration.ts new file mode 100644 index 00000000000..a6cd356a42b --- /dev/null +++ b/cloud/apps/relay/src/cell-admission-migration-registration.ts @@ -0,0 +1,346 @@ +import { + isRelayCellConnectionHardCap, + RELAY_DEFAULT_REGION, + RELAY_REGIONS, + relayCellAdmissionBounds, + type RelayCellConnectionHardCap, + type RelayRegion +} from '@orca-cloud/relay-contract' +import { + decodeMembership, + encodeMembership, + lockedSelector, + normalizeMembership, + requireSelectorMatchesAdmission, + type CellAdmissionMembership, + type CellAdmissionSelector +} from './cell-admission-selector.js' +import type { RelayDatabase, SqlRow } from './database.js' + +export type MigrationCellRegistration = { + id: string + url: string + capacityRequests: number + connectionHardCap: RelayCellConnectionHardCap + connectionUnobservedBound: number + region?: RelayRegion +} + +type NormalizedMigrationCellRegistration = MigrationCellRegistration & { region: RelayRegion } + +type AddMigrationCellsInput = { + attemptId: string + expectedGeneration: number + cells: MigrationCellRegistration[] +} + +const SELECTOR_ID = 'general' + +export class RelayMigrationCellRegistrar { + constructor( + private readonly database: RelayDatabase, + private readonly now: () => number + ) {} + + async add(input: AddMigrationCellsInput): Promise<{ + changed: boolean + selector: CellAdmissionSelector + }> { + validateAttempt(input) + const cells = normalizeMigrationCells(input.cells) + const encodedCells = encodeCells(cells) + return await this.database.transaction(async (transaction) => { + const inventory = await transaction.queryLocked( + `SELECT * FROM relay_cells ORDER BY cell_id ASC` + ) + const selector = await lockedSelector(transaction) + await requireSelectorMatchesAdmission(transaction, selector) + const intent = ( + await transaction.queryLocked( + `SELECT * FROM relay_admission_selector_intents WHERE attempt_id = ?`, + [input.attemptId] + ) + )[0] + const addition = ( + await transaction.queryLocked( + `SELECT * FROM relay_admission_selector_cell_additions WHERE attempt_id = ?`, + [input.attemptId] + ) + )[0] + if (Boolean(intent) !== Boolean(addition)) { + throw new Error('admission_selector_attempt_mismatch') + } + if (intent && addition) { + const membership = decodeMembership(text(intent, 'membership_json')) + if ( + integer(intent, 'expected_generation') !== input.expectedGeneration || + integer(intent, 'intended_generation') !== input.expectedGeneration + 1 || + text(addition, 'cells_json') !== encodedCells || + encodeMembership(membership) !== + encodeMembership( + membershipAfterAddition( + decodeMembership(text(intent, 'previous_membership_json')), + cells + ) + ) + ) { + throw new Error('admission_selector_attempt_mismatch') + } + if ( + selector.generation === input.expectedGeneration + 1 && + selector.attemptId === input.attemptId && + encodeMembership(selector.membership) === encodeMembership(membership) + ) { + await requireExactAddedCells(transaction, inventory, cells) + return { changed: false, selector } + } + throw new Error('admission_selector_generation_mismatch') + } + if (selector.generation < 1) { + throw new Error('admission_selector_boundary_inactive') + } + if (selector.generation !== input.expectedGeneration) { + throw new Error('admission_selector_generation_mismatch') + } + const inventoryIds = new Set(inventory.map((row) => text(row, 'cell_id'))) + const inventoryUrls = new Set(inventory.map((row) => text(row, 'cell_url'))) + if (cells.some((cell) => inventoryIds.has(cell.id) || inventoryUrls.has(cell.url))) { + throw new Error('admission_selector_cell_already_exists') + } + return await this.commit(transaction, input, cells, selector, encodedCells) + }) + } + + private async commit( + transaction: RelayDatabase, + input: AddMigrationCellsInput, + cells: NormalizedMigrationCellRegistration[], + selector: CellAdmissionSelector, + encodedCells: string + ): Promise<{ changed: true; selector: CellAdmissionSelector }> { + const membership = membershipAfterAddition(selector.membership, cells) + const encodedMembership = encodeMembership(membership) + const now = this.now() + await transaction.query( + `INSERT INTO relay_admission_selector_intents + (attempt_id, expected_generation, intended_generation, previous_membership_json, + membership_json, created_at, committed_at) + VALUES (?, ?, ?, ?, ?, ?, NULL)`, + [ + input.attemptId, + input.expectedGeneration, + input.expectedGeneration + 1, + encodeMembership(selector.membership), + encodedMembership, + now + ] + ) + await transaction.query( + `INSERT INTO relay_admission_selector_cell_additions (attempt_id, cells_json) + VALUES (?, ?)`, + [input.attemptId, encodedCells] + ) + for (const cell of cells) await insertCell(transaction, cell, now) + const intendedGeneration = input.expectedGeneration + 1 + const updated = await transaction.query( + `UPDATE relay_admission_selectors + SET generation = ?, attempt_id = ?, membership_json = ?, updated_at = ? + WHERE selector_id = ? AND generation = ?`, + [ + intendedGeneration, + input.attemptId, + encodedMembership, + now, + SELECTOR_ID, + input.expectedGeneration + ] + ) + if (integer(updated[0]!, 'changes') !== 1) { + throw new Error('admission_selector_generation_mismatch') + } + await transaction.query( + `UPDATE relay_admission_selector_intents SET committed_at = ? + WHERE attempt_id = ?`, + [now, input.attemptId] + ) + return { + changed: true, + selector: { + generation: intendedGeneration, + attemptId: input.attemptId, + membership + } + } + } +} + +function encodeCells(cells: NormalizedMigrationCellRegistration[]): string { + return JSON.stringify( + cells.map((cell) => + cell.region === RELAY_DEFAULT_REGION + ? { + id: cell.id, + url: cell.url, + capacityRequests: cell.capacityRequests, + connectionHardCap: cell.connectionHardCap, + connectionUnobservedBound: cell.connectionUnobservedBound + } + : cell + ) + ) +} + +async function insertCell( + transaction: RelayDatabase, + cell: NormalizedMigrationCellRegistration, + now: number +): Promise { + await transaction.query( + `INSERT INTO relay_cells + (cell_id, cell_url, enabled, capacity_requests, reserved_requests, + observed_requests, last_heartbeat_at, updated_at) + VALUES (?, ?, 1, ?, 0, 0, ?, ?)`, + [cell.id, cell.url, cell.capacityRequests, now, now] + ) + await transaction.query( + `INSERT INTO relay_cell_regions (cell_id, region) VALUES (?, ?)`, + [cell.id, cell.region] + ) + await transaction.query( + `INSERT INTO relay_cell_admission (cell_id, admission_state, updated_at) + VALUES (?, 'migration-only', ?)`, + [cell.id, now] + ) + await transaction.query( + `INSERT INTO relay_cell_connection_limits + (cell_id, hard_cap, unobserved_bound, updated_at) + VALUES (?, ?, ?, ?)`, + [cell.id, cell.connectionHardCap, cell.connectionUnobservedBound, now] + ) +} + +function normalizeMigrationCells( + input: MigrationCellRegistration[] +): NormalizedMigrationCellRegistration[] { + if (input.length === 0 || input.length > 128) { + throw new Error('invalid_admission_selector_cells') + } + const cells = input + .map((cell) => ({ ...cell, region: cell.region ?? RELAY_DEFAULT_REGION })) + .sort((left, right) => left.id.localeCompare(right.id)) + if ( + new Set(cells.map(({ id }) => id)).size !== cells.length || + new Set(cells.map(({ url }) => url)).size !== cells.length + ) { + throw new Error('admission_selector_duplicate_cell') + } + for (const cell of cells) validateMigrationCell(cell) + return cells +} + +function validateMigrationCell(cell: NormalizedMigrationCellRegistration): void { + if ( + cell.id.length === 0 || + cell.id.length > 128 || + cell.url.length === 0 || + cell.url.length > 2_048 || + !Number.isSafeInteger(cell.capacityRequests) || + cell.capacityRequests <= 0 || + cell.capacityRequests > 100_000 || + !isRelayCellConnectionHardCap(cell.connectionHardCap) || + !Number.isSafeInteger(cell.connectionUnobservedBound) || + cell.connectionUnobservedBound < 0 || + cell.connectionUnobservedBound > + relayCellAdmissionBounds(cell.connectionHardCap).maxUnobservedBound || + !RELAY_REGIONS.includes(cell.region) + ) { + throw new Error('invalid_admission_selector_cell') + } +} + +function membershipAfterAddition( + membership: CellAdmissionMembership, + cells: NormalizedMigrationCellRegistration[] +): CellAdmissionMembership { + return normalizeMembership({ + existingOnly: membership.existingOnly, + migrationOnly: [...membership.migrationOnly, ...cells.map(({ id }) => id)], + general: membership.general + }) +} + +async function requireExactAddedCells( + database: RelayDatabase, + inventory: SqlRow[], + cells: NormalizedMigrationCellRegistration[] +): Promise { + const byId = new Map(inventory.map((row) => [text(row, 'cell_id'), row])) + for (const cell of cells) { + const row = byId.get(cell.id) + const admission = await admissionState(database, cell.id) + const region = await cellRegion(database, cell.id) + const limit = await connectionLimit(database, cell.id) + if ( + !row || + text(row, 'cell_url') !== cell.url || + integer(row, 'enabled') !== 1 || + integer(row, 'capacity_requests') !== cell.capacityRequests || + region !== cell.region || + admission !== 'migration-only' || + !limit || + integer(limit, 'hard_cap') !== cell.connectionHardCap || + integer(limit, 'unobserved_bound') !== cell.connectionUnobservedBound + ) { + throw new Error('admission_selector_attempt_mismatch') + } + } +} + +async function cellRegion(database: RelayDatabase, cellId: string): Promise { + return ( + await database.query(`SELECT region FROM relay_cell_regions WHERE cell_id = ?`, [cellId]) + )[0]?.region +} + +async function admissionState(database: RelayDatabase, cellId: string): Promise { + return ( + await database.query( + `SELECT admission_state FROM relay_cell_admission WHERE cell_id = ?`, + [cellId] + ) + )[0]?.admission_state +} + +async function connectionLimit( + database: RelayDatabase, + cellId: string +): Promise { + return ( + await database.query( + `SELECT hard_cap, unobserved_bound FROM relay_cell_connection_limits + WHERE cell_id = ?`, + [cellId] + ) + )[0] +} + +function validateAttempt(input: AddMigrationCellsInput): void { + if (!/^[A-Za-z0-9_-]{8,128}$/.test(input.attemptId)) { + throw new Error('invalid_admission_selector_attempt') + } + if (!Number.isSafeInteger(input.expectedGeneration) || input.expectedGeneration < 0) { + throw new Error('invalid_admission_selector_generation') + } +} + +function integer(row: SqlRow, field: string): number { + const value = Number(row[field]) + if (!Number.isSafeInteger(value)) throw new Error(`invalid integer field ${field}`) + return value +} + +function text(row: SqlRow, field: string): string { + const value = row[field] + if (typeof value !== 'string') throw new Error(`invalid text field ${field}`) + return value +} diff --git a/cloud/apps/relay/src/cell-admission-selector.test.ts b/cloud/apps/relay/src/cell-admission-selector.test.ts new file mode 100644 index 00000000000..41b84f91d79 --- /dev/null +++ b/cloud/apps/relay/src/cell-admission-selector.test.ts @@ -0,0 +1,580 @@ +import { ASSIGNMENT_LIMITS } from '@orca-cloud/relay-contract' +import { createHash } from 'node:crypto' +import { afterEach, describe, expect, it } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.js' +import { encodeMembership } from './cell-admission-selector.js' +import { + openInMemoryRelayDatabase, + type RelayDatabase, + type RelayLockOptions, + type SqlRow +} from './database.js' + +const CELLS = [ + { id: 'c1', url: 'https://c1.example.com', capacityRequests: 20 }, + { id: 'c2', url: 'https://c2.example.com', capacityRequests: 20 }, + { id: 'c3', url: 'https://c3.example.com', capacityRequests: 20 } +] + +const INITIAL_MEMBERSHIP = { + existingOnly: ['c1'], + migrationOnly: ['c2'], + general: ['c3'] +} +const BASE_MEMBERSHIP = { + existingOnly: [], + migrationOnly: [], + general: ['c1', 'c2', 'c3'] +} + +class FailAfterIntentDatabase implements RelayDatabase { + private transactionsUntilFailure = Number.POSITIVE_INFINITY + + constructor(private readonly delegate: RelayDatabase) {} + + failSecondTransaction(): void { + this.transactionsUntilFailure = 2 + } + + async query(sql: string, params?: unknown[]): Promise { + return await this.delegate.query(sql, params) + } + + async queryLocked( + sql: string, + params?: unknown[], + options?: RelayLockOptions + ): Promise { + return await this.delegate.queryLocked(sql, params, options) + } + + async transaction(operation: (transaction: RelayDatabase) => Promise): Promise { + this.transactionsUntilFailure-- + if (this.transactionsUntilFailure === 0) { + this.transactionsUntilFailure = Number.POSITIVE_INFINITY + throw new Error('injected_commit_ambiguity') + } + return await this.delegate.transaction(operation) + } + + async close(): Promise { + await this.delegate.close() + } +} + +class AfterTransactionDatabase implements RelayDatabase { + private afterTransaction: (() => Promise) | undefined + + constructor(private readonly delegate: RelayDatabase) {} + + runAfterNextTransaction(operation: () => Promise): void { + this.afterTransaction = operation + } + + async query(sql: string, params?: unknown[]): Promise { + return await this.delegate.query(sql, params) + } + + async queryLocked( + sql: string, + params?: unknown[], + options?: RelayLockOptions + ): Promise { + return await this.delegate.queryLocked(sql, params, options) + } + + async transaction(operation: (transaction: RelayDatabase) => Promise): Promise { + const result = await this.delegate.transaction(operation) + const after = this.afterTransaction + this.afterTransaction = undefined + if (after) await after() + return result + } + + async close(): Promise { + await this.delegate.close() + } +} + +function membershipSha256(membership: typeof INITIAL_MEMBERSHIP): string { + return createHash('sha256').update(encodeMembership(membership)).digest('hex') +} + +const BASE_MEMBERSHIP_SHA256 = membershipSha256(BASE_MEMBERSHIP) + +describe('relay cell admission selector', () => { + let database: RelayDatabase | undefined + + afterEach(async () => await database?.close()) + + async function setup( + now: () => number, + requireLiveCells = false + ): Promise { + database = await openInMemoryRelayDatabase() + const store = new RelayAssignmentStore(database, now, { + requireLiveCells, + heartbeatTtlMs: 45_000 + }) + await store.reconcileCells(CELLS) + return store + } + + async function heartbeat(store: RelayAssignmentStore, cellIndex: number): Promise { + const cell = CELLS[cellIndex]! + await store.recordCellHeartbeat({ + cellId: cell.id, + cellUrl: cell.url, + cellIncarnation: `${cellIndex + 1}1111111-1111-4111-8111-111111111111`, + startedAt: 50, + ready: true, + observedRequests: 0 + }) + } + + it('preserves sticky assignments while reserving ordinary placement for general cells', async () => { + const store = await setup(() => 100) + const existing = { userId: 'existing', relayHostId: 'host000000000001' } + expect(await store.assign(existing)).toMatchObject({ cellId: 'c1' }) + + await store.applyCellAdmissionSelector({ + attemptId: 'cutover_001', + expectedGeneration: 0, + expectedMembershipSha256: BASE_MEMBERSHIP_SHA256, + membership: INITIAL_MEMBERSHIP + }) + + expect(await store.assign(existing)).toMatchObject({ cellId: 'c1' }) + await expect( + store.assign({ userId: 'new', relayHostId: 'host000000000002' }) + ).resolves.toMatchObject({ cellId: 'c3' }) + await expect(store.startEvacuation(existing, 'c2')).resolves.toMatchObject({ + sourceCellId: 'c1', + targetCellId: 'c2' + }) + }) + + it('excludes existing-only and migration-only cells from dormant placement', async () => { + let now = 100 + const store = await setup(() => now) + const identity = { userId: 'dormant', relayHostId: 'host000000000001' } + const assignment = await store.assign(identity) + const control = await store.activateControl(identity, { + cellId: assignment.cellId, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + await store.releaseActivity(identity, control) + now += ASSIGNMENT_LIMITS.dormantTtlMs + 1 + + await store.applyCellAdmissionSelector({ + attemptId: 'cutover_002', + expectedGeneration: 0, + expectedMembershipSha256: BASE_MEMBERSHIP_SHA256, + membership: INITIAL_MEMBERSHIP + }) + + await expect(store.rebalanceDormant(identity, 'c2')).rejects.toThrow( + 'target_cell_unavailable' + ) + await expect(store.rebalanceDormant(identity, 'c3')).resolves.toMatchObject({ + cellId: 'c3' + }) + }) + + it('does not recover an unfenced dead existing-only assignment', async () => { + let now = 100 + const store = await setup(() => now, true) + await heartbeat(store, 0) + await heartbeat(store, 1) + await heartbeat(store, 2) + const identity = { userId: 'dead-source', relayHostId: 'host000000000001' } + expect(await store.assign(identity)).toMatchObject({ cellId: 'c1' }) + await store.applyCellAdmissionSelector({ + attemptId: 'cutover_003', + expectedGeneration: 0, + expectedMembershipSha256: BASE_MEMBERSHIP_SHA256, + membership: INITIAL_MEMBERSHIP + }) + + now += 45_001 + await heartbeat(store, 1) + await heartbeat(store, 2) + expect(await store.evacuateDeadCells()).toBe(0) + expect(await store.resolve(identity)).toBeNull() + }) + + it('commits atomically and resolves an ambiguous response idempotently', async () => { + const store = await setup(() => 100) + const input = { + attemptId: 'cutover_004', + expectedGeneration: 0, + expectedMembershipSha256: BASE_MEMBERSHIP_SHA256, + membership: INITIAL_MEMBERSHIP + } + + await expect(store.applyCellAdmissionSelector(input)).resolves.toMatchObject({ + changed: true, + selector: { generation: 1, membership: INITIAL_MEMBERSHIP } + }) + await expect(store.applyCellAdmissionSelector(input)).resolves.toMatchObject({ + changed: false, + selector: { generation: 1, membership: INITIAL_MEMBERSHIP } + }) + await expect(store.inspectCellAdmissionSelector(input.attemptId)).resolves.toMatchObject({ + intent: { state: 'committed' } + }) + expect(await store.cellDeploymentStatus('c1')).toMatchObject({ + enabled: false, + admissionState: 'existing-only' + }) + expect(await store.cellDeploymentStatus('c2')).toMatchObject({ + enabled: true, + admissionState: 'migration-only' + }) + }) + + it('rejects a generation-zero membership change between intent and commit', async () => { + const delegate = await openInMemoryRelayDatabase() + const hooked = new AfterTransactionDatabase(delegate) + database = hooked + const store = new RelayAssignmentStore(hooked, () => 100) + await store.reconcileCells(CELLS) + const inspected = (await store.inspectCellAdmissionSelector()).selector.membership + const changed = { + existingOnly: ['c2'], + migrationOnly: [], + general: ['c1', 'c3'] + } + hooked.runAfterNextTransaction(async () => { + await delegate.transaction(async (transaction) => { + await transaction.query( + `UPDATE relay_cells SET enabled = 0 WHERE cell_id = ?`, + ['c2'] + ) + await transaction.query( + `UPDATE relay_cell_admission SET admission_state = ? WHERE cell_id = ?`, + ['existing-only', 'c2'] + ) + await transaction.query( + `UPDATE relay_admission_selectors SET membership_json = ? + WHERE selector_id = ? AND generation = 0`, + [encodeMembership(changed), 'general'] + ) + }) + }) + + await expect(store.applyCellAdmissionSelector({ + attemptId: 'cutover_race_001', + expectedGeneration: 0, + expectedMembershipSha256: membershipSha256(inspected), + membership: inspected + })).rejects.toThrow('admission_selector_membership_mismatch') + await expect(store.applyCellAdmissionSelector({ + attemptId: 'cutover_race_001', + expectedGeneration: 0, + membership: inspected + })).rejects.toThrow('admission_selector_membership_fingerprint_required') + await expect(store.inspectCellAdmissionSelector()).resolves.toMatchObject({ + selector: { generation: 0, membership: changed } + }) + }) + + it('preserves admission age for cells unchanged by a selector CAS', async () => { + let now = 100 + const store = await setup(() => now) + await store.applyCellAdmissionSelector({ + attemptId: 'cutover_age_001', + expectedGeneration: 0, + expectedMembershipSha256: BASE_MEMBERSHIP_SHA256, + membership: INITIAL_MEMBERSHIP + }) + + now = 200 + await store.applyCellAdmissionSelector({ + attemptId: 'cutover_age_002', + expectedGeneration: 1, + membership: { + existingOnly: ['c1'], + migrationOnly: [], + general: ['c2', 'c3'] + } + }) + + expect( + await database!.query( + `SELECT cell_id, admission_state, updated_at + FROM relay_cell_admission ORDER BY cell_id` + ) + ).toEqual([ + { cell_id: 'c1', admission_state: 'existing-only', updated_at: 100 }, + { cell_id: 'c2', admission_state: 'general', updated_at: 200 }, + { cell_id: 'c3', admission_state: 'general', updated_at: 100 } + ]) + }) + + it('persists failed CAS intent without mutating current membership', async () => { + const store = await setup(() => 100) + await store.applyCellAdmissionSelector({ + attemptId: 'cutover_005', + expectedGeneration: 0, + expectedMembershipSha256: BASE_MEMBERSHIP_SHA256, + membership: INITIAL_MEMBERSHIP + }) + + await expect( + store.applyCellAdmissionSelector({ + attemptId: 'stale_attempt', + expectedGeneration: 5, + membership: { + existingOnly: ['c1'], + migrationOnly: [], + general: ['c2', 'c3'] + } + }) + ).rejects.toThrow('admission_selector_generation_mismatch') + await expect(store.inspectCellAdmissionSelector('stale_attempt')).resolves.toMatchObject({ + selector: { generation: 1, membership: INITIAL_MEMBERSHIP }, + intent: { state: 'diverged' } + }) + }) + + it('rejects duplicate or incomplete selector membership before mutation', async () => { + const store = await setup(() => 100) + + await expect( + store.applyCellAdmissionSelector({ + attemptId: 'duplicate_001', + expectedGeneration: 0, + expectedMembershipSha256: BASE_MEMBERSHIP_SHA256, + membership: { + existingOnly: ['c1', 'c1'], + migrationOnly: ['c2'], + general: ['c3'] + } + }) + ).rejects.toThrow('admission_selector_duplicate_cell') + await expect( + store.applyCellAdmissionSelector({ + attemptId: 'incomplete_001', + expectedGeneration: 0, + expectedMembershipSha256: BASE_MEMBERSHIP_SHA256, + membership: { + existingOnly: ['c1'], + migrationOnly: ['c2'], + general: [] + } + }) + ).rejects.toThrow('admission_selector_incomplete_membership') + await expect(store.inspectCellAdmissionSelector()).resolves.toMatchObject({ + selector: { + generation: 0, + membership: { existingOnly: [], migrationOnly: [], general: ['c1', 'c2', 'c3'] } + } + }) + }) + + it('inspects an unchanged durable intent after an ambiguous apply failure', async () => { + const underlying = await openInMemoryRelayDatabase() + const failingDatabase = new FailAfterIntentDatabase(underlying) + database = failingDatabase + const store = new RelayAssignmentStore(failingDatabase, () => 100) + await store.reconcileCells(CELLS) + failingDatabase.failSecondTransaction() + + const input = { + attemptId: 'ambiguous_001', + expectedGeneration: 0, + expectedMembershipSha256: BASE_MEMBERSHIP_SHA256, + membership: INITIAL_MEMBERSHIP + } + await expect(store.applyCellAdmissionSelector(input)).rejects.toThrow( + 'injected_commit_ambiguity' + ) + await expect(store.inspectCellAdmissionSelector(input.attemptId)).resolves.toMatchObject({ + selector: { generation: 0 }, + intent: { state: 'unchanged' } + }) + await expect(store.applyCellAdmissionSelector(input)).resolves.toMatchObject({ + changed: true, + selector: { generation: 1, membership: INITIAL_MEMBERSHIP } + }) + }) + + it('allows only one concurrent CAS and never re-enables legacy cells', async () => { + const store = await setup(() => 100) + await store.applyCellAdmissionSelector({ + attemptId: 'cutover_006', + expectedGeneration: 0, + expectedMembershipSha256: BASE_MEMBERSHIP_SHA256, + membership: INITIAL_MEMBERSHIP + }) + await expect(store.setCellEnabled('c1', true)).rejects.toThrow( + 'admission_selector_boundary_active' + ) + + const nextMembership = { + existingOnly: ['c1'], + migrationOnly: [], + general: ['c2', 'c3'] + } + const results = await Promise.allSettled([ + store.applyCellAdmissionSelector({ + attemptId: 'promote_001', + expectedGeneration: 1, + membership: nextMembership + }), + store.applyCellAdmissionSelector({ + attemptId: 'promote_002', + expectedGeneration: 1, + membership: nextMembership + }) + ]) + expect(results.filter(({ status }) => status === 'fulfilled')).toHaveLength(1) + expect(results.filter(({ status }) => status === 'rejected')).toHaveLength(1) + + await expect( + store.applyCellAdmissionSelector({ + attemptId: 'reenable_001', + expectedGeneration: 2, + membership: { + existingOnly: [], + migrationOnly: [], + general: ['c1', 'c2', 'c3'] + } + }) + ).rejects.toThrow('admission_selector_legacy_reenable') + }) + + it('preserves the selector boundary across compatible reconciliation', async () => { + const store = await setup(() => 100) + await store.applyCellAdmissionSelector({ + attemptId: 'cutover_007', + expectedGeneration: 0, + expectedMembershipSha256: BASE_MEMBERSHIP_SHA256, + membership: INITIAL_MEMBERSHIP + }) + + await expect(store.reconcileCells(CELLS, false)).resolves.toBeUndefined() + await expect(store.reconcileCells([CELLS[0]!, CELLS[2]!])).rejects.toThrow( + 'admission_selector_boundary_active' + ) + await expect(store.inspectCellAdmissionSelector()).resolves.toMatchObject({ + selector: { generation: 1, membership: INITIAL_MEMBERSHIP } + }) + }) + + it('atomically adds exact migration-only cells after the selector boundary', async () => { + const store = await setup(() => 100) + await store.applyCellAdmissionSelector({ + attemptId: 'cutover_008', + expectedGeneration: 0, + expectedMembershipSha256: BASE_MEMBERSHIP_SHA256, + membership: INITIAL_MEMBERSHIP + }) + const input = { + attemptId: 'add_cells_001', + expectedGeneration: 1, + cells: [ + { + id: 'c5', + url: 'https://c5.example.com', + capacityRequests: 4_000, + connectionHardCap: 1_000 as const, + connectionUnobservedBound: 60 + }, + { + id: 'c4', + url: 'https://c4.example.com', + capacityRequests: 4_000, + connectionHardCap: 600 as const, + connectionUnobservedBound: 60 + } + ] + } + + await expect(store.addMigrationCells(input)).resolves.toMatchObject({ + changed: true, + selector: { + generation: 2, + attemptId: input.attemptId, + membership: { + existingOnly: ['c1'], + migrationOnly: ['c2', 'c4', 'c5'], + general: ['c3'] + } + } + }) + await expect(store.addMigrationCells(input)).resolves.toMatchObject({ + changed: false, + selector: { generation: 2 } + }) + await expect(store.inspectCellAdmissionSelector(input.attemptId)).resolves.toMatchObject({ + intent: { state: 'committed' } + }) + await expect(store.cellDeploymentStatus('c4')).resolves.toMatchObject({ + enabled: true, + admissionState: 'migration-only', + cellUrl: 'https://c4.example.com', + capacityRequests: 4_000, + connectionCapacity: { + hardCap: 600, + unobservedBound: 60 + } + }) + await expect(store.cellDeploymentStatus('c5')).resolves.toMatchObject({ + connectionCapacity: { + hardCap: 1_000, + ordinaryConnectionLimit: 900, + normalAdmissionPause: 840 + } + }) + }) + + it('rejects additive cells before cutover and exact-attempt config reuse', async () => { + const store = await setup(() => 100) + const cell = { + id: 'c4', + url: 'https://c4.example.com', + capacityRequests: 4_000, + connectionHardCap: 600 as const, + connectionUnobservedBound: 60 + } + await expect( + store.addMigrationCells({ + attemptId: 'add_cells_002', + expectedGeneration: 0, + cells: [cell] + }) + ).rejects.toThrow('admission_selector_boundary_inactive') + await store.applyCellAdmissionSelector({ + attemptId: 'cutover_009', + expectedGeneration: 0, + expectedMembershipSha256: BASE_MEMBERSHIP_SHA256, + membership: INITIAL_MEMBERSHIP + }) + await store.addMigrationCells({ + attemptId: 'add_cells_002', + expectedGeneration: 1, + cells: [cell] + }) + await expect( + store.addMigrationCells({ + attemptId: 'add_cells_002', + expectedGeneration: 1, + cells: [{ ...cell, capacityRequests: 3_999 }] + }) + ).rejects.toThrow('admission_selector_attempt_mismatch') + await expect( + store.applyCellAdmissionSelector({ + attemptId: 'add_cells_002', + expectedGeneration: 2, + membership: { + existingOnly: ['c1'], + migrationOnly: ['c2', 'c4'], + general: ['c3'] + } + }) + ).rejects.toThrow('admission_selector_attempt_mismatch') + }) +}) diff --git a/cloud/apps/relay/src/cell-admission-selector.ts b/cloud/apps/relay/src/cell-admission-selector.ts new file mode 100644 index 00000000000..92764e53e40 --- /dev/null +++ b/cloud/apps/relay/src/cell-admission-selector.ts @@ -0,0 +1,505 @@ +import { createHash } from 'node:crypto' +import type { RelayDatabase, SqlRow } from './database.js' + +export const CELL_ADMISSION_STATES = [ + 'existing-only', + 'migration-only', + 'general' +] as const + +export type CellAdmissionState = (typeof CELL_ADMISSION_STATES)[number] + +export type CellAdmissionMembership = { + existingOnly: string[] + migrationOnly: string[] + general: string[] +} + +export type CellAdmissionSelector = { + generation: number + attemptId: string | null + membership: CellAdmissionMembership +} + +export type CellAdmissionSelectorInspection = { + selector: CellAdmissionSelector + intent: null | { + attemptId: string + expectedGeneration: number + intendedGeneration: number + previousMembership: CellAdmissionMembership + membership: CellAdmissionMembership + state: 'unchanged' | 'committed' | 'diverged' + } +} + +type ApplySelectorInput = { + attemptId: string + expectedGeneration: number + expectedMembershipSha256?: string + membership: CellAdmissionMembership +} + +const SELECTOR_ID = 'general' + +export function stateFromEnabled(enabled: boolean): CellAdmissionState { + return enabled ? 'general' : 'existing-only' +} + +export function enabledForState(state: CellAdmissionState): number { + return state === 'existing-only' ? 0 : 1 +} + +export function parseCellAdmissionState(value: string): CellAdmissionState { + if (!CELL_ADMISSION_STATES.includes(value as CellAdmissionState)) { + throw new Error('invalid_cell_admission_state') + } + return value as CellAdmissionState +} + +export async function ensureCellAdmission( + transaction: RelayDatabase, + cellId: string, + fallback: CellAdmissionState, + now: number +): Promise { + await transaction.query( + `INSERT INTO relay_cell_admission (cell_id, admission_state, updated_at) + VALUES (?, ?, ?) ON CONFLICT (cell_id) DO NOTHING`, + [cellId, fallback, now] + ) +} + +export async function cellAdmissionState( + database: RelayDatabase, + cellId: string +): Promise { + const row = ( + await database.query( + `SELECT admission_state FROM relay_cell_admission WHERE cell_id = ?`, + [cellId] + ) + )[0] + if (!row) throw new Error('cell_admission_missing') + return admissionState(row) +} + +export async function cellAdmissionStates( + database: RelayDatabase +): Promise> { + const rows = await database.query( + `SELECT cell.cell_id, admission.admission_state + FROM relay_cells cell + LEFT JOIN relay_cell_admission admission ON admission.cell_id = cell.cell_id + ORDER BY cell.cell_id ASC` + ) + return new Map(rows.map((row) => [text(row, 'cell_id'), admissionState(row)])) +} + +export async function setCellAdmissionBeforeBoundary( + transaction: RelayDatabase, + cellId: string, + state: CellAdmissionState, + now: number +): Promise { + const cells = await transaction.queryLocked( + `SELECT cell_id FROM relay_cells ORDER BY cell_id ASC` + ) + if (!cells.some((row) => text(row, 'cell_id') === cellId)) { + throw new Error('cell_not_found') + } + if ((await lockedSelector(transaction)).generation > 0) { + throw new Error('admission_selector_boundary_active') + } + const updated = await transaction.query( + `UPDATE relay_cells + SET updated_at = CASE WHEN enabled <> ? THEN ? ELSE updated_at END, enabled = ? + WHERE cell_id = ?`, + [enabledForState(state), now, enabledForState(state), cellId] + ) + if (integer(updated[0]!, 'changes') !== 1) throw new Error('cell_not_found') + await ensureCellAdmission(transaction, cellId, state, now) + await transaction.query( + `UPDATE relay_cell_admission + SET updated_at = CASE WHEN admission_state <> ? THEN ? ELSE updated_at END, + admission_state = ? + WHERE cell_id = ?`, + [state, now, state, cellId] + ) + await synchronizeCellAdmissionBoundary(transaction, now) +} + +export async function synchronizeCellAdmissionBoundary( + transaction: RelayDatabase, + now: number +): Promise { + const selector = await lockedSelector(transaction) + if (selector.generation > 0) { + await requireSelectorMatchesAdmission(transaction, selector) + return + } + await transaction.query( + `UPDATE relay_admission_selectors SET membership_json = ?, updated_at = ? + WHERE selector_id = ? AND generation = 0`, + [encodeMembership(await currentMembership(transaction)), now, SELECTOR_ID] + ) +} + +export class RelayCellAdmissionSelector { + constructor( + private readonly database: RelayDatabase, + private readonly now: () => number + ) {} + + async apply(input: ApplySelectorInput): Promise<{ + changed: boolean + selector: CellAdmissionSelector + }> { + const membership = normalizeMembership(input.membership) + validateAttempt(input) + const encoded = encodeMembership(membership) + await this.persistIntent(input, membership, encoded) + return await this.database.transaction(async (transaction) => { + const cells = await transaction.queryLocked( + `SELECT cell_id FROM relay_cells ORDER BY cell_id ASC` + ) + requireExactMembership(cells, membership) + const selector = await lockedSelector(transaction) + if ( + selector.generation === input.expectedGeneration + 1 && + selector.attemptId === input.attemptId && + encodeMembership(selector.membership) === encoded + ) { + return { changed: false, selector } + } + if (selector.generation !== input.expectedGeneration) { + throw new Error('admission_selector_generation_mismatch') + } + requireExpectedMembership(selector, input.expectedMembershipSha256) + await requireSelectorMatchesAdmission(transaction, selector) + if (selector.generation > 0) { + const nextNonExisting = new Set([...membership.migrationOnly, ...membership.general]) + if (selector.membership.existingOnly.some((cellId) => nextNonExisting.has(cellId))) { + throw new Error('admission_selector_legacy_reenable') + } + } + const now = this.now() + for (const [state, cellIds] of membershipEntries(membership)) { + for (const cellId of cellIds) { + await transaction.query( + `UPDATE relay_cell_admission + SET updated_at = CASE WHEN admission_state <> ? THEN ? ELSE updated_at END, + admission_state = ? + WHERE cell_id = ?`, + [state, now, state, cellId] + ) + await transaction.query( + `UPDATE relay_cells + SET updated_at = CASE WHEN enabled <> ? THEN ? ELSE updated_at END, enabled = ? + WHERE cell_id = ?`, + [enabledForState(state), now, enabledForState(state), cellId] + ) + } + } + const intendedGeneration = input.expectedGeneration + 1 + const updated = await transaction.query( + `UPDATE relay_admission_selectors + SET generation = ?, attempt_id = ?, membership_json = ?, updated_at = ? + WHERE selector_id = ? AND generation = ?`, + [ + intendedGeneration, + input.attemptId, + encoded, + now, + SELECTOR_ID, + input.expectedGeneration + ] + ) + if (integer(updated[0]!, 'changes') !== 1) { + throw new Error('admission_selector_generation_mismatch') + } + await transaction.query( + `UPDATE relay_admission_selector_intents SET committed_at = ? + WHERE attempt_id = ?`, + [now, input.attemptId] + ) + return { + changed: true, + selector: { + generation: intendedGeneration, + attemptId: input.attemptId, + membership + } + } + }) + } + + async inspect(attemptId?: string): Promise { + return await this.database.transaction(async (transaction) => { + await transaction.queryLocked(`SELECT cell_id FROM relay_cells ORDER BY cell_id ASC`) + const selector = await lockedSelector(transaction) + await requireSelectorMatchesAdmission(transaction, selector) + if (!attemptId) return { selector, intent: null } + const row = ( + await transaction.queryLocked( + `SELECT * FROM relay_admission_selector_intents WHERE attempt_id = ?`, + [attemptId] + ) + )[0] + if (!row) return { selector, intent: null } + const membership = decodeMembership(text(row, 'membership_json')) + const previousMembership = decodeMembership(text(row, 'previous_membership_json')) + const intendedGeneration = integer(row, 'intended_generation') + const committed = + selector.generation === intendedGeneration && + selector.attemptId === attemptId && + encodeMembership(selector.membership) === encodeMembership(membership) + const unchanged = + selector.generation === integer(row, 'expected_generation') && + encodeMembership(selector.membership) === text(row, 'previous_membership_json') + return { + selector, + intent: { + attemptId, + expectedGeneration: integer(row, 'expected_generation'), + intendedGeneration, + previousMembership, + membership, + state: committed ? 'committed' : unchanged ? 'unchanged' : 'diverged' + } + } + }) + } + + private async persistIntent( + input: ApplySelectorInput, + membership: CellAdmissionMembership, + encoded: string + ): Promise { + await this.database.transaction(async (transaction) => { + const cells = await transaction.queryLocked( + `SELECT cell_id FROM relay_cells ORDER BY cell_id ASC` + ) + requireExactMembership(cells, membership) + const selector = await lockedSelector(transaction) + const existing = ( + await transaction.queryLocked( + `SELECT * FROM relay_admission_selector_intents WHERE attempt_id = ?`, + [input.attemptId] + ) + )[0] + if (existing) { + const addition = ( + await transaction.queryLocked( + `SELECT attempt_id FROM relay_admission_selector_cell_additions + WHERE attempt_id = ?`, + [input.attemptId] + ) + )[0] + if ( + addition || + integer(existing, 'expected_generation') !== input.expectedGeneration || + integer(existing, 'intended_generation') !== input.expectedGeneration + 1 || + text(existing, 'membership_json') !== encoded || + (input.expectedMembershipSha256 !== undefined && + membershipSha256(text(existing, 'previous_membership_json')) !== + input.expectedMembershipSha256) + ) { + throw new Error('admission_selector_attempt_mismatch') + } + return + } + requireExpectedMembership(selector, input.expectedMembershipSha256) + await transaction.query( + `INSERT INTO relay_admission_selector_intents + (attempt_id, expected_generation, intended_generation, previous_membership_json, + membership_json, created_at, committed_at) + VALUES (?, ?, ?, ?, ?, ?, NULL)`, + [ + input.attemptId, + input.expectedGeneration, + input.expectedGeneration + 1, + encodeMembership(selector.membership), + encoded, + this.now() + ] + ) + }) + } +} + +function requireExpectedMembership( + selector: CellAdmissionSelector, + expectedSha256: string | undefined +): void { + if (expectedSha256 === undefined) return + if ( + !/^[a-f0-9]{64}$/.test(expectedSha256) || + membershipSha256(encodeMembership(selector.membership)) !== expectedSha256 + ) { + throw new Error('admission_selector_membership_mismatch') + } +} + +function membershipSha256(encoded: string): string { + return createHash('sha256').update(encoded).digest('hex') +} + +export async function lockedSelector( + transaction: RelayDatabase +): Promise { + let row = ( + await transaction.queryLocked( + `SELECT * FROM relay_admission_selectors WHERE selector_id = ?`, + [SELECTOR_ID] + ) + )[0] + if (!row) { + const membership = await currentMembership(transaction) + await transaction.query( + `INSERT INTO relay_admission_selectors + (selector_id, generation, attempt_id, membership_json, updated_at) + VALUES (?, 0, NULL, ?, 0) ON CONFLICT (selector_id) DO NOTHING`, + [SELECTOR_ID, encodeMembership(membership)] + ) + row = ( + await transaction.queryLocked( + `SELECT * FROM relay_admission_selectors WHERE selector_id = ?`, + [SELECTOR_ID] + ) + )[0] + } + if (!row) throw new Error('admission_selector_missing') + return { + generation: integer(row, 'generation'), + attemptId: optionalText(row, 'attempt_id'), + membership: decodeMembership(text(row, 'membership_json')) + } +} + +async function currentMembership(database: RelayDatabase): Promise { + const rows = await database.query( + `SELECT cell.cell_id, admission.admission_state + FROM relay_cells cell + LEFT JOIN relay_cell_admission admission ON admission.cell_id = cell.cell_id + ORDER BY cell.cell_id ASC` + ) + const membership: CellAdmissionMembership = { + existingOnly: [], + migrationOnly: [], + general: [] + } + for (const row of rows) membership[keyForState(admissionState(row))].push(text(row, 'cell_id')) + return membership +} + +export async function requireSelectorMatchesAdmission( + database: RelayDatabase, + selector: CellAdmissionSelector +): Promise { + const current = await currentMembership(database) + if (encodeMembership(current) !== encodeMembership(selector.membership)) { + throw new Error('admission_selector_membership_drift') + } +} + +export function normalizeMembership( + input: CellAdmissionMembership +): CellAdmissionMembership { + const membership = { + existingOnly: [...input.existingOnly].sort(), + migrationOnly: [...input.migrationOnly].sort(), + general: [...input.general].sort() + } + const all = [...membership.existingOnly, ...membership.migrationOnly, ...membership.general] + if (new Set(all).size !== all.length) throw new Error('admission_selector_duplicate_cell') + return membership +} + +function requireExactMembership(rows: SqlRow[], membership: CellAdmissionMembership): void { + const expected = rows.map((row) => text(row, 'cell_id')).sort() + const actual = [ + ...membership.existingOnly, + ...membership.migrationOnly, + ...membership.general + ].sort() + if (JSON.stringify(actual) !== JSON.stringify(expected)) { + throw new Error('admission_selector_incomplete_membership') + } +} + +function membershipEntries( + membership: CellAdmissionMembership +): Array<[CellAdmissionState, string[]]> { + return [ + ['existing-only', membership.existingOnly], + ['migration-only', membership.migrationOnly], + ['general', membership.general] + ] +} + +export function encodeMembership(membership: CellAdmissionMembership): string { + return JSON.stringify(normalizeMembership(membership)) +} + +export function decodeMembership(value: string): CellAdmissionMembership { + const parsed = JSON.parse(value) as Partial + if ( + !Array.isArray(parsed.existingOnly) || + !Array.isArray(parsed.migrationOnly) || + !Array.isArray(parsed.general) || + [...parsed.existingOnly, ...parsed.migrationOnly, ...parsed.general].some( + (cellId) => typeof cellId !== 'string' + ) + ) { + throw new Error('admission_selector_invalid_membership') + } + return normalizeMembership({ + existingOnly: parsed.existingOnly as string[], + migrationOnly: parsed.migrationOnly as string[], + general: parsed.general as string[] + }) +} + +function admissionState(row: SqlRow): CellAdmissionState { + return parseCellAdmissionState(text(row, 'admission_state')) +} + +function keyForState(state: CellAdmissionState): keyof CellAdmissionMembership { + if (state === 'existing-only') return 'existingOnly' + if (state === 'migration-only') return 'migrationOnly' + return 'general' +} + +function validateAttempt(input: { + attemptId: string + expectedGeneration: number + expectedMembershipSha256?: string +}): void { + if (!/^[A-Za-z0-9_-]{8,128}$/.test(input.attemptId)) { + throw new Error('invalid_admission_selector_attempt') + } + if (!Number.isSafeInteger(input.expectedGeneration) || input.expectedGeneration < 0) { + throw new Error('invalid_admission_selector_generation') + } + if (input.expectedGeneration === 0 && input.expectedMembershipSha256 === undefined) { + throw new Error('admission_selector_membership_fingerprint_required') + } +} + +function optionalText(row: SqlRow, field: string): string | null { + if (row[field] === null || row[field] === undefined) return null + return text(row, field) +} + +function integer(row: SqlRow, field: string): number { + const value = Number(row[field]) + if (!Number.isSafeInteger(value)) throw new Error(`invalid integer field ${field}`) + return value +} + +function text(row: SqlRow, field: string): string { + const value = row[field] + if (typeof value !== 'string') throw new Error(`invalid text field ${field}`) + return value +} diff --git a/cloud/apps/relay/src/cell-admission-startup-postgres.test.ts b/cloud/apps/relay/src/cell-admission-startup-postgres.test.ts new file mode 100644 index 00000000000..6f93fecce6f --- /dev/null +++ b/cloud/apps/relay/src/cell-admission-startup-postgres.test.ts @@ -0,0 +1,83 @@ +import pg from 'pg' +import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.js' +import { reconcileCellAdmissionAtStartup } from './cell-admission-startup.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 schema = 'relay_startup_retry_test' +const cell: RelayCellConfig = { + id: 'startup-retry-cell', + url: 'https://startup-retry.example.test', + capacityRequests: 4_000 +} + +afterEach(() => vi.restoreAllMocks()) + +describePostgres('PostgreSQL director startup reconciliation', () => { + const databases: RelayDatabase[] = [] + let scopedDatabaseUrl = '' + + beforeAll(async () => { + const client = new pg.Client({ connectionString: databaseUrl }) + await client.connect() + try { + await client.query(`DROP SCHEMA IF EXISTS ${schema} CASCADE`) + await client.query(`CREATE SCHEMA ${schema}`) + } finally { + await client.end() + } + const url = new URL(databaseUrl!) + url.searchParams.set('options', `-c search_path=${schema}`) + scopedDatabaseUrl = url.toString() + databases.push( + await openRelayDatabase({ databaseUrl: scopedDatabaseUrl, dataDir: '' }), + await openRelayDatabase({ databaseUrl: scopedDatabaseUrl, dataDir: '' }) + ) + }) + + afterAll(async () => { + for (const database of databases) await database.close() + const client = new pg.Client({ connectionString: databaseUrl }) + await client.connect() + try { + await client.query(`DROP SCHEMA IF EXISTS ${schema} CASCADE`) + } finally { + await client.end() + } + }) + + it('recovers after the cell inventory lock outlasts transaction retries', async () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const store = new RelayAssignmentStore(databases[0]!, () => 100) + await store.reconcileCells([cell], false) + let releaseLock: () => void = () => undefined + let reportLocked: () => void = () => undefined + const lockReleased = new Promise((resolve) => (releaseLock = resolve)) + const lockAcquired = new Promise((resolve) => (reportLocked = resolve)) + const holder = databases[1]!.transaction(async (transaction) => { + await transaction.queryLocked(`SELECT cell_id FROM relay_cells ORDER BY cell_id ASC`) + reportLocked() + await lockReleased + }) + await lockAcquired + + const releaseTimer = setTimeout(releaseLock, 3_500) + try { + await reconcileCellAdmissionAtStartup({ role: 'director', cells: [cell] }, store) + } finally { + clearTimeout(releaseTimer) + releaseLock() + await holder + } + + expect(await databases[0]!.query(`SELECT cell_id FROM relay_cells`)).toEqual([ + { cell_id: cell.id } + ]) + const warnings = warn.mock.calls.flat().join('\n') + expect(warnings).toContain('orca_relay_startup_reconcile_recovered') + expect(warnings).not.toContain('orca_relay_postgres_transaction_exhausted') + }, 10_000) +}) diff --git a/cloud/apps/relay/src/cell-admission-startup.test.ts b/cloud/apps/relay/src/cell-admission-startup.test.ts new file mode 100644 index 00000000000..d08d9eefc16 --- /dev/null +++ b/cloud/apps/relay/src/cell-admission-startup.test.ts @@ -0,0 +1,119 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { + reconcileCellAdmissionAtStartup, + roleOwnsAssignmentMaintenance +} from './cell-admission-startup.js' +import type { RelayCellConfig } from './config.js' + +const cells: RelayCellConfig[] = [ + { + id: 'candidate', + url: 'https://candidate.relay.example.com', + capacityRequests: 4_000, + initiallyEnabled: false + } +] + +afterEach(() => { + vi.useRealTimers() + vi.restoreAllMocks() +}) + +describe('cell admission startup authority', () => { + it('reserves global assignment maintenance for director-capable roles', () => { + expect(roleOwnsAssignmentMaintenance('cell')).toBe(false) + expect(roleOwnsAssignmentMaintenance('director')).toBe(true) + expect(roleOwnsAssignmentMaintenance('combined')).toBe(true) + }) + + it('does not let a cell process reconcile its own admission', async () => { + const reconcileCellsAtStartup = vi.fn() + await reconcileCellAdmissionAtStartup({ role: 'cell', cells }, { reconcileCellsAtStartup }) + expect(reconcileCellsAtStartup).not.toHaveBeenCalled() + }) + + it.each(['director', 'combined'] as const)( + 'lets the %s process reconcile configured admission without disabling missing cells', + async (role) => { + const reconcileCellsAtStartup = vi.fn() + await reconcileCellAdmissionAtStartup({ role, cells }, { reconcileCellsAtStartup }) + expect(reconcileCellsAtStartup).toHaveBeenCalledWith(cells) + } + ) + + it('waits out transient database pressure during director startup', async () => { + vi.useFakeTimers() + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const lockUnavailable = Object.assign(new Error('lock unavailable'), { code: '55P03' }) + const reconcileCellsAtStartup = vi + .fn() + .mockRejectedValueOnce(lockUnavailable) + .mockRejectedValueOnce(lockUnavailable) + .mockResolvedValue(undefined) + + const startup = reconcileCellAdmissionAtStartup( + { role: 'director', cells }, + { reconcileCellsAtStartup } + ) + await vi.runAllTimersAsync() + await startup + + expect(reconcileCellsAtStartup).toHaveBeenCalledTimes(3) + expect(warn).toHaveBeenCalledWith( + expect.stringContaining('orca_relay_startup_reconcile_recovered') + ) + }) + + it('fails startup immediately for a permanent reconciliation error', async () => { + const failure = new Error('invalid cell configuration') + const reconcileCellsAtStartup = vi.fn().mockRejectedValue(failure) + + await expect( + reconcileCellAdmissionAtStartup({ role: 'director', cells }, { reconcileCellsAtStartup }) + ).rejects.toBe(failure) + expect(reconcileCellsAtStartup).toHaveBeenCalledTimes(1) + }) + + it('bounds transient startup retries', async () => { + vi.useFakeTimers() + vi.spyOn(Math, 'random').mockReturnValue(0) + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const lockUnavailable = Object.assign(new Error('lock unavailable'), { code: '55P03' }) + const reconcileCellsAtStartup = vi.fn().mockRejectedValue(lockUnavailable) + + const startup = reconcileCellAdmissionAtStartup( + { role: 'director', cells }, + { reconcileCellsAtStartup } + ) + const rejection = expect(startup).rejects.toBe(lockUnavailable) + await vi.runAllTimersAsync() + await rejection + + expect(reconcileCellsAtStartup).toHaveBeenCalledTimes(20) + expect(warn).toHaveBeenCalledWith( + expect.stringContaining('orca_relay_startup_reconcile_exhausted') + ) + }) + + it('bounds retry wall time when each reconciliation is slow', async () => { + vi.useFakeTimers() + vi.setSystemTime(0) + vi.spyOn(Math, 'random').mockReturnValue(0) + vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const lockUnavailable = Object.assign(new Error('lock unavailable'), { code: '55P03' }) + const reconcileCellsAtStartup = vi.fn().mockImplementation(async () => { + vi.setSystemTime(Date.now() + 10_000) + throw lockUnavailable + }) + + const startup = reconcileCellAdmissionAtStartup( + { role: 'director', cells }, + { reconcileCellsAtStartup } + ) + const rejection = expect(startup).rejects.toBe(lockUnavailable) + await vi.runAllTimersAsync() + await rejection + + expect(reconcileCellsAtStartup).toHaveBeenCalledTimes(5) + }) +}) diff --git a/cloud/apps/relay/src/cell-admission-startup.ts b/cloud/apps/relay/src/cell-admission-startup.ts new file mode 100644 index 00000000000..384b6aab304 --- /dev/null +++ b/cloud/apps/relay/src/cell-admission-startup.ts @@ -0,0 +1,56 @@ +import type { RelayAssignmentStore } from './assignment-store.js' +import type { RelayConfig } from './config.js' +import { isRelayDatabaseTransientError } from './database.js' + +type CellAdmissionStartupConfig = Pick +type CellAdmissionStore = Pick + +const STARTUP_RECONCILE_ATTEMPTS = 20 +const STARTUP_RECONCILE_RETRY_WINDOW_MS = 45_000 +const STARTUP_RECONCILE_RETRY_BASE_MS = 250 +const STARTUP_RECONCILE_RETRY_JITTER_MS = 250 + +export function roleOwnsAssignmentMaintenance(role: RelayConfig['role']): boolean { + // Cell workers share the database but the director is the sole authority + // for global expiry, evacuation, and dead-cell maintenance. + return role !== 'cell' +} + +export async function reconcileCellAdmissionAtStartup( + config: CellAdmissionStartupConfig, + assignments: CellAdmissionStore +): Promise { + // 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 + const retryDeadline = Date.now() + STARTUP_RECONCILE_RETRY_WINDOW_MS + for (let attempt = 1; attempt <= STARTUP_RECONCILE_ATTEMPTS; attempt += 1) { + try { + await assignments.reconcileCellsAtStartup(config.cells) + if (attempt > 1) { + console.warn( + JSON.stringify({ event: 'orca_relay_startup_reconcile_recovered', attempts: attempt }) + ) + } + return + } catch (error) { + const remainingMs = retryDeadline - Date.now() + if ( + attempt === STARTUP_RECONCILE_ATTEMPTS || + remainingMs <= 0 || + !isRelayDatabaseTransientError(error) + ) { + if (isRelayDatabaseTransientError(error)) { + console.warn( + JSON.stringify({ event: 'orca_relay_startup_reconcile_exhausted', attempts: attempt }) + ) + } + throw error + } + const delayMs = + STARTUP_RECONCILE_RETRY_BASE_MS + + Math.floor(Math.random() * (STARTUP_RECONCILE_RETRY_JITTER_MS + 1)) + await new Promise((resolve) => setTimeout(resolve, Math.min(delayMs, remainingMs))) + } + } +} diff --git a/cloud/apps/relay/src/cell-connection-hard-cap-consistency.test.ts b/cloud/apps/relay/src/cell-connection-hard-cap-consistency.test.ts new file mode 100644 index 00000000000..af062df482d --- /dev/null +++ b/cloud/apps/relay/src/cell-connection-hard-cap-consistency.test.ts @@ -0,0 +1,27 @@ +import { + RELAY_ADMISSION_BUDGETS, + RELAY_CELL_ADMISSION_BOUNDS, + RELAY_CELL_CONNECTION_HARD_CAP, + RELAY_CELL_CONNECTION_HARD_CAPS, + relayCellAdmissionBounds +} from '@orca-cloud/relay-contract' +import { describe, expect, it } from 'vitest' + +describe('cell connection hard cap', () => { + it('derives its own bounds from the cap', () => { + expect(RELAY_CELL_ADMISSION_BOUNDS.hardCap).toBe(RELAY_CELL_CONNECTION_HARD_CAP) + expect(RELAY_CELL_ADMISSION_BOUNDS.socketAdmissionCeiling).toBe( + RELAY_CELL_CONNECTION_HARD_CAP - RELAY_ADMISSION_BUDGETS.reservedHostControls + ) + expect(RELAY_CELL_ADMISSION_BOUNDS.maxUnobservedBound).toBe( + RELAY_CELL_ADMISSION_BOUNDS.socketAdmissionCeiling - 1 + ) + for (const hardCap of RELAY_CELL_CONNECTION_HARD_CAPS) { + const bounds = relayCellAdmissionBounds(hardCap) + expect(bounds.socketAdmissionCeiling).toBe( + hardCap - RELAY_ADMISSION_BUDGETS.reservedHostControls + ) + expect(bounds.maxUnobservedBound).toBe(bounds.socketAdmissionCeiling - 1) + } + }) +}) diff --git a/cloud/apps/relay/src/cell-fence-legacy-adoption-postgres.test.ts b/cloud/apps/relay/src/cell-fence-legacy-adoption-postgres.test.ts new file mode 100644 index 00000000000..9fc1b915481 --- /dev/null +++ b/cloud/apps/relay/src/cell-fence-legacy-adoption-postgres.test.ts @@ -0,0 +1,122 @@ +import { afterAll, beforeAll, describe, expect, it } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.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 = { + id: 'legacy-fence-adoption-postgres', + url: 'https://legacy-fence-adoption-postgres.example.com', + capacityRequests: 100 +} +const incarnation = '11111111-1111-4111-8111-111111111111' +const attemptId = '22222222-2222-4222-8222-222222222222' + +describePostgres('PostgreSQL legacy fence adoption', () => { + const databases: RelayDatabase[] = [] + + beforeAll(async () => { + databases.push( + await openRelayDatabase({ databaseUrl, dataDir: '' }), + await openRelayDatabase({ databaseUrl, dataDir: '' }) + ) + await cleanup() + }) + + afterAll(async () => { + await cleanup() + for (const database of databases) await database.close() + }) + + async function cleanup(): Promise { + const database = databases[0] + if (!database) return + await database.query( + `DELETE FROM relay_cell_fence_plan_bindings WHERE attempt_id = ?`, + [attemptId] + ) + await database.query( + `DELETE FROM relay_cell_fence_apply_invocations WHERE attempt_id = ?`, + [attemptId] + ) + await database.query(`DELETE FROM relay_cell_committed_fences WHERE cell_id = ?`, [ + cell.id + ]) + await database.query( + `DELETE FROM relay_cell_legacy_fence_adoptions WHERE cell_id = ?`, + [cell.id] + ) + await database.query(`DELETE FROM relay_cell_fence_attempts WHERE cell_id = ?`, [cell.id]) + await database.query(`DELETE FROM relay_cell_fences WHERE cell_id = ?`, [cell.id]) + await database.query(`DELETE FROM relay_cell_runtime WHERE cell_id = ?`, [cell.id]) + await database.query(`DELETE FROM relay_cell_admission WHERE cell_id = ?`, [cell.id]) + await database.query(`DELETE FROM relay_cells WHERE cell_id = ?`, [cell.id]) + } + + it('allows either adoption or attempt preparation, never both', async () => { + let now = 100 + const stores = databases.map( + (database) => + new RelayAssignmentStore(database, () => now, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + ) + await stores[0]!.reconcileCells([cell], false) + await stores[0]!.recordCellHeartbeat({ + cellId: cell.id, + cellUrl: cell.url, + cellIncarnation: incarnation, + startedAt: 50, + ready: true, + observedRequests: 0 + }) + await stores[0]!.setCellEnabled(cell.id, false) + now += 45_001 + const evidence = { + attemptId, + environment: 'production' as const, + cellId: cell.id, + cellIncarnation: incarnation, + migName: 'orca-relay-c3', + instanceGroup: 'https://compute.example/instanceGroups/orca-relay-c3', + generationIdentity: 'https://compute.example/instanceTemplates/orca-relay-c3-abc', + fenceCommit: 'a'.repeat(40), + planSha256: 'b'.repeat(64), + planObjectName: + 'terraform/state/relay-fence-plans/production/22222222-2222-4222-8222-222222222222.tfplan', + planObjectGeneration: '123456789', + varFileSha256: 'c'.repeat(64), + terraformStateLineage: '33333333-3333-4333-8333-333333333333', + terraformStateSerial: 7, + terraformStateObjectGeneration: '987654321', + terraformStateObjectSha256: 'd'.repeat(64), + requestReason: + 'orca-relay-fence/22222222-2222-4222-8222-222222222222' + } + + const results = await Promise.allSettled([ + stores[0]!.adoptLegacyCellFence(cell.id, incarnation), + stores[1]!.prepareCellFenceAttempt(evidence) + ]) + expect(results.filter(({ status }) => status === 'fulfilled')).toHaveLength(1) + expect(results.filter(({ status }) => status === 'rejected')).toHaveLength(1) + if (results[0]!.status === 'fulfilled') { + await stores[0]!.commitLegacyCellFenceAdoption(cell.id, incarnation) + } + const attempts = await databases[0]!.query( + `SELECT COUNT(*) AS count FROM relay_cell_fence_attempts WHERE cell_id = ?`, + [cell.id] + ) + const fences = await databases[0]!.query( + `SELECT COUNT(*) AS count FROM relay_cell_fences WHERE cell_id = ?`, + [cell.id] + ) + const adoptions = await databases[0]!.query( + `SELECT COUNT(*) AS count FROM relay_cell_legacy_fence_adoptions WHERE cell_id = ?`, + [cell.id] + ) + expect(Number(attempts[0]!.count) + Number(fences[0]!.count)).toBe(1) + expect(Number(adoptions[0]!.count)).toBe(Number(fences[0]!.count)) + }) +}) diff --git a/cloud/apps/relay/src/cell-heartbeat-client.test.ts b/cloud/apps/relay/src/cell-heartbeat-client.test.ts new file mode 100644 index 00000000000..2aa708bed1b --- /dev/null +++ b/cloud/apps/relay/src/cell-heartbeat-client.test.ts @@ -0,0 +1,176 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import type { RelayConfig } from './config.js' +import { startCellHeartbeat } from './cell-heartbeat-client.js' + +const CONFIG = { + role: 'cell', + cellId: 'cell-a', + cellUrl: 'https://relay-a.example.com', + directorUrl: 'https://relay.example.com', + heartbeatAudience: 'https://relay.example.com/v1/admin/cell-heartbeat' +} as RelayConfig + +describe('cell heartbeat client', () => { + afterEach(() => vi.restoreAllMocks()) + + it('sends an immediate authenticated heartbeat without putting credentials in the URL', async () => { + const requests: Array<{ url: string; init?: RequestInit }> = [] + const fetchImpl = vi.fn(async (url: string | URL | Request, init?: RequestInit) => { + requests.push({ url: String(url), init }) + return new Response('{}', { status: 200 }) + }) as typeof fetch + const client = startCellHeartbeat(CONFIG, { + ready: async () => true, + observedRequests: () => 7, + connectionCounts: () => ({ + totalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0 + }), + fetch: fetchImpl, + identityToken: async () => 'secret-token', + now: () => 123, + incarnation: '11111111-1111-4111-8111-111111111111', + intervalMs: 60_000 + })! + await vi.waitFor(() => expect(requests).toHaveLength(1)) + client.stop() + + expect(requests[0]!.url).toBe('https://relay.example.com/v1/admin/cell-heartbeat') + expect(requests[0]!.url).not.toContain('secret-token') + expect(requests[0]!.init?.headers).toMatchObject({ authorization: 'Bearer secret-token' }) + expect(JSON.parse(String(requests[0]!.init?.body))).toEqual({ + v: 1, + cellId: 'cell-a', + cellUrl: 'https://relay-a.example.com', + region: 'us-central1', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 123, + ready: true, + observedRequests: 7 + }) + }) + + it('advertises per-host drain only when its dedicated trust boundary is configured', async () => { + const requests: RequestInit[] = [] + const client = startCellHeartbeat( + { + ...CONFIG, + rehomeAudience: 'https://relay.example.com/v1/admin/host-drain', + rehomeDirectorServiceAccount: 'relay-director@example.com' + }, + { + ready: async () => true, + observedRequests: () => 0, + connectionCounts: () => ({ + totalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0 + }), + regionalRehomeSafety: () => ({ + observedAt: 120, + sqlFailures: 0, + reconnects: 2, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + }), + fetch: async (_input, init) => { + requests.push(init ?? {}) + return new Response('{}', { status: 200 }) + }, + identityToken: async () => 'secret-token', + now: () => 123, + incarnation: '11111111-1111-4111-8111-111111111111', + intervalMs: 60_000 + } + )! + await vi.waitFor(() => expect(requests).toHaveLength(2)) + client.stop() + + expect(JSON.parse(String(requests[1]!.body))).toMatchObject({ + regionalRehomeProtocol: 1, + safety: { + observedAt: 120, + sqlFailures: 0, + reconnects: 2, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + } + }) + }) + + it('reports the complete enforced connection ledger for a limited cell', async () => { + const requests: RequestInit[] = [] + const client = startCellHeartbeat( + { + ...CONFIG, + connectionHardCap: 600, + connectionUnobservedBound: 60 + }, + { + ready: async () => true, + observedRequests: () => 9, + connectionCounts: () => ({ + totalConnections: 10, + inFlightConnections: 2, + reservedConnectionUnits: 3, + enforcedConnectionUnits: 15, + inclusionWatermark: 42 + }), + fetch: async (_input, init) => { + requests.push(init ?? {}) + return new Response('{}', { status: 200 }) + }, + identityToken: async () => 'secret-token', + now: () => 123, + incarnation: '11111111-1111-4111-8111-111111111111', + intervalMs: 60_000 + } + )! + await vi.waitFor(() => expect(requests).toHaveLength(1)) + client.stop() + + expect(JSON.parse(String(requests[0]!.body))).toMatchObject({ + totalConnections: 10, + inFlightConnections: 2, + reservedConnectionUnits: 3, + enforcedConnectionUnits: 15, + connectionInclusionWatermark: 42, + connectionHardCap: 600, + connectionUnobservedBound: 60 + }) + }) + + it('does not start outside an explicitly configured cell role', () => { + expect( + startCellHeartbeat({ ...CONFIG, role: 'director' }, { + ready: async () => true, + observedRequests: () => 0, + connectionCounts: () => ({ + totalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0 + }) + }) + ).toBeNull() + expect( + startCellHeartbeat({ ...CONFIG, directorUrl: undefined }, { + ready: async () => true, + observedRequests: () => 0, + connectionCounts: () => ({ + totalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0 + }) + }) + ).toBeNull() + }) +}) diff --git a/cloud/apps/relay/src/cell-heartbeat-client.ts b/cloud/apps/relay/src/cell-heartbeat-client.ts new file mode 100644 index 00000000000..3bbcd08ecd6 --- /dev/null +++ b/cloud/apps/relay/src/cell-heartbeat-client.ts @@ -0,0 +1,123 @@ +import { randomUUID } from 'node:crypto' +import type { RelayConfig } from './config.js' +import { googleMetadataIdentityToken } from './google-metadata-identity-token.js' +import type { RegionalRehomeSafetySnapshot } from './relay-observability.js' + +type ConnectionCounts = { + totalConnections: number + inFlightConnections: number + reservedConnectionUnits: number + enforcedConnectionUnits: number + inclusionWatermark?: number +} + +type HeartbeatClientOptions = { + ready: () => Promise + observedRequests: () => number + connectionCounts: () => ConnectionCounts + regionalRehomeSafety?: () => RegionalRehomeSafetySnapshot + fetch?: typeof fetch + identityToken?: (audience: string) => Promise + now?: () => number + incarnation?: string + intervalMs?: number +} + +export type CellHeartbeatClient = { stop: () => void; send: () => Promise } + +export function startCellHeartbeat( + config: RelayConfig, + options: HeartbeatClientOptions +): CellHeartbeatClient | null { + if (config.role !== 'cell' || !config.directorUrl || !config.heartbeatAudience) return null + const fetchImpl = options.fetch ?? fetch + const tokenProvider = + options.identityToken ?? ((audience) => googleMetadataIdentityToken(audience, fetchImpl)) + const startedAt = (options.now ?? Date.now)() + const cellIncarnation = options.incarnation ?? randomUUID() + let stopped = false + let inFlight = false + + const send = async (): Promise => { + if (stopped || inFlight) return + inFlight = true + try { + const [token, ready] = await Promise.all([ + tokenProvider(config.heartbeatAudience!), + options.ready() + ]) + const connectionCounts = + config.connectionHardCap === undefined ? null : options.connectionCounts() + const response = await fetchImpl( + new URL('/v1/admin/cell-heartbeat', config.directorUrl).toString(), + { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: JSON.stringify({ + v: 1, + cellId: config.cellId, + cellUrl: config.cellUrl, + region: config.region ?? 'us-central1', + cellIncarnation, + startedAt, + ready, + observedRequests: options.observedRequests(), + ...(config.connectionHardCap === undefined + ? {} + : { + totalConnections: connectionCounts!.totalConnections, + inFlightConnections: connectionCounts!.inFlightConnections, + reservedConnectionUnits: connectionCounts!.reservedConnectionUnits, + enforcedConnectionUnits: connectionCounts!.enforcedConnectionUnits, + connectionInclusionWatermark: + connectionCounts!.inclusionWatermark, + connectionHardCap: config.connectionHardCap, + connectionUnobservedBound: config.connectionUnobservedBound + }) + }), + signal: AbortSignal.timeout(10_000) + } + ) + if (!response.ok) throw new Error(`director_heartbeat_${response.status}`) + if (options.regionalRehomeSafety) { + const statusResponse = await fetchImpl( + new URL('/v1/admin/cell-rehome-status', config.directorUrl).toString(), + { + method: 'POST', + headers: { + authorization: `Bearer ${token}`, + 'content-type': 'application/json' + }, + body: JSON.stringify({ + v: 1, + cellId: config.cellId, + cellIncarnation, + regionalRehomeProtocol: + config.rehomeAudience && config.rehomeDirectorServiceAccount ? 1 : 0, + safety: options.regionalRehomeSafety() + }), + signal: AbortSignal.timeout(10_000) + } + ) + if (!statusResponse.ok && statusResponse.status !== 404) { + throw new Error(`director_rehome_status_${statusResponse.status}`) + } + } + } catch (error) { + // A heartbeat must fail closed without ever logging its bearer token. + console.warn('[orca-relay] cell heartbeat failed', error instanceof Error ? error.message : '') + } finally { + inFlight = false + } + } + const timer = setInterval(() => void send(), options.intervalMs ?? 15_000) + timer.unref() + void send() + return { + send, + stop: () => { + stopped = true + clearInterval(timer) + } + } +} diff --git a/cloud/apps/relay/src/config.test.ts b/cloud/apps/relay/src/config.test.ts new file mode 100644 index 00000000000..bb0522dcbd3 --- /dev/null +++ b/cloud/apps/relay/src/config.test.ts @@ -0,0 +1,264 @@ +import { describe, expect, it } from 'vitest' +import { + loadRelayConfig, + RELAY_CELL_CONNECTION_HARD_CAP, + RELAY_DATABASE_POOL_MAX, + RELAY_DIRECTOR_DATABASE_POOL_MAX, + RELAY_MAX_CELL_CAPACITY_REQUESTS, + RELAY_PUBLIC_RESOLVE_CONCURRENCY, + RELAY_PUBLIC_RESOLVE_WAIT_MS +} from './config.js' + +function cellEnvironment(capacity: number): NodeJS.ProcessEnv { + return { + ORCA_RELAY_PUBLIC_URL: 'https://c1.relay.example.com', + ORCA_RELAY_CELL_URL: 'https://c1.relay.example.com', + ORCA_RELAY_AUTH_ISSUER: 'https://auth.example.com', + ORCA_RELAY_JWKS_URL: 'https://auth.example.com/.well-known/jwks.json', + ORCA_RELAY_ASSIGNMENT_SIGNING_KEY: 'assignment-key-with-at-least-thirty-two-bytes', + ORCA_RELAY_ROLE: 'cell', + ORCA_RELAY_CELL_ID: 'gce-c1', + ORCA_RELAY_CELL_CAPACITY: String(capacity), + ORCA_RELAY_ADMIN_AUDIENCE: 'https://relay.example.com/v1/admin/drain', + ORCA_RELAY_DEPLOY_SERVICE_ACCOUNT: 'deploy@example.iam.gserviceaccount.com' + } +} + +describe('GCE relay capacity configuration', () => { + it('requires distinct dedicated admin identities and accepts omitted values', () => { + const env = cellEnvironment(4_000) + expect(loadRelayConfig(env)).toMatchObject({ + capacityServiceAccount: undefined, + asiaProofServiceAccount: undefined, + monitorServiceAccount: undefined, + fenceServiceAccount: undefined + }) + env.ORCA_RELAY_MONITOR_SERVICE_ACCOUNT = '' + env.ORCA_RELAY_CAPACITY_SERVICE_ACCOUNT = 'capacity@example.iam.gserviceaccount.com' + env.ORCA_RELAY_ASIA_PROOF_SERVICE_ACCOUNT = 'proof@example.iam.gserviceaccount.com' + env.ORCA_RELAY_FENCE_SERVICE_ACCOUNT = 'fence@example.iam.gserviceaccount.com' + expect(loadRelayConfig(env)).toMatchObject({ + capacityServiceAccount: 'capacity@example.iam.gserviceaccount.com', + asiaProofServiceAccount: 'proof@example.iam.gserviceaccount.com', + monitorServiceAccount: undefined, + fenceServiceAccount: 'fence@example.iam.gserviceaccount.com' + }) + env.ORCA_RELAY_MONITOR_SERVICE_ACCOUNT = env.ORCA_RELAY_DEPLOY_SERVICE_ACCOUNT + expect(() => loadRelayConfig(env)).toThrow('relay admin service accounts must be distinct') + env.ORCA_RELAY_MONITOR_SERVICE_ACCOUNT = undefined + env.ORCA_RELAY_ASIA_PROOF_SERVICE_ACCOUNT = env.ORCA_RELAY_CAPACITY_SERVICE_ACCOUNT + expect(() => loadRelayConfig(env)).toThrow('relay admin service accounts must be distinct') + }) + + it('requires a distinct paired identity and exact audience for regional host drain', () => { + const env = cellEnvironment(4_000) + env.ORCA_RELAY_REHOME_DIRECTOR_SERVICE_ACCOUNT = + 'relay-director@example.iam.gserviceaccount.com' + expect(() => loadRelayConfig(env)).toThrow( + 'relay rehome identity and audience must be configured together' + ) + env.ORCA_RELAY_REHOME_AUDIENCE = 'https://relay.example.com/v1/admin/host-drain' + expect(loadRelayConfig(env)).toMatchObject({ + rehomeDirectorServiceAccount: 'relay-director@example.iam.gserviceaccount.com', + rehomeAudience: 'https://relay.example.com/v1/admin/host-drain' + }) + env.ORCA_RELAY_REHOME_AUDIENCE = 'https://relay.example.com/v1/admin/drain' + expect(() => loadRelayConfig(env)).toThrow( + 'relay rehome audience must target the host drain route' + ) + env.ORCA_RELAY_REHOME_AUDIENCE = 'https://relay.example.com/v1/admin/host-drain' + env.ORCA_RELAY_RUNTIME_SERVICE_ACCOUNT = + env.ORCA_RELAY_REHOME_DIRECTOR_SERVICE_ACCOUNT + expect(() => loadRelayConfig(env)).toThrow( + 'relay rehome director identity must differ from the cell runtime identity' + ) + }) + + it('accepts a measured capacity above the Cloud Run request ceiling', () => { + expect(loadRelayConfig(cellEnvironment(4_000)).cells[0]?.capacityRequests).toBe(4_000) + }) + + it('retains a finite process-wide capacity bound', () => { + expect(() => loadRelayConfig(cellEnvironment(RELAY_MAX_CELL_CAPACITY_REQUESTS + 1))).toThrow() + }) + + it('keeps legacy cells uncapped and requires a supported hard-cap pair', () => { + const env = cellEnvironment(4_000) + expect(loadRelayConfig(env)).toMatchObject({ + connectionHardCap: undefined, + connectionUnobservedBound: undefined + }) + + env.ORCA_RELAY_CELL_CONNECTION_HARD_CAP = String(RELAY_CELL_CONNECTION_HARD_CAP) + expect(() => loadRelayConfig(env)).toThrow( + 'connection hard cap and unobserved bound must be configured together' + ) + env.ORCA_RELAY_CELL_CONNECTION_UNOBSERVED_BOUND = '60' + expect(loadRelayConfig(env)).toMatchObject({ + connectionHardCap: 600, + connectionUnobservedBound: 60 + }) + env.ORCA_RELAY_CELL_CONNECTION_HARD_CAP = '599' + expect(() => loadRelayConfig(env)).toThrow() + env.ORCA_RELAY_CELL_CONNECTION_HARD_CAP = '600' + env.ORCA_RELAY_CELL_CONNECTION_UNOBSERVED_BOUND = '500' + expect(() => loadRelayConfig(env)).toThrow() + env.ORCA_RELAY_CELL_CONNECTION_UNOBSERVED_BOUND = '600' + expect(() => loadRelayConfig(env)).toThrow() + env.ORCA_RELAY_CELL_CONNECTION_HARD_CAP = '1000' + env.ORCA_RELAY_CELL_CONNECTION_UNOBSERVED_BOUND = '60' + expect(loadRelayConfig(env)).toMatchObject({ + connectionHardCap: 1_000, + connectionUnobservedBound: 60 + }) + env.ORCA_RELAY_CELL_CONNECTION_UNOBSERVED_BOUND = '900' + expect(() => loadRelayConfig(env)).toThrow() + }) + + it('accepts hard-cap metadata in director cell inventory', () => { + const env = cellEnvironment(4_000) + env.ORCA_RELAY_ROLE = 'director' + env.ORCA_RELAY_PUBLIC_URL = 'https://relay.example.com' + env.ORCA_RELAY_CELL_URL = 'https://relay.example.com' + env.ORCA_RELAY_CELLS_JSON = JSON.stringify([ + { + id: 'gce-c7', + url: 'https://c7.relay.example.com', + capacityRequests: 4_000, + connectionHardCap: 1_000, + connectionUnobservedBound: 60 + } + ]) + + expect(loadRelayConfig(env).cells[0]).toMatchObject({ + connectionHardCap: 1_000, + connectionUnobservedBound: 60 + }) + }) + + it('accepts mixed 600- and 1000-cap director inventory', () => { + const env = cellEnvironment(4_000) + env.ORCA_RELAY_ROLE = 'director' + env.ORCA_RELAY_PUBLIC_URL = 'https://relay.example.com' + env.ORCA_RELAY_CELL_URL = 'https://relay.example.com' + env.ORCA_RELAY_CELLS_JSON = JSON.stringify([ + { + id: 'gce-c1', + url: 'https://c1.relay.example.com', + capacityRequests: 4_000, + connectionHardCap: 600, + connectionUnobservedBound: 60 + }, + { + id: 'gce-c2', + url: 'https://c2.relay.example.com', + capacityRequests: 4_000, + connectionHardCap: 1_000, + connectionUnobservedBound: 60 + } + ]) + + expect(loadRelayConfig(env).cells.map((cell) => cell.connectionHardCap)).toEqual([ + 600, 1_000 + ]) + }) + + it('accepts only a canonical served image digest', () => { + const env = cellEnvironment(4_000) + env.ORCA_RELAY_IMAGE_DIGEST = `sha256:${'a'.repeat(64)}` + expect(loadRelayConfig(env).imageDigest).toBe(env.ORCA_RELAY_IMAGE_DIGEST) + env.ORCA_RELAY_IMAGE_DIGEST = 'relay:latest' + expect(() => loadRelayConfig(env)).toThrow() + }) + + it('accepts a statically declared candidate that starts disabled', () => { + const env = cellEnvironment(4_000) + env.ORCA_RELAY_ROLE = 'director' + env.ORCA_RELAY_PUBLIC_URL = 'https://relay.example.com' + env.ORCA_RELAY_CELL_URL = 'https://relay.example.com' + env.ORCA_RELAY_CELLS_JSON = JSON.stringify([ + { + id: 'gce-candidate', + url: 'https://candidate.relay.example.com', + capacityRequests: 4_000, + region: 'us-central1', + initiallyEnabled: false + } + ]) + expect(loadRelayConfig(env).cells).toEqual([ + { + id: 'gce-candidate', + url: 'https://candidate.relay.example.com', + capacityRequests: 4_000, + region: 'us-central1', + initiallyEnabled: false + } + ]) + }) + + it('reserves a smaller PostgreSQL connection budget for directors', () => { + const cellEnv = cellEnvironment(4_000) + expect(loadRelayConfig(cellEnv).databasePoolMax).toBe(RELAY_DATABASE_POOL_MAX) + + const directorEnv: NodeJS.ProcessEnv = { ...cellEnv, ORCA_RELAY_ROLE: 'director' } + directorEnv.ORCA_RELAY_CELLS_JSON = JSON.stringify([ + { + id: 'gce-c1', + url: 'https://c1.relay.example.com', + capacityRequests: 4_000 + } + ]) + expect(loadRelayConfig(directorEnv).databasePoolMax).toBe(RELAY_DIRECTOR_DATABASE_POOL_MAX) + + directorEnv.ORCA_RELAY_DATABASE_POOL_MAX = '7' + expect(loadRelayConfig(directorEnv).databasePoolMax).toBe(7) + }) + + it('defaults to bounded public assignment admission and supports an emergency stop', () => { + const env = cellEnvironment(4_000) + expect(loadRelayConfig(env)).toMatchObject({ + publicAssignmentsEnabled: true, + regionalPlacementEnabled: true, + publicAssignmentConcurrency: 2, + publicAssignmentQueueMax: 128, + publicAssignmentWaitMs: 4_000, + publicResolveConcurrency: RELAY_PUBLIC_RESOLVE_CONCURRENCY, + publicResolveWaitMs: RELAY_PUBLIC_RESOLVE_WAIT_MS, + publicAssignmentRetryAfterSeconds: 5 + }) + + env.ORCA_RELAY_PUBLIC_ASSIGNMENTS_ENABLED = 'false' + env.ORCA_RELAY_REGIONAL_PLACEMENT_ENABLED = 'false' + env.ORCA_RELAY_PUBLIC_ASSIGNMENT_CONCURRENCY = '4' + env.ORCA_RELAY_PUBLIC_ASSIGNMENT_QUEUE_MAX = '256' + env.ORCA_RELAY_PUBLIC_ASSIGNMENT_WAIT_MS = '8000' + env.ORCA_RELAY_PUBLIC_ASSIGNMENT_RETRY_AFTER_SECONDS = '30' + expect(loadRelayConfig(env)).toMatchObject({ + publicAssignmentsEnabled: false, + regionalPlacementEnabled: false, + publicAssignmentConcurrency: 4, + publicAssignmentQueueMax: 256, + publicAssignmentWaitMs: 8_000, + publicResolveConcurrency: RELAY_PUBLIC_RESOLVE_CONCURRENCY, + publicResolveWaitMs: RELAY_PUBLIC_RESOLVE_WAIT_MS, + publicAssignmentRetryAfterSeconds: 30 + }) + }) + + it('fails closed when public request lanes exceed the director database pool', () => { + const env = cellEnvironment(4_000) + env.ORCA_RELAY_ROLE = 'director' + env.ORCA_RELAY_CELLS_JSON = JSON.stringify([ + { + id: 'gce-c1', + url: 'https://c1.relay.example.com', + capacityRequests: 4_000 + } + ]) + env.ORCA_RELAY_DATABASE_POOL_MAX = '2' + + expect(() => loadRelayConfig(env)).toThrow( + 'public relay admission must leave database pool headroom' + ) + }) +}) diff --git a/cloud/apps/relay/src/config.ts b/cloud/apps/relay/src/config.ts new file mode 100644 index 00000000000..2bf23444a71 --- /dev/null +++ b/cloud/apps/relay/src/config.ts @@ -0,0 +1,348 @@ +import { z } from 'zod' +import { + isRelayCellConnectionHardCap, + RELAY_CELL_CONNECTION_HARD_CAP, + RELAY_MAX_CELL_CONNECTION_UNOBSERVED_BOUND, + RELAY_DEFAULT_REGION, + RelayRegionSchema, + relayCellAdmissionBounds, + type RelayCellConnectionHardCap, + type RelayRegion +} from '@orca-cloud/relay-contract' + +export const RELAY_MAX_CELL_CAPACITY_REQUESTS = 100_000 +export const RELAY_DATABASE_POOL_MAX = 10 +export const RELAY_DIRECTOR_DATABASE_POOL_MAX = 3 +export const RELAY_PUBLIC_RESOLVE_CONCURRENCY = 1 +export const RELAY_PUBLIC_RESOLVE_WAIT_MS = 5_000 +export { RELAY_CELL_CONNECTION_HARD_CAP } + +const RelayCellConnectionHardCapSchema = z.custom( + isRelayCellConnectionHardCap +) + +const EnvironmentBooleanSchema = z + .enum(['true', 'false']) + .default('true') + .transform((value) => value === 'true') + +const OptionalServiceAccountSchema = z.preprocess( + (value) => (value === '' ? undefined : value), + z.string().email().optional() +) + +const EnvSchema = z.object({ + PORT: z.coerce.number().int().positive().default(8080), + ORCA_RELAY_PUBLIC_URL: z.string().url(), + ORCA_RELAY_CELL_URL: z.string().url(), + ORCA_RELAY_AUTH_ISSUER: z.string().url(), + ORCA_RELAY_AUTH_AUDIENCE: z.literal('orca-relay').default('orca-relay'), + ORCA_RELAY_JWKS_URL: z.string().url(), + ORCA_RELAY_ASSIGNMENT_SIGNING_KEY: z.string().min(32), + ORCA_RELAY_ROLE: z.enum(['combined', 'director', 'cell']).default('combined'), + ORCA_RELAY_CELL_ID: z.string().min(1).max(128).default('combined'), + ORCA_RELAY_REGION: RelayRegionSchema.default(RELAY_DEFAULT_REGION), + ORCA_RELAY_CELL_CAPACITY: z.coerce + .number() + .int() + .positive() + .max(RELAY_MAX_CELL_CAPACITY_REQUESTS) + .default(900), + ORCA_RELAY_CELL_CONNECTION_HARD_CAP: z.coerce + .number() + .int() + .pipe(RelayCellConnectionHardCapSchema) + .optional(), + ORCA_RELAY_CELL_CONNECTION_UNOBSERVED_BOUND: z.coerce + .number() + .int() + .nonnegative() + .max(RELAY_MAX_CELL_CONNECTION_UNOBSERVED_BOUND) + .optional(), + ORCA_RELAY_CELLS_JSON: z.string().default('[]'), + ORCA_RELAY_ADMIN_AUDIENCE: z.string().url(), + ORCA_RELAY_DEPLOY_SERVICE_ACCOUNT: z.string().email(), + ORCA_RELAY_CAPACITY_SERVICE_ACCOUNT: OptionalServiceAccountSchema, + ORCA_RELAY_ASIA_PROOF_SERVICE_ACCOUNT: OptionalServiceAccountSchema, + ORCA_RELAY_MONITOR_SERVICE_ACCOUNT: OptionalServiceAccountSchema, + ORCA_RELAY_FENCE_SERVICE_ACCOUNT: OptionalServiceAccountSchema, + ORCA_RELAY_FENCE_BROKER_SERVICE_ACCOUNT: OptionalServiceAccountSchema, + ORCA_RELAY_REHOME_DIRECTOR_SERVICE_ACCOUNT: OptionalServiceAccountSchema, + ORCA_RELAY_REHOME_AUDIENCE: z.preprocess( + (value) => (value === '' ? undefined : value), + z.string().url().optional() + ), + ORCA_RELAY_RUNTIME_SERVICE_ACCOUNT: z.string().email().optional(), + ORCA_RELAY_DIRECTOR_URL: z.string().url().optional(), + ORCA_RELAY_HEARTBEAT_AUDIENCE: z.string().url().optional(), + ORCA_RELAY_IMAGE_DIGEST: z.string().regex(/^sha256:[a-f0-9]{64}$/).optional(), + ORCA_RELAY_ADMIN_JWKS_URL: z.string().url().default('https://www.googleapis.com/oauth2/v3/certs'), + ORCA_RELAY_DATABASE_POOL_MAX: z.coerce.number().int().positive().max(100).optional(), + ORCA_RELAY_PUBLIC_ASSIGNMENTS_ENABLED: EnvironmentBooleanSchema, + ORCA_RELAY_REGIONAL_PLACEMENT_ENABLED: EnvironmentBooleanSchema, + ORCA_RELAY_PUBLIC_ASSIGNMENT_CONCURRENCY: z.coerce.number().int().positive().max(100).default(2), + ORCA_RELAY_PUBLIC_STICKY_CONCURRENCY: z.coerce.number().int().positive().max(100).default(1), + ORCA_RELAY_PUBLIC_STICKY_QUEUE_MAX: z.coerce.number().int().positive().max(4_096).default(64), + ORCA_RELAY_PUBLIC_STICKY_WAIT_MS: z.coerce.number().int().positive().max(30_000).default(2_000), + ORCA_RELAY_PUBLIC_STICKY_RETRY_AFTER_SECONDS: z.coerce + .number() + .int() + .positive() + .max(60) + .default(2), + ORCA_RELAY_PUBLIC_ASSIGNMENT_QUEUE_MAX: z.coerce + .number() + .int() + .positive() + .max(4_096) + .default(128), + ORCA_RELAY_PUBLIC_ASSIGNMENT_WAIT_MS: z.coerce + .number() + .int() + .positive() + .max(30_000) + .default(4_000), + ORCA_RELAY_PUBLIC_ASSIGNMENT_RETRY_AFTER_SECONDS: z.coerce + .number() + .int() + .positive() + .max(300) + .default(5), + DATABASE_URL: z.string().optional(), + ORCA_RELAY_DATA_DIR: z.string().default('./data/relay') +}) + +const RelayCellConfigSchema = z + .object({ + id: z.string().min(1).max(128), + url: z.string().url(), + capacityRequests: z.number().int().positive().max(RELAY_MAX_CELL_CAPACITY_REQUESTS), + region: RelayRegionSchema.default(RELAY_DEFAULT_REGION), + initiallyEnabled: z.boolean().optional(), + connectionHardCap: RelayCellConnectionHardCapSchema.optional(), + connectionUnobservedBound: z + .number() + .int() + .nonnegative() + .max(RELAY_MAX_CELL_CONNECTION_UNOBSERVED_BOUND) + .optional() + }) + .strict() + .superRefine((value, context) => { + if ( + (value.connectionHardCap === undefined) !== + (value.connectionUnobservedBound === undefined) + ) { + context.addIssue({ + code: 'custom', + message: 'connection hard cap and unobserved bound must be configured together' + }) + } else if ( + value.connectionHardCap !== undefined && + value.connectionUnobservedBound! > + relayCellAdmissionBounds(value.connectionHardCap).maxUnobservedBound + ) { + context.addIssue({ + code: 'custom', + path: ['connectionUnobservedBound'], + message: 'connection unobserved bound must leave ordinary admission capacity' + }) + } + }) + +export type RelayCellConfig = Omit, 'region'> & { + region?: RelayRegion +} + +export type RelayConfig = { + port: number + publicUrl: string + cellUrl: string + authIssuer: string + authAudience: 'orca-relay' + jwksUrl: string + assignmentSigningKey: Uint8Array + role: 'combined' | 'director' | 'cell' + cellId: string + region?: RelayRegion + cells: RelayCellConfig[] + adminAudience: string + deployServiceAccount: string + capacityServiceAccount?: string + asiaProofServiceAccount?: string + monitorServiceAccount?: string + fenceServiceAccount?: string + fenceBrokerServiceAccount?: string + rehomeDirectorServiceAccount?: string + rehomeAudience?: string + runtimeServiceAccount: string + directorUrl?: string + heartbeatAudience?: string + imageDigest?: string + connectionHardCap?: RelayCellConnectionHardCap + connectionUnobservedBound?: number + adminJwksUrl: string + databasePoolMax: number + publicAssignmentsEnabled: boolean + regionalPlacementEnabled?: boolean + publicAssignmentConcurrency: number + publicAssignmentQueueMax: number + publicAssignmentWaitMs: number + publicResolveConcurrency: number + publicResolveWaitMs: number + publicAssignmentRetryAfterSeconds: number + publicStickyConcurrency?: number + publicStickyQueueMax?: number + publicStickyWaitMs?: number + publicStickyRetryAfterSeconds?: number + databaseUrl?: string + dataDir: string +} + +function canonicalOrigin(value: string, name: string): string { + const url = new URL(value) + if (url.origin !== value || url.pathname !== '/') throw new Error(`${name} must be an origin`) + const loopback = ['127.0.0.1', 'localhost', '::1', '[::1]'].includes(url.hostname) + if (url.protocol !== 'https:' && !(loopback && url.protocol === 'http:')) { + throw new Error(`${name} must use HTTPS outside loopback development`) + } + return value +} + +export function loadRelayConfig(env: NodeJS.ProcessEnv = process.env): RelayConfig { + const parsed = EnvSchema.parse(env) + const adminServiceAccounts = [ + parsed.ORCA_RELAY_DEPLOY_SERVICE_ACCOUNT, + parsed.ORCA_RELAY_CAPACITY_SERVICE_ACCOUNT, + parsed.ORCA_RELAY_ASIA_PROOF_SERVICE_ACCOUNT, + parsed.ORCA_RELAY_MONITOR_SERVICE_ACCOUNT, + parsed.ORCA_RELAY_FENCE_SERVICE_ACCOUNT, + parsed.ORCA_RELAY_FENCE_BROKER_SERVICE_ACCOUNT, + parsed.ORCA_RELAY_REHOME_DIRECTOR_SERVICE_ACCOUNT + ].filter((value): value is string => value !== undefined) + if (new Set(adminServiceAccounts).size !== adminServiceAccounts.length) { + throw new Error('relay admin service accounts must be distinct') + } + if ( + (parsed.ORCA_RELAY_REHOME_DIRECTOR_SERVICE_ACCOUNT === undefined) !== + (parsed.ORCA_RELAY_REHOME_AUDIENCE === undefined) + ) { + throw new Error('relay rehome identity and audience must be configured together') + } + if ( + parsed.ORCA_RELAY_REHOME_DIRECTOR_SERVICE_ACCOUNT && + parsed.ORCA_RELAY_REHOME_DIRECTOR_SERVICE_ACCOUNT === + (parsed.ORCA_RELAY_RUNTIME_SERVICE_ACCOUNT ?? parsed.ORCA_RELAY_DEPLOY_SERVICE_ACCOUNT) + ) { + throw new Error('relay rehome director identity must differ from the cell runtime identity') + } + if ( + parsed.ORCA_RELAY_REHOME_AUDIENCE && + new URL(parsed.ORCA_RELAY_REHOME_AUDIENCE).pathname !== '/v1/admin/host-drain' + ) { + throw new Error('relay rehome audience must target the host drain route') + } + const directorUrl = parsed.ORCA_RELAY_DIRECTOR_URL + ? canonicalOrigin(parsed.ORCA_RELAY_DIRECTOR_URL, 'ORCA_RELAY_DIRECTOR_URL') + : undefined + const publicUrl = canonicalOrigin(parsed.ORCA_RELAY_PUBLIC_URL, 'ORCA_RELAY_PUBLIC_URL') + const configuredCells = z + .array(RelayCellConfigSchema) + .max(128) + .parse(JSON.parse(parsed.ORCA_RELAY_CELLS_JSON) as unknown) + .map((cell) => ({ ...cell, url: canonicalOrigin(cell.url, `cell ${cell.id}`) })) + const ownCell = { + id: parsed.ORCA_RELAY_CELL_ID, + region: parsed.ORCA_RELAY_REGION, + url: canonicalOrigin(parsed.ORCA_RELAY_CELL_URL, 'ORCA_RELAY_CELL_URL'), + capacityRequests: parsed.ORCA_RELAY_CELL_CAPACITY, + connectionHardCap: parsed.ORCA_RELAY_CELL_CONNECTION_HARD_CAP, + connectionUnobservedBound: parsed.ORCA_RELAY_CELL_CONNECTION_UNOBSERVED_BOUND, + initiallyEnabled: true + } + if ( + (ownCell.connectionHardCap === undefined) !== + (ownCell.connectionUnobservedBound === undefined) + ) { + throw new Error('connection hard cap and unobserved bound must be configured together') + } + if ( + ownCell.connectionHardCap !== undefined && + ownCell.connectionUnobservedBound! > + relayCellAdmissionBounds(ownCell.connectionHardCap).maxUnobservedBound + ) { + throw new Error('connection unobserved bound must leave ordinary admission capacity') + } + const cells = parsed.ORCA_RELAY_ROLE === 'director' ? configuredCells : [ownCell] + if (cells.length === 0) throw new Error('director requires at least one configured cell') + if (new Set(cells.map(({ id }) => id)).size !== cells.length) { + throw new Error('relay cell ids must be unique') + } + const databasePoolMax = + parsed.ORCA_RELAY_DATABASE_POOL_MAX ?? + (parsed.ORCA_RELAY_ROLE === 'director' + ? RELAY_DIRECTOR_DATABASE_POOL_MAX + : RELAY_DATABASE_POOL_MAX) + if ( + parsed.ORCA_RELAY_ROLE !== 'cell' && + parsed.ORCA_RELAY_PUBLIC_ASSIGNMENT_CONCURRENCY >= databasePoolMax + ) { + throw new Error('public relay admission must leave database pool headroom') + } + if ( + parsed.ORCA_RELAY_ROLE !== 'cell' && + parsed.ORCA_RELAY_PUBLIC_ASSIGNMENT_CONCURRENCY + parsed.ORCA_RELAY_PUBLIC_STICKY_CONCURRENCY > + databasePoolMax + ) { + throw new Error('sticky and placement admission together must fit the database pool') + } + return { + port: parsed.PORT, + publicUrl, + cellUrl: ownCell.url, + authIssuer: canonicalOrigin(parsed.ORCA_RELAY_AUTH_ISSUER, 'ORCA_RELAY_AUTH_ISSUER'), + authAudience: parsed.ORCA_RELAY_AUTH_AUDIENCE, + jwksUrl: parsed.ORCA_RELAY_JWKS_URL, + assignmentSigningKey: new TextEncoder().encode(parsed.ORCA_RELAY_ASSIGNMENT_SIGNING_KEY), + role: parsed.ORCA_RELAY_ROLE, + cellId: ownCell.id, + region: ownCell.region, + cells, + adminAudience: parsed.ORCA_RELAY_ADMIN_AUDIENCE, + deployServiceAccount: parsed.ORCA_RELAY_DEPLOY_SERVICE_ACCOUNT, + capacityServiceAccount: parsed.ORCA_RELAY_CAPACITY_SERVICE_ACCOUNT, + asiaProofServiceAccount: parsed.ORCA_RELAY_ASIA_PROOF_SERVICE_ACCOUNT, + monitorServiceAccount: parsed.ORCA_RELAY_MONITOR_SERVICE_ACCOUNT, + fenceServiceAccount: parsed.ORCA_RELAY_FENCE_SERVICE_ACCOUNT, + fenceBrokerServiceAccount: parsed.ORCA_RELAY_FENCE_BROKER_SERVICE_ACCOUNT, + rehomeDirectorServiceAccount: parsed.ORCA_RELAY_REHOME_DIRECTOR_SERVICE_ACCOUNT, + rehomeAudience: parsed.ORCA_RELAY_REHOME_AUDIENCE, + runtimeServiceAccount: + parsed.ORCA_RELAY_RUNTIME_SERVICE_ACCOUNT ?? parsed.ORCA_RELAY_DEPLOY_SERVICE_ACCOUNT, + directorUrl, + heartbeatAudience: + parsed.ORCA_RELAY_HEARTBEAT_AUDIENCE ?? + (directorUrl || parsed.ORCA_RELAY_ROLE === 'director' + ? new URL('/v1/admin/cell-heartbeat', directorUrl ?? publicUrl).toString() + : undefined), + imageDigest: parsed.ORCA_RELAY_IMAGE_DIGEST, + connectionHardCap: ownCell.connectionHardCap, + connectionUnobservedBound: ownCell.connectionUnobservedBound, + adminJwksUrl: parsed.ORCA_RELAY_ADMIN_JWKS_URL, + databasePoolMax, + publicAssignmentsEnabled: parsed.ORCA_RELAY_PUBLIC_ASSIGNMENTS_ENABLED, + regionalPlacementEnabled: parsed.ORCA_RELAY_REGIONAL_PLACEMENT_ENABLED, + publicAssignmentConcurrency: parsed.ORCA_RELAY_PUBLIC_ASSIGNMENT_CONCURRENCY, + publicAssignmentQueueMax: parsed.ORCA_RELAY_PUBLIC_ASSIGNMENT_QUEUE_MAX, + publicAssignmentWaitMs: parsed.ORCA_RELAY_PUBLIC_ASSIGNMENT_WAIT_MS, + publicResolveConcurrency: RELAY_PUBLIC_RESOLVE_CONCURRENCY, + publicResolveWaitMs: RELAY_PUBLIC_RESOLVE_WAIT_MS, + publicAssignmentRetryAfterSeconds: parsed.ORCA_RELAY_PUBLIC_ASSIGNMENT_RETRY_AFTER_SECONDS, + publicStickyConcurrency: parsed.ORCA_RELAY_PUBLIC_STICKY_CONCURRENCY, + publicStickyQueueMax: parsed.ORCA_RELAY_PUBLIC_STICKY_QUEUE_MAX, + publicStickyWaitMs: parsed.ORCA_RELAY_PUBLIC_STICKY_WAIT_MS, + publicStickyRetryAfterSeconds: parsed.ORCA_RELAY_PUBLIC_STICKY_RETRY_AFTER_SECONDS, + databaseUrl: parsed.DATABASE_URL, + dataDir: parsed.ORCA_RELAY_DATA_DIR + } +} diff --git a/cloud/apps/relay/src/connection-reservation-debt-retention.test.ts b/cloud/apps/relay/src/connection-reservation-debt-retention.test.ts new file mode 100644 index 00000000000..fd50f9d6747 --- /dev/null +++ b/cloud/apps/relay/src/connection-reservation-debt-retention.test.ts @@ -0,0 +1,130 @@ +import { afterEach, describe, expect, it } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.js' +import type { RelayCellConfig } from './config.js' +import { openInMemoryRelayDatabase, type RelayDatabase } from './database.js' + +const DEBT_RETENTION_MS = 10 * 60 * 1_000 + +const LIMITED_CELL: RelayCellConfig = { + id: 'limited', + url: 'https://limited.example.com', + capacityRequests: 1_000, + connectionHardCap: 600, + connectionUnobservedBound: 50 +} + +const databases: RelayDatabase[] = [] + +afterEach(async () => { + for (const database of databases.splice(0)) await database.close() +}) + +async function setup(now: () => number): Promise<{ + database: RelayDatabase + store: RelayAssignmentStore +}> { + const database = await openInMemoryRelayDatabase() + databases.push(database) + const store = new RelayAssignmentStore(database, now, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + await store.reconcileCells([LIMITED_CELL], true) + await store.recordCellHeartbeat({ + cellId: LIMITED_CELL.id, + cellUrl: LIMITED_CELL.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0, + connectionHardCap: 600, + connectionUnobservedBound: 50 + }) + return { database, store } +} + +async function insertDebt( + database: RelayDatabase, + reservationId: string, + relayHostId: string, + timeoutAt: number, + claimActivityId: string | null = null +): Promise { + await database.query( + `INSERT INTO relay_control_connection_reservations + (reservation_id, idempotency_key, user_id, relay_host_id, assignment_epoch, + cell_id, state, claim_activity_id, created_at, timeout_at, updated_at) + VALUES (?, ?, 'user-1', ?, 1, ?, 'late-arrival-debt', ?, ?, ?, ?)`, + [ + reservationId, + reservationId, + relayHostId, + LIMITED_CELL.id, + claimActivityId, + timeoutAt - 10_000, + timeoutAt, + timeoutAt + ] + ) +} + +async function reservationStates(database: RelayDatabase): Promise> { + const rows = await database.query( + `SELECT reservation_id, state FROM relay_control_connection_reservations` + ) + return new Map(rows.map((row) => [String(row['reservation_id']), String(row['state'])])) +} + +describe('late-arrival debt retention', () => { + it('releases unclaimed debt past retention and keeps fresh or claimed debt', async () => { + const now = 100_000_000 + const { database, store } = await setup(() => now) + await insertDebt(database, 'stale-debt', 'host000000000001', now - DEBT_RETENTION_MS - 1) + await insertDebt(database, 'fresh-debt', 'host000000000002', now - 30_000) + await insertDebt( + database, + 'claimed-debt', + 'host000000000003', + now - DEBT_RETENTION_MS - 1, + 'control:1' + ) + + await store.releaseExpiredActivityLeases() + + expect(await reservationStates(database)).toEqual( + new Map([ + ['stale-debt', 'released'], + ['fresh-debt', 'late-arrival-debt'], + ['claimed-debt', 'late-arrival-debt'] + ]) + ) + }) + + it('restores placement once stale debt no longer consumes connection headroom', async () => { + const now = 100_000_000 + const { database, store } = await setup(() => now) + // Headroom needs enforced + outstanding + unobserved(50) < hardCap(600) - + // rebindReserve(100); 450 stale debt rows are exactly enough to block. + for (let index = 0; index < 450; index++) { + await insertDebt( + database, + `stale-${index}`, + `host${String(index).padStart(12, '0')}`, + now - DEBT_RETENTION_MS - 1 + ) + } + await expect( + store.assign({ userId: 'user-9', relayHostId: 'hostfffffffffff9' }) + ).rejects.toThrow('relay_capacity_exhausted') + + await store.releaseExpiredActivityLeases() + + await expect( + store.assign({ userId: 'user-9', relayHostId: 'hostfffffffffff9' }) + ).resolves.toMatchObject({ cellId: LIMITED_CELL.id }) + }) +}) diff --git a/cloud/apps/relay/src/control-lease-recovery-postgres.test.ts b/cloud/apps/relay/src/control-lease-recovery-postgres.test.ts new file mode 100644 index 00000000000..2fa01c27df8 --- /dev/null +++ b/cloud/apps/relay/src/control-lease-recovery-postgres.test.ts @@ -0,0 +1,223 @@ +import { EventEmitter } from 'node:events' +import { ASSIGNMENT_LIMITS, RELAY_CLOSE_CODE } from '@orca-cloud/relay-contract' +import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' +import type WebSocket from 'ws' +import { RelayAssignmentStore } from './assignment-store.js' +import type { RelayConfig } from './config.js' +import type { RelayCredentialStore } from './credential-store.js' +import { openRelayDatabase, type RelayDatabase } from './database.js' +import { HostSessionRegistry, type HostSession } from './host-session-registry.js' +import type { RelayRuntimeObserver } from './relay-observability.js' +import type { RelayTokenClaims } from './relay-token-verifier.js' +import { ProcessQueuedByteBudget } from './splice-forwarder.js' + +const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL +const describePostgres = databaseUrl ? describe : describe.skip + +const sourceCell = { + id: 'control-recovery-source', + url: 'https://control-recovery-source.example.com', + capacityRequests: 100 +} +const targetCell = { + id: 'control-recovery-target', + url: 'https://control-recovery-target.example.com', + capacityRequests: 100 +} +const userId = 'control-recovery-user' +const identities = ['controlrecovery1', 'controlrecovery2'].map((relayHostId) => ({ + userId, + relayHostId +})) + +class FakeSocket extends EventEmitter { + readonly OPEN = 1 + readonly CLOSED = 3 + readyState = this.OPEN + readonly send = vi.fn() + readonly close = vi.fn((code?: number, reason?: string) => { + this.readyState = this.CLOSED + this.emit('close', code, Buffer.from(reason ?? '')) + }) +} + +const observer = { + recordAuth: vi.fn(), + recordForwardedBytes: vi.fn(), + recordHttp: vi.fn(), + recordReconnect: vi.fn(), + recordSql: vi.fn() +} satisfies RelayRuntimeObserver + +type RegistryInternals = { + activate( + socket: WebSocket, + identity: RelayTokenClaims, + existing: HostSession | null, + generation: number, + rebind: boolean, + assignmentEpoch: number, + appVersion: string + ): Promise + heartbeat(session: HostSession): void +} + +describePostgres('expired control lease after a database outage', () => { + let database: RelayDatabase + let now = 1_900_000_000_000 + + const removeFixtureRows = async (): Promise => { + await database.query(`DELETE FROM relay_control_connection_reservations WHERE user_id = ?`, [ + userId + ]) + await database.query(`DELETE FROM relay_assignment_activity_leases WHERE user_id = ?`, [userId]) + await database.query(`DELETE FROM relay_assignment_migrations WHERE user_id = ?`, [userId]) + await database.query(`DELETE FROM relay_assignments WHERE user_id = ?`, [userId]) + for (const cell of [sourceCell, targetCell]) { + await database.query(`DELETE FROM relay_cell_connection_snapshots WHERE cell_id = ?`, [cell.id]) + await database.query(`DELETE FROM relay_cell_connection_runtime WHERE cell_id = ?`, [cell.id]) + await database.query(`DELETE FROM relay_cell_connection_limits WHERE cell_id = ?`, [cell.id]) + await database.query(`DELETE FROM relay_cell_runtime WHERE cell_id = ?`, [cell.id]) + await database.query(`DELETE FROM relay_cells WHERE cell_id = ?`, [cell.id]) + } + } + + const createRegistry = (store: RelayAssignmentStore): HostSessionRegistry => + new HostSessionRegistry( + { + role: 'cell', + cellId: sourceCell.id + } as RelayConfig, + vi.fn(), + {} as RelayCredentialStore, + store, + new ProcessQueuedByteBudget(), + observer, + () => now + ) + + const activate = async ( + registry: HostSessionRegistry, + store: RelayAssignmentStore, + index: number + ): Promise<{ activityId: string; session: HostSession; socket: FakeSocket }> => { + const identity = identities[index]! + const assignment = await store.assign(identity) + const socket = new FakeSocket() + const token = { + sub: identity.userId, + prof: 'control-recovery-profile', + relayHostId: identity.relayHostId, + purpose: 'host-control', + exp: 4_102_444_800 + } satisfies RelayTokenClaims + await (registry as unknown as RegistryInternals).activate( + socket as unknown as WebSocket, + token, + null, + 1, + false, + assignment.assignmentEpoch, + 'control-recovery-test' + ) + const session = registry.get(identity)! + clearInterval(session.heartbeatTimer!) + session.heartbeatTimer = null + return { activityId: session.controlActivityId!, session, socket } + } + + const heartbeat = (registry: HostSessionRegistry, session: HostSession): void => { + session.activityRenewalDueAt = now + session.lastPongAt = now + const internals = registry as unknown as RegistryInternals + internals.heartbeat(session) + } + + const leaseRows = async (relayHostId: string) => + await database.query( + `SELECT activity_id, cell_id FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ?`, + [userId, relayHostId] + ) + + const waitForRenewal = async ( + socket: FakeSocket, + relayHostId: string, + expectedRows: number + ): Promise => { + await expect + .poll( + async () => + socket.close.mock.calls.length > 0 || + (await leaseRows(relayHostId)).length === expectedRows + ) + .toBe(true) + } + + beforeAll(async () => { + database = await openRelayDatabase({ databaseUrl, dataDir: '' }) + }) + + beforeEach(async () => { + now = 1_900_000_000_000 + await removeFixtureRows() + }) + + afterEach(async () => { + await removeFixtureRows() + }) + + afterAll(async () => { + if (database) await database.close() + }) + + it('re-acquires a reaped lease while the assignment remains on this cell', async () => { + const store = new RelayAssignmentStore(database, () => now) + await store.reconcileCells([sourceCell]) + const registry = createRegistry(store) + const { activityId, session, socket } = await activate(registry, store, 0) + + now += ASSIGNMENT_LIMITS.activityLeaseMs + 1 + await store.releaseExpiredActivityLeases() + expect(await leaseRows(identities[0]!.relayHostId)).toHaveLength(0) + + heartbeat(registry, session) + await waitForRenewal(socket, identities[0]!.relayHostId, 1) + + expect(socket.close).not.toHaveBeenCalled() + expect(await leaseRows(identities[0]!.relayHostId)).toEqual([ + expect.objectContaining({ activity_id: activityId, cell_id: sourceCell.id }) + ]) + registry.drain(0) + }) + + it('closes without stealing back a lease that moved to another cell', async () => { + const store = new RelayAssignmentStore(database, () => now) + await store.reconcileCells([sourceCell, targetCell]) + const registry = createRegistry(store) + const { activityId, session, socket } = await activate(registry, store, 1) + const identity = identities[1]! + await database.query( + `UPDATE relay_assignment_activity_leases SET cell_id = ? + WHERE user_id = ? AND relay_host_id = ? AND activity_id = ?`, + [targetCell.id, identity.userId, identity.relayHostId, activityId] + ) + + await expect( + store.renewControlActivity(identity, { + activityId, + cellId: sourceCell.id, + expiresAt: now + ASSIGNMENT_LIMITS.activityLeaseMs + }) + ).rejects.toThrow('control_activity_moved') + + heartbeat(registry, session) + await expect.poll(() => socket.close.mock.calls.length).toBe(1) + + expect(socket.close).toHaveBeenCalledWith(RELAY_CLOSE_CODE.DRAINING, 'control activity moved') + expect(await leaseRows(identity.relayHostId)).toEqual([ + expect.objectContaining({ activity_id: activityId, cell_id: targetCell.id }) + ]) + registry.drain(0) + }) +}) diff --git a/cloud/apps/relay/src/control-renewal-postgres.test.ts b/cloud/apps/relay/src/control-renewal-postgres.test.ts new file mode 100644 index 00000000000..fb49e3af3a2 --- /dev/null +++ b/cloud/apps/relay/src/control-renewal-postgres.test.ts @@ -0,0 +1,349 @@ +import { ASSIGNMENT_LIMITS, RELAY_PROTOCOL_LIMITS } from '@orca-cloud/relay-contract' +import { afterAll, beforeAll, describe, expect, it } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.js' +import { + openRelayDatabase, + type RelayDatabase, + type RelayLockOptions, + type SqlRow +} from './database.js' + +const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL +const describePostgres = databaseUrl ? describe : describe.skip + +const sourceCell = { + id: 'control-renewal-source', + url: 'https://control-renewal-source.example.com', + capacityRequests: 100 +} +const targetCell = { + id: 'control-renewal-target', + url: 'https://control-renewal-target.example.com', + capacityRequests: 100 +} +const userId = 'control-renewal-postgres-user' +const identities = Array.from({ length: 6 }, (_, index) => ({ + userId, + relayHostId: `controlrenewal${index + 1}` +})) + +function signal(): { promise: Promise; resolve: () => void } { + let resolve!: () => void + return { promise: new Promise((done) => (resolve = done)), resolve } +} + +class StallFirstTransactionDatabase implements RelayDatabase { + readonly stalled = signal() + readonly continue = signal() + private stallNext = true + + constructor(private readonly database: RelayDatabase) {} + + async query(sql: string, params?: unknown[]): Promise { + return await this.database.query(sql, params) + } + + async queryLocked( + sql: string, + params?: unknown[], + options?: RelayLockOptions + ): Promise { + return await this.database.queryLocked(sql, params, options) + } + + async transaction(operation: (transaction: RelayDatabase) => Promise): Promise { + if (this.stallNext) { + this.stallNext = false + this.stalled.resolve() + await this.continue.promise + } + return await this.database.transaction(operation) + } + + async close(): Promise {} +} + +class StallFirstRenewalQueryDatabase implements RelayDatabase { + readonly dialect = 'postgres' as const + readonly stalled = signal() + readonly continue = signal() + private stallNext = true + + constructor(private readonly database: RelayDatabase) {} + + async query(sql: string, params?: unknown[]): Promise { + if (this.stallNext && sql.includes('WITH assignment_state AS MATERIALIZED')) { + this.stallNext = false + this.stalled.resolve() + await this.continue.promise + } + return await this.database.query(sql, params) + } + + async queryLocked( + sql: string, + params?: unknown[], + options?: RelayLockOptions + ): Promise { + return await this.database.queryLocked(sql, params, options) + } + + async transaction(operation: (transaction: RelayDatabase) => Promise): Promise { + return await this.database.transaction(operation) + } + + async close(): Promise {} +} + +class RenewalQueryProbeDatabase implements RelayDatabase { + readonly dialect = 'postgres' as const + renewalQueries = 0 + transactions = 0 + + constructor(private readonly database: RelayDatabase) {} + + async query(sql: string, params?: unknown[]): Promise { + if (sql.includes('WITH assignment_state AS MATERIALIZED')) this.renewalQueries++ + return await this.database.query(sql, params) + } + + async queryLocked( + sql: string, + params?: unknown[], + options?: RelayLockOptions + ): Promise { + return await this.database.queryLocked(sql, params, options) + } + + async transaction(operation: (transaction: RelayDatabase) => Promise): Promise { + this.transactions++ + return await this.database.transaction(operation) + } + + async close(): Promise {} +} + +describePostgres('PostgreSQL control renewal', () => { + let database: RelayDatabase + let now = 1_900_000_000_000 + const controlId = (cellId: string): string => `control:${cellId}:1` + + const removeFixtureRows = async (): Promise => { + await database.query(`DELETE FROM relay_control_connection_reservations WHERE user_id = ?`, [ + userId + ]) + await database.query(`DELETE FROM relay_assignment_activity_leases WHERE user_id = ?`, [userId]) + await database.query(`DELETE FROM relay_assignment_migrations WHERE user_id = ?`, [userId]) + await database.query(`DELETE FROM relay_assignments WHERE user_id = ?`, [userId]) + for (const cell of [sourceCell, targetCell]) { + await database.query(`DELETE FROM relay_cell_connection_snapshots WHERE cell_id = ?`, [cell.id]) + await database.query(`DELETE FROM relay_cell_connection_runtime WHERE cell_id = ?`, [cell.id]) + await database.query(`DELETE FROM relay_cell_connection_limits WHERE cell_id = ?`, [cell.id]) + await database.query(`DELETE FROM relay_cell_runtime WHERE cell_id = ?`, [cell.id]) + await database.query(`DELETE FROM relay_cells WHERE cell_id = ?`, [cell.id]) + } + } + + beforeAll(async () => { + database = await openRelayDatabase({ databaseUrl, dataDir: '' }) + await removeFixtureRows() + const store = new RelayAssignmentStore(database, () => now) + await store.reconcileCells([sourceCell]) + for (const identity of identities) { + const assignment = await store.assign(identity) + expect(assignment.cellId).toBe(sourceCell.id) + await store.activateControl(identity, { + cellId: sourceCell.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + } + await store.reconcileCells([sourceCell, targetCell]) + }) + + afterAll(async () => { + if (!database) return + await removeFixtureRows() + await database.close() + }) + + it('reaches PostgreSQL past an acquireActivity stalled in the identity queue', async () => { + const probe = new StallFirstTransactionDatabase(database) + const store = new RelayAssignmentStore(probe, () => now) + const identity = identities[0]! + const queued = store.acquireActivity(identity, { + activityId: 'splice:queue-blocker', + kind: 'splice', + cellId: sourceCell.id + }) + await probe.stalled.promise + const expiresAt = now + 105_000 + + try { + await Promise.race([ + store.renewControlActivity(identity, { + activityId: controlId(sourceCell.id), + cellId: sourceCell.id, + expiresAt + }), + new Promise((_resolve, reject) => + setTimeout(() => reject(new Error('renewal_waited_for_identity_queue')), 2_000) + ) + ]) + const row = ( + await database.query( + `SELECT expires_at FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? AND activity_id = ?`, + [identity.userId, identity.relayHostId, controlId(sourceCell.id)] + ) + )[0] + expect(Number(row!.expires_at)).toBe(expiresAt) + } finally { + probe.continue.resolve() + await queued + } + }) + + it('does not let a late older renewal rewind lease or activity timestamps', async () => { + const probe = new StallFirstRenewalQueryDatabase(database) + const store = new RelayAssignmentStore(probe, () => now) + const identity = identities[1]! + const earlierNow = now + const earlierExpiry = earlierNow + 105_000 + const earlier = store.renewControlActivity(identity, { + activityId: controlId(sourceCell.id), + cellId: sourceCell.id, + expiresAt: earlierExpiry + }) + await probe.stalled.promise + + now += 30_000 + const laterNow = now + const laterExpiry = laterNow + 105_000 + await store.renewControlActivity(identity, { + activityId: controlId(sourceCell.id), + cellId: sourceCell.id, + expiresAt: laterExpiry + }) + probe.continue.resolve() + await earlier + + const lease = ( + await database.query( + `SELECT expires_at, updated_at FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? AND activity_id = ?`, + [identity.userId, identity.relayHostId, controlId(sourceCell.id)] + ) + )[0] + const assignment = ( + await database.query( + `SELECT lease_expires_at, last_activity_at FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + )[0] + expect(lease).toMatchObject({ + expires_at: String(laterExpiry), + updated_at: String(laterNow) + }) + expect(assignment).toMatchObject({ + lease_expires_at: String(laterExpiry), + last_activity_at: String(laterNow) + }) + }) + + it('allows the source only while its exact forward migration remains active', async () => { + const identity = identities[2]! + const store = new RelayAssignmentStore(database, () => now) + const migration = await store.startEvacuation(identity, targetCell.id) + const activeExpiry = now + 105_000 + + await expect( + store.renewControlActivity(identity, { + activityId: controlId(sourceCell.id), + cellId: sourceCell.id, + expiresAt: activeExpiry + }) + ).resolves.toBeUndefined() + + await database.query( + `UPDATE relay_assignment_migrations SET completed_at = ? + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [now, identity.userId, identity.relayHostId, migration.assignmentEpoch] + ) + now += 30_000 + await expect( + store.renewControlActivity(identity, { + activityId: controlId(sourceCell.id), + cellId: sourceCell.id, + expiresAt: now + 105_000 + }) + ).rejects.toThrow('activity_cell_not_authoritative') + + const lease = ( + await database.query( + `SELECT expires_at FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? AND activity_id = ?`, + [identity.userId, identity.relayHostId, controlId(sourceCell.id)] + ) + )[0] + expect(Number(lease!.expires_at)).toBe(activeExpiry) + }) + + it('does not resurrect a control released while its renewal was stalled', async () => { + const probe = new StallFirstRenewalQueryDatabase(database) + const renewalStore = new RelayAssignmentStore(probe, () => now) + const releaseStore = new RelayAssignmentStore(database, () => now) + const identity = identities[3]! + const renewal = renewalStore.renewControlActivity(identity, { + activityId: controlId(sourceCell.id), + cellId: sourceCell.id, + expiresAt: now + 105_000 + }) + await probe.stalled.promise + + await expect(releaseStore.releaseActivity(identity, controlId(sourceCell.id))).resolves.toBe( + true + ) + probe.continue.resolve() + await expect(renewal).rejects.toThrow('control_activity_not_found') + const rows = await database.query( + `SELECT activity_id FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? AND activity_id = ?`, + [identity.userId, identity.relayHostId, controlId(sourceCell.id)] + ) + expect(rows).toHaveLength(0) + }) + + it('rejects a renewal expiry beyond the maximum control lease horizon', async () => { + const identity = identities[4]! + const store = new RelayAssignmentStore(database, () => now) + const maximumExpiry = + now + + ASSIGNMENT_LIMITS.activityLeaseMs + + RELAY_PROTOCOL_LIMITS.controlPingIntervalMs * 2 + + await expect( + store.renewControlActivity(identity, { + activityId: controlId(sourceCell.id), + cellId: sourceCell.id, + expiresAt: maximumExpiry + 1 + }) + ).rejects.toThrow('invalid_activity_expiry') + }) + + it('uses one autocommitted PostgreSQL statement for a steady renewal', async () => { + const probe = new RenewalQueryProbeDatabase(database) + const store = new RelayAssignmentStore(probe, () => now) + const identity = identities[5]! + + await store.renewControlActivity(identity, { + activityId: controlId(sourceCell.id), + cellId: sourceCell.id, + expiresAt: now + 105_000 + }) + + expect(probe.renewalQueries).toBe(1) + expect(probe.transactions).toBe(0) + }) +}) diff --git a/cloud/apps/relay/src/control-reservation-reconnect-claim.test.ts b/cloud/apps/relay/src/control-reservation-reconnect-claim.test.ts new file mode 100644 index 00000000000..e6fb44e336c --- /dev/null +++ b/cloud/apps/relay/src/control-reservation-reconnect-claim.test.ts @@ -0,0 +1,201 @@ +import { afterEach, describe, expect, it } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.js' +import type { RelayCellConfig } from './config.js' +import { + openInMemoryRelayDatabase, + openRelayDatabase, + type RelayDatabase +} from './database.js' + +const KEY_PREFIX = 'reconnect-claim' +const USER_ID = `${KEY_PREFIX}-user-1` +const RELAY_HOST_ID = 'reconnectclaim01' +const CELL: RelayCellConfig = { + id: `${KEY_PREFIX}-cell`, + url: `https://${KEY_PREFIX}-cell.example.com`, + capacityRequests: 1_000, + connectionHardCap: 600, + connectionUnobservedBound: 50 +} +const CELL_INCARNATION = '11111111-1111-4111-8111-111111111111' +const IDENTITY = { userId: USER_ID, relayHostId: RELAY_HOST_ID } +const CONTROL_ACTIVITY_ID = `control:${CELL.id}:1` + +const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL +const backends: { name: string; open: () => Promise }[] = [ + { name: 'sqlite', open: openInMemoryRelayDatabase }, + ...(databaseUrl + ? [ + { + name: 'postgres', + open: () => openRelayDatabase({ databaseUrl, dataDir: '' }) + } + ] + : []) +] + +const databases: RelayDatabase[] = [] + +async function removeScopedRows(database: RelayDatabase): Promise { + await database.query( + `DELETE FROM relay_control_connection_reservations WHERE user_id LIKE '${KEY_PREFIX}-%'` + ) + await database.query( + `DELETE FROM relay_assignment_activity_leases WHERE user_id LIKE '${KEY_PREFIX}-%'` + ) + await database.query(`DELETE FROM relay_assignments WHERE user_id LIKE '${KEY_PREFIX}-%'`) + for (const table of [ + 'relay_cell_connection_snapshots', + 'relay_cell_connection_runtime', + 'relay_cell_connection_limits', + 'relay_cell_runtime', + 'relay_cells' + ]) { + await database.query(`DELETE FROM ${table} WHERE cell_id = ?`, [CELL.id]) + } +} + +afterEach(async () => { + for (const database of databases.splice(0)) { + await removeScopedRows(database) + await database.close() + } +}) + +async function setup( + open: () => Promise, + now: () => number +): Promise<{ database: RelayDatabase; store: RelayAssignmentStore }> { + const database = await open() + databases.push(database) + await removeScopedRows(database) + const store = new RelayAssignmentStore(database, now, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + await store.reconcileCells([CELL], false) + await heartbeat(store, 0) + return { database, store } +} + +async function heartbeat(store: RelayAssignmentStore, watermark: number): Promise { + await store.recordCellHeartbeat({ + cellId: CELL.id, + cellUrl: CELL.url, + cellIncarnation: CELL_INCARNATION, + startedAt: 50, + ready: true, + observedRequests: 0, + totalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0, + connectionHardCap: CELL.connectionHardCap, + connectionUnobservedBound: CELL.connectionUnobservedBound, + connectionInclusionWatermark: watermark + }) +} + +async function reservations( + database: RelayDatabase +): Promise<{ state: string; claimActivityId: string | null }[]> { + const rows = await database.query( + `SELECT state, claim_activity_id FROM relay_control_connection_reservations + WHERE user_id = ? ORDER BY created_at ASC, reservation_id ASC`, + [USER_ID] + ) + return rows.map((row) => ({ + state: String(row['state']), + claimActivityId: + row['claim_activity_id'] === null || row['claim_activity_id'] === undefined + ? null + : String(row['claim_activity_id']) + })) +} + +describe.each(backends)('control reservation claim on reconnect ($name)', ({ open }) => { + it('claims the fresh reservation when the same generation reconnects', async () => { + let now = 100_000_000 + const { database, store } = await setup(open, () => now) + + const assignment = await store.assign(IDENTITY) + await store.activateControl(IDENTITY, { + cellId: CELL.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1, + connectionInclusionWatermark: 1 + }) + // Cell telemetry now includes the control, so its reservation is released. + now += 1_000 + await heartbeat(store, 2) + expect(await reservations(database)).toEqual([ + { state: 'released', claimActivityId: CONTROL_ACTIVITY_ID } + ]) + + // Control drops and the host is re-granted the same cell and epoch. + now += 1_000 + await store.releaseActivity(IDENTITY, CONTROL_ACTIVITY_ID) + now += 1_000 + await store.assign(IDENTITY) + expect(await reservations(database)).toEqual([ + { state: 'released', claimActivityId: CONTROL_ACTIVITY_ID }, + { state: 'reserved', claimActivityId: null } + ]) + + now += 1_000 + await store.activateControl(IDENTITY, { + cellId: CELL.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1, + connectionInclusionWatermark: 3 + }) + + expect(await reservations(database)).toEqual([ + { state: 'released', claimActivityId: CONTROL_ACTIVITY_ID }, + { state: 'claimed', claimActivityId: CONTROL_ACTIVITY_ID } + ]) + }) + + it('still refuses to claim a second reservation while the first is outstanding', async () => { + let now = 100_000_000 + const { database, store } = await setup(open, () => now) + + const assignment = await store.assign(IDENTITY) + await store.activateControl(IDENTITY, { + cellId: CELL.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1, + connectionInclusionWatermark: 1 + }) + await database.query( + `INSERT INTO relay_control_connection_reservations + (reservation_id, idempotency_key, user_id, relay_host_id, assignment_epoch, + cell_id, state, claim_activity_id, created_at, timeout_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, 'reserved', NULL, ?, ?, ?)`, + [ + `${KEY_PREFIX}-extra`, + `${KEY_PREFIX}-extra`, + USER_ID, + RELAY_HOST_ID, + assignment.assignmentEpoch, + CELL.id, + now + 1, + now + 90_000, + now + 1 + ] + ) + + now += 1_000 + await store.activateControl(IDENTITY, { + cellId: CELL.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1, + connectionInclusionWatermark: 2 + }) + + expect(await reservations(database)).toEqual([ + { state: 'claimed', claimActivityId: CONTROL_ACTIVITY_ID }, + { state: 'reserved', claimActivityId: null } + ]) + }) +}) diff --git a/cloud/apps/relay/src/credential-store.test.ts b/cloud/apps/relay/src/credential-store.test.ts new file mode 100644 index 00000000000..9f279b93791 --- /dev/null +++ b/cloud/apps/relay/src/credential-store.test.ts @@ -0,0 +1,301 @@ +import { describe, expect, it } from 'vitest' +import { + hashCredential, + RelayCredentialStore, + type CredentialReservation, + type RelayIdentity +} from './credential-store.js' +import { openInMemoryRelayDatabase, type RelayDatabase } from './database.js' + +const identity: RelayIdentity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } +const relayDeviceId = 'device-1' + +async function inviteBasis( + store: RelayCredentialStore, + generation = 1 +): Promise<{ reservation: CredentialReservation; basisConnId: string }> { + const invite = await store.createInvite(identity, relayDeviceId) + const reservation = await store.reserveCredential(identity.relayHostId, invite.inviteToken) + if (!reservation) throw new Error('invite did not reserve') + const basisConnId = `basis-${generation}` + await store.recordConnectionBasis({ + ...reservation, + basisConnId, + owningControlGeneration: generation, + deadline: reservation.leaseExpiresAt + }) + return { reservation, basisConnId } +} + +describe('relay credential store', () => { + it('issues invites with a skew margin under the client-side TTL ceiling', async () => { + const database = await openInMemoryRelayDatabase() + const store = new RelayCredentialStore(database, () => 1_000_000) + const invite = await store.createInvite(identity, relayDeviceId) + // Zero-tolerance released desktops reject expiry past local now + 10min; + // issuing 30s under the ceiling absorbs that much desktop clock lag. + expect(invite.expiresAt).toBe(1_000_000 + 10 * 60 * 1_000 - 30 * 1_000) + }) + + it('persists one invite reservation with bounded attempts and cooldown', async () => { + let now = 100 + const database = await openInMemoryRelayDatabase() + const store = new RelayCredentialStore(database, () => now) + const invite = await store.createInvite(identity, relayDeviceId) + const first = await store.reserveCredential(identity.relayHostId, invite.inviteToken, 'first') + expect(first).toMatchObject({ credentialKind: 'invite', reservationId: 'first' }) + expect(await store.reserveCredential(identity.relayHostId, invite.inviteToken, 'second')).toBeNull() + now = first!.leaseExpiresAt + expect(await store.reserveCredential(identity.relayHostId, invite.inviteToken, 'second')).toBeNull() + await database.query(`UPDATE relay_invites SET cooldown_until = ? WHERE token_hash = ?`, [ + now, + hashCredential(invite.inviteToken) + ]) + expect(await store.reserveCredential(identity.relayHostId, invite.inviteToken, 'second')).toMatchObject({ + reservationId: 'second' + }) + await database.close() + }) + + it('validates director moves without reservation and enforces account-global mint rate', async () => { + const database = await openInMemoryRelayDatabase() + const store = new RelayCredentialStore(database, () => 100) + const invite = await store.createInvite(identity, relayDeviceId) + expect(await store.validateInviteForMove(identity.relayHostId, invite.inviteToken)).toBe(true) + const rows = await database.query(`SELECT state, attempt_count FROM relay_invites WHERE token_hash = ?`, [ + hashCredential(invite.inviteToken) + ]) + expect(rows[0]).toMatchObject({ state: 'available', attempt_count: 0 }) + for (let index = 1; index < 30; index++) { + await store.createInvite(identity, `device-${index + 1}`) + } + await expect(store.createInvite(identity, 'device-over-limit')).rejects.toMatchObject({ + code: 'rate_limit_exceeded' + }) + expect(await database.query(`SELECT * FROM relay_audit_events`)).toHaveLength(30) + await database.close() + }) + + it('serializes relay-basis and authenticated-direct under one global result', async () => { + const database = await openInMemoryRelayDatabase() + const store = new RelayCredentialStore(database, () => 100) + const { basisConnId } = await inviteBasis(store) + await store.recordDirectAuthorization({ + ...identity, + relayDeviceId, + directAuthId: 'direct-1', + owningControlGeneration: 1, + deadline: 1_000 + }) + const base = { + ...identity, + relayDeviceId, + reqId: 'req-1', + newResumeTokenHash: hashCredential('pending-resume'), + owningControlGeneration: 1 + } + const [relay, direct] = await Promise.all([ + store.installCredential({ + ...base, + authorization: { mode: 'relay-basis' as const, basisConnId } + }), + store.installCredential({ + ...base, + authorization: { mode: 'authenticated-direct' as const, directAuthId: 'direct-1' } + }) + ]) + expect(relay).toEqual(direct) + expect(relay.currentVersion).toBe(1) + expect(await store.installStatus(base)).toEqual(relay) + const devices = await database.query(`SELECT * FROM relay_devices`) + expect(devices).toHaveLength(1) + await database.close() + }) + + it('rolls back token and invite effects when result persistence fails', async () => { + const database = await openInMemoryRelayDatabase() + const normalStore = new RelayCredentialStore(database, () => 100) + const { basisConnId, reservation } = await inviteBasis(normalStore) + const fault: RelayDatabase = { + query: (sql, params) => database.query(sql, params), + queryLocked: (sql, params) => database.queryLocked(sql, params), + close: () => database.close(), + transaction: async (operation) => + await database.transaction(async (transaction) => + await operation({ + query: async (sql, params) => { + if (sql.includes('INSERT INTO relay_install_results')) throw new Error('injected SQL failure') + return await transaction.query(sql, params) + }, + queryLocked: (sql, params) => transaction.queryLocked(sql, params), + transaction: (nested) => transaction.transaction(nested), + close: () => transaction.close() + }) + ) + } + const faultStore = new RelayCredentialStore(fault, () => 100) + const input = { + ...identity, + relayDeviceId, + reqId: 'req-fault', + newResumeTokenHash: hashCredential('pending-resume'), + owningControlGeneration: 1, + authorization: { mode: 'relay-basis' as const, basisConnId } + } + await expect(faultStore.installCredential(input)).rejects.toThrow('injected SQL failure') + expect(await normalStore.installStatus(input)).toBeNull() + expect(await database.query(`SELECT * FROM relay_devices`)).toEqual([]) + const invites = await database.query(`SELECT state FROM relay_invites WHERE token_hash = ?`, [ + reservation.tokenHash + ]) + expect(invites[0]?.state).toBe('reserved') + await database.close() + }) + + it('renews only a tuple-bound current credential and replays its committed result', async () => { + let now = 100 + const database = await openInMemoryRelayDatabase() + const store = new RelayCredentialStore(database, () => now) + const { basisConnId } = await inviteBasis(store) + const resumeToken = 'resume-token-1' + await store.installCredential({ + ...identity, + relayDeviceId, + reqId: 'install-1', + newResumeTokenHash: hashCredential(resumeToken), + owningControlGeneration: 1, + authorization: { mode: 'relay-basis', basisConnId } + }) + const reservation = await store.reserveCredential(identity.relayHostId, resumeToken) + if (!reservation) throw new Error('resume did not reserve') + expect(await store.resolveResume(identity.relayHostId, resumeToken)).toEqual({ + userId: identity.userId, + relayDeviceId + }) + await store.recordConnectionBasis({ + ...reservation, + basisConnId: 'resume-basis', + owningControlGeneration: 1, + deadline: 1_000 + }) + now = 200 + const confirmed = await store.confirmResume({ + ...identity, + reqId: 'confirm-1', + basisConnId: 'resume-basis', + owningControlGeneration: 1 + }) + expect(confirmed).toMatchObject({ renewed: true, acceptedAs: 'current' }) + await store.deactivateBasis('resume-basis') + expect( + await store.confirmResume({ + ...identity, + reqId: 'confirm-1', + basisConnId: 'resume-basis', + owningControlGeneration: 1 + }) + ).toEqual(confirmed) + await expect( + store.confirmResume({ + ...identity, + reqId: 'confirm-1', + basisConnId: 'other-basis', + owningControlGeneration: 1 + }) + ).rejects.toMatchObject({ code: 'confirmation_tuple_mismatch' }) + await database.close() + }) + + it('leaves a now-grace version unchanged and rejects late, expired, and revoked tuples', async () => { + let now = 100 + const database = await openInMemoryRelayDatabase() + const store = new RelayCredentialStore(database, () => now) + const { basisConnId } = await inviteBasis(store) + const firstToken = 'resume-token-1' + const first = await store.installCredential({ + ...identity, + relayDeviceId, + reqId: 'install-1', + newResumeTokenHash: hashCredential(firstToken), + owningControlGeneration: 1, + authorization: { mode: 'relay-basis', basisConnId } + }) + const firstReservation = await store.reserveCredential(identity.relayHostId, firstToken) + if (!firstReservation) throw new Error('first resume did not reserve') + await store.recordConnectionBasis({ + ...firstReservation, + basisConnId: 'grace-basis', + owningControlGeneration: 1, + deadline: 100_000_000 + }) + await store.recordDirectAuthorization({ + ...identity, + relayDeviceId, + directAuthId: 'rotate-direct', + owningControlGeneration: 1, + deadline: 1_000 + }) + now = 200 + await store.installCredential({ + ...identity, + relayDeviceId, + reqId: 'install-2', + newResumeTokenHash: hashCredential('resume-token-2'), + expectedCurrentHash: hashCredential(firstToken), + owningControlGeneration: 1, + authorization: { mode: 'authenticated-direct', directAuthId: 'rotate-direct' } + }) + const grace = await store.confirmResume({ + ...identity, + reqId: 'confirm-grace', + basisConnId: 'grace-basis', + owningControlGeneration: 1 + }) + expect(grace).toMatchObject({ acceptedAs: 'current', renewed: false, currentVersion: 2 }) + expect(grace.resumeExpiresAt).toBeGreaterThan(first.resumeExpiresAt) + + now += 24 * 60 * 60 * 1000 + 1 + await expect( + store.confirmResume({ + ...identity, + reqId: 'confirm-expired-grace', + basisConnId: 'grace-basis', + owningControlGeneration: 1 + }) + ).rejects.toMatchObject({ code: 'reject-expired' }) + + const currentReservation = await store.reserveCredential(identity.relayHostId, 'resume-token-2') + if (!currentReservation) throw new Error('current resume did not reserve') + await store.recordConnectionBasis({ + ...currentReservation, + basisConnId: 'revoked-basis', + owningControlGeneration: 1, + deadline: now + 1_000 + }) + await store.revoke(identity, relayDeviceId) + await expect( + store.confirmResume({ + ...identity, + reqId: 'confirm-revoked', + basisConnId: 'revoked-basis', + owningControlGeneration: 1 + }) + ).rejects.toMatchObject({ code: 'reject-revoked' }) + + await store.recordConnectionBasis({ + ...currentReservation, + basisConnId: 'late-basis', + owningControlGeneration: 1, + deadline: now - 1 + }) + await expect( + store.confirmResume({ + ...identity, + reqId: 'confirm-late', + basisConnId: 'late-basis', + owningControlGeneration: 1 + }) + ).rejects.toMatchObject({ code: 'confirmation_not_active' }) + await database.close() + }) +}) diff --git a/cloud/apps/relay/src/credential-store.ts b/cloud/apps/relay/src/credential-store.ts new file mode 100644 index 00000000000..a5697e8c699 --- /dev/null +++ b/cloud/apps/relay/src/credential-store.ts @@ -0,0 +1,745 @@ +import { createHash, randomBytes, randomUUID, timingSafeEqual } from 'node:crypto' +import { + decideResumeCommit, + RELAY_PROTOCOL_LIMITS, + type DeviceCredentialInstalled, + type DeviceResumeConfirmed +} from '@orca-cloud/relay-contract' +import type { RelayDatabase, SqlRow } from './database.js' + +const CREDENTIAL_GRACE_MS = 24 * 60 * 60 * 1000 +// Released desktops validate invite expiry against their own clock with zero +// tolerance at exactly inviteTtlMs; issuing under the ceiling keeps pairing +// working for clients whose clocks trail the cell by up to this margin. +const INVITE_ISSUE_SKEW_MARGIN_MS = 30 * 1000 + +export type RelayIdentity = { userId: string; relayHostId: string } +export type CredentialReservation = RelayIdentity & { + credentialKind: 'invite' | 'resume' + relayDeviceId: string + tokenHash: string + reservationId: string + leaseExpiresAt: number + acceptedCredentialVersion?: number + acceptedAs?: 'current' | 'grace' + resumeExpiresAt?: number + graceExpiresAt?: number +} + +export type InstallInput = RelayIdentity & { + relayDeviceId: string + reqId: string + newResumeTokenHash: string + expectedCurrentHash?: string + owningControlGeneration: number + authorization: + | { mode: 'relay-basis'; basisConnId: string } + | { mode: 'authenticated-direct'; directAuthId: string } +} + +export class RelayStoreError extends Error { + constructor(readonly code: string) { + super(code) + } +} + +export function hashCredential(token: string): string { + return createHash('sha256').update(token).digest('base64url') +} + +function equalHash(left: string, right: string): boolean { + const leftBytes = Buffer.from(left) + const rightBytes = Buffer.from(right) + return leftBytes.length === rightBytes.length && timingSafeEqual(leftBytes, rightBytes) +} + +function number(row: SqlRow, field: string): number { + const value = Number(row[field]) + if (!Number.isSafeInteger(value)) throw new RelayStoreError(`invalid_${field}`) + return value +} + +function optionalNumber(row: SqlRow, field: string): number | undefined { + return row[field] === null || row[field] === undefined ? undefined : number(row, field) +} + +function string(row: SqlRow, field: string): string { + const value = row[field] + if (typeof value !== 'string') throw new RelayStoreError(`invalid_${field}`) + return value +} + +export class RelayCredentialStore { + constructor( + private readonly database: RelayDatabase, + private readonly now: () => number = Date.now + ) {} + + async createInvite(identity: RelayIdentity, relayDeviceId: string): Promise<{ + inviteToken: string + expiresAt: number + maxAttempts: number + }> { + const inviteToken = randomBytes(32).toString('base64url') + const tokenHash = hashCredential(inviteToken) + const now = this.now() + const expiresAt = now + RELAY_PROTOCOL_LIMITS.inviteTtlMs - INVITE_ISSUE_SKEW_MARGIN_MS + await this.database.transaction(async (transaction) => { + await this.consumeRateWith(transaction, `account:${identity.userId}`, 'invite-mint', 30, 60_000, now) + await transaction.query( + `UPDATE relay_invites SET state = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND relay_device_id = ? + AND state IN (?, ?, ?)`, + [ + 'invalidated', now, identity.userId, identity.relayHostId, relayDeviceId, + 'available', 'reserved', 'cooldown' + ] + ) + await this.auditWith(transaction, { + type: 'invite-created', + ...identity, + relayDeviceId, + detail: { expiresAt } + }) + await transaction.query( + `INSERT INTO relay_invites + (user_id, relay_host_id, relay_device_id, token_hash, state, attempt_count, + max_attempts, expires_at, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + [ + identity.userId, identity.relayHostId, relayDeviceId, tokenHash, 'available', 0, + RELAY_PROTOCOL_LIMITS.inviteMaxAttempts, expiresAt, now, now + ] + ) + }) + return { inviteToken, expiresAt, maxAttempts: RELAY_PROTOCOL_LIMITS.inviteMaxAttempts } + } + + async reserveCredential( + relayHostId: string, + token: string, + reservationId: string = randomUUID() + ): Promise { + const tokenHash = hashCredential(token) + const now = this.now() + return await this.database.transaction(async (transaction) => { + const inviteRows = await transaction.query( + `SELECT * FROM relay_invites WHERE relay_host_id = ? AND token_hash = ?`, + [relayHostId, tokenHash] + ) + const invite = inviteRows[0] + if (invite && equalHash(string(invite, 'token_hash'), tokenHash)) { + const expiresAt = number(invite, 'expires_at') + const attempts = number(invite, 'attempt_count') + const maxAttempts = number(invite, 'max_attempts') + let state = string(invite, 'state') + const reservationExpiresAt = optionalNumber(invite, 'reservation_expires_at') + if (expiresAt <= now) state = 'expired' + else if (state === 'reserved' && (reservationExpiresAt ?? 0) <= now) { + await transaction.query( + `UPDATE relay_invites SET state = ?, reservation_id = NULL, + reservation_expires_at = NULL, cooldown_until = ?, updated_at = ? + WHERE token_hash = ?`, + [ + 'cooldown', + now + RELAY_PROTOCOL_LIMITS.inviteAttemptCooldownMs, + now, + tokenHash + ] + ) + return null + } + const cooldownUntil = optionalNumber(invite, 'cooldown_until') ?? 0 + if ( + (state !== 'available' && !(state === 'cooldown' && cooldownUntil <= now)) || + attempts >= maxAttempts + ) { + return null + } + const leaseExpiresAt = Math.min( + expiresAt, + now + RELAY_PROTOCOL_LIMITS.inviteReservationLeaseMs + ) + await transaction.query( + `UPDATE relay_invites SET state = ?, attempt_count = ?, reservation_id = ?, + reservation_expires_at = ?, cooldown_until = NULL, updated_at = ? + WHERE token_hash = ?`, + ['reserved', attempts + 1, reservationId, leaseExpiresAt, now, tokenHash] + ) + return { + userId: string(invite, 'user_id'), + relayHostId, + relayDeviceId: string(invite, 'relay_device_id'), + credentialKind: 'invite', + tokenHash, + reservationId, + leaseExpiresAt + } + } + + const deviceRows = await transaction.query( + `SELECT * FROM relay_devices + WHERE relay_host_id = ? AND (current_hash = ? OR grace_hash = ?)`, + [relayHostId, tokenHash, tokenHash] + ) + const device = deviceRows[0] + if (!device || optionalNumber(device, 'revoked_at') !== undefined) return null + const currentHash = string(device, 'current_hash') + const currentExpiresAt = number(device, 'current_expires_at') + const graceHash = device.grace_hash + const graceExpiresAt = optionalNumber(device, 'grace_expires_at') + let acceptedAs: 'current' | 'grace' + let acceptedCredentialVersion: number + if (equalHash(currentHash, tokenHash) && currentExpiresAt > now) { + acceptedAs = 'current' + acceptedCredentialVersion = number(device, 'current_version') + } else if ( + typeof graceHash === 'string' && + equalHash(graceHash, tokenHash) && + (graceExpiresAt ?? 0) > now + ) { + acceptedAs = 'grace' + acceptedCredentialVersion = number(device, 'grace_version') + } else { + return null + } + return { + userId: string(device, 'user_id'), + relayHostId, + relayDeviceId: string(device, 'relay_device_id'), + credentialKind: 'resume', + tokenHash, + reservationId, + leaseExpiresAt: now + RELAY_PROTOCOL_LIMITS.hostAttachDeadlineMs, + acceptedCredentialVersion, + acceptedAs, + resumeExpiresAt: currentExpiresAt, + graceExpiresAt + } + }) + } + + async recordConnectionBasis(input: CredentialReservation & { + basisConnId: string + owningControlGeneration: number + deadline: number + }): Promise { + const now = this.now() + await this.database.query( + `INSERT INTO relay_connection_bases + (basis_conn_id, user_id, relay_host_id, relay_device_id, owning_control_generation, + credential_kind, invite_token_hash, accepted_credential_version, accepted_as, + deadline, active, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + [ + input.basisConnId, input.userId, input.relayHostId, input.relayDeviceId, + input.owningControlGeneration, input.credentialKind, + input.credentialKind === 'invite' ? input.tokenHash : null, + input.acceptedCredentialVersion, input.acceptedAs, input.deadline, 1, now + ] + ) + } + + async failReservation(reservation: CredentialReservation): Promise { + if (reservation.credentialKind !== 'invite') return + const now = this.now() + await this.database.transaction(async (transaction) => { + const rows = await transaction.query( + `SELECT attempt_count, max_attempts FROM relay_invites + WHERE token_hash = ? AND reservation_id = ? AND state = ?`, + [reservation.tokenHash, reservation.reservationId, 'reserved'] + ) + const invite = rows[0] + if (!invite) return + const exhausted = number(invite, 'attempt_count') >= number(invite, 'max_attempts') + await transaction.query( + `UPDATE relay_invites SET state = ?, reservation_id = NULL, + reservation_expires_at = NULL, cooldown_until = ?, updated_at = ? + WHERE token_hash = ? AND reservation_id = ?`, + [ + exhausted ? 'invalidated' : 'cooldown', + exhausted ? null : now + RELAY_PROTOCOL_LIMITS.inviteAttemptCooldownMs, + now, + reservation.tokenHash, + reservation.reservationId + ] + ) + }) + } + + async recordDirectAuthorization(input: RelayIdentity & { + relayDeviceId: string + directAuthId: string + owningControlGeneration: number + deadline: number + }): Promise { + await this.database.query( + `INSERT INTO relay_direct_authorizations + (direct_auth_id, user_id, relay_host_id, relay_device_id, + owning_control_generation, deadline) + VALUES (?, ?, ?, ?, ?, ?)`, + [ + input.directAuthId, input.userId, input.relayHostId, input.relayDeviceId, + input.owningControlGeneration, input.deadline + ] + ) + } + + async installCredential(input: InstallInput): Promise { + let lastError: unknown + for (let attempt = 0; attempt <= 3; attempt++) { + try { + return await this.installCredentialOnce(input) + } catch (error) { + if (error instanceof RelayStoreError) throw error + const committed = await this.installStatus(input) + if (committed) return committed + const code = (error as { code?: unknown }).code + const retryable = ['23505', '40P01', '55P03', '40001'].includes(String(code)) + if (!retryable || attempt === 3) throw error + lastError = error + } + } + throw lastError + } + + private async installCredentialOnce(input: InstallInput): Promise { + return await this.database.transaction(async (transaction) => { + const existing = await this.installStatusWith(transaction, input) + if (existing) return existing + const now = this.now() + const basisInviteHash = await this.validateInstallAuthorization(transaction, input, now) + const devices = await transaction.query( + `SELECT * FROM relay_devices + WHERE user_id = ? AND relay_host_id = ? AND relay_device_id = ?`, + [input.userId, input.relayHostId, input.relayDeviceId] + ) + const current = devices[0] + if (input.expectedCurrentHash) { + if (!current || !equalHash(string(current, 'current_hash'), input.expectedCurrentHash)) { + throw new RelayStoreError('current_hash_mismatch') + } + } + const currentVersion = current ? number(current, 'current_version') + 1 : 1 + const resumeExpiresAt = now + RELAY_PROTOCOL_LIMITS.resumeTtlMs + const graceExpiresAt = current ? now + CREDENTIAL_GRACE_MS : undefined + await transaction.query( + `INSERT INTO relay_devices + (user_id, relay_host_id, relay_device_id, current_hash, current_version, + current_expires_at, grace_hash, grace_version, grace_expires_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT (user_id, relay_host_id, relay_device_id) DO UPDATE SET + current_hash = excluded.current_hash, + current_version = excluded.current_version, + current_expires_at = excluded.current_expires_at, + grace_hash = excluded.grace_hash, + grace_version = excluded.grace_version, + grace_expires_at = excluded.grace_expires_at, + revoked_at = NULL, + updated_at = excluded.updated_at`, + [ + input.userId, input.relayHostId, input.relayDeviceId, input.newResumeTokenHash, + currentVersion, resumeExpiresAt, current ? string(current, 'current_hash') : null, + current ? number(current, 'current_version') : null, graceExpiresAt, now + ] + ) + if (basisInviteHash) { + await transaction.query( + `UPDATE relay_invites SET state = ?, reservation_id = NULL, + reservation_expires_at = NULL, updated_at = ? WHERE token_hash = ?`, + ['consumed', now, basisInviteHash] + ) + } else { + await transaction.query( + `UPDATE relay_invites SET state = ?, reservation_id = NULL, + reservation_expires_at = NULL, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND relay_device_id = ? + AND state IN (?, ?, ?)`, + [ + 'invalidated', now, input.userId, input.relayHostId, input.relayDeviceId, + 'available', 'reserved', 'cooldown' + ] + ) + } + await transaction.query( + `UPDATE relay_direct_authorizations SET consumed_at = ? + WHERE user_id = ? AND relay_host_id = ? AND relay_device_id = ? + AND consumed_at IS NULL`, + [now, input.userId, input.relayHostId, input.relayDeviceId] + ) + const result: DeviceCredentialInstalled = { + v: 1, + reqId: input.reqId, + authorizationMode: input.authorization.mode, + currentVersion, + resumeExpiresAt, + ...(graceExpiresAt === undefined ? {} : { graceExpiresAt }) + } + await transaction.query( + `INSERT INTO relay_install_results + (user_id, relay_host_id, relay_device_id, req_id, authorization_mode, + result_json, committed_at) VALUES (?, ?, ?, ?, ?, ?, ?)`, + [ + input.userId, input.relayHostId, input.relayDeviceId, input.reqId, + input.authorization.mode, JSON.stringify(result), now + ] + ) + await this.auditWith(transaction, { + type: 'credential-installed', + userId: input.userId, + relayHostId: input.relayHostId, + relayDeviceId: input.relayDeviceId, + detail: { reqId: input.reqId, authorizationMode: input.authorization.mode, currentVersion } + }) + return result + }) + } + + async installStatus(input: RelayIdentity & { + relayDeviceId: string + reqId: string + }): Promise { + return await this.installStatusWith(this.database, input) + } + + async confirmResume(input: RelayIdentity & { + reqId: string + basisConnId: string + owningControlGeneration: number + }): Promise { + return await this.database.transaction(async (transaction) => { + const prior = await transaction.query( + `SELECT basis_conn_id, result_json FROM relay_confirm_results + WHERE user_id = ? AND relay_host_id = ? AND req_id = ?`, + [input.userId, input.relayHostId, input.reqId] + ) + if (prior[0]) { + if (string(prior[0], 'basis_conn_id') !== input.basisConnId) { + throw new RelayStoreError('confirmation_tuple_mismatch') + } + return JSON.parse(string(prior[0], 'result_json')) as DeviceResumeConfirmed + } + const rows = await transaction.query( + `SELECT * FROM relay_connection_bases WHERE basis_conn_id = ?`, + [input.basisConnId] + ) + const basis = rows[0] + const now = this.now() + if ( + !basis || + string(basis, 'user_id') !== input.userId || + string(basis, 'relay_host_id') !== input.relayHostId || + string(basis, 'credential_kind') !== 'resume' || + number(basis, 'owning_control_generation') !== input.owningControlGeneration || + number(basis, 'active') !== 1 || + number(basis, 'deadline') < now + ) { + throw new RelayStoreError('confirmation_not_active') + } + const relayDeviceId = string(basis, 'relay_device_id') + await transaction.query( + `UPDATE relay_devices SET updated_at = updated_at + WHERE user_id = ? AND relay_host_id = ? AND relay_device_id = ?`, + [input.userId, input.relayHostId, relayDeviceId] + ) + const devices = await transaction.query( + `SELECT * FROM relay_devices + WHERE user_id = ? AND relay_host_id = ? AND relay_device_id = ?`, + [input.userId, input.relayHostId, relayDeviceId] + ) + const device = devices[0] + if (!device) throw new RelayStoreError('credential_not_found') + const acceptedVersion = number(basis, 'accepted_credential_version') + const decision = decideResumeCommit( + { + currentVersion: number(device, 'current_version'), + currentHash: string(device, 'current_hash'), + currentExpiresAt: number(device, 'current_expires_at'), + graceVersion: optionalNumber(device, 'grace_version'), + graceHash: typeof device.grace_hash === 'string' ? device.grace_hash : undefined, + graceExpiresAt: optionalNumber(device, 'grace_expires_at'), + revokedAt: optionalNumber(device, 'revoked_at') + }, + acceptedVersion, + now + ) + if (decision.startsWith('reject-')) throw new RelayStoreError(decision) + const renewed = decision === 'renew-current' + const resumeExpiresAt = renewed + ? now + RELAY_PROTOCOL_LIMITS.resumeTtlMs + : number(device, 'current_expires_at') + if (renewed) { + await transaction.query( + `UPDATE relay_devices SET current_expires_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND relay_device_id = ?`, + [resumeExpiresAt, now, input.userId, input.relayHostId, relayDeviceId] + ) + } + const result: DeviceResumeConfirmed = { + v: 1, + reqId: input.reqId, + currentVersion: number(device, 'current_version'), + acceptedAs: string(basis, 'accepted_as') as 'current' | 'grace', + renewed, + resumeExpiresAt, + ...(optionalNumber(device, 'grace_expires_at') === undefined + ? {} + : { graceExpiresAt: optionalNumber(device, 'grace_expires_at') }) + } + await transaction.query( + `INSERT INTO relay_confirm_results + (user_id, relay_host_id, req_id, basis_conn_id, tuple_json, result_json, committed_at) + VALUES (?, ?, ?, ?, ?, ?, ?)`, + [ + input.userId, input.relayHostId, input.reqId, input.basisConnId, + JSON.stringify({ + relayDeviceId, + acceptedCredentialVersion: acceptedVersion, + acceptedAs: result.acceptedAs, + confirmDeadline: number(basis, 'deadline'), + owningControlGeneration: input.owningControlGeneration + }), + JSON.stringify(result), + now + ] + ) + await this.auditWith(transaction, { + type: 'resume-confirmed', + userId: input.userId, + relayHostId: input.relayHostId, + relayDeviceId, + detail: { reqId: input.reqId, basisConnId: input.basisConnId, renewed } + }) + return result + }) + } + + async deactivateBasis(basisConnId: string): Promise { + await this.database.query(`UPDATE relay_connection_bases SET active = ? WHERE basis_conn_id = ?`, [ + 0, + basisConnId + ]) + } + + async revoke(identity: RelayIdentity, relayDeviceId: string): Promise { + await this.database.transaction(async (transaction) => { + const now = this.now() + await transaction.query( + `UPDATE relay_devices SET revoked_at = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND relay_device_id = ?`, + [now, now, identity.userId, identity.relayHostId, relayDeviceId] + ) + await transaction.query( + `UPDATE relay_invites SET state = ?, updated_at = ? + WHERE user_id = ? AND relay_host_id = ? AND relay_device_id = ?`, + ['invalidated', now, identity.userId, identity.relayHostId, relayDeviceId] + ) + await this.auditWith(transaction, { + type: 'device-revoked', + ...identity, + relayDeviceId, + detail: {} + }) + }) + } + + async resolveResume( + relayHostId: string, + token: string + ): Promise<{ userId: string; relayDeviceId: string } | null> { + const tokenHash = hashCredential(token) + const rows = await this.database.query( + `SELECT * FROM relay_devices + WHERE relay_host_id = ? AND (current_hash = ? OR grace_hash = ?)`, + [relayHostId, tokenHash, tokenHash] + ) + const device = rows[0] + if (!device || optionalNumber(device, 'revoked_at') !== undefined) return null + const now = this.now() + const currentValid = + equalHash(string(device, 'current_hash'), tokenHash) && + number(device, 'current_expires_at') > now + const graceHash = device.grace_hash + const graceValid = + typeof graceHash === 'string' && + equalHash(graceHash, tokenHash) && + (optionalNumber(device, 'grace_expires_at') ?? 0) > now + return currentValid || graceValid + ? { userId: string(device, 'user_id'), relayDeviceId: string(device, 'relay_device_id') } + : null + } + + async validateInviteForMove(relayHostId: string, token: string): Promise { + return Boolean(await this.resolveInviteForMove(relayHostId, token)) + } + + async resolveInviteForMove( + relayHostId: string, + token: string + ): Promise<{ userId: string; relayDeviceId: string } | null> { + const tokenHash = hashCredential(token) + const rows = await this.database.query( + `SELECT user_id, relay_device_id, token_hash, state, attempt_count, max_attempts, expires_at + FROM relay_invites WHERE relay_host_id = ? AND token_hash = ?`, + [relayHostId, tokenHash] + ) + const invite = rows[0] + const valid = Boolean( + invite && + equalHash(string(invite, 'token_hash'), tokenHash) && + ['available', 'reserved', 'cooldown'].includes(string(invite, 'state')) && + number(invite, 'attempt_count') < number(invite, 'max_attempts') && + number(invite, 'expires_at') > this.now() + ) + return valid && invite + ? { userId: string(invite, 'user_id'), relayDeviceId: string(invite, 'relay_device_id') } + : null + } + + async consumeRate(input: { + scopeKey: string + kind: string + limit: number + windowMs: number + }): Promise { + await this.database.transaction(async (transaction) => { + await this.consumeRateWith( + transaction, + input.scopeKey, + input.kind, + input.limit, + input.windowMs, + this.now() + ) + }) + } + + async cleanup(): Promise { + const now = this.now() + await this.database.transaction(async (transaction) => { + await transaction.query( + `UPDATE relay_invites SET state = ?, reservation_id = NULL, + reservation_expires_at = NULL, updated_at = ? + WHERE expires_at <= ? AND state IN (?, ?, ?)`, + ['expired', now, now, 'available', 'reserved', 'cooldown'] + ) + await transaction.query( + `UPDATE relay_invites SET state = ?, reservation_id = NULL, + reservation_expires_at = NULL, cooldown_until = ?, updated_at = ? + WHERE state = ? AND reservation_expires_at <= ? AND expires_at > ?`, + ['cooldown', now + RELAY_PROTOCOL_LIMITS.inviteAttemptCooldownMs, now, 'reserved', now, now] + ) + await transaction.query( + `UPDATE relay_connection_bases SET active = ? WHERE active = ? AND deadline <= ?`, + [0, 1, now] + ) + await transaction.query( + `UPDATE relay_direct_authorizations SET consumed_at = ? + WHERE consumed_at IS NULL AND deadline <= ?`, + [now, now] + ) + await transaction.query( + `DELETE FROM relay_rate_windows WHERE window_started_at < ?`, + [now - 24 * 60 * 60 * 1000] + ) + }) + } + + private async installStatusWith( + database: RelayDatabase, + input: RelayIdentity & { relayDeviceId: string; reqId: string } + ): Promise { + const rows = await database.query( + `SELECT result_json FROM relay_install_results + WHERE user_id = ? AND relay_host_id = ? AND relay_device_id = ? AND req_id = ?`, + [input.userId, input.relayHostId, input.relayDeviceId, input.reqId] + ) + return rows[0] ? (JSON.parse(string(rows[0], 'result_json')) as DeviceCredentialInstalled) : null + } + + private async validateInstallAuthorization( + transaction: RelayDatabase, + input: InstallInput, + now: number + ): Promise { + if (input.authorization.mode === 'relay-basis') { + const rows = await transaction.query( + `SELECT * FROM relay_connection_bases WHERE basis_conn_id = ?`, + [input.authorization.basisConnId] + ) + const basis = rows[0] + if ( + !basis || + string(basis, 'user_id') !== input.userId || + string(basis, 'relay_host_id') !== input.relayHostId || + string(basis, 'relay_device_id') !== input.relayDeviceId || + string(basis, 'credential_kind') !== 'invite' || + number(basis, 'owning_control_generation') !== input.owningControlGeneration || + number(basis, 'active') !== 1 || + number(basis, 'deadline') < now + ) { + throw new RelayStoreError('invalid_relay_basis') + } + return string(basis, 'invite_token_hash') + } + const rows = await transaction.query( + `SELECT * FROM relay_direct_authorizations WHERE direct_auth_id = ?`, + [input.authorization.directAuthId] + ) + const direct = rows[0] + if ( + !direct || + string(direct, 'user_id') !== input.userId || + string(direct, 'relay_host_id') !== input.relayHostId || + string(direct, 'relay_device_id') !== input.relayDeviceId || + number(direct, 'owning_control_generation') !== input.owningControlGeneration || + number(direct, 'deadline') < now || + optionalNumber(direct, 'consumed_at') !== undefined + ) { + throw new RelayStoreError('invalid_direct_authorization') + } + await transaction.query( + `UPDATE relay_direct_authorizations SET consumed_at = ? WHERE direct_auth_id = ?`, + [now, input.authorization.directAuthId] + ) + return null + } + + private async consumeRateWith( + transaction: RelayDatabase, + scopeKey: string, + kind: string, + limit: number, + windowMs: number, + now: number + ): Promise { + const windowStartedAt = Math.floor(now / windowMs) * windowMs + const rows = await transaction.query( + `INSERT INTO relay_rate_windows + (scope_key, window_kind, window_started_at, count) VALUES (?, ?, ?, ?) + ON CONFLICT (scope_key, window_kind, window_started_at) DO UPDATE + SET count = relay_rate_windows.count + 1 RETURNING count`, + [scopeKey, kind, windowStartedAt, 1] + ) + if (number(rows[0]!, 'count') > limit) throw new RelayStoreError('rate_limit_exceeded') + } + + private async auditWith( + transaction: RelayDatabase, + event: RelayIdentity & { + type: string + relayDeviceId?: string + detail: Record + } + ): Promise { + await transaction.query( + `INSERT INTO relay_audit_events + (id, at, type, user_id, relay_host_id, relay_device_id, detail_json) + VALUES (?, ?, ?, ?, ?, ?, ?)`, + [ + randomUUID(), this.now(), event.type, event.userId, event.relayHostId, + event.relayDeviceId, JSON.stringify(event.detail) + ] + ) + } +} diff --git a/cloud/apps/relay/src/database-postgres-timeout.test.ts b/cloud/apps/relay/src/database-postgres-timeout.test.ts new file mode 100644 index 00000000000..a04a9eea042 --- /dev/null +++ b/cloud/apps/relay/src/database-postgres-timeout.test.ts @@ -0,0 +1,201 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +const fakes = vi.hoisted(() => ({ + configs: [] as Array>, + query: vi.fn(async () => ({ rows: [], rowCount: 0 })), + release: vi.fn(), + end: vi.fn(async () => undefined) +})) + +vi.mock('pg', () => ({ + default: { + Pool: class { + totalCount = 1 + idleCount = 1 + waitingCount = 0 + end = fakes.end + on = vi.fn() + connect = vi.fn(async () => ({ query: fakes.query, release: fakes.release })) + + constructor(config: Record) { + fakes.configs.push(config) + } + } + } +})) + +import { openRelayDatabase } from './database.js' +import { applyPostgresSchema } from './postgres-schema-startup.js' + +afterEach(() => { + vi.restoreAllMocks() +}) + +describe('PostgreSQL relay deadlines', () => { + beforeEach(() => { + fakes.configs.length = 0 + fakes.query.mockClear() + fakes.release.mockClear() + fakes.end.mockClear() + }) + + it('bounds pool acquisition, statements, locks, and abandoned transactions', async () => { + const database = await openRelayDatabase({ + databaseUrl: 'postgresql://relay:secret@127.0.0.1:5432/relay', + dataDir: './unused', + poolMax: 3, + applicationName: 'orca-relay/director/director' + }) + + expect(fakes.configs).toEqual([ + expect.objectContaining({ + max: 3, + application_name: 'orca-relay/director/director', + connectionTimeoutMillis: 2_000, + statement_timeout: 5_000, + lock_timeout: 1_000, + idle_in_transaction_session_timeout: 5_000 + }) + ]) + await database.close() + }) +}) + +describe('PostgreSQL schema startup', () => { + it('retries lock and statement timeouts with bounded backoff', async () => { + vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const query = vi + .fn<(statement: string) => Promise>() + .mockRejectedValueOnce(Object.assign(new Error('lock timeout'), { code: '55P03' })) + .mockRejectedValueOnce(Object.assign(new Error('statement timeout'), { code: '57014' })) + .mockResolvedValue(undefined) + const delays: number[] = [] + + await applyPostgresSchema(['CREATE TABLE test'], query, { + random: () => 0, + wait: async (delayMs) => { + delays.push(delayMs) + } + }) + + expect(query).toHaveBeenCalledTimes(3) + expect(delays).toEqual([125, 250]) + }) + + it('retries only the PostgreSQL concurrent type-creation collision', async () => { + vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const collision = Object.assign(new Error('duplicate type'), { + code: '23505', + constraint: 'pg_type_typname_nsp_index' + }) + const query = vi + .fn<(statement: string) => Promise>() + .mockRejectedValueOnce(collision) + .mockResolvedValue(undefined) + + await applyPostgresSchema(['CREATE TABLE IF NOT EXISTS test'], query, { + wait: async () => undefined + }) + + expect(query).toHaveBeenCalledTimes(2) + }) + + it('retries only the PostgreSQL concurrent index-creation collision', async () => { + vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const collision = Object.assign(new Error('duplicate index'), { + code: '23505', + constraint: 'pg_class_relname_nsp_index' + }) + const query = vi + .fn<(statement: string) => Promise>() + .mockRejectedValueOnce(collision) + .mockResolvedValue(undefined) + + await applyPostgresSchema(['CREATE INDEX IF NOT EXISTS test_index ON test(id)'], query, { + wait: async () => undefined + }) + + expect(query).toHaveBeenCalledTimes(2) + }) + + it.each([ + ['pg_type_typname_nsp_index', 'CREATE TABLE test'], + ['pg_class_relname_nsp_index', 'CREATE INDEX test_index ON test(id)'] + ])('does not retry %s for non-idempotent DDL', async (constraint, statement) => { + const error = Object.assign(new Error('duplicate catalog object'), { + code: '23505', + constraint + }) + const query = vi.fn<(statement: string) => Promise>().mockRejectedValue(error) + const pause = vi.fn(async () => undefined) + + await expect( + applyPostgresSchema([statement], query, { wait: pause }) + ).rejects.toBe(error) + + expect(pause).not.toHaveBeenCalled() + }) + + it('does not retry unrelated unique violations', async () => { + const error = Object.assign(new Error('duplicate row'), { + code: '23505', + constraint: 'application_key' + }) + const query = vi.fn<(statement: string) => Promise>().mockRejectedValue(error) + const pause = vi.fn(async () => undefined) + + await expect( + applyPostgresSchema(['CREATE TABLE test'], query, { wait: pause }) + ).rejects.toBe(error) + + expect(pause).not.toHaveBeenCalled() + }) + + it('fails immediately for non-timeout schema errors', async () => { + const error = Object.assign(new Error('permission denied'), { code: '42501' }) + const query = vi.fn<(statement: string) => Promise>().mockRejectedValue(error) + const pause = vi.fn(async () => undefined) + + await expect( + applyPostgresSchema(['CREATE TABLE test'], query, { wait: pause }) + ).rejects.toBe(error) + + expect(query).toHaveBeenCalledTimes(1) + expect(pause).not.toHaveBeenCalled() + }) + + it('stops retrying at the shared startup deadline', async () => { + vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const error = Object.assign(new Error('lock timeout'), { code: '55P03' }) + const delays: number[] = [] + let now = 0 + const query = vi + .fn<(statement: string) => Promise>() + .mockImplementationOnce(async () => { + now = 200 + }) + .mockRejectedValue(error) + + await expect( + applyPostgresSchema(['CREATE TABLE first', 'CREATE TABLE second'], query, { + now: () => now, + random: () => 1, + retryDeadlineMs: 300, + wait: async (delayMs) => { + delays.push(delayMs) + now += delayMs + } + }) + ).rejects.toBe(error) + + expect(query).toHaveBeenCalledTimes(3) + expect(delays).toEqual([100]) + expect(console.warn).toHaveBeenLastCalledWith( + JSON.stringify({ + event: 'orca_relay_postgres_schema_retry_exhausted', + code: '55P03', + attempts: 2 + }) + ) + }) +}) diff --git a/cloud/apps/relay/src/database-transient-error.test.ts b/cloud/apps/relay/src/database-transient-error.test.ts new file mode 100644 index 00000000000..b29ff519328 --- /dev/null +++ b/cloud/apps/relay/src/database-transient-error.test.ts @@ -0,0 +1,18 @@ +import { describe, expect, it } from 'vitest' +import { isRelayDatabaseTransientError } from './database.js' + +describe('relay database transient errors', () => { + it.each(['40P01', '40001', '55P03', '57014', '53300', '57P03', '08001', '08006'])( + 'classifies PostgreSQL code %s as retryable overload', + (code) => { + expect(isRelayDatabaseTransientError({ code })).toBe(true) + } + ) + + it('classifies pool acquisition timeout without hiding programming failures', () => { + expect( + isRelayDatabaseTransientError(new Error('timeout exceeded when trying to connect')) + ).toBe(true) + expect(isRelayDatabaseTransientError(new TypeError('broken invariant'))).toBe(false) + }) +}) diff --git a/cloud/apps/relay/src/database.test.ts b/cloud/apps/relay/src/database.test.ts new file mode 100644 index 00000000000..32e50a7bc6a --- /dev/null +++ b/cloud/apps/relay/src/database.test.ts @@ -0,0 +1,154 @@ +import { mkdtempSync, rmSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterEach, describe, expect, it } from 'vitest' +import { openInMemoryRelayDatabase, openRelayDatabase } from './database.js' + +const temporaryDirectories: string[] = [] + +afterEach(() => { + for (const directory of temporaryDirectories.splice(0)) { + rmSync(directory, { recursive: true, force: true }) + } +}) + +describe('relay database', () => { + it('creates every durable relay state table', async () => { + const database = await openInMemoryRelayDatabase() + const rows = await database.query( + `SELECT name FROM sqlite_master WHERE type = 'table' AND name LIKE 'relay_%' ORDER BY name` + ) + expect(rows.map((row) => row.name)).toEqual([ + 'relay_admission_selector_cell_additions', + 'relay_admission_selector_intents', + 'relay_admission_selectors', + 'relay_assignment_activity_leases', + 'relay_assignment_migration_incarnations', + 'relay_assignment_migrations', + 'relay_assignment_region_preferences', + 'relay_assignments', + 'relay_audit_events', + 'relay_cell_admission', + 'relay_cell_capabilities', + 'relay_cell_committed_fences', + 'relay_cell_connection_limits', + 'relay_cell_connection_runtime', + 'relay_cell_connection_snapshots', + 'relay_cell_drain_attempt_states', + 'relay_cell_drain_attempts', + 'relay_cell_drain_recovery_attempts', + 'relay_cell_fence_apply_invocations', + 'relay_cell_fence_attempts', + 'relay_cell_fence_plan_bindings', + 'relay_cell_fences', + 'relay_cell_legacy_fence_adoptions', + 'relay_cell_regions', + 'relay_cell_rehome_safety', + 'relay_cell_runtime', + 'relay_cells', + 'relay_confirm_results', + 'relay_confirmable_splices', + 'relay_connection_bases', + 'relay_control_connection_reservations', + 'relay_devices', + 'relay_direct_authorizations', + 'relay_install_results', + 'relay_invites', + 'relay_migration_leases', + 'relay_post_drain_migration_pins', + 'relay_rate_windows', + 'relay_region_rehome_attempts', + 'relay_region_rehome_control', + 'relay_region_rehome_worker_state' + ]) + await database.close() + }) + + it('rolls back every effect and serializes concurrent transactions', async () => { + const database = await openInMemoryRelayDatabase() + await expect( + database.transaction(async (transaction) => { + await transaction.query( + `INSERT INTO relay_rate_windows + (scope_key, window_kind, window_started_at, count) VALUES (?, ?, ?, ?)`, + ['scope', 'invite', 1, 1] + ) + throw new Error('injected failure') + }) + ).rejects.toThrow('injected failure') + expect(await database.query(`SELECT * FROM relay_rate_windows`)).toEqual([]) + + await Promise.all([ + database.transaction(async (transaction) => { + await transaction.query( + `INSERT INTO relay_rate_windows + (scope_key, window_kind, window_started_at, count) VALUES (?, ?, ?, ?)`, + ['scope', 'invite', 1, 1] + ) + }), + database.transaction(async (transaction) => { + const rows = await transaction.query( + `SELECT count FROM relay_rate_windows + WHERE scope_key = ? AND window_kind = ? AND window_started_at = ?`, + ['scope', 'invite', 1] + ) + await transaction.query( + `UPDATE relay_rate_windows SET count = ? + WHERE scope_key = ? AND window_kind = ? AND window_started_at = ?`, + [Number(rows[0]?.count ?? 0) + 1, 'scope', 'invite', 1] + ) + }) + ]) + const rows = await database.query(`SELECT count FROM relay_rate_windows`) + expect(Number(rows[0]?.count)).toBe(2) + await database.close() + }) + + it('persists SQLite state across process-style reopen', async () => { + const dataDir = mkdtempSync(join(tmpdir(), 'orca-relay-db-')) + temporaryDirectories.push(dataDir) + const first = await openRelayDatabase({ dataDir }) + await first.query( + `INSERT INTO relay_rate_windows + (scope_key, window_kind, window_started_at, count) VALUES (?, ?, ?, ?)`, + ['scope', 'connect', 1, 7] + ) + await first.close() + const second = await openRelayDatabase({ dataDir }) + const rows = await second.query(`SELECT count FROM relay_rate_windows`) + expect(Number(rows[0]?.count)).toBe(7) + await second.close() + }) + + it('defaults cells created by an older schema user to the US on reopen', async () => { + const dataDir = mkdtempSync(join(tmpdir(), 'orca-relay-region-db-')) + temporaryDirectories.push(dataDir) + const first = await openRelayDatabase({ dataDir }) + await first.query( + `INSERT INTO relay_cells + (cell_id, cell_url, enabled, capacity_requests, reserved_requests, + observed_requests, last_heartbeat_at, updated_at) + VALUES (?, ?, 1, 10, 0, 0, 1, 1)`, + ['legacy-cell', 'https://legacy.relay.example.com'] + ) + await first.close() + + const second = await openRelayDatabase({ dataDir }) + expect( + await second.query(`SELECT region FROM relay_cell_regions WHERE cell_id = ?`, [ + 'legacy-cell' + ]) + ).toEqual([{ region: 'us-central1' }]) + await second.close() + }) + + it('indexes region preference expiry by observation time', async () => { + const database = await openInMemoryRelayDatabase() + const rows = await database.query( + `SELECT sql FROM sqlite_master + WHERE type = 'index' AND name = 'relay_assignment_region_preferences_observed'` + ) + expect(rows[0]?.sql).toContain('(observed_at)') + await database.close() + }) +}) diff --git a/cloud/apps/relay/src/database.ts b/cloud/apps/relay/src/database.ts new file mode 100644 index 00000000000..f7208863f42 --- /dev/null +++ b/cloud/apps/relay/src/database.ts @@ -0,0 +1,935 @@ +import { mkdirSync } from 'node:fs' +import { join } from 'node:path' +import { DatabaseSync } from 'node:sqlite' +import pg from 'pg' +import { + emptyPostgresPoolPressureCounts, + PostgresPoolPressure, + type PostgresPoolPressureCounts +} from './postgres-pool-pressure.js' +import { applyPostgresSchema } from './postgres-schema-startup.js' + +export type SqlRow = Record +export type RelayLockOptions = { failIfUnavailable?: boolean } +export type RelayTransactionOptions = { reportRetries?: boolean } + +export interface RelayDatabase { + readonly dialect?: 'sqlite' | 'postgres' + query(sql: string, params?: unknown[]): Promise + queryLocked( + sql: string, + params?: unknown[], + options?: RelayLockOptions + ): Promise + transaction( + operation: (transaction: RelayDatabase) => Promise, + options?: RelayTransactionOptions + ): Promise + close(): Promise +} + +const SCHEMA = ` +CREATE TABLE IF NOT EXISTS relay_invites ( + user_id TEXT NOT NULL, + relay_host_id TEXT NOT NULL, + relay_device_id TEXT NOT NULL, + token_hash TEXT NOT NULL UNIQUE, + state TEXT NOT NULL, + attempt_count BIGINT NOT NULL, + max_attempts BIGINT NOT NULL, + expires_at BIGINT NOT NULL, + reservation_id TEXT, + reservation_expires_at BIGINT, + cooldown_until BIGINT, + created_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL, + PRIMARY KEY (user_id, relay_host_id, relay_device_id, token_hash) +); +CREATE INDEX IF NOT EXISTS relay_invites_device + ON relay_invites(user_id, relay_host_id, relay_device_id); + +CREATE TABLE IF NOT EXISTS relay_devices ( + user_id TEXT NOT NULL, + relay_host_id TEXT NOT NULL, + relay_device_id TEXT NOT NULL, + current_hash TEXT NOT NULL, + current_version BIGINT NOT NULL, + current_expires_at BIGINT NOT NULL, + grace_hash TEXT, + grace_version BIGINT, + grace_expires_at BIGINT, + revoked_at BIGINT, + updated_at BIGINT NOT NULL, + PRIMARY KEY (user_id, relay_host_id, relay_device_id) +); +CREATE INDEX IF NOT EXISTS relay_devices_current_hash ON relay_devices(relay_host_id, current_hash); +CREATE INDEX IF NOT EXISTS relay_devices_grace_hash ON relay_devices(relay_host_id, grace_hash); + +CREATE TABLE IF NOT EXISTS relay_install_results ( + user_id TEXT NOT NULL, + relay_host_id TEXT NOT NULL, + relay_device_id TEXT NOT NULL, + req_id TEXT NOT NULL, + authorization_mode TEXT NOT NULL, + result_json TEXT NOT NULL, + committed_at BIGINT NOT NULL, + PRIMARY KEY (user_id, relay_host_id, relay_device_id, req_id) +); + +CREATE TABLE IF NOT EXISTS relay_confirmable_splices ( + basis_conn_id TEXT PRIMARY KEY, + user_id TEXT NOT NULL, + relay_host_id TEXT NOT NULL, + owning_control_generation BIGINT NOT NULL, + relay_device_id TEXT NOT NULL, + accepted_credential_version BIGINT NOT NULL, + accepted_as TEXT NOT NULL, + confirm_deadline BIGINT NOT NULL, + active BIGINT NOT NULL, + created_at BIGINT NOT NULL +); + +CREATE TABLE IF NOT EXISTS relay_connection_bases ( + basis_conn_id TEXT PRIMARY KEY, + user_id TEXT NOT NULL, + relay_host_id TEXT NOT NULL, + relay_device_id TEXT NOT NULL, + owning_control_generation BIGINT NOT NULL, + credential_kind TEXT NOT NULL, + invite_token_hash TEXT, + accepted_credential_version BIGINT, + accepted_as TEXT, + deadline BIGINT NOT NULL, + active BIGINT NOT NULL, + created_at BIGINT NOT NULL +); + +-- Why: the maintenance sweep matches (active, deadline) while inactive bases +-- accumulate unboundedly. Unindexed it seq-scans millions of rows every cycle +-- and holds the maintenance transaction open long enough to time out +-- assignment lock waits. +CREATE INDEX IF NOT EXISTS relay_connection_bases_active_deadline + ON relay_connection_bases(active, deadline); + +CREATE TABLE IF NOT EXISTS relay_direct_authorizations ( + direct_auth_id TEXT PRIMARY KEY, + user_id TEXT NOT NULL, + relay_host_id TEXT NOT NULL, + relay_device_id TEXT NOT NULL, + owning_control_generation BIGINT NOT NULL, + deadline BIGINT NOT NULL, + consumed_at BIGINT +); + +CREATE TABLE IF NOT EXISTS relay_confirm_results ( + user_id TEXT NOT NULL, + relay_host_id TEXT NOT NULL, + req_id TEXT NOT NULL, + basis_conn_id TEXT NOT NULL, + tuple_json TEXT NOT NULL, + result_json TEXT NOT NULL, + committed_at BIGINT NOT NULL, + PRIMARY KEY (user_id, relay_host_id, req_id) +); + +CREATE TABLE IF NOT EXISTS relay_assignments ( + user_id TEXT NOT NULL, + relay_host_id TEXT NOT NULL, + cell_id TEXT NOT NULL, + assignment_epoch BIGINT NOT NULL, + lease_expires_at BIGINT NOT NULL, + last_activity_at BIGINT NOT NULL, + reserved_controls BIGINT NOT NULL, + reserved_splices BIGINT NOT NULL, + reserved_invites BIGINT NOT NULL, + pending_installs BIGINT NOT NULL, + pending_confirmations BIGINT NOT NULL, + migration_leases BIGINT NOT NULL, + PRIMARY KEY (user_id, relay_host_id) +); + +CREATE TABLE IF NOT EXISTS relay_assignment_region_preferences ( + user_id TEXT NOT NULL, + relay_host_id TEXT NOT NULL, + preferred_region TEXT NOT NULL + CHECK (preferred_region IN ('us-central1', 'asia-east2')), + observed_at BIGINT NOT NULL, + PRIMARY KEY (user_id, relay_host_id) +); +CREATE INDEX IF NOT EXISTS relay_assignment_region_preferences_observed + ON relay_assignment_region_preferences(observed_at); + +CREATE TABLE IF NOT EXISTS relay_region_rehome_worker_state ( + worker_id TEXT PRIMARY KEY, + next_dispatch_at BIGINT NOT NULL, + paused_until BIGINT NOT NULL, + consecutive_failures BIGINT NOT NULL, + updated_at BIGINT NOT NULL +); + +CREATE TABLE IF NOT EXISTS relay_region_rehome_control ( + control_id TEXT PRIMARY KEY, + generation BIGINT NOT NULL, + enabled BIGINT NOT NULL, + observation_started_at BIGINT NOT NULL, + not_before BIGINT NOT NULL, + rate_per_minute BIGINT NOT NULL, + preference_max_age_ms BIGINT NOT NULL, + drain_grace_ms BIGINT NOT NULL, + updated_at BIGINT NOT NULL +); + +CREATE TABLE IF NOT EXISTS relay_region_rehome_attempts ( + attempt_id TEXT PRIMARY KEY, + user_id TEXT NOT NULL, + relay_host_id TEXT NOT NULL, + preferred_region TEXT NOT NULL CHECK (preferred_region = 'asia-east2'), + source_cell_id TEXT NOT NULL, + source_cell_incarnation TEXT NOT NULL, + target_cell_id TEXT NOT NULL, + target_cell_incarnation TEXT NOT NULL, + previous_epoch BIGINT NOT NULL, + assignment_epoch BIGINT NOT NULL, + drain_grace_ms BIGINT NOT NULL, + send_attempts BIGINT NOT NULL, + last_send_attempt_at BIGINT, + drain_receipt_at BIGINT, + drain_outcome TEXT CHECK ( + drain_outcome IN ('accepted', 'already-accepted', 'host-not-connected') + ), + completed_at BIGINT, + aborted_at BIGINT, + created_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL, + UNIQUE (user_id, relay_host_id, assignment_epoch) +); +CREATE INDEX IF NOT EXISTS relay_region_rehome_attempts_pending + ON relay_region_rehome_attempts(drain_receipt_at, last_send_attempt_at, completed_at, aborted_at); + +CREATE TABLE IF NOT EXISTS relay_cells ( + cell_id TEXT PRIMARY KEY, + cell_url TEXT NOT NULL UNIQUE, + enabled BIGINT NOT NULL, + capacity_requests BIGINT NOT NULL, + reserved_requests BIGINT NOT NULL, + observed_requests BIGINT NOT NULL, + last_heartbeat_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL +); + +CREATE TABLE IF NOT EXISTS relay_cell_regions ( + cell_id TEXT PRIMARY KEY, + region TEXT NOT NULL CHECK (region IN ('us-central1', 'asia-east2')) +); + +CREATE TABLE IF NOT EXISTS relay_cell_admission ( + cell_id TEXT PRIMARY KEY, + admission_state TEXT NOT NULL + CHECK (admission_state IN ('existing-only', 'migration-only', 'general')), + updated_at BIGINT NOT NULL +); + +CREATE TABLE IF NOT EXISTS relay_admission_selectors ( + selector_id TEXT PRIMARY KEY, + generation BIGINT NOT NULL, + attempt_id TEXT, + membership_json TEXT NOT NULL, + updated_at BIGINT NOT NULL +); + +CREATE TABLE IF NOT EXISTS relay_admission_selector_intents ( + attempt_id TEXT PRIMARY KEY, + expected_generation BIGINT NOT NULL, + intended_generation BIGINT NOT NULL, + previous_membership_json TEXT NOT NULL, + membership_json TEXT NOT NULL, + created_at BIGINT NOT NULL, + committed_at BIGINT +); + +CREATE TABLE IF NOT EXISTS relay_admission_selector_cell_additions ( + attempt_id TEXT PRIMARY KEY, + cells_json TEXT NOT NULL +); + +CREATE TABLE IF NOT EXISTS relay_cell_runtime ( + cell_id TEXT PRIMARY KEY, + cell_url TEXT NOT NULL, + cell_incarnation TEXT NOT NULL, + started_at BIGINT NOT NULL, + ready BIGINT NOT NULL, + observed_requests BIGINT NOT NULL, + last_heartbeat_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL +); +CREATE INDEX IF NOT EXISTS relay_cell_runtime_heartbeat + ON relay_cell_runtime(ready, last_heartbeat_at); + +CREATE TABLE IF NOT EXISTS relay_cell_capabilities ( + cell_id TEXT PRIMARY KEY, + cell_incarnation TEXT NOT NULL, + regional_rehome_protocol BIGINT NOT NULL, + last_heartbeat_at BIGINT NOT NULL +); + +CREATE TABLE IF NOT EXISTS relay_cell_rehome_safety ( + cell_id TEXT PRIMARY KEY, + cell_incarnation TEXT NOT NULL, + observed_at BIGINT NOT NULL, + sql_failures BIGINT NOT NULL, + reconnects BIGINT NOT NULL, + control_activity_recovery_failures BIGINT NOT NULL, + database_pool_waiting BIGINT NOT NULL, + database_pool_waiters_max BIGINT NOT NULL, + database_pool_wait_ms_max BIGINT NOT NULL +); + +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 +); + +CREATE TABLE IF NOT EXISTS relay_cell_connection_runtime ( + cell_id TEXT PRIMARY KEY, + cell_incarnation TEXT NOT NULL, + total_connections BIGINT NOT NULL, + in_flight_connections BIGINT NOT NULL, + reserved_connection_units BIGINT NOT NULL, + enforced_connection_units BIGINT NOT NULL, + last_heartbeat_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL +); +CREATE INDEX IF NOT EXISTS relay_cell_connection_runtime_heartbeat + ON relay_cell_connection_runtime(last_heartbeat_at); + +CREATE TABLE IF NOT EXISTS relay_cell_connection_snapshots ( + cell_id TEXT PRIMARY KEY, + cell_incarnation TEXT NOT NULL, + inclusion_watermark BIGINT NOT NULL, + total_connections BIGINT NOT NULL, + in_flight_connections BIGINT NOT NULL, + reserved_connection_units BIGINT NOT NULL, + enforced_connection_units BIGINT NOT NULL, + snapshot_at BIGINT NOT NULL +); +CREATE INDEX IF NOT EXISTS relay_cell_connection_snapshot_freshness + ON relay_cell_connection_snapshots(snapshot_at); + +CREATE TABLE IF NOT EXISTS relay_cell_fences ( + cell_id TEXT PRIMARY KEY, + cell_incarnation TEXT NOT NULL, + attested_at BIGINT NOT NULL, + expires_at BIGINT NOT NULL +); +CREATE INDEX IF NOT EXISTS relay_cell_fences_expiry + ON relay_cell_fences(expires_at); + +CREATE TABLE IF NOT EXISTS relay_cell_committed_fences ( + cell_id TEXT PRIMARY KEY, + attempt_id TEXT NOT NULL UNIQUE, + cell_incarnation TEXT NOT NULL, + attested_at BIGINT NOT NULL, + expires_at BIGINT NOT NULL +); +CREATE INDEX IF NOT EXISTS relay_cell_committed_fences_expiry + ON relay_cell_committed_fences(expires_at); + +CREATE TABLE IF NOT EXISTS relay_cell_legacy_fence_adoptions ( + cell_id TEXT PRIMARY KEY, + cell_incarnation TEXT NOT NULL, + attested_at BIGINT NOT NULL, + expires_at BIGINT NOT NULL +); +CREATE INDEX IF NOT EXISTS relay_cell_legacy_fence_adoptions_expiry + ON relay_cell_legacy_fence_adoptions(expires_at); + +CREATE TABLE IF NOT EXISTS relay_cell_fence_attempts ( + attempt_id TEXT PRIMARY KEY, + environment TEXT NOT NULL, + cell_id TEXT NOT NULL, + cell_incarnation TEXT NOT NULL, + mig_name TEXT NOT NULL, + instance_group TEXT NOT NULL, + generation_identity TEXT NOT NULL, + fence_commit TEXT NOT NULL, + plan_sha256 TEXT NOT NULL, + gce_operation TEXT, + created_at BIGINT NOT NULL, + expires_at BIGINT NOT NULL, + apply_started_at BIGINT, + completed_at BIGINT, + aborted_at BIGINT +); +CREATE INDEX IF NOT EXISTS relay_cell_fence_attempts_expiry + ON relay_cell_fence_attempts(expires_at); +CREATE INDEX IF NOT EXISTS relay_cell_fence_attempts_cell + ON relay_cell_fence_attempts(cell_id, created_at); + +CREATE TABLE IF NOT EXISTS relay_cell_fence_plan_bindings ( + attempt_id TEXT PRIMARY KEY, + plan_object_name TEXT NOT NULL, + plan_object_generation TEXT, + var_file_sha256 TEXT NOT NULL, + terraform_state_lineage TEXT NOT NULL, + terraform_state_serial BIGINT NOT NULL, + terraform_state_object_generation TEXT NOT NULL, + terraform_state_object_sha256 TEXT NOT NULL, + request_reason TEXT NOT NULL, + FOREIGN KEY (attempt_id) REFERENCES relay_cell_fence_attempts(attempt_id) +); + +CREATE TABLE IF NOT EXISTS relay_cell_fence_apply_invocations ( + invocation_id TEXT PRIMARY KEY, + attempt_id TEXT NOT NULL, + request_reason TEXT NOT NULL UNIQUE, + started_at BIGINT NOT NULL, + gce_operation TEXT, + FOREIGN KEY (attempt_id) REFERENCES relay_cell_fence_attempts(attempt_id) +); +CREATE INDEX IF NOT EXISTS relay_cell_fence_apply_invocations_attempt + ON relay_cell_fence_apply_invocations(attempt_id, started_at); + +CREATE TABLE IF NOT EXISTS relay_cell_drain_attempts ( + cell_id TEXT PRIMARY KEY, + cell_incarnation TEXT NOT NULL, + planned_grace_ms BIGINT NOT NULL, + attempted_at BIGINT NOT NULL, + retry_after BIGINT NOT NULL, + recover_forward_attempted_at BIGINT +); + +CREATE TABLE IF NOT EXISTS relay_cell_drain_attempt_states ( + attempt_id TEXT PRIMARY KEY, + cell_id TEXT NOT NULL, + cell_incarnation TEXT NOT NULL, + trace_value TEXT NOT NULL UNIQUE, + planned_grace_ms BIGINT NOT NULL, + state TEXT NOT NULL CHECK ( + state IN ( + 'prepared', + 'send-may-have-started', + 'application-receipt', + 'proven-not-delivered' + ) + ), + prepared_at BIGINT NOT NULL, + send_may_have_started_at BIGINT, + send_permit_expires_at BIGINT, + application_receipt_at BIGINT, + backend_success_status BIGINT, + backend_instance TEXT, + receipt_cell_incarnation TEXT, + retry_after BIGINT, + recover_forward_attempted_at BIGINT, + proven_not_delivered_at BIGINT +); +CREATE INDEX IF NOT EXISTS relay_cell_drain_attempt_states_cell + ON relay_cell_drain_attempt_states(cell_id, prepared_at); + +CREATE TABLE IF NOT EXISTS relay_cell_drain_recovery_attempts ( + drain_attempt_id TEXT NOT NULL, + cell_incarnation TEXT NOT NULL, + attempted_at BIGINT NOT NULL, + PRIMARY KEY (drain_attempt_id, cell_incarnation), + FOREIGN KEY (drain_attempt_id) REFERENCES relay_cell_drain_attempt_states(attempt_id) +); + +CREATE TABLE IF NOT EXISTS relay_assignment_activity_leases ( + user_id TEXT NOT NULL, + relay_host_id TEXT NOT NULL, + activity_id TEXT NOT NULL, + activity_kind TEXT NOT NULL, + cell_id TEXT NOT NULL, + request_units BIGINT NOT NULL, + expires_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL, + PRIMARY KEY (user_id, relay_host_id, activity_id) +); +CREATE INDEX IF NOT EXISTS relay_assignment_activity_expiry + ON relay_assignment_activity_leases(expires_at); + +CREATE TABLE IF NOT EXISTS relay_control_connection_reservations ( + reservation_id TEXT PRIMARY KEY, + idempotency_key TEXT NOT NULL UNIQUE, + user_id TEXT NOT NULL, + relay_host_id TEXT NOT NULL, + assignment_epoch BIGINT NOT NULL, + cell_id TEXT NOT NULL, + state TEXT NOT NULL + CHECK (state IN ('reserved', 'late-arrival-debt', 'claimed', 'released')), + inclusion_watermark BIGINT, + claim_activity_id TEXT, + created_at BIGINT NOT NULL, + timeout_at BIGINT NOT NULL, + claimed_at BIGINT, + released_at BIGINT, + updated_at BIGINT NOT NULL +); +CREATE INDEX IF NOT EXISTS relay_control_connection_reservation_headroom + ON relay_control_connection_reservations(cell_id, state); +CREATE INDEX IF NOT EXISTS relay_control_connection_reservation_assignment + ON relay_control_connection_reservations( + user_id, relay_host_id, assignment_epoch, cell_id, created_at + ); + +CREATE TABLE IF NOT EXISTS relay_rate_windows ( + scope_key TEXT NOT NULL, + window_kind TEXT NOT NULL, + window_started_at BIGINT NOT NULL, + count BIGINT NOT NULL, + PRIMARY KEY (scope_key, window_kind, window_started_at) +); + +CREATE TABLE IF NOT EXISTS relay_migration_leases ( + user_id TEXT NOT NULL, + relay_host_id TEXT NOT NULL, + source_cell_id TEXT NOT NULL, + target_cell_id TEXT NOT NULL, + assignment_epoch BIGINT NOT NULL, + expires_at BIGINT NOT NULL, + completed_at BIGINT, + PRIMARY KEY (user_id, relay_host_id, assignment_epoch) +); + +CREATE TABLE IF NOT EXISTS relay_assignment_migrations ( + user_id TEXT NOT NULL, + relay_host_id TEXT NOT NULL, + source_cell_id TEXT NOT NULL, + target_cell_id TEXT NOT NULL, + previous_epoch BIGINT NOT NULL, + assignment_epoch BIGINT NOT NULL, + source_request_units BIGINT NOT NULL, + target_reserved_units BIGINT NOT NULL, + expires_at BIGINT NOT NULL, + target_registered_at BIGINT, + completed_at BIGINT, + aborted_at BIGINT, + created_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL, + PRIMARY KEY (user_id, relay_host_id, assignment_epoch) +); +CREATE INDEX IF NOT EXISTS relay_assignment_migrations_active + ON relay_assignment_migrations(expires_at, completed_at, aborted_at); + +CREATE TABLE IF NOT EXISTS relay_assignment_migration_incarnations ( + user_id TEXT NOT NULL, + relay_host_id TEXT NOT NULL, + assignment_epoch BIGINT NOT NULL, + source_cell_incarnation TEXT NOT NULL, + target_cell_incarnation TEXT NOT NULL, + PRIMARY KEY (user_id, relay_host_id, assignment_epoch) +); + +CREATE TABLE IF NOT EXISTS relay_post_drain_migration_pins ( + user_id TEXT NOT NULL, + relay_host_id TEXT NOT NULL, + assignment_epoch BIGINT NOT NULL, + drain_attempt_id TEXT NOT NULL, + source_cell_id TEXT NOT NULL, + source_cell_incarnation TEXT NOT NULL, + target_cell_id TEXT NOT NULL, + target_cell_incarnation TEXT NOT NULL, + source_request_units BIGINT NOT NULL, + target_reserved_units BIGINT NOT NULL, + pinned_at BIGINT NOT NULL, + PRIMARY KEY (user_id, relay_host_id, assignment_epoch) +); +CREATE INDEX IF NOT EXISTS relay_post_drain_migration_pins_attempt + ON relay_post_drain_migration_pins(drain_attempt_id); + +CREATE TABLE IF NOT EXISTS relay_audit_events ( + id TEXT PRIMARY KEY, + at BIGINT NOT NULL, + type TEXT NOT NULL, + user_id TEXT, + relay_host_id TEXT, + relay_device_id TEXT, + detail_json TEXT NOT NULL +); +CREATE INDEX IF NOT EXISTS relay_audit_events_at ON relay_audit_events(at); +` + +function postgresSql(sql: string): string { + let index = 0 + return sql.replace(/\?/g, () => `$${++index}`) +} + +function returnsRows(sql: string): boolean { + return /^\s*(select|with)/i.test(sql) || /returning/i.test(sql) +} + +const POSTGRES_TRANSACTION_PHASES = [ + ['relay_region_rehome_', 'regional-rehome'], + ['relay_assignment_activity_leases', 'activity-lease'], + ['relay_assignment_migration', 'migration'], + ['relay_migration_leases', 'migration'], + ['relay_post_drain_migration_pins', 'migration'], + ['relay_cell_connection_runtime', 'cell-runtime'], + ['relay_cell_connection_snapshots', 'cell-runtime'], + ['relay_cell_runtime', 'cell-runtime'], + ['relay_cell_drain_', 'cell-operation'], + ['relay_cell_fence', 'cell-operation'], + ['relay_cell_committed_fences', 'cell-operation'], + ['relay_cell_legacy_fence_adoptions', 'cell-operation'], + ['relay_admission_selector', 'admission'], + ['relay_cell_admission', 'admission'], + ['relay_control_connection_reservations', 'connection'], + ['relay_confirmable_splices', 'connection'], + ['relay_connection_bases', 'connection'], + ['relay_direct_authorizations', 'connection'], + ['relay_confirm_results', 'connection'], + ['relay_assignment_region_preferences', 'assignment'], + ['relay_assignments', 'assignment'], + ['relay_cell_', 'cell-inventory'], + ['relay_cells', 'cell-inventory'], + ['relay_invites', 'credential'], + ['relay_devices', 'credential'], + ['relay_install_results', 'credential'], + ['relay_rate_windows', 'rate-limit'], + ['relay_audit_events', 'audit'] +] as const + +const postgresTransactionPhaseByError = new WeakMap() + +function postgresTransactionPhase(sql: string): string { + const normalized = sql.toLowerCase() + return POSTGRES_TRANSACTION_PHASES.find(([table]) => normalized.includes(table))?.[1] ?? 'other' +} + +function rememberPostgresTransactionPhase(error: unknown, sql: string): void { + if (typeof error === 'object' && error !== null) { + postgresTransactionPhaseByError.set(error, postgresTransactionPhase(sql)) + } +} + +function postgresTransactionErrorPhase(error: unknown): string { + return typeof error === 'object' && error !== null + ? (postgresTransactionPhaseByError.get(error) ?? 'transaction') + : 'transaction' +} + +class SqliteTransaction implements RelayDatabase { + readonly dialect = 'sqlite' as const + + constructor(protected readonly database: DatabaseSync) {} + + async query(sql: string, params: unknown[] = []): Promise { + const statement = this.database.prepare(sql) + const bound = params.map((value) => (value === undefined ? null : value)) as never[] + if (returnsRows(sql)) return statement.all(...bound) as SqlRow[] + const result = statement.run(...bound) + return [{ changes: Number(result.changes) }] + } + + async queryLocked( + sql: string, + params: unknown[] = [], + _options: RelayLockOptions = {} + ): Promise { + return await this.query(sql, params) + } + + async transaction( + operation: (transaction: RelayDatabase) => Promise, + _options: RelayTransactionOptions = {} + ): Promise { + return await operation(this) + } + + async close(): Promise {} +} + +class SqliteDatabase extends SqliteTransaction { + private tail: Promise = Promise.resolve() + + override async query(sql: string, params: unknown[] = []): Promise { + await this.tail + return await super.query(sql, params) + } + + override async transaction(operation: (transaction: RelayDatabase) => Promise): Promise { + const previous = this.tail + let release!: () => void + this.tail = new Promise((resolve) => (release = resolve)) + await previous + this.database.exec('BEGIN IMMEDIATE') + try { + const result = await operation(new SqliteTransaction(this.database)) + this.database.exec('COMMIT') + return result + } catch (error) { + this.database.exec('ROLLBACK') + throw error + } finally { + release() + } + } + + override async close(): Promise { + await this.tail + this.database.close() + } +} + +class PostgresTransaction implements RelayDatabase { + readonly dialect = 'postgres' as const + + constructor(protected readonly client: pg.PoolClient) {} + + async query(sql: string, params: unknown[] = []): Promise { + try { + const result = await this.client.query(postgresSql(sql), params) + return returnsRows(sql) ? (result.rows as SqlRow[]) : [{ changes: result.rowCount ?? 0 }] + } catch (error) { + rememberPostgresTransactionPhase(error, sql) + throw error + } + } + + async queryLocked( + sql: string, + params: unknown[] = [], + options: RelayLockOptions = {} + ): Promise { + try { + return await this.query( + `${sql} FOR UPDATE${options.failIfUnavailable ? ' NOWAIT' : ''}`, + params + ) + } catch (error) { + if ( + options.failIfUnavailable && + String((error as { code?: unknown }).code) === '55P03' + ) { + throw new Error('database_lock_unavailable') + } + throw error + } + } + + async transaction( + operation: (transaction: RelayDatabase) => Promise, + _options: RelayTransactionOptions = {} + ): Promise { + return await operation(this) + } + + async close(): Promise {} +} + +const POSTGRES_TRANSACTION_ATTEMPTS = 3 +const POSTGRES_RETRY_MAX_DELAY_MS = 25 +const POSTGRES_CONNECTION_TIMEOUT_MS = 2_000 +const POSTGRES_STATEMENT_TIMEOUT_MS = 5_000 +const POSTGRES_LOCK_TIMEOUT_MS = 1_000 +const POSTGRES_IDLE_TRANSACTION_TIMEOUT_MS = 5_000 + +function retryablePostgresTransactionError(error: unknown): boolean { + const code = String((error as { code?: unknown }).code) + return code === '40P01' || code === '40001' || code === '55P03' +} + +export function isRelayDatabaseTransientError(error: unknown): boolean { + const code = String((error as { code?: unknown }).code) + if (['40P01', '40001', '55P03', '57014', '53300', '57P03', '08001', '08006'].includes(code)) { + return true + } + return String((error as { message?: unknown }).message).includes( + 'timeout exceeded when trying to connect' + ) +} + +async function waitForPostgresRetry(random: () => number = Math.random): Promise { + const delayMs = Math.floor(random() * (POSTGRES_RETRY_MAX_DELAY_MS + 1)) + await new Promise((resolve) => setTimeout(resolve, delayMs)) +} + +class PostgresDatabase implements RelayDatabase { + readonly dialect = 'postgres' as const + private readonly pressure: PostgresPoolPressure + + constructor(private readonly pool: pg.Pool) { + this.pressure = new PostgresPoolPressure(pool) + } + + async query(sql: string, params: unknown[] = []): Promise { + const client = await this.pressure.connect() + try { + const result = await client.query(postgresSql(sql), params) + return returnsRows(sql) ? (result.rows as SqlRow[]) : [{ changes: result.rowCount ?? 0 }] + } finally { + client.release() + } + } + + async queryLocked( + sql: string, + params: unknown[] = [], + options: RelayLockOptions = {} + ): Promise { + try { + return await this.query( + `${sql} FOR UPDATE${options.failIfUnavailable ? ' NOWAIT' : ''}`, + params + ) + } catch (error) { + if ( + options.failIfUnavailable && + String((error as { code?: unknown }).code) === '55P03' + ) { + throw new Error('database_lock_unavailable') + } + throw error + } + } + + async transaction( + operation: (transaction: RelayDatabase) => Promise, + options: RelayTransactionOptions = {} + ): Promise { + for (let attempt = 1; attempt <= POSTGRES_TRANSACTION_ATTEMPTS; attempt++) { + const client = await this.pressure.connect() + try { + await client.query('BEGIN') + const result = await operation(new PostgresTransaction(client)) + await client.query('COMMIT') + return result + } catch (error) { + await client.query('ROLLBACK').catch(() => undefined) + if (!retryablePostgresTransactionError(error) || attempt === POSTGRES_TRANSACTION_ATTEMPTS) { + if (retryablePostgresTransactionError(error) && options.reportRetries !== false) { + console.warn( + JSON.stringify({ + event: 'orca_relay_postgres_transaction_exhausted', + code: String((error as { code?: unknown }).code), + attempts: attempt, + phase: postgresTransactionErrorPhase(error) + }) + ) + } + throw error + } + if (options.reportRetries !== false) { + console.warn( + JSON.stringify({ + event: 'orca_relay_postgres_transaction_retry', + code: String((error as { code?: unknown }).code), + attempt, + phase: postgresTransactionErrorPhase(error) + }) + ) + } + } finally { + client.release() + } + // A PostgreSQL transaction is unusable after an abort, so retry all work + // on a fresh pooled client with a small full-jitter delay. + await waitForPostgresRetry() + } + throw new Error('postgres_transaction_retry_exhausted') + } + + async close(): Promise { + await this.pool.end() + } + + consumePoolPressure(): PostgresPoolPressureCounts { + return this.pressure.consumeCounts() + } + + peekPoolPressure(): PostgresPoolPressureCounts { + return this.pressure.peekCounts() + } +} + +export function consumeRelayDatabasePoolPressure( + database: RelayDatabase +): PostgresPoolPressureCounts { + return database instanceof PostgresDatabase + ? database.consumePoolPressure() + : emptyPostgresPoolPressureCounts() +} + +export function readRelayDatabasePoolPressure( + database: RelayDatabase +): PostgresPoolPressureCounts { + return database instanceof PostgresDatabase + ? database.peekPoolPressure() + : emptyPostgresPoolPressureCounts() +} + +export function absorbPostgresIdleClientErrors(pool: Pick): void { + pool.on('error', () => { + // Why: node-postgres removes failed idle clients itself; leaving `error` + // unhandled would crash the cell and turn a SQL outage into autoheal churn. + console.warn('[orca-relay] idle PostgreSQL client failed') + }) +} + +async function applySchema(database: RelayDatabase): Promise { + for (const statement of SCHEMA.split(';')) { + if (statement.trim()) await database.query(statement) + } +} + +async function applySchemaWithPostgresRetries(database: RelayDatabase): Promise { + await applyPostgresSchema( + SCHEMA.split(';').filter((statement) => statement.trim()), + async (statement) => await database.query(statement) + ) +} + +async function backfillRelayCellRegions(database: RelayDatabase): Promise { + await database.query( + `INSERT INTO relay_cell_regions (cell_id, region) + SELECT cell_id, 'us-central1' FROM relay_cells WHERE true + ON CONFLICT (cell_id) DO NOTHING` + ) +} + +export async function openRelayDatabase(input: { + databaseUrl?: string + dataDir: string + poolMax?: number + applicationName?: string +}): Promise { + let database: RelayDatabase + if (input.databaseUrl) { + const pool = new pg.Pool({ + connectionString: input.databaseUrl, + max: input.poolMax ?? 10, + application_name: input.applicationName, + connectionTimeoutMillis: POSTGRES_CONNECTION_TIMEOUT_MS, + statement_timeout: POSTGRES_STATEMENT_TIMEOUT_MS, + lock_timeout: POSTGRES_LOCK_TIMEOUT_MS, + idle_in_transaction_session_timeout: POSTGRES_IDLE_TRANSACTION_TIMEOUT_MS + }) + absorbPostgresIdleClientErrors(pool) + database = new PostgresDatabase(pool) + } else { + mkdirSync(input.dataDir, { recursive: true }) + const sqlite = new DatabaseSync(join(input.dataDir, 'orca-relay.sqlite')) + sqlite.exec('PRAGMA journal_mode = WAL; PRAGMA foreign_keys = ON;') + database = new SqliteDatabase(sqlite) + } + try { + if (input.databaseUrl) await applySchemaWithPostgresRetries(database) + else await applySchema(database) + await backfillRelayCellRegions(database) + return database + } catch (error) { + await database.close().catch(() => undefined) + throw error + } +} + +export async function openInMemoryRelayDatabase(): Promise { + const sqlite = new DatabaseSync(':memory:') + sqlite.exec('PRAGMA foreign_keys = ON;') + const database = new SqliteDatabase(sqlite) + await applySchema(database) + await backfillRelayCellRegions(database) + return database +} diff --git a/cloud/apps/relay/src/fault-injection-test-entry.ts b/cloud/apps/relay/src/fault-injection-test-entry.ts new file mode 100644 index 00000000000..3f2cc6b6246 --- /dev/null +++ b/cloud/apps/relay/src/fault-injection-test-entry.ts @@ -0,0 +1,58 @@ +import { loadRelayConfig } from './config.js' +import { reconcileCellAdmissionAtStartup } from './cell-admission-startup.js' +import { openRelayDatabase, type RelayDatabase } from './database.js' +import { createRelayServer } from './relay-server.js' +import { readFileSync } from 'node:fs' + +if (process.env.NODE_ENV !== 'test') { + throw new Error('fault injection entry is test-only') +} + +const pattern = process.env.ORCA_RELAY_TEST_FAULT_SQL +const clockFile = process.env.ORCA_RELAY_TEST_CLOCK_FILE +if (!pattern && !clockFile) throw new Error('a test fault configuration is required') + +const config = loadRelayConfig() +const realDatabase = await openRelayDatabase({ + databaseUrl: config.databaseUrl, + dataDir: config.dataDir +}) +let faulted = false + +function wrap(transaction: RelayDatabase): RelayDatabase { + return { + query: async (sql, params) => { + if (pattern && !faulted && sql.includes(pattern)) { + faulted = true + throw new Error('injected SQL failure') + } + return await transaction.query(sql, params) + }, + queryLocked: async (sql, params, options) => + await transaction.queryLocked(sql, params, options), + transaction: async (operation) => + await transaction.transaction(async (nested) => await operation(wrap(nested))), + close: async () => await transaction.close() + } +} + +const database = wrap(realDatabase) +const now = clockFile + ? (): number => { + const offset = Number(readFileSync(clockFile, 'utf8')) + if (!Number.isFinite(offset)) throw new Error('invalid test clock offset') + return Date.now() + offset + } + : Date.now +const { server, sessions, assignments } = createRelayServer(config, database, { now }) +await reconcileCellAdmissionAtStartup(config, assignments) +server.listen(config.port, () => { + console.log(`[orca-relay] listening on ${config.publicUrl} (port ${config.port})`) +}) + +const shutdown = (): void => { + sessions.drain(0) + server.close(() => void realDatabase.close()) +} +process.once('SIGTERM', shutdown) +process.once('SIGINT', shutdown) diff --git a/cloud/apps/relay/src/google-metadata-identity-token.ts b/cloud/apps/relay/src/google-metadata-identity-token.ts new file mode 100644 index 00000000000..f51f6282a65 --- /dev/null +++ b/cloud/apps/relay/src/google-metadata-identity-token.ts @@ -0,0 +1,22 @@ +const METADATA_IDENTITY_ENDPOINT = + 'http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity' +const METADATA_IDENTITY_TIMEOUT_MS = 5_000 + +export async function googleMetadataIdentityToken( + audience: string, + fetchImpl: typeof fetch = fetch +): Promise { + const endpoint = new URL(METADATA_IDENTITY_ENDPOINT) + endpoint.searchParams.set('audience', audience) + endpoint.searchParams.set('format', 'full') + const response = await fetchImpl(endpoint, { + headers: { 'Metadata-Flavor': 'Google' }, + signal: AbortSignal.timeout(METADATA_IDENTITY_TIMEOUT_MS) + }) + if (!response.ok) throw new Error(`metadata_identity_${response.status}`) + const token = (await response.text()).trim() + if (token.length === 0 || token.length > 16 * 1024) { + throw new Error('metadata_identity_invalid') + } + return token +} diff --git a/cloud/apps/relay/src/host-session-registry.test.ts b/cloud/apps/relay/src/host-session-registry.test.ts new file mode 100644 index 00000000000..f632d5d4358 --- /dev/null +++ b/cloud/apps/relay/src/host-session-registry.test.ts @@ -0,0 +1,1007 @@ +import { EventEmitter } from 'node:events' +import { + ASSIGNMENT_LIMITS, + CONTROL_CONTINUITY_LIMITS, + RELAY_CLOSE_CODE, + RELAY_PROTOCOL_LIMITS +} from '@orca-cloud/relay-contract' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import type WebSocket from 'ws' +import type { RelayAssignmentStore } from './assignment-store.js' +import type { RelayConfig } from './config.js' +import type { RelayCredentialStore } from './credential-store.js' +import { HostSessionRegistry, type HostSession } from './host-session-registry.js' +import { relayHostLogDigest } from './relay-host-log-digest.js' +import type { RelayRuntimeObserver } from './relay-observability.js' +import { + REGIONAL_REHOME_TRUST_PROBE_ATTEMPT_ID, + REGIONAL_REHOME_TRUST_PROBE_HOST_ID, + REGIONAL_REHOME_TRUST_PROBE_USER_ID +} from './regional-rehome-trust-probe.js' +import type { RelayTokenClaims } from './relay-token-verifier.js' +import { ProcessQueuedByteBudget } from './splice-forwarder.js' + +class FakeSocket extends EventEmitter { + readonly OPEN = 1 + readonly CLOSING = 2 + readonly CLOSED = 3 + readyState = this.OPEN + readonly send = vi.fn() + readonly close = vi.fn((code?: number, reason?: string) => { + this.readyState = this.CLOSED + this.emit('close', code, Buffer.from(reason ?? '')) + }) + readonly terminate = vi.fn(() => { + this.readyState = this.CLOSED + this.emit('close') + }) +} + +type ActivateSession = ( + socket: WebSocket, + identity: RelayTokenClaims, + existing: HostSession | null, + generation: number, + rebind: boolean, + assignmentEpoch: number, + appVersion?: string +) => Promise + +const config = { + port: 8080, + publicUrl: 'https://relay-c3.example.com', + cellUrl: 'https://relay-c3.example.com', + authIssuer: 'https://auth.example.com', + authAudience: 'orca-relay', + jwksUrl: 'https://auth.example.com/jwks', + assignmentSigningKey: new Uint8Array(32), + role: 'cell', + cellId: 'production-gce-c3', + cells: [ + { + id: 'production-gce-c3', + url: 'https://relay-c3.example.com', + capacityRequests: 4_000 + } + ], + adminAudience: 'https://relay-c3.example.com/v1/admin/drain', + deployServiceAccount: 'deploy@example.com', + runtimeServiceAccount: 'runtime@example.com', + adminJwksUrl: 'https://auth.example.com/admin-jwks', + databasePoolMax: 10, + publicAssignmentsEnabled: true, + publicAssignmentConcurrency: 2, + publicAssignmentQueueMax: 128, + publicAssignmentWaitMs: 4_000, + publicResolveConcurrency: 1, + publicResolveWaitMs: 5_000, + publicAssignmentRetryAfterSeconds: 5, + dataDir: './test-data' +} satisfies RelayConfig + +const identity = { + sub: 'user-1', + prof: 'profile-1', + relayHostId: 'abcdefghijklmnop', + purpose: 'host-control', + exp: 4_102_444_800 +} satisfies RelayTokenClaims + +function deferred(): { + promise: Promise + resolve(value: T): void +} { + let resolve!: (value: T) => void + const promise = new Promise((complete) => { + resolve = complete + }) + return { promise, resolve } +} + +function createRegistry( + activateControl: RelayAssignmentStore['activateControl'], + store: Partial = {}, + verifyRelayToken: (token: string) => Promise = vi.fn() +): { + registry: HostSessionRegistry + activate: ActivateSession + acquireActivity: ReturnType + renewControlActivity: ReturnType + releaseActivity: ReturnType + observer: { + recordControlClose: ReturnType + recordSpliceClose: ReturnType + } +} { + const acquireActivity = vi.fn().mockResolvedValue(undefined) + const renewControlActivity = vi.fn().mockResolvedValue(undefined) + const releaseActivity = vi.fn().mockResolvedValue(true) + const assignments = { + activateControl, + markMigrationTargetRegistered: vi.fn().mockResolvedValue(undefined), + resolve: vi.fn().mockResolvedValue({ cellId: config.cellId }), + acquireActivity, + renewControlActivity, + releaseActivity + } as unknown as RelayAssignmentStore + const observer = { + recordAuth: vi.fn(), + recordForwardedBytes: vi.fn(), + recordHttp: vi.fn(), + recordReconnect: vi.fn(), + recordSql: vi.fn(), + recordControlClose: vi.fn(), + recordSpliceClose: vi.fn() + } satisfies RelayRuntimeObserver + const registry = new HostSessionRegistry( + config, + verifyRelayToken, + store as RelayCredentialStore, + assignments, + new ProcessQueuedByteBudget(), + observer + ) + // Mirrors the production signature exactly so a future positional shift fails to compile. + const bound = ( + registry as unknown as { + activate: ( + socket: WebSocket, + identity: RelayTokenClaims, + existing: HostSession | null, + generation: number, + rebind: boolean, + assignmentEpoch: number, + appVersion: string, + connectionInclusionWatermark?: number + ) => Promise + } + ).activate.bind(registry) + const activate: ActivateSession = ( + socket, + identity, + existing, + generation, + rebind, + assignmentEpoch, + appVersion = '1.4.173' + ) => bound(socket, identity, existing, generation, rebind, assignmentEpoch, appVersion) + return { registry, activate, acquireActivity, renewControlActivity, releaseActivity, observer } +} + +describe('host session cleanup races', () => { + beforeEach(() => vi.useFakeTimers()) + afterEach(() => { + vi.clearAllTimers() + vi.useRealTimers() + }) + + it('logs control closes with a host digest and counts them, never the raw id', async () => { + const activateControl = vi + .fn() + .mockResolvedValue('control:production-gce-c3:1') + const { activate, observer } = createRegistry(activateControl) + const socket = new FakeSocket() + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + try { + await activate(socket as unknown as WebSocket, identity, null, 1, false, 1) + socket.emit('error', new RangeError('Max payload size exceeded')) + socket.close(1006, 'network reset') + + expect(observer.recordControlClose).toHaveBeenCalledWith(1006) + const line = warn.mock.calls + .map((call) => String(call[0])) + .find((entry) => entry.includes('control closed')) + expect(line).toContain(`host=${relayHostLogDigest(identity.relayHostId)}`) + expect(line).toContain('code=1006') + expect(line).toContain('Max payload size exceeded') + expect(line).not.toContain(identity.relayHostId) + } finally { + warn.mockRestore() + } + }) + + it('contains a dependency failure to one socket instead of the process', async () => { + const activateControl = vi + .fn() + .mockResolvedValue('control:production-gce-c3:1') + // The same rejection shape as a pg-pool connect timeout; unguarded, it + // became an unhandled rejection that crashed whole production cells. + const verifyRelayToken = vi.fn(async (): Promise => { + throw new Error('Connection terminated due to connection timeout') + }) + const { activate } = createRegistry(activateControl, {}, verifyRelayToken) + const socket = new FakeSocket() + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + try { + await activate(socket as unknown as WebSocket, identity, null, 1, false, 7) + socket.emit( + 'message', + Buffer.from(JSON.stringify({ type: 'auth-refresh', relayJwt: 'refreshed' })), + false + ) + await vi.waitFor(() => + expect(socket.close).toHaveBeenCalledWith( + RELAY_CLOSE_CODE.LIMIT_EXCEEDED, + 'relay temporarily unavailable' + ) + ) + expect(warn).toHaveBeenCalledWith( + '[orca-relay] auth refresh failed: Connection terminated due to connection timeout' + ) + } finally { + warn.mockRestore() + } + }) + + it('attributes a control close to its client build without trusting the version string', async () => { + const activateControl = vi + .fn() + .mockResolvedValue('control:production-gce-c3:1') + const { activate } = createRegistry(activateControl) + const socket = new FakeSocket() + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + try { + // A client controls this string, so it must be bounded and stripped like any close reason. + await activate( + socket as unknown as WebSocket, + identity, + null, + 1, + false, + 1, + `1.4.173\n${'x'.repeat(200)}` + ) + socket.close(4408, 'replaced by a newer generation') + + const line = warn.mock.calls + .map((call) => String(call[0])) + .find((entry) => entry.includes('control closed')) + expect(line).toContain('app="1.4.173') + // The raw newline is escaped by JSON.stringify either way; only its escaped + // form proves the strip ran, so assert on that. + expect(line).not.toContain('\\n') + expect(line).toMatch(/app="[^"]{1,80}"/) + } finally { + warn.mockRestore() + } + }) + + it('reports the work a generation replacement destroyed, not the drained aftermath', async () => { + const activateControl = vi + .fn() + .mockResolvedValue('control:production-gce-c3:1') + const { registry, activate } = createRegistry(activateControl, { + failReservation: vi.fn().mockResolvedValue(undefined) + }) + const firstSocket = new FakeSocket() + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + try { + await activate(firstSocket as unknown as WebSocket, identity, null, 1, false, 1) + const session = registry.get({ + userId: identity.sub, + relayHostId: identity.relayHostId + })! + // Teardown drains both maps before closing the socket, so a close handler that + // reads them live always reports zero regardless of what was actually killed. + session.activeSplices.set('conn-a', () => session.activeSplices.delete('conn-a')) + session.activeSplices.set('conn-b', () => session.activeSplices.delete('conn-b')) + const clientSocket = new FakeSocket() + session.pendingConns.set('conn-c', { + connId: 'conn-c', + connTicket: 'ticket', + reservation: { userId: identity.sub, relayHostId: identity.relayHostId }, + client: clientSocket as unknown as WebSocket, + attachTimer: setTimeout(() => {}, 60_000), + credentialActivityId: null + } as unknown as Parameters[1]) + + const secondSocket = new FakeSocket() + await activate(secondSocket as unknown as WebSocket, identity, session, 2, false, 1) + + const line = warn.mock.calls + .map((call) => String(call[0])) + .find((entry) => entry.includes('replaced by a newer generation')) + expect(line).toContain('splices=2') + expect(line).toContain('pending=1') + } finally { + warn.mockRestore() + } + }) + + it('keeps the first drain snapshot when a drain is retried', async () => { + const activateControl = vi + .fn() + .mockResolvedValue('control:production-gce-c3:1') + const { registry, activate } = createRegistry(activateControl) + const socket = new FakeSocket() + await activate(socket as unknown as WebSocket, identity, null, 1, false, 1) + // Captured before draining, because teardown removes the session from the map. + const session = registry.get({ + userId: identity.sub, + relayHostId: identity.relayHostId + })! + session.activeSplices.set('conn-a', () => session.activeSplices.delete('conn-a')) + + // POST /v1/admin/drain has no idempotency guard, and SIGTERM then SIGINT both + // reach drain(), so a second teardown can be scheduled for the same session. + registry.drain(0) + const scheduled = vi.getTimerCount() + registry.drain(0) + // Pin the premise: if drain ever gains an idempotency guard, the retry schedules no + // second teardown and the assertion below stops defending the write-once snapshot + // while still passing. Compare against the count before the retry rather than an + // absolute, since the session's heartbeat interval is also pending. + expect(vi.getTimerCount()).toBe(scheduled + 1) + vi.advanceTimersByTime(1) + + // Asserting registry state, not the log line: FakeSocket closes synchronously, so + // the line is already emitted before the second teardown runs and would pass either way. + expect(session.closingCounts).toEqual({ splices: 1, pending: 0 }) + }) + + it('drains only the incarnation-bound host and makes replay idempotent', async () => { + const activateControl = vi + .fn() + .mockResolvedValue('control:production-gce-c3:1') + const { registry, activate } = createRegistry( + activateControl, + {}, + vi.fn(async () => identity) + ) + const firstSocket = new FakeSocket() + const secondSocket = new FakeSocket() + const secondIdentity = { + ...identity, + sub: 'user-2', + relayHostId: 'ponmlkjihgfedcba' + } + await activate(firstSocket as unknown as WebSocket, identity, null, 1, false, 7) + await activate(secondSocket as unknown as WebSocket, secondIdentity, null, 1, false, 3) + const trustProbe = { + attemptId: REGIONAL_REHOME_TRUST_PROBE_ATTEMPT_ID, + userId: REGIONAL_REHOME_TRUST_PROBE_USER_ID, + relayHostId: REGIONAL_REHOME_TRUST_PROBE_HOST_ID, + sourceAssignmentEpoch: 1, + graceMs: 0 + } + expect(registry.get(trustProbe)).toBeNull() + expect(registry.drainHost(trustProbe)).toBe('host-not-connected') + expect(registry.drainHost(trustProbe)).toBe('host-not-connected') + expect(firstSocket.send).not.toHaveBeenCalledWith(expect.stringContaining('"drain"')) + expect(secondSocket.send).not.toHaveBeenCalledWith(expect.stringContaining('"drain"')) + const request = { + attemptId: '11111111-1111-4111-8111-111111111111', + userId: identity.sub, + relayHostId: identity.relayHostId, + sourceAssignmentEpoch: 7, + graceMs: 30_000 + } + + expect(registry.drainHost(request)).toBe('accepted') + expect(firstSocket.send).toHaveBeenCalledWith( + JSON.stringify({ type: 'drain', graceMs: 30_000, recovery: 'resolve-director' }) + ) + expect(secondSocket.send).not.toHaveBeenCalledWith(expect.stringContaining('"drain"')) + firstSocket.emit( + 'message', + Buffer.from(JSON.stringify({ type: 'auth-refresh', relayJwt: 'refreshed' })), + false + ) + await vi.waitFor(() => expect(registry.get(request)?.state).toBe('drain-only')) + const timers = vi.getTimerCount() + expect(registry.drainHost(request)).toBe('already-accepted') + expect(vi.getTimerCount()).toBe(timers) + expect(() => + registry.drainHost({ + ...request, + attemptId: '22222222-2222-4222-8222-222222222222' + }) + ).toThrow('regional_rehome_attempt_conflict') + expect(() => + registry.drainHost({ ...request, sourceAssignmentEpoch: 8 }) + ).toThrow('regional_rehome_assignment_epoch_mismatch') + + const rebound = new FakeSocket() + await activate( + rebound as unknown as WebSocket, + identity, + registry.get(request), + 1, + true, + 7 + ) + expect(registry.get(request)?.state).toBe('drain-only') + expect(rebound.send).toHaveBeenCalledWith(expect.stringContaining('"type":"drain"')) + + await vi.advanceTimersByTimeAsync(30_000) + expect(registry.get(request)).toBeNull() + expect(registry.get({ userId: secondIdentity.sub, relayHostId: secondIdentity.relayHostId })) + .not.toBeNull() + expect(secondSocket.close).not.toHaveBeenCalled() + }) + + it('refreshes the logged client build when a control rebinds', async () => { + const activateControl = vi + .fn() + .mockResolvedValue('control:production-gce-c3:1') + const { registry, activate } = createRegistry(activateControl) + const firstSocket = new FakeSocket() + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + try { + await activate(firstSocket as unknown as WebSocket, identity, null, 1, false, 1, '1.4.100') + const session = registry.get({ + userId: identity.sub, + relayHostId: identity.relayHostId + })! + const rebindSocket = new FakeSocket() + await activate(rebindSocket as unknown as WebSocket, identity, session, 1, true, 1, '1.4.200') + + // The rebind closes the predecessor. That line is a churn line, so it must carry the + // build that socket ran, not the successor's — the refresh above lands before it closes. + const rebound = warn.mock.calls + .map((call) => String(call[0])) + .find((entry) => entry.includes('control rebound')) + expect(rebound).toContain('app="1.4.100"') + + rebindSocket.close(1006, 'network reset') + const line = warn.mock.calls + .map((call) => String(call[0])) + .find((entry) => entry.includes('code=1006')) + expect(line).toContain('app="1.4.200"') + } finally { + warn.mockRestore() + } + }) + + it('keeps every live control socket indexed when first activations overlap', async () => { + const firstControl = deferred() + const secondControl = deferred() + const activateControl = vi + .fn() + .mockReturnValueOnce(firstControl.promise) + .mockReturnValueOnce(secondControl.promise) + const { registry, activate } = createRegistry(activateControl) + const firstSocket = new FakeSocket() + const secondSocket = new FakeSocket() + + const first = activate(firstSocket as unknown as WebSocket, identity, null, 1, false, 1) + const second = activate(secondSocket as unknown as WebSocket, identity, null, 1, false, 1) + secondControl.resolve('control:production-gce-c3:1') + await Promise.resolve() + firstControl.resolve('control:production-gce-c3:1') + await Promise.all([first, second]) + + const liveSockets = [firstSocket, secondSocket].filter( + (socket) => socket.readyState === socket.OPEN + ) + const session = registry.get({ userId: identity.sub, relayHostId: identity.relayHostId }) + expect(liveSockets).toHaveLength(1) + expect(session?.socket).toBe(liveSockets[0]) + }) + + it('does not publish a control that closes during activation', async () => { + const blocked = deferred() + const activateControl = vi + .fn() + .mockReturnValueOnce(blocked.promise) + const { registry, activate, releaseActivity } = createRegistry(activateControl) + const socket = new FakeSocket() + + const activation = activate(socket as unknown as WebSocket, identity, null, 1, false, 1) + socket.close() + blocked.resolve('control:production-gce-c3:1') + await activation + + expect(registry.get({ userId: identity.sub, relayHostId: identity.relayHostId })).toBeNull() + expect(releaseActivity).toHaveBeenCalledWith( + { userId: identity.sub, relayHostId: identity.relayHostId }, + 'control:production-gce-c3:1' + ) + }) + + it('does not rebind a control that closes during activation', async () => { + const blocked = deferred() + const activateControl = vi + .fn() + .mockResolvedValueOnce('control:production-gce-c3:1') + .mockReturnValueOnce(blocked.promise) + const { registry, activate, releaseActivity } = createRegistry(activateControl) + const originalSocket = new FakeSocket() + await activate(originalSocket as unknown as WebSocket, identity, null, 1, false, 1) + const original = registry.get({ userId: identity.sub, relayHostId: identity.relayHostId }) + expect(original).not.toBeNull() + + const rebindSocket = new FakeSocket() + const rebinding = activate( + rebindSocket as unknown as WebSocket, + identity, + original, + 1, + true, + 1 + ) + rebindSocket.close() + blocked.resolve('control:production-gce-c3:1') + await rebinding + + const session = registry.get({ userId: identity.sub, relayHostId: identity.relayHostId }) + expect(session).toBe(original) + expect(session?.socket).toBe(originalSocket) + expect(session?.state).toBe('active') + expect(releaseActivity).toHaveBeenCalledWith( + { userId: identity.sub, relayHostId: identity.relayHostId }, + 'control:production-gce-c3:1' + ) + }) + + it('rejects client lookup when the indexed control socket is not open', async () => { + const reservation = { + userId: identity.sub, + relayHostId: identity.relayHostId, + credentialKind: 'resume', + relayDeviceId: 'device-1', + leaseExpiresAt: Date.now() + 60_000 + } + const store = { + resolveResume: vi.fn().mockResolvedValue({ userId: identity.sub }), + reserveCredential: vi.fn().mockResolvedValue(reservation), + failReservation: vi.fn().mockResolvedValue(undefined) + } + const activateControl = vi + .fn() + .mockResolvedValueOnce('control:production-gce-c3:1') + const { registry, activate } = createRegistry(activateControl, store) + const controlSocket = new FakeSocket() + await activate(controlSocket as unknown as WebSocket, identity, null, 1, false, 1) + const session = registry.get({ userId: identity.sub, relayHostId: identity.relayHostId }) + expect(session?.state).toBe('active') + // A dead socket that never delivered its close event: state stays active, + // so only the readyState guard can protect the lookup. + controlSocket.readyState = controlSocket.CLOSED + + const client = new FakeSocket() + await registry.acceptClient(client as unknown as WebSocket, identity.relayHostId, 'credential') + + expect(store.failReservation).toHaveBeenCalledWith(reservation) + expect(client.send).toHaveBeenCalledWith( + JSON.stringify({ type: 'relay-hello', ok: false, code: RELAY_CLOSE_CODE.HOST_OFFLINE }) + ) + expect(session?.pendingConns.size).toBe(0) + }) + + it('fails a control waiting behind a stalled activation without breaking serialization', async () => { + const stalled = deferred() + const activateControl = vi + .fn() + .mockReturnValueOnce(stalled.promise) + const { registry, activate } = createRegistry(activateControl) + const stalledSocket = new FakeSocket() + const first = activate(stalledSocket as unknown as WebSocket, identity, null, 1, false, 1) + const waitingSocket = new FakeSocket() + const second = activate(waitingSocket as unknown as WebSocket, identity, null, 1, false, 1) + + // Let the first activation reach the store (clearing its own queue timer) + // before the waiting control's deadline elapses. + await Promise.resolve() + await Promise.resolve() + vi.advanceTimersByTime(30_000) + expect(waitingSocket.close).toHaveBeenCalledWith( + RELAY_CLOSE_CODE.LIMIT_EXCEEDED, + 'control activation queue stalled' + ) + // The waiting control never reached the store; serialization held. + expect(activateControl).toHaveBeenCalledOnce() + + stalled.resolve('control:production-gce-c3:1') + await Promise.all([first, second]) + const session = registry.get({ userId: identity.sub, relayHostId: identity.relayHostId }) + expect(session?.socket).toBe(stalledSocket) + expect(session?.state).toBe('active') + }) + + it('rejects an activation that returns after drain begins', async () => { + const blocked = deferred() + const activateControl = vi + .fn() + .mockResolvedValueOnce('control:production-gce-c3:1') + .mockReturnValueOnce(blocked.promise) + const { registry, activate, renewControlActivity, releaseActivity } = + createRegistry(activateControl) + const originalSocket = new FakeSocket() + await activate(originalSocket as unknown as WebSocket, identity, null, 1, false, 1) + const original = registry.get({ + userId: identity.sub, + relayHostId: identity.relayHostId + }) + expect(original).not.toBeNull() + + const replacementSocket = new FakeSocket() + const replacement = activate( + replacementSocket as unknown as WebSocket, + identity, + original, + 2, + false, + 1 + ) + registry.drain(100) + blocked.resolve('control:production-gce-c3:2') + await replacement + vi.advanceTimersByTime(100) + vi.advanceTimersByTime(15_000) + + expect(replacementSocket.close).toHaveBeenCalledWith( + RELAY_CLOSE_CODE.DRAINING, + 'relay draining' + ) + expect(registry.get({ userId: identity.sub, relayHostId: identity.relayHostId })).toBeNull() + expect(renewControlActivity).not.toHaveBeenCalled() + expect(releaseActivity).toHaveBeenCalledWith( + { userId: identity.sub, relayHostId: identity.relayHostId }, + 'control:production-gce-c3:2' + ) + }) + + it('keeps a replacement mapped after stale orphan cleanup', async () => { + const activateControl = vi + .fn() + .mockResolvedValueOnce('control:production-gce-c3:1') + .mockResolvedValueOnce('control:production-gce-c3:2') + const { registry, activate } = createRegistry(activateControl) + const originalSocket = new FakeSocket() + await activate(originalSocket as unknown as WebSocket, identity, null, 1, false, 1) + const original = registry.get({ + userId: identity.sub, + relayHostId: identity.relayHostId + }) + expect(original).not.toBeNull() + originalSocket.close() + + const replacementSocket = new FakeSocket() + await activate( + replacementSocket as unknown as WebSocket, + identity, + original, + 2, + false, + 1 + ) + const replacement = registry.get({ + userId: identity.sub, + relayHostId: identity.relayHostId + }) + vi.advanceTimersByTime(CONTROL_CONTINUITY_LIMITS.orphanGraceMs) + + expect(replacement).not.toBeNull() + expect(registry.get({ userId: identity.sub, relayHostId: identity.relayHostId })).toBe( + replacement + ) + expect(registry.runtimeCounts().controls).toBe(1) + registry.drain(0) + vi.advanceTimersByTime(0) + }) + + it('stops the heartbeat of an actively replaced generation', async () => { + const activateControl = vi + .fn() + .mockResolvedValueOnce('control:production-gce-c3:1') + .mockResolvedValueOnce('control:production-gce-c3:2') + const { registry, activate, renewControlActivity } = createRegistry(activateControl) + const originalSocket = new FakeSocket() + await activate(originalSocket as unknown as WebSocket, identity, null, 1, false, 1) + const original = registry.get({ + userId: identity.sub, + relayHostId: identity.relayHostId + }) + expect(original).not.toBeNull() + + await activate( + new FakeSocket() as unknown as WebSocket, + identity, + original, + 2, + false, + 1 + ) + vi.advanceTimersByTime(15_000) + + expect(renewControlActivity).toHaveBeenCalledOnce() + expect(renewControlActivity).toHaveBeenCalledWith( + { userId: identity.sub, relayHostId: identity.relayHostId }, + expect.objectContaining({ + activityId: 'control:production-gce-c3:2', + cellId: 'production-gce-c3' + }) + ) + registry.drain(0) + vi.advanceTimersByTime(0) + }) + + it('keeps 15s pings while halving steady-state control renewals', async () => { + const activateControl = vi + .fn() + .mockResolvedValue('control:production-gce-c3:1') + const { registry, activate, renewControlActivity } = createRegistry(activateControl) + const socket = new FakeSocket() + const activatedAt = Date.now() + await activate(socket as unknown as WebSocket, identity, null, 1, false, 1) + + for (let interval = 0; interval < 4; interval++) { + await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs) + socket.emit('message', Buffer.from(JSON.stringify({ type: 'pong' })), false) + } + + const pings = socket.send.mock.calls.filter((call) => + String(call[0]).includes('"ping"') + ) + expect(pings).toHaveLength(4) + expect(renewControlActivity).toHaveBeenCalledTimes(2) + const firstExpiry = Number(renewControlActivity.mock.calls[0]![1].expiresAt) + const secondExpiry = Number(renewControlActivity.mock.calls[1]![1].expiresAt) + expect(firstExpiry).toBe( + activatedAt + + RELAY_PROTOCOL_LIMITS.controlPingIntervalMs + + ASSIGNMENT_LIMITS.activityLeaseMs + + RELAY_PROTOCOL_LIMITS.controlPingIntervalMs + ) + expect(secondExpiry - firstExpiry).toBe(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs * 2) + registry.drain(0) + vi.advanceTimersByTime(0) + }) + + it('retries a failed control renewal on the next ping', async () => { + const activateControl = vi + .fn() + .mockResolvedValue('control:production-gce-c3:1') + const { registry, activate, renewControlActivity } = createRegistry(activateControl) + renewControlActivity.mockRejectedValueOnce(new Error('pool timeout')) + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const socket = new FakeSocket() + try { + await activate(socket as unknown as WebSocket, identity, null, 1, false, 1) + await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs) + expect(renewControlActivity).toHaveBeenCalledOnce() + await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs) + expect(renewControlActivity).toHaveBeenCalledTimes(2) + } finally { + warn.mockRestore() + registry.drain(0) + vi.advanceTimersByTime(0) + } + }) + + it('retries past a stalled control renewal without waiting for it', async () => { + const activateControl = vi + .fn() + .mockResolvedValue('control:production-gce-c3:1') + const { registry, activate, renewControlActivity } = createRegistry(activateControl) + const stalled = deferred() + renewControlActivity.mockReturnValueOnce(stalled.promise) + const socket = new FakeSocket() + await activate(socket as unknown as WebSocket, identity, null, 1, false, 1) + + await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs * 2) + + expect(renewControlActivity).toHaveBeenCalledTimes(2) + stalled.resolve(undefined) + await vi.advanceTimersByTimeAsync(0) + registry.drain(0) + vi.advanceTimersByTime(0) + }) + + it('ignores a superseded renewal resolving after a fresher success', async () => { + const activateControl = vi + .fn() + .mockResolvedValue('control:production-gce-c3:1') + const { registry, activate, renewControlActivity } = createRegistry(activateControl) + const stalled = deferred() + renewControlActivity.mockReturnValueOnce(stalled.promise) + const socket = new FakeSocket() + await activate(socket as unknown as WebSocket, identity, null, 1, false, 1) + + await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs * 2) + expect(renewControlActivity).toHaveBeenCalledTimes(2) + stalled.resolve(undefined) + await vi.advanceTimersByTimeAsync(0) + await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs) + + expect(renewControlActivity).toHaveBeenCalledTimes(2) + await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs) + + expect(renewControlActivity).toHaveBeenCalledTimes(3) + registry.drain(0) + vi.advanceTimersByTime(0) + }) + + it('re-acquires a missing control activity lease', async () => { + const activateControl = vi + .fn() + .mockResolvedValue('control:production-gce-c3:1') + const { registry, activate, acquireActivity, renewControlActivity } = + createRegistry(activateControl) + renewControlActivity.mockRejectedValueOnce(new Error('control_activity_not_found')) + const socket = new FakeSocket() + await activate(socket as unknown as WebSocket, identity, null, 1, false, 1) + + await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs) + + expect(acquireActivity).toHaveBeenCalledWith( + { userId: identity.sub, relayHostId: identity.relayHostId }, + { + activityId: 'control:production-gce-c3:1', + kind: 'control', + cellId: config.cellId + } + ) + expect(socket.close).not.toHaveBeenCalled() + registry.drain(0) + vi.advanceTimersByTime(0) + }) + + it('closes a control whose activity lease moved to another cell', async () => { + const activateControl = vi + .fn() + .mockResolvedValue('control:production-gce-c3:1') + const { registry, activate, acquireActivity, renewControlActivity } = + createRegistry(activateControl) + renewControlActivity.mockRejectedValueOnce(new Error('control_activity_moved')) + const socket = new FakeSocket() + await activate(socket as unknown as WebSocket, identity, null, 1, false, 1) + + await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs) + + expect(acquireActivity).not.toHaveBeenCalled() + expect(socket.close).toHaveBeenCalledWith(RELAY_CLOSE_CODE.DRAINING, 'control activity moved') + registry.drain(0) + vi.advanceTimersByTime(0) + }) + + it('closes when a missing control activity cannot be re-acquired on this cell', async () => { + const activateControl = vi + .fn() + .mockResolvedValue('control:production-gce-c3:1') + const { registry, activate, acquireActivity, renewControlActivity } = + createRegistry(activateControl) + renewControlActivity.mockRejectedValueOnce(new Error('control_activity_not_found')) + acquireActivity.mockRejectedValueOnce(new Error('activity_cell_not_authoritative')) + const socket = new FakeSocket() + await activate(socket as unknown as WebSocket, identity, null, 1, false, 1) + + await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs) + + expect(socket.close).toHaveBeenCalledWith( + RELAY_CLOSE_CODE.DRAINING, + 'control migration completed' + ) + registry.drain(0) + vi.advanceTimersByTime(0) + }) + + it('closes a control when renewal finds its cell is no longer authoritative', async () => { + const activateControl = vi + .fn() + .mockResolvedValue('control:production-gce-c3:1') + const { registry, activate, renewControlActivity } = createRegistry(activateControl) + renewControlActivity.mockRejectedValueOnce(new Error('activity_cell_not_authoritative')) + const socket = new FakeSocket() + await activate(socket as unknown as WebSocket, identity, null, 1, false, 1) + + await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs) + + expect(socket.close).toHaveBeenCalledWith( + RELAY_CLOSE_CODE.DRAINING, + 'control migration completed' + ) + registry.drain(0) + vi.advanceTimersByTime(0) + }) + + it('continues renewing throughout the drain grace period', async () => { + const activateControl = vi + .fn() + .mockResolvedValue('control:production-gce-c3:1') + const { registry, activate, renewControlActivity } = createRegistry(activateControl) + const socket = new FakeSocket() + await activate(socket as unknown as WebSocket, identity, null, 1, false, 1) + registry.drain(60_000) + + for (let interval = 0; interval < 3; interval++) { + await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs) + socket.emit('message', Buffer.from(JSON.stringify({ type: 'pong' })), false) + } + + expect(renewControlActivity).toHaveBeenCalledTimes(2) + expect(socket.readyState).toBe(socket.OPEN) + vi.advanceTimersByTime(15_000) + }) +}) + +describe('control renewal cadence across a rebind', () => { + beforeEach(() => vi.useFakeTimers()) + afterEach(() => { + vi.clearAllTimers() + vi.useRealTimers() + }) + + it('keeps the halved cadence after a rebind lands under a stalled renewal', async () => { + const activateControl = vi + .fn() + .mockResolvedValue('control:production-gce-c3:1') + const { registry, activate, renewControlActivity } = createRegistry(activateControl) + const ping = RELAY_PROTOCOL_LIMITS.controlPingIntervalMs + const socket = new FakeSocket() + await activate(socket as unknown as WebSocket, identity, null, 1, false, 1) + + const beat = async (target: FakeSocket): Promise => { + await vi.advanceTimersByTimeAsync(ping) + target.emit('message', Buffer.from(JSON.stringify({ type: 'pong' })), false) + } + + // Age the session so its attempt counter is well above zero. + for (let tick = 0; tick < 5; tick++) await beat(socket) + expect(renewControlActivity).toHaveBeenCalledTimes(3) + + // The next renewal stalls and is still in flight when the control rebinds. + const stalled = deferred() + renewControlActivity.mockReturnValueOnce(stalled.promise) + for (let tick = 0; tick < 2; tick++) await beat(socket) + expect(renewControlActivity).toHaveBeenCalledTimes(4) + + const session = registry.get({ userId: identity.sub, relayHostId: identity.relayHostId })! + const rebindSocket = new FakeSocket() + await activate(rebindSocket as unknown as WebSocket, identity, session, 1, true, 1) + stalled.resolve(undefined) + await vi.advanceTimersByTimeAsync(0) + + const before = renewControlActivity.mock.calls.length + for (let tick = 0; tick < 4; tick++) await beat(rebindSocket) + expect(renewControlActivity.mock.calls.length - before).toBe(2) + + registry.drain(0) + vi.advanceTimersByTime(0) + }) +}) + +describe('control lease recovery after the session is gone', () => { + beforeEach(() => vi.useFakeTimers()) + afterEach(() => { + vi.clearAllTimers() + vi.useRealTimers() + }) + + it('does not re-acquire a lease for a session a newer generation already replaced', async () => { + const activateControl = vi + .fn() + .mockResolvedValueOnce('control:production-gce-c3:1') + .mockResolvedValueOnce('control:production-gce-c3:2') + const { registry, activate, acquireActivity, renewControlActivity } = + createRegistry(activateControl) + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + try { + const socket = new FakeSocket() + await activate(socket as unknown as WebSocket, identity, null, 1, false, 1) + const session = registry.get({ userId: identity.sub, relayHostId: identity.relayHostId })! + + // The renewal is still in flight when a newer generation takes over. + let failRenewal!: (error: Error) => void + renewControlActivity.mockReturnValueOnce( + new Promise((_resolve, reject) => (failRenewal = reject)) + ) + await vi.advanceTimersByTimeAsync(RELAY_PROTOCOL_LIMITS.controlPingIntervalMs) + expect(renewControlActivity).toHaveBeenCalledOnce() + + const newer = new FakeSocket() + await activate(newer as unknown as WebSocket, identity, session, 2, false, 1) + + // Its release already removed the lease, so the renewal reports it missing. + failRenewal(new Error('control_activity_not_found')) + await vi.advanceTimersByTimeAsync(0) + + expect(acquireActivity).not.toHaveBeenCalled() + } finally { + warn.mockRestore() + registry.drain(0) + vi.advanceTimersByTime(0) + } + }) +}) diff --git a/cloud/apps/relay/src/host-session-registry.ts b/cloud/apps/relay/src/host-session-registry.ts new file mode 100644 index 00000000000..11b7d1de030 --- /dev/null +++ b/cloud/apps/relay/src/host-session-registry.ts @@ -0,0 +1,1226 @@ +import { createHash, createHmac, randomBytes, randomUUID, timingSafeEqual } from 'node:crypto' +import { + ASSIGNMENT_LIMITS, + AuthRefreshSchema, + buildHostChallengePlaintext, + buildHostProofMacInput, + buildHostProofTranscript, + CONTROL_CONTINUITY_LIMITS, + DeviceCredentialInstallSchema, + DeviceCredentialInstallStatusSchema, + DeviceResumeConfirmSchema, + DeviceRevokeSchema, + HostChallengeAckSchema, + HostHelloSchema, + InviteCreateSchema, + RELAY_PROTOCOL_LIMITS, + RELAY_CLOSE_CODE +} from '@orca-cloud/relay-contract' +import nacl from 'tweetnacl' +import type WebSocket from 'ws' +import type { RawData } from 'ws' +import type { RelayConfig } from './config.js' +import type { RelayAssignmentStore } from './assignment-store.js' +import { + RelayCredentialStore, + type CredentialReservation +} from './credential-store.js' +import { relayHostLogDigest } from './relay-host-log-digest.js' +import type { RelayTokenClaims } from './relay-token-verifier.js' +import type { RelayRuntimeObserver } from './relay-observability.js' +import type { PendingHostDataReservation } from './relay-connection-ledger.js' +import { closeRelayWebSocket } from './relay-websocket-close.js' +import { ProcessQueuedByteBudget, wireSplice } from './splice-forwarder.js' + +// Peer-supplied close reasons are logged; keep them printable and short. +function printableCloseReason(reason: Buffer | string): string { + return reason + .toString() + .replace(/[^\x20-\x7e]/g, '') + .slice(0, 80) +} + +type VerifyRelayToken = (token: string) => Promise +type HostState = 'proving' | 'active' | 'orphaned' | 'drain-only' | 'closed' + +const CONTROL_ACTIVITY_RENEWAL_INTERVAL_MS = RELAY_PROTOCOL_LIMITS.controlPingIntervalMs * 2 +// Preserve the existing 75s renewal runway after doubling the successful-call interval. +const CONTROL_ACTIVITY_LEASE_MS = + ASSIGNMENT_LIMITS.activityLeaseMs + + CONTROL_ACTIVITY_RENEWAL_INTERVAL_MS - + RELAY_PROTOCOL_LIMITS.controlPingIntervalMs + +export type HostSession = { + identity: RelayTokenClaims + readonly relayHostId: string + readonly generation: number + readonly assignmentEpoch: number + readonly controlActivityId: string | null + readonly controlResumeSecret: string + // Why: reconnect churn is only actionable once it can be pinned to a client build. + // Refreshed on rebind so it describes the socket that closed, not the first one. + appVersion: string + state: HostState + socket: WebSocket | null + leaseExpiresAt: number + orphanTimer: ReturnType | null + heartbeatTimer: ReturnType | null + lastPongAt: number + activityRenewalDueAt: number + activityRenewalAttempt: number + activityRenewalCompletedAttempt: number + activeConnIds: Set + activeSplices: Map void> + pendingConns: Map + // Why: relay-initiated teardown drains these maps before closing the control + // socket, so the close handler would otherwise always report zero destroyed work. + closingCounts: { splices: number; pending: number } | null + regionalDrainAttemptId: string | null + regionalDrainTimer: ReturnType | null + regionalDrainExpiresAt: number | null +} + +export type RegionalHostDrainOutcome = + | 'accepted' + | 'already-accepted' + | 'host-not-connected' + +type PendingConnection = { + connId: string + connTicket: string + reservation: CredentialReservation + client: WebSocket + attachTimer: ReturnType + credentialActivityId: string | null + capacityReservation?: PendingHostDataReservation +} + +function decodeCanonicalBase64(value: string, bytes: number): Uint8Array | null { + if (!/^(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?$/.test(value)) { + return null + } + const decoded = Buffer.from(value, 'base64') + return decoded.length === bytes && decoded.toString('base64') === value ? decoded : null +} + +function relayHostId(publicKey: Uint8Array): string { + return createHash('sha256').update(publicKey).digest('base64url').slice(0, 16) +} + +function payload(raw: RawData, expectedType: string): unknown { + if (typeof raw !== 'string' && !Buffer.isBuffer(raw)) return null + try { + const parsed = JSON.parse(raw.toString()) as Record + if (parsed.type !== expectedType) return null + const { type: _type, ...rest } = parsed + return rest + } catch { + return null + } +} + +function send(socket: WebSocket, type: string, message: object): void { + socket.send(JSON.stringify({ type, ...message })) +} + +// Hosts abandon connects after 15s; waiting much longer than that behind a +// stalled predecessor only accumulates doomed sockets. +const ACTIVATION_QUEUE_WAIT_MS = 30_000 + +export class HostSessionRegistry { + private readonly sessions = new Map() + private readonly activationQueues = new Map>() + private draining = false + + constructor( + private readonly config: RelayConfig, + private readonly verifyRelayToken: VerifyRelayToken, + private readonly store: RelayCredentialStore, + private readonly assignments: RelayAssignmentStore, + private readonly queuedByteBudget: ProcessQueuedByteBudget, + private readonly observer: RelayRuntimeObserver, + private readonly now: () => number = Date.now + ) {} + + async acceptClient( + socket: WebSocket, + hostId: string, + credential: string, + capacityReservation?: PendingHostDataReservation + ): Promise { + if (this.draining) { + capacityReservation?.release() + this.rejectClient(socket, RELAY_CLOSE_CODE.DRAINING) + return + } + if (this.config.role === 'cell') { + const outerIdentity = + (await this.store.resolveResume(hostId, credential)) ?? + (await this.store.resolveInviteForMove(hostId, credential)) + const assignment = outerIdentity + ? await this.assignments.resolve({ userId: outerIdentity.userId, relayHostId: hostId }) + : null + if (!assignment || assignment.cellId !== this.config.cellId) { + capacityReservation?.release() + this.observer.recordAuth(false) + this.rejectClient(socket, RELAY_CLOSE_CODE.WRONG_CELL) + return + } + } + const reservation = await this.store.reserveCredential(hostId, credential) + if (!reservation) { + capacityReservation?.release() + this.observer.recordAuth(false) + this.rejectClient(socket, RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL) + return + } + this.observer.recordAuth(true) + const session = this.sessions.get(this.key(reservation.userId, hostId)) + if ( + !session || + session.state !== 'active' || + !session.socket || + session.socket.readyState !== session.socket.OPEN + ) { + capacityReservation?.release() + await this.store.failReservation(reservation) + this.rejectClient(socket, RELAY_CLOSE_CODE.HOST_OFFLINE) + return + } + if (session.activeConnIds.size + session.pendingConns.size >= 8) { + capacityReservation?.release() + await this.store.failReservation(reservation) + this.rejectClient(socket, RELAY_CLOSE_CODE.LIMIT_EXCEEDED) + return + } + const connId = randomUUID() + const connTicket = randomBytes(32).toString('base64url') + const identity = { userId: reservation.userId, relayHostId: hostId } + const credentialActivityId = + this.config.role === 'cell' + ? `${reservation.credentialKind === 'invite' ? 'invite' : 'confirmation'}:${connId}` + : null + if (credentialActivityId) { + try { + await this.assignments.acquireActivity(identity, { + activityId: credentialActivityId, + kind: reservation.credentialKind === 'invite' ? 'invite' : 'confirmation', + cellId: this.config.cellId + }) + } catch { + capacityReservation?.release() + await this.store.failReservation(reservation) + this.rejectClient(socket, RELAY_CLOSE_CODE.LIMIT_EXCEEDED) + return + } + } + const attachTimer = setTimeout(() => { + session.pendingConns.delete(connId) + capacityReservation?.release() + this.failReservationBestEffort(reservation) + if (credentialActivityId) this.releaseActivityBestEffort(identity, credentialActivityId) + this.rejectClient(socket, RELAY_CLOSE_CODE.HOST_OFFLINE) + }, RELAY_PROTOCOL_LIMITS.hostAttachDeadlineMs) + const pending: PendingConnection = { + connId, + connTicket, + reservation, + client: socket, + attachTimer, + credentialActivityId, + capacityReservation + } + capacityReservation?.bind(connId) + session.pendingConns.set(connId, pending) + send(session.socket, 'conn-open', { + connId, + connTicket, + kind: reservation.credentialKind, + relayDeviceId: reservation.relayDeviceId, + attachDeadlineMs: RELAY_PROTOCOL_LIMITS.hostAttachDeadlineMs + }) + socket.once('close', () => { + const current = session.pendingConns.get(connId) + if (current?.client === socket) { + clearTimeout(current.attachTimer) + session.pendingConns.delete(connId) + current.capacityReservation?.release() + this.failReservationBestEffort(current.reservation) + if (current.credentialActivityId) { + this.releaseActivityBestEffort(identity, current.credentialActivityId) + } + } + }) + } + + async acceptHostData( + socket: WebSocket, + connId: string, + connTicket: string, + generation: number + ): Promise { + const session = [...this.sessions.values()].find((candidate) => + candidate.pendingConns.has(connId) + ) + const pending = session?.pendingConns.get(connId) + if ( + !session || + !pending || + pending.connTicket !== connTicket || + session.generation !== generation || + (session.state !== 'active' && session.state !== 'drain-only') + ) { + this.observer.recordAuth(false) + socket.close(RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL, 'invalid host data ticket') + return false + } + this.observer.recordAuth(true) + clearTimeout(pending.attachTimer) + session.pendingConns.delete(connId) + session.activeConnIds.add(connId) + const basisDeadline = + pending.reservation.credentialKind === 'resume' + ? this.now() + RELAY_PROTOCOL_LIMITS.resumeConfirmationDeadlineMs + : pending.reservation.leaseExpiresAt + const identity = { + userId: pending.reservation.userId, + relayHostId: pending.reservation.relayHostId + } + const spliceActivityId = this.config.role === 'cell' ? `splice:${connId}` : null + try { + if (spliceActivityId) { + await this.assignments.acquireActivity(identity, { + activityId: spliceActivityId, + kind: 'splice', + cellId: this.config.cellId + }) + } + await this.store.recordConnectionBasis({ + ...pending.reservation, + basisConnId: connId, + owningControlGeneration: session.generation, + deadline: basisDeadline + }) + } catch { + session.activeConnIds.delete(connId) + await this.store.failReservation(pending.reservation) + if (spliceActivityId) this.releaseActivityBestEffort(identity, spliceActivityId) + if (pending.credentialActivityId) { + this.releaseActivityBestEffort(identity, pending.credentialActivityId) + } + this.rejectClient(pending.client, RELAY_CLOSE_CODE.LIMIT_EXCEEDED) + socket.close(RELAY_CLOSE_CODE.LIMIT_EXCEEDED, 'basis persistence failed') + return false + } + const close = wireSplice({ + client: pending.client, + host: socket, + budget: this.queuedByteBudget, + onClose: () => { + session.activeConnIds.delete(connId) + session.activeSplices.delete(connId) + this.deactivateBasisBestEffort(connId) + if (spliceActivityId) this.releaseActivityBestEffort(identity, spliceActivityId) + if (pending.credentialActivityId) { + this.releaseActivityBestEffort(identity, pending.credentialActivityId) + } + }, + onForwardedBytes: (bytes) => this.observer.recordForwardedBytes(bytes), + onClosed: (closeInfo) => { + this.observer.recordSpliceClose?.(closeInfo.trigger) + // Only abnormal closes are logged; routine peer disconnects would be + // one line per phone backgrounding. + if ( + closeInfo.code === RELAY_CLOSE_CODE.LIMIT_EXCEEDED || + closeInfo.trigger.includes('error') || + closeInfo.trigger.includes('oversize') + ) { + console.warn( + `[orca-relay] splice closed host=${relayHostLogDigest(session.relayHostId)}` + + ` trigger=${closeInfo.trigger} code=${closeInfo.code}` + + ` reason=${JSON.stringify(closeInfo.reason)}` + ) + } + } + }) + session.activeSplices.set(connId, close) + if (pending.client.readyState !== pending.client.OPEN || socket.readyState !== socket.OPEN) { + close() + return false + } + send(pending.client, 'relay-hello', { + ok: true, + credentialKind: pending.reservation.credentialKind, + leaseExpiresAt: pending.reservation.leaseExpiresAt, + ...(pending.reservation.credentialKind === 'resume' + ? { + acceptedCredentialVersion: pending.reservation.acceptedCredentialVersion, + acceptedAs: pending.reservation.acceptedAs, + resumeExpiresAt: pending.reservation.resumeExpiresAt, + ...(pending.reservation.graceExpiresAt === undefined + ? {} + : { graceExpiresAt: pending.reservation.graceExpiresAt }) + } + : {}) + }) + return true + } + + acceptControl( + socket: WebSocket, + identity: RelayTokenClaims, + connectionInclusionWatermark?: number + ): void { + if (this.draining) { + socket.close(RELAY_CLOSE_CODE.DRAINING, 'relay draining') + return + } + let firstFrameTimer: ReturnType | null = setTimeout(() => { + socket.close(RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL, 'host hello timeout') + }, 2_000) + socket.once('message', (raw, isBinary) => { + if (firstFrameTimer) clearTimeout(firstFrameTimer) + firstFrameTimer = null + if (isBinary) { + socket.close(RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL, 'host hello must be text') + return + } + this.guardSessionTask( + () => + this.beginProof( + socket, + identity, + payload(raw, 'host-hello'), + connectionInclusionWatermark + ), + socket, + 'host hello proof' + ) + }) + } + + // A dependency failure (e.g. a database connect timeout) must cost one + // handshake or command, not the process: an unhandled rejection here has + // crashed whole cells and wiped their in-memory draining flag. 4429 is the + // endpoint-scoped close in the contract, so the client retries this cell. + private guardSessionTask( + task: () => Promise, + socket: WebSocket | null, + context: string + ): void { + void Promise.resolve() + .then(task) + .catch((error: unknown) => { + const message = (error instanceof Error ? error.message : 'unknown') + // Untruncated, unlike peer-supplied close reasons: this is the + // primary diagnostic for the next rejection class. + .replace(/[^\x20-\x7e]/g, '') + console.warn(`[orca-relay] ${context} failed: ${message}`) + socket?.close(RELAY_CLOSE_CODE.LIMIT_EXCEEDED, 'relay temporarily unavailable') + }) + // Terminal: a throw in the handler above must not itself crash the process. + .catch(() => {}) + } + + get(identity: RelayIdentityKey): HostSession | null { + return this.sessions.get(this.key(identity.userId, identity.relayHostId)) ?? null + } + + hasActiveControl(identity: RelayIdentityKey): boolean { + const session = this.get(identity) + return ( + session !== null && + session.state === 'active' && + session.socket !== null && + session.socket.readyState === session.socket.OPEN + ) + } + + runtimeCounts(): { controls: number; splices: number; pendingSplices: number } { + let controls = 0 + let splices = 0 + let pendingSplices = 0 + for (const session of this.sessions.values()) { + const socket = session.socket + if ( + socket !== null && + socket.readyState === socket.OPEN && + (session.state === 'active' || session.state === 'drain-only') + ) { + controls++ + } + splices += session.activeSplices.size + pendingSplices += session.pendingConns.size + } + return { controls, splices, pendingSplices } + } + + drain(graceMs: number): void { + this.draining = true + for (const session of this.sessions.values()) { + if (session.state === 'closed') continue + session.state = 'drain-only' + if (session.socket) send(session.socket, 'drain', { graceMs, recovery: 'resolve-director' }) + setTimeout(() => this.closeDrainedSession(session), graceMs) + } + } + + drainHost(input: { + attemptId: string + userId: string + relayHostId: string + sourceAssignmentEpoch: number + graceMs: number + }): RegionalHostDrainOutcome { + const session = this.get(input) + if (!session || session.state === 'closed') return 'host-not-connected' + if (session.assignmentEpoch !== input.sourceAssignmentEpoch) { + throw new Error('regional_rehome_assignment_epoch_mismatch') + } + if (session.regionalDrainAttemptId) { + if (session.regionalDrainAttemptId !== input.attemptId) { + throw new Error('regional_rehome_attempt_conflict') + } + this.reassertRegionalDrain(session) + return 'already-accepted' + } + session.regionalDrainAttemptId = input.attemptId + session.regionalDrainExpiresAt = this.now() + input.graceMs + this.reassertRegionalDrain(session) + session.regionalDrainTimer = setTimeout( + () => this.closeDrainedSession(session), + input.graceMs + ) + return 'accepted' + } + + isDraining(): boolean { + return this.draining + } + + private async beginProof( + socket: WebSocket, + identity: RelayTokenClaims, + candidate: unknown, + connectionInclusionWatermark?: number + ): Promise { + const hello = HostHelloSchema.safeParse(candidate) + const hostPublicKey = hello.success + ? decodeCanonicalBase64(hello.data.hostPublicKeyB64, 32) + : null + if (!hello.success) { + this.observer.recordAuth(false) + socket.close(RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL, 'invalid host hello') + return + } + if (!hostPublicKey) { + this.observer.recordAuth(false) + socket.close(RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL, 'invalid host public key') + return + } + if ( + hello.data.relayHostId !== identity.relayHostId || + relayHostId(hostPublicKey) !== identity.relayHostId + ) { + this.observer.recordAuth(false) + socket.close(RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL, 'host key binding mismatch') + return + } + // Combined is staging-only compatibility; stamped cells require the durable director epoch. + const assignmentValid = + 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 + }) + if (!assignmentValid) { + this.observer.recordAuth(false) + socket.close(RELAY_CLOSE_CODE.WRONG_CELL, 'wrong assignment epoch') + return + } + + const key = this.key(identity.sub, identity.relayHostId) + const existing = this.sessions.get(key) + const rebind = Boolean( + existing && + hello.data.controlResumeSecret && + hello.data.controlResumeSecret === existing.controlResumeSecret && + (existing.state === 'orphaned' || existing.state === 'active') + ) + const generation = rebind ? existing!.generation : (existing?.generation ?? 0) + 1 + const ephemeral = nacl.box.keyPair() + const challengeNonce = randomBytes(nacl.box.nonceLength) + const challengeSecret = randomBytes(32) + const challengeId = randomUUID() + const issuedAt = this.now() + const expiresAt = issuedAt + 10_000 + const transcript = buildHostProofTranscript({ + relayOrigin: this.config.publicUrl, + relayEphemeralPublicKey: ephemeral.publicKey, + challengeNonce, + challengeId, + issuedAt, + expiresAt, + userId: identity.sub, + profileId: identity.prof, + organizationId: identity.org ?? '', + relayHostId: identity.relayHostId, + hostPublicKey, + assignmentEpoch: hello.data.assignmentEpoch, + previousGeneration: hello.data.previousGeneration, + resumeRequested: rebind + }) + const plaintext = buildHostChallengePlaintext(transcript, challengeSecret) + const ciphertext = nacl.box(plaintext, challengeNonce, hostPublicKey, ephemeral.secretKey) + const expectedProof = createHmac('sha256', challengeSecret) + .update(buildHostProofMacInput(transcript)) + .digest() + send(socket, 'host-challenge', { + challengeId, + relayEphemeralPublicKeyB64: Buffer.from(ephemeral.publicKey).toString('base64'), + nonceB64: Buffer.from(challengeNonce).toString('base64'), + ciphertextB64: Buffer.from(ciphertext).toString('base64'), + expiresAt + }) + const proofTimer = setTimeout(() => { + socket.close(RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL, 'host proof timeout') + }, 10_000) + socket.once('message', (raw, isBinary) => { + clearTimeout(proofTimer) + const ack = isBinary ? null : HostChallengeAckSchema.safeParse(payload(raw, 'host-challenge-ack')) + const proof = ack?.success ? decodeCanonicalBase64(ack.data.proofB64, 32) : null + if ( + !ack?.success || + ack.data.challengeId !== challengeId || + !proof || + this.now() > expiresAt || + !timingSafeEqual(proof, expectedProof) + ) { + this.observer.recordAuth(false) + socket.close(RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL, 'invalid host proof') + return + } + this.observer.recordAuth(true) + this.guardSessionTask( + () => + this.activate( + socket, + identity, + existing ?? null, + generation, + rebind, + hello.data.assignmentEpoch, + hello.data.appVersion, + connectionInclusionWatermark + ), + socket, + 'host activation' + ) + }) + } + + private activate( + socket: WebSocket, + identity: RelayTokenClaims, + existing: HostSession | null, + generation: number, + rebind: boolean, + assignmentEpoch: number, + appVersion: string, + connectionInclusionWatermark?: number + ): Promise { + const key = this.key(identity.sub, identity.relayHostId) + const previous = this.activationQueues.get(key) ?? Promise.resolve() + // The timeout only fails this waiting socket; the queue entry still chains + // behind the stalled predecessor so activations never run concurrently. + let queueWaitExpired = false + const queueWaitTimer = setTimeout(() => { + queueWaitExpired = true + socket.close(RELAY_CLOSE_CODE.LIMIT_EXCEEDED, 'control activation queue stalled') + }, ACTIVATION_QUEUE_WAIT_MS) + queueWaitTimer.unref?.() + const activation = previous.catch(() => undefined).then(async () => { + clearTimeout(queueWaitTimer) + if (queueWaitExpired) return + if ((this.sessions.get(key) ?? null) !== existing) { + socket.close(RELAY_CLOSE_CODE.PEER_DROPPED, 'control activation superseded') + return + } + await this.activateCurrent( + socket, + identity, + existing, + generation, + rebind, + assignmentEpoch, + appVersion, + connectionInclusionWatermark + ) + }) + this.activationQueues.set(key, activation) + const cleanup = (): void => { + if (this.activationQueues.get(key) === activation) this.activationQueues.delete(key) + } + void activation.then(cleanup, cleanup) + return activation + } + + private async activateCurrent( + socket: WebSocket, + identity: RelayTokenClaims, + existing: HostSession | null, + generation: number, + rebind: boolean, + assignmentEpoch: number, + appVersion: string, + connectionInclusionWatermark?: number + ): Promise { + let controlActivityId: string | null = null + if (this.config.role === 'cell') { + try { + controlActivityId = await this.assignments.activateControl( + { userId: identity.sub, relayHostId: identity.relayHostId }, + { + cellId: this.config.cellId, + assignmentEpoch, + generation, + connectionInclusionWatermark + } + ) + await this.assignments.markMigrationTargetRegistered( + { userId: identity.sub, relayHostId: identity.relayHostId }, + { cellId: this.config.cellId, assignmentEpoch } + ) + } catch { + // A failure after activateControl succeeded must not strand the + // acquired control activity until lease expiry. + if (controlActivityId) { + this.releaseActivityBestEffort( + { userId: identity.sub, relayHostId: identity.relayHostId }, + controlActivityId + ) + } + socket.close(RELAY_CLOSE_CODE.WRONG_CELL, 'assignment changed during host proof') + return + } + } + if (this.draining || socket.readyState !== socket.OPEN) { + if (controlActivityId) { + this.releaseActivityBestEffort( + { userId: identity.sub, relayHostId: identity.relayHostId }, + controlActivityId + ) + } + if (socket.readyState === socket.OPEN) { + socket.close(RELAY_CLOSE_CODE.DRAINING, 'relay draining') + } + return + } + if (existing) this.observer.recordReconnect() + if (rebind && existing) { + const previousSocket = existing.socket + if (existing.orphanTimer) clearTimeout(existing.orphanTimer) + existing.orphanTimer = null + existing.socket = socket + existing.state = existing.regionalDrainAttemptId ? 'drain-only' : 'active' + existing.appVersion = appVersion + existing.leaseExpiresAt = this.now() + 55 * 60 * 1000 + existing.lastPongAt = this.now() + existing.activityRenewalDueAt = + this.now() + RELAY_PROTOCOL_LIMITS.controlPingIntervalMs + this.wireActiveControl(existing) + this.sendHelloAck(existing) + if (existing.regionalDrainAttemptId) this.reassertRegionalDrain(existing) + previousSocket?.close(RELAY_CLOSE_CODE.PEER_DROPPED, 'control rebound') + return + } + if (existing) { + existing.state = 'closed' + if (existing.heartbeatTimer) clearInterval(existing.heartbeatTimer) + if (existing.orphanTimer) clearTimeout(existing.orphanTimer) + if (existing.regionalDrainTimer) clearTimeout(existing.regionalDrainTimer) + existing.heartbeatTimer = null + existing.orphanTimer = null + existing.regionalDrainTimer = null + existing.closingCounts ??= { + splices: existing.activeSplices.size, + pending: existing.pendingConns.size + } + for (const close of existing.activeSplices.values()) close() + for (const pending of existing.pendingConns.values()) { + clearTimeout(pending.attachTimer) + pending.capacityReservation?.release() + this.failReservationBestEffort(pending.reservation) + if (pending.credentialActivityId) { + this.releaseActivityBestEffort( + { + userId: pending.reservation.userId, + relayHostId: pending.reservation.relayHostId + }, + pending.credentialActivityId + ) + } + this.rejectClient(pending.client, RELAY_CLOSE_CODE.PEER_DROPPED) + } + existing.socket?.close(RELAY_CLOSE_CODE.PEER_DROPPED, 'replaced by a newer generation') + this.releaseControlActivity(existing) + } + const session: HostSession = { + identity, + relayHostId: identity.relayHostId, + generation, + assignmentEpoch, + controlActivityId, + controlResumeSecret: randomBytes(32).toString('base64url'), + appVersion, + state: 'active', + socket, + leaseExpiresAt: this.now() + 55 * 60 * 1000, + orphanTimer: null, + heartbeatTimer: null, + lastPongAt: this.now(), + activityRenewalDueAt: this.now() + RELAY_PROTOCOL_LIMITS.controlPingIntervalMs, + activityRenewalAttempt: 0, + activityRenewalCompletedAttempt: 0, + activeConnIds: new Set(), + activeSplices: new Map(), + pendingConns: new Map(), + closingCounts: null, + regionalDrainAttemptId: null, + regionalDrainTimer: null, + regionalDrainExpiresAt: null + } + this.sessions.set(this.key(identity.sub, identity.relayHostId), session) + this.wireActiveControl(session) + this.sendHelloAck(session) + } + + private wireActiveControl(session: HostSession): void { + const socket = session.socket! + const wiredAt = this.now() + // Why: pin the build to THIS socket. A rebind refreshes session.appVersion and + // only then closes the predecessor, whose close event always lands after that + // write, so reading it at log time would stamp the successor's build. + const appVersion = session.appVersion + let socketError: string | null = null + // Why: an unhandled ws 'error' (e.g. an oversize control frame) would + // otherwise throw process-wide; the message also explains the close below. + socket.on('error', (error) => { + socketError ??= error.message + }) + socket.once('close', (code, reason) => { + this.observer.recordControlClose?.(code) + // One line per control close makes reconnect churners attributable by + // host digest without exposing the raw relay host id. + console.warn( + `[orca-relay] control closed host=${relayHostLogDigest(session.relayHostId)}` + + ` gen=${session.generation} state=${session.state} ageMs=${this.now() - wiredAt}` + + ` app=${JSON.stringify(printableCloseReason(appVersion))}` + + ` splices=${session.closingCounts?.splices ?? session.activeSplices.size}` + + ` pending=${session.closingCounts?.pending ?? session.pendingConns.size}` + + ` code=${code} reason=${JSON.stringify(printableCloseReason(reason))}` + + (socketError === null ? '' : ` error=${JSON.stringify(printableCloseReason(socketError))}`) + ) + }) + socket.on('message', (raw, isBinary) => { + if (isBinary || (session.state !== 'active' && session.state !== 'drain-only')) return + try { + const parsed = JSON.parse(raw.toString()) as Record + if (parsed.type === 'pong') { + session.lastPongAt = this.now() + return + } + if (parsed.type === 'auth-refresh') { + // Close the socket the message arrived on: after a rebind, + // session.socket already points at the successor. + this.guardSessionTask(() => this.acceptRefresh(session, raw), socket, 'auth refresh') + return + } + this.guardSessionTask( + () => this.acceptControlCommand(session, parsed.type, raw), + socket, + 'control command' + ) + } catch { + socket.close(RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL, 'invalid control JSON') + } + }) + socket.once('close', () => { + if (session.socket !== socket || session.state === 'closed') return + session.socket = null + session.state = session.regionalDrainAttemptId ? 'drain-only' : 'orphaned' + if (session.heartbeatTimer) clearInterval(session.heartbeatTimer) + session.heartbeatTimer = null + session.orphanTimer = setTimeout(() => { + session.state = 'closed' + for (const close of session.activeSplices.values()) close() + const key = this.key(session.identity.sub, session.relayHostId) + if (this.sessions.get(key) === session) this.sessions.delete(key) + this.releaseControlActivity(session) + }, CONTROL_CONTINUITY_LIMITS.orphanGraceMs) + }) + if (session.heartbeatTimer) clearInterval(session.heartbeatTimer) + session.heartbeatTimer = setInterval( + () => this.heartbeat(session), + RELAY_PROTOCOL_LIMITS.controlPingIntervalMs + ) + } + + private async acceptRefresh(session: HostSession, raw: RawData): Promise { + const parsed = AuthRefreshSchema.safeParse(payload(raw, 'auth-refresh')) + if (!parsed.success) return + const refreshed = await this.verifyRelayToken(parsed.data.relayJwt) + const sameIdentity = + refreshed && + refreshed.sub === session.identity.sub && + refreshed.prof === session.identity.prof && + refreshed.org === session.identity.org && + refreshed.relayHostId === session.identity.relayHostId + if (!sameIdentity) { + session.socket?.close(RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL, 'refresh identity changed') + return + } + session.identity = refreshed + if (!session.regionalDrainAttemptId) session.state = 'active' + } + + private heartbeat(session: HostSession): void { + const key = this.key(session.identity.sub, session.relayHostId) + if (this.sessions.get(key) !== session || session.state === 'closed') { + if (session.heartbeatTimer) clearInterval(session.heartbeatTimer) + session.heartbeatTimer = null + return + } + const now = this.now() + if (!session.socket) return + const controlActivityId = session.controlActivityId + if (controlActivityId && now >= session.activityRenewalDueAt) { + const attempt = ++session.activityRenewalAttempt + const startedAt = now + void this.assignments + .renewControlActivity( + { userId: session.identity.sub, relayHostId: session.relayHostId }, + { + activityId: controlActivityId, + cellId: this.config.cellId, + expiresAt: startedAt + CONTROL_ACTIVITY_LEASE_MS + } + ) + .then(() => { + if (attempt <= session.activityRenewalCompletedAttempt) return + session.activityRenewalCompletedAttempt = attempt + session.activityRenewalDueAt = startedAt + CONTROL_ACTIVITY_RENEWAL_INTERVAL_MS + }) + .catch(async (error: unknown) => { + if (error instanceof Error && error.message === 'activity_cell_not_authoritative') { + // Completion fences a late drain-only heartbeat after all source work is gone. + session.socket?.close(RELAY_CLOSE_CODE.DRAINING, 'control migration completed') + return + } + if (error instanceof Error && error.message === 'control_activity_not_found') { + if ( + this.sessions.get(key) !== session || + session.state === 'closed' || + !session.socket + ) { + return + } + try { + await this.assignments.acquireActivity( + { userId: session.identity.sub, relayHostId: session.relayHostId }, + { + activityId: controlActivityId, + kind: 'control', + cellId: this.config.cellId + } + ) + this.observer.recordControlActivityRecovery?.(true) + } catch (acquireError: unknown) { + this.observer.recordControlActivityRecovery?.(false) + if ( + acquireError instanceof Error && + acquireError.message === 'activity_cell_not_authoritative' + ) { + session.socket?.close(RELAY_CLOSE_CODE.DRAINING, 'control migration completed') + return + } + console.warn('[orca-relay] control activity recovery failed') + } + return + } + if (error instanceof Error && error.message === 'control_activity_moved') { + session.socket?.close(RELAY_CLOSE_CODE.DRAINING, 'control activity moved') + return + } + console.warn('[orca-relay] control activity renewal failed') + }) + // Terminal handler: a throw inside the async catch above (e.g. a + // future await) must not become a process-killing rejection. + .catch(() => { + console.warn('[orca-relay] control activity renewal handling failed') + }) + } + if (now - session.lastPongAt > 75_000) { + session.socket.close(RELAY_CLOSE_CODE.PEER_DROPPED, 'control silence timeout') + return + } + const expiresAt = session.identity.exp * 1000 + if (now > expiresAt + CONTROL_CONTINUITY_LIMITS.expiredAuthExistingSpliceGraceMs) { + session.socket.close(RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL, 'relay authorization expired') + return + } + if (now > expiresAt) session.state = 'drain-only' + if (now > session.leaseExpiresAt) { + send(session.socket, 'drain', { graceMs: 0, recovery: 'resolve-director' }) + session.socket.close(RELAY_CLOSE_CODE.DRAINING, 'control lease expired') + return + } + send(session.socket, 'ping', { t: now }) + } + + private sendHelloAck(session: HostSession): void { + if (!session.socket) return + send(session.socket, 'host-hello-ack', { + v: 1, + generation: session.generation, + controlResumeSecret: session.controlResumeSecret, + leaseExpiresAt: session.leaseExpiresAt, + activeConnIds: [...session.activeConnIds], + pendingConns: [...session.pendingConns.values()].map((pending) => ({ + connId: pending.connId, + connTicket: pending.connTicket + })) + }) + } + + private closeDrainedSession(session: HostSession): void { + if (session.state === 'closed') return + if (session.heartbeatTimer) clearInterval(session.heartbeatTimer) + if (session.orphanTimer) clearTimeout(session.orphanTimer) + if (session.regionalDrainTimer) clearTimeout(session.regionalDrainTimer) + session.heartbeatTimer = null + session.orphanTimer = null + session.regionalDrainTimer = null + session.regionalDrainExpiresAt = null + session.closingCounts ??= { + splices: session.activeSplices.size, + pending: session.pendingConns.size + } + for (const close of session.activeSplices.values()) { + close(RELAY_CLOSE_CODE.DRAINING, 'relay draining') + } + for (const pending of session.pendingConns.values()) { + clearTimeout(pending.attachTimer) + pending.capacityReservation?.release() + this.failReservationBestEffort(pending.reservation) + if (pending.credentialActivityId) { + this.releaseActivityBestEffort( + { + userId: pending.reservation.userId, + relayHostId: pending.reservation.relayHostId + }, + pending.credentialActivityId + ) + } + this.rejectClient(pending.client, RELAY_CLOSE_CODE.DRAINING) + } + session.pendingConns.clear() + session.state = 'closed' + if (session.socket) { + closeRelayWebSocket( + session.socket, + RELAY_CLOSE_CODE.DRAINING, + 'resolve configured director' + ) + } + const key = this.key(session.identity.sub, session.relayHostId) + if (this.sessions.get(key) === session) this.sessions.delete(key) + this.releaseControlActivity(session) + } + + private reassertRegionalDrain(session: HostSession): void { + session.state = 'drain-only' + if (!session.socket) return + send(session.socket, 'drain', { + graceMs: Math.max(0, (session.regionalDrainExpiresAt ?? this.now()) - this.now()), + recovery: 'resolve-director' + }) + } + + private key(userId: string, hostId: string): string { + return `${userId}\0${hostId}` + } + + private async acceptControlCommand( + session: HostSession, + type: unknown, + raw: RawData + ): Promise { + if (typeof type !== 'string' || !session.socket) return + try { + const identity = { userId: session.identity.sub, relayHostId: session.relayHostId } + if (type === 'invite-create') { + if (session.state !== 'active') throw new Error('authorization_expired') + const request = InviteCreateSchema.parse(payload(raw, type)) + const activityId = `invite-offer:${request.reqId}` + if (this.config.role === 'cell') { + await this.assignments.acquireActivity(identity, { + activityId, + kind: 'invite', + cellId: this.config.cellId + }) + } + let invite + try { + invite = await this.store.createInvite(identity, request.relayDeviceId) + if (this.config.role === 'cell') { + await this.assignments.acquireActivity(identity, { + activityId, + kind: 'invite', + cellId: this.config.cellId, + expiresAt: invite.expiresAt + }) + } + } catch (error) { + if (this.config.role === 'cell') { + await this.assignments.releaseActivity(identity, activityId) + } + throw error + } + send(session.socket, 'invite-created', { reqId: request.reqId, ...invite }) + return + } + if (type === 'device-credential-install') { + const request = DeviceCredentialInstallSchema.parse(payload(raw, type)) + if ( + session.state !== 'active' && + request.authorization.mode === 'authenticated-direct' + ) { + throw new Error('authorization_expired') + } + const installActivityId = `install:${request.reqId}` + if (this.config.role === 'cell') { + await this.assignments.acquireActivity(identity, { + activityId: installActivityId, + kind: 'install', + cellId: this.config.cellId + }) + } + let result + try { + if (request.authorization.mode === 'authenticated-direct') { + await this.store.recordDirectAuthorization({ + ...identity, + relayDeviceId: request.relayDeviceId, + directAuthId: request.authorization.directAuthId, + owningControlGeneration: session.generation, + deadline: this.now() + RELAY_PROTOCOL_LIMITS.resumeConfirmationDeadlineMs + }) + } + result = await this.store.installCredential({ + ...identity, + ...request, + owningControlGeneration: session.generation + }) + } finally { + if (this.config.role === 'cell') { + await this.assignments.releaseActivity(identity, installActivityId) + } + } + if (request.authorization.mode === 'relay-basis') { + await this.assignments.releaseActivity( + identity, + `invite:${request.authorization.basisConnId}` + ) + } + send(session.socket, 'device-credential-installed', result) + return + } + if (type === 'device-credential-install-status') { + const request = DeviceCredentialInstallStatusSchema.parse(payload(raw, type)) + const result = await this.store.installStatus({ ...identity, ...request }) + send(session.socket, 'device-credential-install-status-result', { + v: 1, + reqId: request.reqId, + state: result ? 'committed' : 'not-found', + ...(result ? { result } : {}) + }) + return + } + if (type === 'device-resume-confirm') { + const request = DeviceResumeConfirmSchema.parse(payload(raw, type)) + const result = await this.store.confirmResume({ + ...identity, + ...request, + owningControlGeneration: session.generation + }) + await this.assignments.releaseActivity(identity, `confirmation:${request.basisConnId}`) + send(session.socket, 'device-resume-confirmed', result) + return + } + if (type === 'device-revoke') { + const request = DeviceRevokeSchema.parse(payload(raw, type)) + await this.store.revoke(identity, request.relayDeviceId) + send(session.socket, 'device-revoked', { reqId: request.reqId }) + return + } + this.sendControlError(session, undefined, 'unknown_control_message') + } catch (error) { + const reqId = (() => { + const candidate = payload(raw, type) + return typeof candidate === 'object' && candidate && 'reqId' in candidate + ? String(candidate.reqId) + : undefined + })() + this.sendControlError( + session, + reqId, + error instanceof Error ? error.message : 'control_operation_failed' + ) + } + } + + private sendControlError(session: HostSession, reqId: string | undefined, code: string): void { + if (session.socket) send(session.socket, 'control-error', { ...(reqId ? { reqId } : {}), code }) + } + + private rejectClient(socket: WebSocket, code: number): void { + send(socket, 'relay-hello', { ok: false, code }) + closeRelayWebSocket(socket, code, 'relay connection rejected') + } + + private releaseControlActivity(session: HostSession): void { + if (!session.controlActivityId) return + this.releaseActivityBestEffort( + { userId: session.identity.sub, relayHostId: session.relayHostId }, + session.controlActivityId + ) + } + + private releaseActivityBestEffort(identity: RelayIdentityKey, activityId: string): void { + void this.assignments.releaseActivity(identity, activityId).catch(() => { + // Why: expiry cleanup is the durable fallback; a transient SQL failure while + // closing a socket must not become an unhandled rejection that kills the cell. + console.warn('[orca-relay] activity release deferred to lease cleanup') + }) + } + + private failReservationBestEffort(reservation: CredentialReservation): void { + // Why: reservation deadlines and basis cleanup are durable recovery paths; + // transient SQL errors during socket callbacks must stay process-contained. + void this.store.failReservation(reservation).catch(() => { + console.warn('[orca-relay] reservation release deferred to credential cleanup') + }) + } + + private deactivateBasisBestEffort(connId: string): void { + void this.store.deactivateBasis(connId).catch(() => { + console.warn('[orca-relay] basis deactivation deferred to credential cleanup') + }) + } +} + +export type RelayIdentityKey = { userId: string; relayHostId: string } diff --git a/cloud/apps/relay/src/index.ts b/cloud/apps/relay/src/index.ts new file mode 100644 index 00000000000..635f3ae9b38 --- /dev/null +++ b/cloud/apps/relay/src/index.ts @@ -0,0 +1,131 @@ +import { + formatAssignmentInventorySnapshot, + readAssignmentInventorySnapshot +} from './assignment-inventory-snapshot.js' +import { RelayAssignmentStore } from './assignment-store.js' +import { loadRelayConfig } from './config.js' +import { startCellHeartbeat } from './cell-heartbeat-client.js' +import { + reconcileCellAdmissionAtStartup, + roleOwnsAssignmentMaintenance +} from './cell-admission-startup.js' +import { + consumeRelayDatabasePoolPressure, + openRelayDatabase, + readRelayDatabasePoolPressure +} from './database.js' +import { runAssignmentCleanup } from './assignment-cleanup-steps.js' +import { runRelayBackgroundOperation } from './relay-background-operation.js' +import { observedRelayRequests } from './relay-observability.js' +import { startRegionalRehomeWorker } from './regional-rehome-worker.js' +import { createRelayServer } from './relay-server.js' +import { + formatRegisteredMigrationInventory, + readRegisteredMigrationInventory +} from './registered-migration-inventory.js' + +const config = loadRelayConfig() +const database = await openRelayDatabase({ + databaseUrl: config.databaseUrl, + dataDir: config.dataDir, + poolMax: config.databasePoolMax, + applicationName: `orca-relay/${config.role}/${config.cellId}` +}) +await reconcileCellAdmissionAtStartup(config, new RelayAssignmentStore(database)) +const { + server, + sessions, + store, + assignments, + observability, + runtimeCounts, + connectionSnapshot, + ready, + cellIncarnation +} = createRelayServer(config, database) +const cleanupTimer = setInterval( + () => + void runRelayBackgroundOperation( + () => store.cleanup(), + '[orca-relay] credential cleanup failed' + ), + 30_000 +) +const assignmentCleanupTimer = roleOwnsAssignmentMaintenance(config.role) + ? setInterval(() => { + void runAssignmentCleanup(assignments) + }, 30_000) + : null +const inventorySnapshotTimer = roleOwnsAssignmentMaintenance(config.role) + ? setInterval(() => { + void runRelayBackgroundOperation(async () => { + const snapshot = await readAssignmentInventorySnapshot(database, Date.now()) + for (const line of formatAssignmentInventorySnapshot(snapshot)) console.warn(line) + }, '[orca-relay] inventory snapshot failed') + }, 60_000) + : null +const migrationInventoryTimer = roleOwnsAssignmentMaintenance(config.role) + ? setInterval(() => { + void runRelayBackgroundOperation(async () => { + const inventory = await readRegisteredMigrationInventory(database, Date.now()) + for (const line of formatRegisteredMigrationInventory(inventory)) console.warn(line) + }, '[orca-relay] migration inventory failed') + }, 5 * 60_000) + : null +cleanupTimer.unref() +assignmentCleanupTimer?.unref() +inventorySnapshotTimer?.unref() +migrationInventoryTimer?.unref() +observability.start(() => ({ + ...runtimeCounts(), + ...consumeRelayDatabasePoolPressure(database) +})) +const regionalRehomeWorker = startRegionalRehomeWorker(config, assignments, { + safetySnapshot: () => ({ + ...observability.regionalRehomeRuntimeSafety(), + ...readRelayDatabasePoolPressure(database) + }) +}) +const heartbeat = startCellHeartbeat(config, { + ready, + incarnation: cellIncarnation, + observedRequests: () => observedRelayRequests(runtimeCounts()), + connectionCounts: () => { + const snapshot = connectionSnapshot() + const counts = runtimeCounts() + return { + totalConnections: snapshot?.physicalConnections ?? counts.totalConnections, + inFlightConnections: + snapshot?.inFlightConnections ?? counts.inFlightConnections ?? 0, + reservedConnectionUnits: + snapshot?.reservedConnectionUnits ?? counts.reservedConnectionUnits ?? 0, + enforcedConnectionUnits: + snapshot?.enforcedConnectionUnits ?? + counts.enforcedConnectionUnits ?? + counts.totalConnections, + inclusionWatermark: snapshot?.inclusionWatermark + } + }, + regionalRehomeSafety: () => ({ + ...observability.regionalRehomeRuntimeSafety(), + ...readRelayDatabasePoolPressure(database) + }) +}) + +server.listen(config.port, () => { + console.log(`[orca-relay] listening on ${config.publicUrl} (port ${config.port})`) +}) + +const shutdown = (): void => { + clearInterval(cleanupTimer) + if (assignmentCleanupTimer) clearInterval(assignmentCleanupTimer) + if (inventorySnapshotTimer) clearInterval(inventorySnapshotTimer) + if (migrationInventoryTimer) clearInterval(migrationInventoryTimer) + observability.stop() + heartbeat?.stop() + regionalRehomeWorker?.stop() + sessions.drain(0) + server.close(() => void database.close()) +} +process.once('SIGTERM', shutdown) +process.once('SIGINT', shutdown) diff --git a/cloud/apps/relay/src/migration-recovery-postgres.test.ts b/cloud/apps/relay/src/migration-recovery-postgres.test.ts new file mode 100644 index 00000000000..aee41278d49 --- /dev/null +++ b/cloud/apps/relay/src/migration-recovery-postgres.test.ts @@ -0,0 +1,1376 @@ +import { ASSIGNMENT_LIMITS } from '@orca-cloud/relay-contract' +import { afterAll, beforeAll, describe, expect, it } from 'vitest' +import { + RelayAssignmentStore, + STRANDED_MIGRATION_ABANDON_MS, + type RelayAssignmentMigration +} from './assignment-store.js' +import { openRelayDatabase, type RelayDatabase } from './database.js' +import { readRegisteredMigrationInventory } from './registered-migration-inventory.js' + +const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL +const describePostgres = databaseUrl ? describe : describe.skip + +type RecoveryFixture = { + store: RelayAssignmentStore + identity: { userId: string; relayHostId: string } + migration: RelayAssignmentMigration + sourceControlId: string + targetControlId: string + cellIncarnation: string + cells: { + source: { id: string; url: string; capacityRequests: number } + failed: { id: string; url: string; capacityRequests: number } + replacement: { id: string; url: string; capacityRequests: number } + } + advancePastHeartbeat: () => void + advance: (milliseconds: number) => void + heartbeat: (cell: { id: string; url: string }) => Promise +} + +describePostgres('PostgreSQL migration recovery', () => { + let database: RelayDatabase + let sequence = 0 + + beforeAll(async () => { + database = await openRelayDatabase({ databaseUrl, dataDir: '' }) + }) + + afterAll(async () => { + await database.query( + `DELETE FROM relay_control_connection_reservations + WHERE user_id LIKE 'recovery-user-%'` + ) + await database.query( + `DELETE FROM relay_post_drain_migration_pins WHERE user_id LIKE 'recovery-user-%'` + ) + await database.query( + `DELETE FROM relay_assignment_migration_incarnations WHERE user_id LIKE 'recovery-user-%'` + ) + await database.query( + `DELETE FROM relay_assignment_activity_leases WHERE user_id LIKE 'recovery-user-%'` + ) + await database.query( + `DELETE FROM relay_assignment_migrations WHERE user_id LIKE 'recovery-user-%'` + ) + await database.query(`DELETE FROM relay_assignments WHERE user_id LIKE 'recovery-user-%'`) + await database.query( + `DELETE FROM relay_cell_fence_apply_invocations WHERE attempt_id IN ( + SELECT attempt_id FROM relay_cell_fence_attempts + WHERE cell_id LIKE 'recovery-cell-%' + )` + ) + await database.query( + `DELETE FROM relay_cell_committed_fences WHERE cell_id LIKE 'recovery-cell-%'` + ) + await database.query( + `DELETE FROM relay_cell_legacy_fence_adoptions + WHERE cell_id LIKE 'recovery-cell-%'` + ) + await database.query( + `DELETE FROM relay_cell_fence_plan_bindings WHERE attempt_id IN ( + SELECT attempt_id FROM relay_cell_fence_attempts + WHERE cell_id LIKE 'recovery-cell-%' + )` + ) + await database.query( + `DELETE FROM relay_cell_fence_attempts WHERE cell_id LIKE 'recovery-cell-%'` + ) + await database.query(`DELETE FROM relay_cell_fences WHERE cell_id LIKE 'recovery-cell-%'`) + await database.query(`DELETE FROM relay_cell_runtime WHERE cell_id LIKE 'recovery-cell-%'`) + await database.query(`DELETE FROM relay_cell_admission WHERE cell_id LIKE 'recovery-cell-%'`) + await database.query(`DELETE FROM relay_cells WHERE cell_id LIKE 'recovery-cell-%'`) + await database.close() + }) + + async function fixture(replacementCapacity = 20): Promise { + sequence++ + let now = 100 + const suffix = String(sequence) + const cellIncarnation = `11111111-1111-4111-8111-${suffix.padStart(12, '0')}` + const cells = { + source: { + id: `recovery-cell-${suffix}-source`, + url: `https://recovery-${suffix}-source.example.com`, + capacityRequests: 20 + }, + failed: { + id: `recovery-cell-${suffix}-failed`, + url: `https://recovery-${suffix}-failed.example.com`, + capacityRequests: 20 + }, + replacement: { + id: `recovery-cell-${suffix}-replacement`, + url: `https://recovery-${suffix}-replacement.example.com`, + capacityRequests: replacementCapacity + } + } + const store = new RelayAssignmentStore(database, () => now, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + await store.reconcileCells(Object.values(cells), false) + const heartbeat = async (cell: { id: string; url: string }): Promise => { + await store.recordCellHeartbeat({ + cellId: cell.id, + cellUrl: cell.url, + cellIncarnation, + startedAt: 50, + ready: true, + observedRequests: 0 + }) + } + for (const cell of Object.values(cells)) await heartbeat(cell) + const identity = { + userId: `recovery-user-${suffix}`, + relayHostId: `recoveryhost${suffix.padStart(4, '0')}` + } + const sourceControlId = `control:${cells.source.id}:1` + await database.query( + `INSERT INTO relay_assignments + (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, + last_activity_at, reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + cells.source.id, + 1, + 90_100, + now, + 1, + 0, + 0, + 0, + 0, + 0 + ] + ) + await database.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, + [ + identity.userId, + identity.relayHostId, + sourceControlId, + 'control', + cells.source.id, + 1, + 90_100, + now + ] + ) + await database.query( + `UPDATE relay_cells SET reserved_requests = 1 WHERE cell_id = ?`, + [cells.source.id] + ) + await store.setCellEnabled(cells.source.id, false) + const migration = await store.startEvacuation(identity, cells.failed.id) + const targetControlId = await store.activateControl(identity, { + cellId: cells.failed.id, + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await store.markMigrationTargetRegistered(identity, { + cellId: cells.failed.id, + assignmentEpoch: migration.assignmentEpoch + }) + return { + store, + identity, + migration, + sourceControlId, + targetControlId, + cellIncarnation, + cells, + advancePastHeartbeat: () => { + now += 45_001 + }, + advance: (milliseconds) => { + now += milliseconds + }, + heartbeat + } + } + + async function completeSourceFence(input: RecoveryFixture): Promise { + const attemptId = input.cellIncarnation + const invocationId = input.cellIncarnation + const requestReason = `relay-recovery-test/${attemptId}` + const evidence = { + attemptId, + environment: 'production' as const, + cellId: input.cells.source.id, + cellIncarnation: input.cellIncarnation, + migName: input.cells.source.id, + instanceGroup: `https://compute.example/instanceGroups/${input.cells.source.id}`, + generationIdentity: `https://compute.example/instanceTemplates/${input.cells.source.id}`, + fenceCommit: 'a'.repeat(40), + planSha256: 'b'.repeat(64), + planObjectName: `relay-fence-plans/${attemptId}.tfplan`, + planObjectGeneration: '1', + varFileSha256: 'c'.repeat(64), + terraformStateLineage: input.cellIncarnation, + terraformStateSerial: 1, + terraformStateObjectGeneration: '1', + terraformStateObjectSha256: 'd'.repeat(64), + requestReason + } + await input.store.prepareCellFenceAttempt(evidence) + await input.store.bindCellFencePlanGeneration(evidence, evidence.planObjectGeneration) + await input.store.startCellFenceApply( + evidence, + invocationId, + `${requestReason}/${invocationId}` + ) + await input.store.recordCellFenceOperation( + evidence, + invocationId, + `${requestReason}/${invocationId}`, + `operation-${input.cells.source.id}` + ) + await input.store.attestCellFenceAttempt( + evidence, + `operation-${input.cells.source.id}` + ) + } + + it('reaps an abandoned registered migration onto its healthy target', async () => { + const input = await fixture() + await input.store.releaseActivity(input.identity, input.sourceControlId) + await input.store.releaseActivity(input.identity, input.targetControlId) + input.advance( + ASSIGNMENT_LIMITS.migrationLeaseMs + STRANDED_MIGRATION_ABANDON_MS + 1 + ) + + await expect(input.store.abortExpiredEvacuations()).resolves.toBe(1) + expect( + await database.query( + `SELECT assignment.cell_id, assignment.assignment_epoch, + migration.completed_at, migration.aborted_at + FROM relay_assignments assignment + JOIN relay_assignment_migrations migration + ON migration.user_id = assignment.user_id + AND migration.relay_host_id = assignment.relay_host_id + AND migration.assignment_epoch = assignment.assignment_epoch + WHERE assignment.user_id = ? AND assignment.relay_host_id = ?`, + [input.identity.userId, input.identity.relayHostId] + ) + ).toEqual([ + { + cell_id: input.cells.failed.id, + assignment_epoch: '2', + completed_at: String( + 100 + ASSIGNMENT_LIMITS.migrationLeaseMs + STRANDED_MIGRATION_ABANDON_MS + 1 + ), + aborted_at: null + } + ]) + }) + + it('reaps immediately after the retired source has a durable completed fence', async () => { + const input = await fixture() + await input.store.releaseActivity(input.identity, input.sourceControlId) + await input.store.releaseActivity(input.identity, input.targetControlId) + input.advancePastHeartbeat() + await completeSourceFence(input) + input.advance(ASSIGNMENT_LIMITS.migrationLeaseMs + 1) + + await expect(input.store.abortExpiredEvacuations()).resolves.toBe(1) + expect( + await database.query( + `SELECT assignment.cell_id, assignment.assignment_epoch, + migration.completed_at, migration.aborted_at + FROM relay_assignments assignment + JOIN relay_assignment_migrations migration + ON migration.user_id = assignment.user_id + AND migration.relay_host_id = assignment.relay_host_id + AND migration.assignment_epoch = assignment.assignment_epoch + WHERE assignment.user_id = ? AND assignment.relay_host_id = ?`, + [input.identity.userId, input.identity.relayHostId] + ) + ).toEqual([ + { + cell_id: input.cells.failed.id, + assignment_epoch: '2', + completed_at: String(100 + 45_001 + ASSIGNMENT_LIMITS.migrationLeaseMs + 1), + aborted_at: null + } + ]) + }) + + it('reaps immediately after the retired source has an adopted legacy fence', async () => { + const input = await fixture() + await input.store.releaseActivity(input.identity, input.sourceControlId) + await input.store.releaseActivity(input.identity, input.targetControlId) + input.advancePastHeartbeat() + await input.store.adoptLegacyCellFence( + input.cells.source.id, + input.cellIncarnation + ) + await input.store.commitLegacyCellFenceAdoption( + input.cells.source.id, + input.cellIncarnation + ) + input.advance(ASSIGNMENT_LIMITS.migrationLeaseMs + 1) + + expect( + ( + await readRegisteredMigrationInventory( + database, + 100 + 45_001 + ASSIGNMENT_LIMITS.migrationLeaseMs + 1 + ) + ).abandoned + ).toBe(1) + + await expect(input.store.abortExpiredEvacuations()).resolves.toBe(1) + expect( + await database.query( + `SELECT assignment.cell_id, assignment.assignment_epoch, + migration.completed_at, migration.aborted_at + FROM relay_assignments assignment + JOIN relay_assignment_migrations migration + ON migration.user_id = assignment.user_id + AND migration.relay_host_id = assignment.relay_host_id + AND migration.assignment_epoch = assignment.assignment_epoch + WHERE assignment.user_id = ? AND assignment.relay_host_id = ?`, + [input.identity.userId, input.identity.relayHostId] + ) + ).toEqual([ + { + cell_id: input.cells.failed.id, + assignment_epoch: '2', + completed_at: String(100 + 45_001 + ASSIGNMENT_LIMITS.migrationLeaseMs + 1), + aborted_at: null + } + ]) + }) + + it('does not reap after temporary legacy adoption without durable commit', async () => { + const input = await fixture() + await input.store.releaseActivity(input.identity, input.sourceControlId) + await input.store.releaseActivity(input.identity, input.targetControlId) + input.advancePastHeartbeat() + await input.store.adoptLegacyCellFence( + input.cells.source.id, + input.cellIncarnation + ) + input.advance(ASSIGNMENT_LIMITS.migrationLeaseMs + 1) + + expect( + ( + await readRegisteredMigrationInventory( + database, + 100 + 45_001 + ASSIGNMENT_LIMITS.migrationLeaseMs + 1 + ) + ).abandoned + ).toBe(0) + await expect(input.store.abortExpiredEvacuations()).resolves.toBe(0) + input.advance(STRANDED_MIGRATION_ABANDON_MS) + await expect(input.store.abortExpiredEvacuations()).resolves.toBe(1) + }) + + it('invalidates an adopted legacy fence when the source heartbeats', async () => { + const input = await fixture() + await input.store.releaseActivity(input.identity, input.sourceControlId) + await input.store.releaseActivity(input.identity, input.targetControlId) + input.advancePastHeartbeat() + await input.store.adoptLegacyCellFence( + input.cells.source.id, + input.cellIncarnation + ) + await input.store.commitLegacyCellFenceAdoption( + input.cells.source.id, + input.cellIncarnation + ) + input.advance(ASSIGNMENT_LIMITS.migrationLeaseMs + 1) + await input.heartbeat(input.cells.source) + + expect( + await database.query( + `SELECT cell_id FROM relay_cell_legacy_fence_adoptions WHERE cell_id = ?`, + [input.cells.source.id] + ) + ).toEqual([]) + await expect(input.store.abortExpiredEvacuations()).resolves.toBe(0) + input.advance(STRANDED_MIGRATION_ABANDON_MS) + await expect(input.store.abortExpiredEvacuations()).resolves.toBe(1) + }) + + it('invalidates an adopted legacy fence when the source restarts', async () => { + const input = await fixture() + await input.store.releaseActivity(input.identity, input.sourceControlId) + await input.store.releaseActivity(input.identity, input.targetControlId) + input.advancePastHeartbeat() + await input.store.adoptLegacyCellFence( + input.cells.source.id, + input.cellIncarnation + ) + await input.store.commitLegacyCellFenceAdoption( + input.cells.source.id, + input.cellIncarnation + ) + input.advance(ASSIGNMENT_LIMITS.migrationLeaseMs + 1) + await input.store.recordCellHeartbeat({ + cellId: input.cells.source.id, + cellUrl: input.cells.source.url, + cellIncarnation: '99999999-9999-4999-8999-999999999999', + startedAt: 51, + ready: true, + observedRequests: 0 + }) + + expect( + await database.query( + `SELECT cell_id FROM relay_cell_legacy_fence_adoptions WHERE cell_id = ?`, + [input.cells.source.id] + ) + ).toEqual([]) + await expect(input.store.abortExpiredEvacuations()).resolves.toBe(0) + input.advance(STRANDED_MIGRATION_ABANDON_MS) + await expect(input.store.abortExpiredEvacuations()).resolves.toBe(1) + }) + + it('serializes legacy adoption with a returning source heartbeat', async () => { + const input = await fixture() + await input.store.releaseActivity(input.identity, input.sourceControlId) + await input.store.releaseActivity(input.identity, input.targetControlId) + input.advancePastHeartbeat() + await input.store.adoptLegacyCellFence( + input.cells.source.id, + input.cellIncarnation + ) + + const [commit, heartbeatResult] = await Promise.allSettled([ + input.store.commitLegacyCellFenceAdoption( + input.cells.source.id, + input.cellIncarnation + ), + input.heartbeat(input.cells.source) + ]) + expect(heartbeatResult.status).toBe('fulfilled') + expect(['fulfilled', 'rejected']).toContain(commit.status) + expect( + await database.query( + `SELECT cell_id FROM relay_cell_legacy_fence_adoptions WHERE cell_id = ?`, + [input.cells.source.id] + ) + ).toEqual([]) + input.advance( + ASSIGNMENT_LIMITS.migrationLeaseMs + STRANDED_MIGRATION_ABANDON_MS + 1 + ) + await expect(input.store.abortExpiredEvacuations()).resolves.toBe(1) + }) + + it('restores the 24-hour guard when the fenced source heartbeats again', async () => { + const input = await fixture() + await input.store.releaseActivity(input.identity, input.sourceControlId) + await input.store.releaseActivity(input.identity, input.targetControlId) + input.advancePastHeartbeat() + await completeSourceFence(input) + input.advance(ASSIGNMENT_LIMITS.migrationLeaseMs + 1) + await input.heartbeat(input.cells.source) + + await expect(input.store.abortExpiredEvacuations()).resolves.toBe(0) + expect( + await database.query( + `SELECT completed_at, aborted_at FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ?`, + [input.identity.userId, input.identity.relayHostId] + ) + ).toEqual([{ completed_at: null, aborted_at: null }]) + input.advance(STRANDED_MIGRATION_ABANDON_MS) + await expect(input.store.abortExpiredEvacuations()).resolves.toBe(1) + }) + + it('restores the 24-hour guard when the fenced source restarts', async () => { + const input = await fixture() + await input.store.releaseActivity(input.identity, input.sourceControlId) + await input.store.releaseActivity(input.identity, input.targetControlId) + input.advancePastHeartbeat() + await completeSourceFence(input) + input.advance(ASSIGNMENT_LIMITS.migrationLeaseMs + 1) + await input.store.recordCellHeartbeat({ + cellId: input.cells.source.id, + cellUrl: input.cells.source.url, + cellIncarnation: '99999999-9999-4999-8999-999999999999', + startedAt: 51, + ready: true, + observedRequests: 0 + }) + + await expect(input.store.abortExpiredEvacuations()).resolves.toBe(0) + expect( + await database.query( + `SELECT completed_at, aborted_at FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ?`, + [input.identity.userId, input.identity.relayHostId] + ) + ).toEqual([{ completed_at: null, aborted_at: null }]) + input.advance(STRANDED_MIGRATION_ABANDON_MS) + await expect(input.store.abortExpiredEvacuations()).resolves.toBe(1) + }) + + it('serializes fenced cleanup with a returning source heartbeat', async () => { + const input = await fixture() + await input.store.releaseActivity(input.identity, input.sourceControlId) + await input.store.releaseActivity(input.identity, input.targetControlId) + input.advancePastHeartbeat() + await completeSourceFence(input) + input.advance(ASSIGNMENT_LIMITS.migrationLeaseMs + 1) + + const [cleanup, heartbeatResult] = await Promise.allSettled([ + input.store.abortExpiredEvacuations(), + input.heartbeat(input.cells.source) + ]) + expect(cleanup.status).toBe('fulfilled') + expect(heartbeatResult.status).toBe('fulfilled') + if (cleanup.status !== 'fulfilled') throw cleanup.reason + if (heartbeatResult.status !== 'fulfilled') throw heartbeatResult.reason + expect([0, 1]).toContain(cleanup.value) + const migrations = await database.query( + `SELECT completed_at, aborted_at FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ?`, + [input.identity.userId, input.identity.relayHostId] + ) + if (cleanup.value === 1) { + expect(migrations[0]?.completed_at).not.toBeNull() + } else { + expect(migrations).toEqual([{ completed_at: null, aborted_at: null }]) + input.advance(STRANDED_MIGRATION_ABANDON_MS) + await expect(input.store.abortExpiredEvacuations()).resolves.toBe(1) + } + }) + + it('keeps a freshly retired source protected without a durable fence', async () => { + const input = await fixture() + await input.store.releaseActivity(input.identity, input.sourceControlId) + await input.store.releaseActivity(input.identity, input.targetControlId) + input.advance(ASSIGNMENT_LIMITS.migrationLeaseMs + 1) + + await expect(input.store.abortExpiredEvacuations()).resolves.toBe(0) + expect( + await database.query( + `SELECT completed_at, aborted_at FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ?`, + [input.identity.userId, input.identity.relayHostId] + ) + ).toEqual([{ completed_at: null, aborted_at: null }]) + input.advance(STRANDED_MIGRATION_ABANDON_MS) + await expect(input.store.abortExpiredEvacuations()).resolves.toBe(1) + }) + + it('rolls an abandoned disabled target back to its source', async () => { + const input = await fixture() + await input.store.releaseActivity(input.identity, input.sourceControlId) + await input.store.releaseActivity(input.identity, input.targetControlId) + await input.store.setCellEnabled(input.cells.failed.id, false) + input.advance( + ASSIGNMENT_LIMITS.migrationLeaseMs + STRANDED_MIGRATION_ABANDON_MS + 1 + ) + + await expect(input.store.abortExpiredEvacuations()).resolves.toBe(1) + expect( + await database.query( + `SELECT assignment.cell_id, assignment.assignment_epoch, + migration.completed_at, migration.aborted_at + FROM relay_assignments assignment + JOIN relay_assignment_migrations migration + ON migration.user_id = assignment.user_id + AND migration.relay_host_id = assignment.relay_host_id + WHERE assignment.user_id = ? AND assignment.relay_host_id = ? + ORDER BY migration.assignment_epoch`, + [input.identity.userId, input.identity.relayHostId] + ) + ).toEqual([ + { + cell_id: input.cells.source.id, + assignment_epoch: '3', + completed_at: null, + aborted_at: String( + 100 + ASSIGNMENT_LIMITS.migrationLeaseMs + STRANDED_MIGRATION_ABANDON_MS + 1 + ) + } + ]) + }) + + it('serializes cleanup against supersession without reporting a false target', async () => { + const input = await fixture() + await input.store.releaseActivity(input.identity, input.sourceControlId) + await input.store.releaseActivity(input.identity, input.targetControlId) + await input.store.setCellEnabled(input.cells.failed.id, false) + input.advance( + ASSIGNMENT_LIMITS.migrationLeaseMs + STRANDED_MIGRATION_ABANDON_MS + 1 + ) + await input.heartbeat(input.cells.source) + await input.heartbeat(input.cells.replacement) + await input.store.attestCellFence(input.cells.failed.id, input.cellIncarnation) + + const [cleanup, supersession] = await Promise.allSettled([ + input.store.abortExpiredEvacuations(), + input.store.supersedeRegisteredCellEvacuations( + input.cells.source.id, + input.cells.failed.id, + input.cells.replacement.id, + 100 + ) + ]) + expect(cleanup.status).toBe('fulfilled') + if (cleanup.status !== 'fulfilled') throw cleanup.reason + if (supersession.status === 'fulfilled') { + expect([0, 1]).toContain(supersession.value) + const assignment = await database.query( + `SELECT cell_id FROM relay_assignments WHERE user_id = ?`, + [input.identity.userId] + ) + expect(assignment).toEqual([{ + cell_id: + supersession.value === 1 + ? input.cells.replacement.id + : input.cells.source.id + }]) + if (supersession.value === 0) expect(cleanup.value).toBeGreaterThanOrEqual(1) + } else { + expect(cleanup.value).toBeGreaterThanOrEqual(1) + expect(supersession.reason).toMatchObject({ message: 'migration_already_superseded' }) + expect( + await database.query( + `SELECT cell_id FROM relay_assignments WHERE user_id = ?`, + [input.identity.userId] + ) + ).toEqual([{ cell_id: input.cells.source.id }]) + } + }) + + it('completes a registered migration from a dead source idempotently', async () => { + const input = await fixture() + await input.store.releaseActivity(input.identity, input.sourceControlId) + input.advancePastHeartbeat() + await input.heartbeat(input.cells.failed) + await input.store.attestCellFence( + input.cells.source.id, + input.cellIncarnation + ) + const completion = { + assignmentEpoch: input.migration.assignmentEpoch, + sourceCellId: input.cells.source.id, + targetCellId: input.cells.failed.id + } + + await expect( + input.store.cellEvacuationStatus( + input.cells.source.id, + input.cells.failed.id, + true + ) + ).resolves.toMatchObject({ inProgress: 0, completed: 1 }) + await expect( + input.store.completeEvacuationFromDeadSource(input.identity, completion) + ).resolves.toMatchObject({ changed: false }) + expect(await reservations(input.cells)).toEqual({ source: 0, failed: 1, replacement: 0 }) + }) + + it('fails dead-source completion without a stale source heartbeat', async () => { + const input = await fixture() + await input.store.releaseActivity(input.identity, input.sourceControlId) + + await expect( + input.store.completeEvacuationFromDeadSource(input.identity, { + assignmentEpoch: input.migration.assignmentEpoch, + sourceCellId: input.cells.source.id, + targetCellId: input.cells.failed.id + }) + ).rejects.toThrow('cell_fence_attestation_missing') + expect(await reservations(input.cells)).toEqual({ source: 0, failed: 2, replacement: 0 }) + }) + + it('supersedes a registered migration with exact accounting and idempotency', async () => { + const input = await fixture() + await database.query( + `INSERT INTO relay_control_connection_reservations + (reservation_id, idempotency_key, user_id, relay_host_id, + assignment_epoch, cell_id, state, inclusion_watermark, + claim_activity_id, created_at, timeout_at, claimed_at, released_at, + updated_at) + VALUES (?, ?, ?, ?, ?, ?, 'late-arrival-debt', NULL, NULL, 100, 100, NULL, NULL, 100)`, + [ + `superseded-${input.identity.userId}`, + `superseded-${input.identity.userId}`, + input.identity.userId, + input.identity.relayHostId, + input.migration.assignmentEpoch, + input.cells.failed.id + ] + ) + await input.store.setCellEnabled(input.cells.failed.id, false) + input.advancePastHeartbeat() + await input.heartbeat(input.cells.source) + await input.heartbeat(input.cells.replacement) + await input.store.attestCellFence( + input.cells.failed.id, + input.cellIncarnation + ) + const supersession = { + assignmentEpoch: input.migration.assignmentEpoch, + sourceCellId: input.cells.source.id, + currentTargetCellId: input.cells.failed.id, + replacementTargetCellId: input.cells.replacement.id + } + + const first = await input.store.supersedeRegisteredEvacuation( + input.identity, + supersession + ) + const retry = await input.store.supersedeRegisteredEvacuation( + input.identity, + supersession + ) + expect(retry).toEqual(first) + expect(first).toMatchObject({ previousEpoch: 2, assignmentEpoch: 3 }) + expect( + await database.query( + `SELECT state FROM relay_control_connection_reservations + WHERE reservation_id = ?`, + [`superseded-${input.identity.userId}`] + ) + ).toEqual([{ state: 'released' }]) + expect(await reservations(input.cells)).toEqual({ source: 1, failed: 0, replacement: 2 }) + expect( + await database.query( + `SELECT assignment_epoch, completed_at, aborted_at + FROM relay_assignment_migrations WHERE user_id = ? + ORDER BY assignment_epoch`, + [input.identity.userId] + ) + ).toEqual([ + { assignment_epoch: '2', completed_at: null, aborted_at: '45101' }, + { assignment_epoch: '3', completed_at: null, aborted_at: null } + ]) + }) + + it('reconciles durable cell accounting before aggregate supersession', async () => { + const input = await fixture() + await input.store.setCellEnabled(input.cells.failed.id, false) + input.advancePastHeartbeat() + await input.heartbeat(input.cells.source) + await input.heartbeat(input.cells.replacement) + await input.store.attestCellFence( + input.cells.failed.id, + input.cellIncarnation + ) + await database.query( + `UPDATE relay_cells + SET reserved_requests = CASE WHEN cell_id = ? THEN 1 ELSE 0 END + WHERE cell_id IN (?, ?)`, + [input.cells.replacement.id, input.cells.source.id, input.cells.replacement.id] + ) + + await expect( + input.store.supersedeRegisteredCellEvacuations( + input.cells.source.id, + input.cells.failed.id, + input.cells.replacement.id, + 100 + ) + ).resolves.toBe(1) + expect(await reservations(input.cells)).toEqual({ + source: 1, + failed: 0, + replacement: 2 + }) + }) + + it('rebuilds an expired registered migration lease before supersession', async () => { + const input = await fixture() + await pinMigration(input, input.migration.assignmentEpoch) + await clearActivities(input) + await database.query( + `UPDATE relay_assignment_migrations SET expires_at = 0 + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [input.identity.userId, input.identity.relayHostId, input.migration.assignmentEpoch] + ) + await fenceFailedCell(input) + + await expect(supersedeAggregate(input)).resolves.toBe(1) + expect(await reservations(input.cells)).toEqual({ + source: 0, + failed: 0, + replacement: 1 + }) + expect( + await database.query( + `SELECT assignment_epoch, source_request_units, target_reserved_units, aborted_at + FROM relay_assignment_migrations WHERE user_id = ? ORDER BY assignment_epoch`, + [input.identity.userId] + ) + ).toEqual([ + { + assignment_epoch: '2', + source_request_units: '1', + target_reserved_units: '2', + aborted_at: '45101' + }, + { + assignment_epoch: '3', + source_request_units: '0', + target_reserved_units: '1', + aborted_at: null + } + ]) + }) + + it('retires an obsolete row without discarding replacement activity', async () => { + const input = await fixture() + await pinMigration(input, input.migration.assignmentEpoch) + await fenceFailedCell(input) + await input.store.supersedeRegisteredEvacuation(input.identity, { + assignmentEpoch: input.migration.assignmentEpoch, + sourceCellId: input.cells.source.id, + currentTargetCellId: input.cells.failed.id, + replacementTargetCellId: input.cells.replacement.id + }) + await database.query( + `UPDATE relay_assignment_migrations SET aborted_at = NULL + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [input.identity.userId, input.identity.relayHostId, input.migration.assignmentEpoch] + ) + const activityBefore = await assignmentActivities(input) + + await expect(supersedeAggregate(input)).resolves.toBe(1) + expect(await assignmentActivities(input)).toEqual(activityBefore) + expect(await reservations(input.cells)).toEqual({ + source: 1, + failed: 0, + replacement: 2 + }) + expect( + await database.query( + `SELECT assignment_epoch, aborted_at FROM relay_assignment_migrations + WHERE user_id = ? ORDER BY assignment_epoch`, + [input.identity.userId] + ) + ).toEqual([ + { assignment_epoch: '2', aborted_at: '45101' }, + { assignment_epoch: '3', aborted_at: null } + ]) + }) + + it('anchors a newer dormant failed-cell epoch and preserves drain evidence', async () => { + const input = await fixture() + const drainAttemptId = await pinMigration(input, input.migration.assignmentEpoch) + await clearActivities(input) + await database.query( + `UPDATE relay_assignments SET assignment_epoch = 4 + WHERE user_id = ? AND relay_host_id = ?`, + [input.identity.userId, input.identity.relayHostId] + ) + await database.query( + `UPDATE relay_assignment_migrations SET expires_at = 0 + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [input.identity.userId, input.identity.relayHostId, input.migration.assignmentEpoch] + ) + await fenceFailedCell(input) + + await expect(supersedeAggregate(input)).resolves.toBe(1) + expect(await reservations(input.cells)).toEqual({ + source: 0, + failed: 0, + replacement: 1 + }) + expect( + await database.query( + `SELECT assignment_epoch, source_request_units, target_reserved_units, aborted_at + FROM relay_assignment_migrations WHERE user_id = ? ORDER BY assignment_epoch`, + [input.identity.userId] + ) + ).toEqual([ + { + assignment_epoch: '2', + source_request_units: '1', + target_reserved_units: '2', + aborted_at: '45101' + }, + { + assignment_epoch: '4', + source_request_units: '0', + target_reserved_units: '1', + aborted_at: '45101' + }, + { + assignment_epoch: '5', + source_request_units: '0', + target_reserved_units: '1', + aborted_at: null + } + ]) + expect( + await database.query( + `SELECT assignment_epoch, drain_attempt_id, source_request_units, + target_reserved_units + FROM relay_post_drain_migration_pins + WHERE user_id = ? ORDER BY assignment_epoch`, + [input.identity.userId] + ) + ).toEqual([ + { + assignment_epoch: '2', + drain_attempt_id: drainAttemptId, + source_request_units: '1', + target_reserved_units: '2' + }, + { + assignment_epoch: '4', + drain_attempt_id: drainAttemptId, + source_request_units: '0', + target_reserved_units: '1' + }, + { + assignment_epoch: '5', + drain_attempt_id: drainAttemptId, + source_request_units: '0', + target_reserved_units: '1' + } + ]) + }) + + it('uses an existing current-epoch migration once when retiring an older row', async () => { + const input = await fixture() + await pinMigration(input, input.migration.assignmentEpoch) + await clearActivities(input) + await database.query( + `UPDATE relay_assignments SET assignment_epoch = 4 + WHERE user_id = ? AND relay_host_id = ?`, + [input.identity.userId, input.identity.relayHostId] + ) + await database.query( + `INSERT INTO relay_assignment_migrations + (user_id, relay_host_id, source_cell_id, target_cell_id, + previous_epoch, assignment_epoch, source_request_units, + target_reserved_units, expires_at, target_registered_at, + completed_at, aborted_at, created_at, updated_at) + VALUES (?, ?, ?, ?, 2, 4, 0, 1, 0, 100, NULL, NULL, 100, 100)`, + [ + input.identity.userId, + input.identity.relayHostId, + input.cells.source.id, + input.cells.failed.id + ] + ) + await database.query( + `INSERT INTO relay_assignment_migration_incarnations + (user_id, relay_host_id, assignment_epoch, source_cell_incarnation, + target_cell_incarnation) + VALUES (?, ?, 4, ?, ?)`, + [ + input.identity.userId, + input.identity.relayHostId, + input.cellIncarnation, + input.cellIncarnation + ] + ) + await pinMigration(input, 4) + await fenceFailedCell(input) + + await expect(supersedeAggregate(input)).resolves.toBe(1) + expect( + await database.query( + `SELECT assignment_epoch, aborted_at FROM relay_assignment_migrations + WHERE user_id = ? ORDER BY assignment_epoch`, + [input.identity.userId] + ) + ).toEqual([ + { assignment_epoch: '2', aborted_at: '45101' }, + { assignment_epoch: '4', aborted_at: '45101' }, + { assignment_epoch: '5', aborted_at: null } + ]) + expect(await reservations(input.cells)).toEqual({ + source: 0, + failed: 0, + replacement: 1 + }) + }) + + it('rejects a newer failed-cell epoch with ambiguous activity', async () => { + const input = await fixture() + await database.query( + `UPDATE relay_assignments SET assignment_epoch = 4 + WHERE user_id = ? AND relay_host_id = ?`, + [input.identity.userId, input.identity.relayHostId] + ) + await fenceFailedCell(input) + + await expect(supersedeAggregate(input)).rejects.toThrow( + 'migration_activity_topology_mismatch' + ) + expect( + await database.query( + `SELECT assignment_epoch, aborted_at FROM relay_assignment_migrations + WHERE user_id = ?`, + [input.identity.userId] + ) + ).toEqual([{ assignment_epoch: '2', aborted_at: null }]) + expect(await reservations(input.cells)).toEqual({ + source: 1, + failed: 2, + replacement: 0 + }) + }) + + it('leaves lease repair retryable when supersession later fails', async () => { + const input = await fixture(1) + await database.query( + `DELETE FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? AND activity_kind = 'migration'`, + [input.identity.userId, input.identity.relayHostId] + ) + await database.query( + `UPDATE relay_assignments SET migration_leases = 0 + WHERE user_id = ? AND relay_host_id = ?`, + [input.identity.userId, input.identity.relayHostId] + ) + await database.query( + `UPDATE relay_cells SET reserved_requests = reserved_requests - 1 + WHERE cell_id = ?`, + [input.cells.failed.id] + ) + await database.query( + `UPDATE relay_assignment_migrations SET expires_at = 0 + WHERE user_id = ? AND relay_host_id = ? AND assignment_epoch = ?`, + [input.identity.userId, input.identity.relayHostId, input.migration.assignmentEpoch] + ) + await fenceFailedCell(input) + + await expect(supersedeAggregate(input)).rejects.toThrow( + 'relay_capacity_exhausted' + ) + expect( + await database.query( + `SELECT activity_id FROM relay_assignment_activity_leases + WHERE user_id = ? AND activity_kind = 'migration'`, + [input.identity.userId] + ) + ).toEqual([{ activity_id: 'migration:2' }]) + expect( + await database.query( + `SELECT aborted_at FROM relay_assignment_migrations + WHERE user_id = ? AND assignment_epoch = 2`, + [input.identity.userId] + ) + ).toEqual([{ aborted_at: null }]) + + await database.query( + `UPDATE relay_cells SET capacity_requests = 2 WHERE cell_id = ?`, + [input.cells.replacement.id] + ) + await expect(supersedeAggregate(input)).resolves.toBe(1) + expect(await reservations(input.cells)).toEqual({ + source: 1, + failed: 0, + replacement: 2 + }) + }) + + it('serializes concurrent supersession retries without duplicating capacity', async () => { + const input = await fixture() + await input.store.setCellEnabled(input.cells.failed.id, false) + input.advancePastHeartbeat() + await input.heartbeat(input.cells.source) + await input.heartbeat(input.cells.replacement) + await input.store.attestCellFence( + input.cells.failed.id, + input.cellIncarnation + ) + const supersession = { + assignmentEpoch: input.migration.assignmentEpoch, + sourceCellId: input.cells.source.id, + currentTargetCellId: input.cells.failed.id, + replacementTargetCellId: input.cells.replacement.id + } + + const results = await Promise.all([ + input.store.supersedeRegisteredEvacuation(input.identity, supersession), + input.store.supersedeRegisteredEvacuation(input.identity, supersession) + ]) + expect(results[0]).toEqual(results[1]) + expect(await reservations(input.cells)).toEqual({ source: 1, failed: 0, replacement: 2 }) + expect( + await database.query( + `SELECT COUNT(*) AS count FROM relay_assignment_migrations WHERE user_id = ?`, + [input.identity.userId] + ) + ).toEqual([{ count: '2' }]) + }) + + it('fails a competing aggregate supersession instead of reporting the wrong target', async () => { + const input = await fixture() + const alternate = { + id: `recovery-cell-${sequence}-alternate`, + url: `https://recovery-${sequence}-alternate.example.com`, + capacityRequests: 20 + } + await input.store.reconcileCells([...Object.values(input.cells), alternate], false) + await input.heartbeat(alternate) + await input.store.setCellEnabled(input.cells.failed.id, false) + input.advancePastHeartbeat() + await input.heartbeat(input.cells.source) + await input.heartbeat(input.cells.replacement) + await input.heartbeat(alternate) + await input.store.attestCellFence(input.cells.failed.id, input.cellIncarnation) + + const results = await Promise.allSettled([ + input.store.supersedeRegisteredCellEvacuations( + input.cells.source.id, + input.cells.failed.id, + input.cells.replacement.id, + 100 + ), + input.store.supersedeRegisteredCellEvacuations( + input.cells.source.id, + input.cells.failed.id, + alternate.id, + 100 + ) + ]) + const fulfilled = results.filter( + (result): result is PromiseFulfilledResult => result.status === 'fulfilled' + ) + const rejected = results.filter( + (result): result is PromiseRejectedResult => result.status === 'rejected' + ) + expect(fulfilled).toHaveLength(1) + expect(fulfilled[0]!.value).toBe(1) + expect(rejected).toHaveLength(1) + expect(rejected[0]!.reason).toMatchObject({ message: 'migration_already_superseded' }) + const successor = await database.query( + `SELECT assignment.cell_id, migration.target_cell_id + FROM relay_assignments assignment + JOIN relay_assignment_migrations migration + ON migration.user_id = assignment.user_id + AND migration.relay_host_id = assignment.relay_host_id + AND migration.assignment_epoch = assignment.assignment_epoch + WHERE assignment.user_id = ? AND migration.previous_epoch = ?`, + [input.identity.userId, input.migration.assignmentEpoch] + ) + expect(successor).toHaveLength(1) + expect(successor[0]!.cell_id).toBe(successor[0]!.target_cell_id) + expect([input.cells.replacement.id, alternate.id]).toContain( + successor[0]!.target_cell_id + ) + }) + + it('rolls back every supersession change when replacement capacity is insufficient', async () => { + const input = await fixture(1) + await input.store.setCellEnabled(input.cells.failed.id, false) + input.advancePastHeartbeat() + await input.heartbeat(input.cells.source) + await input.heartbeat(input.cells.replacement) + await input.store.attestCellFence( + input.cells.failed.id, + input.cellIncarnation + ) + + await expect( + input.store.supersedeRegisteredEvacuation(input.identity, { + assignmentEpoch: input.migration.assignmentEpoch, + sourceCellId: input.cells.source.id, + currentTargetCellId: input.cells.failed.id, + replacementTargetCellId: input.cells.replacement.id + }) + ).rejects.toThrow('relay_capacity_exhausted') + expect(await reservations(input.cells)).toEqual({ source: 1, failed: 2, replacement: 0 }) + expect( + await database.query( + `SELECT cell_id, assignment_epoch FROM relay_assignments WHERE user_id = ?`, + [input.identity.userId] + ) + ).toEqual([{ cell_id: input.cells.failed.id, assignment_epoch: '2' }]) + expect( + await database.query( + `SELECT assignment_epoch, aborted_at FROM relay_assignment_migrations + WHERE user_id = ?`, + [input.identity.userId] + ) + ).toEqual([{ assignment_epoch: '2', aborted_at: null }]) + }) + + it('rolls back supersession when migration request-unit shape drifted', async () => { + const input = await fixture() + await input.store.setCellEnabled(input.cells.failed.id, false) + input.advancePastHeartbeat() + await input.heartbeat(input.cells.source) + await input.heartbeat(input.cells.replacement) + await input.store.attestCellFence(input.cells.failed.id, input.cellIncarnation) + await database.query( + `UPDATE relay_assignment_activity_leases + SET request_units = request_units + 1 + WHERE user_id = ? AND activity_kind = 'migration'`, + [input.identity.userId] + ) + await database.query( + `UPDATE relay_cells SET reserved_requests = reserved_requests + 1 WHERE cell_id = ?`, + [input.cells.failed.id] + ) + + await expect( + input.store.supersedeRegisteredEvacuation(input.identity, { + assignmentEpoch: input.migration.assignmentEpoch, + sourceCellId: input.cells.source.id, + currentTargetCellId: input.cells.failed.id, + replacementTargetCellId: input.cells.replacement.id + }) + ).rejects.toThrow('migration_activity_lease_shape_mismatch') + expect(await reservations(input.cells)).toEqual({ source: 1, failed: 3, replacement: 0 }) + expect( + await database.query( + `SELECT cell_id, assignment_epoch FROM relay_assignments WHERE user_id = ?`, + [input.identity.userId] + ) + ).toEqual([{ cell_id: input.cells.failed.id, assignment_epoch: '2' }]) + }) + + it('rolls back supersession when locked cell reservation accounting drifted', async () => { + const input = await fixture() + await input.store.setCellEnabled(input.cells.failed.id, false) + input.advancePastHeartbeat() + await input.heartbeat(input.cells.source) + await input.heartbeat(input.cells.replacement) + await input.store.attestCellFence(input.cells.failed.id, input.cellIncarnation) + await database.query( + `UPDATE relay_cells SET reserved_requests = 1 WHERE cell_id = ?`, + [input.cells.replacement.id] + ) + + await expect( + input.store.supersedeRegisteredEvacuation(input.identity, { + assignmentEpoch: input.migration.assignmentEpoch, + sourceCellId: input.cells.source.id, + currentTargetCellId: input.cells.failed.id, + replacementTargetCellId: input.cells.replacement.id + }) + ).rejects.toThrow('migration_cell_reservation_accounting_mismatch') + expect(await reservations(input.cells)).toEqual({ source: 1, failed: 2, replacement: 1 }) + expect( + await database.query( + `SELECT assignment_epoch, aborted_at FROM relay_assignment_migrations + WHERE user_id = ?`, + [input.identity.userId] + ) + ).toEqual([{ assignment_epoch: '2', aborted_at: null }]) + }) + + it('rejects supersession before incrementing the maximum safe epoch', async () => { + const input = await fixture() + await expect( + input.store.supersedeRegisteredEvacuation(input.identity, { + assignmentEpoch: Number.MAX_SAFE_INTEGER, + sourceCellId: input.cells.source.id, + currentTargetCellId: input.cells.failed.id, + replacementTargetCellId: input.cells.replacement.id + }) + ).rejects.toThrow('assignment_epoch_exhausted') + }) + + async function fenceFailedCell(input: RecoveryFixture): Promise { + await input.store.setCellEnabled(input.cells.failed.id, false) + input.advancePastHeartbeat() + await input.heartbeat(input.cells.source) + await input.heartbeat(input.cells.replacement) + await input.store.attestCellFence( + input.cells.failed.id, + input.cellIncarnation + ) + } + + async function supersedeAggregate(input: RecoveryFixture): Promise { + return await input.store.supersedeRegisteredCellEvacuations( + input.cells.source.id, + input.cells.failed.id, + input.cells.replacement.id, + 100 + ) + } + + async function clearActivities(input: RecoveryFixture): Promise { + await database.query( + `DELETE FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ?`, + [input.identity.userId, input.identity.relayHostId] + ) + await database.query( + `UPDATE relay_assignments SET reserved_controls = 0, reserved_splices = 0, + reserved_invites = 0, pending_installs = 0, pending_confirmations = 0, + migration_leases = 0 + WHERE user_id = ? AND relay_host_id = ?`, + [input.identity.userId, input.identity.relayHostId] + ) + await database.query( + `UPDATE relay_cells SET reserved_requests = 0 + WHERE cell_id IN (?, ?, ?)`, + [input.cells.source.id, input.cells.failed.id, input.cells.replacement.id] + ) + } + + async function pinMigration( + input: RecoveryFixture, + assignmentEpoch: number + ): Promise { + const drainAttemptId = `${input.identity.userId}-drain-${assignmentEpoch}` + await database.query( + `INSERT INTO relay_post_drain_migration_pins + (user_id, relay_host_id, assignment_epoch, drain_attempt_id, + source_cell_id, source_cell_incarnation, target_cell_id, + target_cell_incarnation, source_request_units, target_reserved_units, + pinned_at) + SELECT migration.user_id, migration.relay_host_id, + migration.assignment_epoch, ?, migration.source_cell_id, + incarnation.source_cell_incarnation, migration.target_cell_id, + incarnation.target_cell_incarnation, migration.source_request_units, + migration.target_reserved_units, 100 + FROM relay_assignment_migrations migration + JOIN relay_assignment_migration_incarnations incarnation + ON incarnation.user_id = migration.user_id + AND incarnation.relay_host_id = migration.relay_host_id + AND incarnation.assignment_epoch = migration.assignment_epoch + WHERE migration.user_id = ? AND migration.relay_host_id = ? + AND migration.assignment_epoch = ?`, + [ + drainAttemptId, + input.identity.userId, + input.identity.relayHostId, + assignmentEpoch + ] + ) + return drainAttemptId + } + + async function assignmentActivities(input: RecoveryFixture): Promise { + return await database.query( + `SELECT activity_id, activity_kind, cell_id, request_units + FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? ORDER BY activity_id`, + [input.identity.userId, input.identity.relayHostId] + ) + } + + async function reservations(cells: RecoveryFixture['cells']): Promise> { + const rows = await database.query( + `SELECT cell_id, reserved_requests FROM relay_cells + WHERE cell_id IN (?, ?, ?) ORDER BY cell_id`, + [cells.source.id, cells.failed.id, cells.replacement.id] + ) + return Object.fromEntries( + rows.map((row) => { + const cellId = String(row.cell_id) + const name = cellId.endsWith('-source') + ? 'source' + : cellId.endsWith('-failed') + ? 'failed' + : 'replacement' + return [name, Number(row.reserved_requests)] + }) + ) + } +}) diff --git a/cloud/apps/relay/src/observed-relay-database.ts b/cloud/apps/relay/src/observed-relay-database.ts new file mode 100644 index 00000000000..4720cae6c98 --- /dev/null +++ b/cloud/apps/relay/src/observed-relay-database.ts @@ -0,0 +1,46 @@ +import type { + RelayDatabase, + RelayLockOptions, + RelayTransactionOptions, + SqlRow +} from './database.js' +import { timedRelayOperation, type RelayRuntimeObserver } from './relay-observability.js' + +export function observeRelayDatabase( + database: RelayDatabase, + observer: RelayRuntimeObserver +): RelayDatabase { + const query = (sql: string, params?: unknown[]): Promise => + timedRelayOperation( + () => database.query(sql, params), + (durationMs, success) => observer.recordSql(durationMs, success) + ) + const queryLocked = ( + sql: string, + params?: unknown[], + options?: RelayLockOptions + ): Promise => + timedRelayOperation( + () => database.queryLocked(sql, params, options), + (durationMs, success) => observer.recordSql(durationMs, success), + (error) => + // NOWAIT contention is an intentional sweep deferral, not a SQL-health failure. + options?.failIfUnavailable === true && + error instanceof Error && + error.message === 'database_lock_unavailable' + ) + return { + dialect: database.dialect, + query, + queryLocked, + transaction: async ( + operation: (transaction: RelayDatabase) => Promise, + options?: RelayTransactionOptions + ): Promise => + await database.transaction( + async (transaction) => await operation(observeRelayDatabase(transaction, observer)), + options + ), + close: async () => await database.close() + } +} diff --git a/cloud/apps/relay/src/postgres-drain-send-locking.test.ts b/cloud/apps/relay/src/postgres-drain-send-locking.test.ts new file mode 100644 index 00000000000..024bf46db5e --- /dev/null +++ b/cloud/apps/relay/src/postgres-drain-send-locking.test.ts @@ -0,0 +1,215 @@ +import { afterAll, afterEach, beforeAll, describe, expect, it } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.js' +import { openRelayDatabase, type RelayDatabase } from './database.js' + +const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL +const describePostgres = databaseUrl ? describe : describe.skip +const source = { + id: 'drain-send-lock-source', + url: 'https://drain-send-lock-source.example.com', + capacityRequests: 100 +} +const target = { + id: 'drain-send-lock-target', + url: 'https://drain-send-lock-target.example.com', + capacityRequests: 100 +} +const identity = { + userId: 'drain-send-lock-user', + relayHostId: 'drainsendlock01' +} +const sourceIncarnation = '11111111-1111-4111-8111-111111111111' +const targetIncarnation = '22222222-2222-4222-8222-222222222222' +const attemptId = '33333333-3333-4333-8333-333333333333' + +describePostgres('PostgreSQL drain-send locking', () => { + let database: RelayDatabase + + beforeAll(async () => { + database = await openRelayDatabase({ databaseUrl, dataDir: '' }) + }) + + afterAll(async () => await database.close()) + + afterEach(async () => { + await database.query(`DELETE FROM relay_post_drain_migration_pins WHERE user_id = ?`, [ + identity.userId + ]) + await database.query(`DELETE FROM relay_cell_drain_attempt_states WHERE attempt_id = ?`, [ + attemptId + ]) + await database.query( + `DELETE FROM relay_control_connection_reservations WHERE user_id = ?`, + [identity.userId] + ) + await database.query(`DELETE FROM relay_migration_leases WHERE user_id = ?`, [ + identity.userId + ]) + await database.query(`DELETE FROM relay_assignment_activity_leases WHERE user_id = ?`, [ + identity.userId + ]) + await database.query( + `DELETE FROM relay_assignment_migration_incarnations WHERE user_id = ?`, + [identity.userId] + ) + await database.query(`DELETE FROM relay_assignment_migrations WHERE user_id = ?`, [ + identity.userId + ]) + await database.query(`DELETE FROM relay_assignments WHERE user_id = ?`, [identity.userId]) + await database.query(`DELETE FROM relay_cell_runtime WHERE cell_id IN (?, ?)`, [ + source.id, + target.id + ]) + await database.query(`DELETE FROM relay_cell_admission WHERE cell_id IN (?, ?)`, [ + source.id, + target.id + ]) + await database.query(`DELETE FROM relay_cells WHERE cell_id IN (?, ?)`, [ + source.id, + target.id + ]) + }) + + it('locks active migrations without locking the nullable incarnation lookup', async () => { + const now = 1_700_000_000_000 + const store = new RelayAssignmentStore(database, () => now, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + await store.reconcileCells([source, target]) + await store.recordCellHeartbeat({ + cellId: source.id, + cellUrl: source.url, + cellIncarnation: sourceIncarnation, + startedAt: now - 1, + ready: true, + observedRequests: 0 + }) + await store.recordCellHeartbeat({ + cellId: target.id, + cellUrl: target.url, + cellIncarnation: targetIncarnation, + startedAt: now - 1, + ready: true, + observedRequests: 0 + }) + await store.assign(identity) + await store.setCellEnabled(source.id, false) + await store.startEvacuation(identity, target.id) + await store.prepareCellDrainAttempt({ + attemptId, + cellId: source.id, + cellIncarnation: sourceIncarnation, + traceValue: '44444444-4444-4444-8444-444444444444', + plannedGraceMs: 120_000 + }) + await expect( + store.prepareCellDrainRecovery({ + attemptId, + cellId: source.id, + cellIncarnation: sourceIncarnation + }) + ).resolves.toMatchObject({ + shouldSend: false, + preparedAttempt: { attemptId, state: 'prepared' } + }) + + await expect( + store.beginCellDrainSend({ + attemptId, + cellId: source.id, + cellIncarnation: sourceIncarnation + }) + ).resolves.toMatchObject({ state: 'send-may-have-started', shouldSend: true }) + await expect( + store.beginCellDrainSend({ + attemptId, + cellId: source.id, + cellIncarnation: sourceIncarnation + }) + ).resolves.toMatchObject({ state: 'send-may-have-started', shouldSend: false }) + await expect( + store.prepareCellDrainRecovery({ + attemptId, + cellId: source.id, + cellIncarnation: sourceIncarnation + }) + ).rejects.toThrow('drain_application_receipt_missing') + await expect( + database.query(`SELECT drain_attempt_id FROM relay_post_drain_migration_pins`) + ).resolves.toEqual([{ drain_attempt_id: attemptId }]) + }) + + it('restores an expired registered migration lease with PostgreSQL locks', async () => { + let now = 1_700_000_000_000 + const startedAt = now - 1 + const store = new RelayAssignmentStore(database, () => now, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + await store.reconcileCells([source, target]) + await store.recordCellHeartbeat({ + cellId: source.id, + cellUrl: source.url, + cellIncarnation: sourceIncarnation, + startedAt, + ready: true, + observedRequests: 0 + }) + await store.recordCellHeartbeat({ + cellId: target.id, + cellUrl: target.url, + cellIncarnation: targetIncarnation, + startedAt, + ready: true, + observedRequests: 0 + }) + await store.assign(identity) + await store.setCellEnabled(source.id, false) + const migration = await store.startEvacuation(identity, target.id) + await store.markMigrationTargetRegistered(identity, { + cellId: target.id, + assignmentEpoch: migration.assignmentEpoch + }) + await store.prepareCellDrainAttempt({ + attemptId, + cellId: source.id, + cellIncarnation: sourceIncarnation, + traceValue: '44444444-4444-4444-8444-444444444444', + plannedGraceMs: 120_000 + }) + + now = migration.expiresAt + 1 + expect(await store.releaseExpiredActivityLeases()).toBeGreaterThan(0) + await store.recordCellHeartbeat({ + cellId: source.id, + cellUrl: source.url, + cellIncarnation: sourceIncarnation, + startedAt, + ready: true, + observedRequests: 0 + }) + await store.recordCellHeartbeat({ + cellId: target.id, + cellUrl: target.url, + cellIncarnation: targetIncarnation, + startedAt, + ready: true, + observedRequests: 0 + }) + await expect( + store.beginCellDrainSend({ + attemptId, + cellId: source.id, + cellIncarnation: sourceIncarnation + }) + ).resolves.toMatchObject({ state: 'send-may-have-started', shouldSend: true }) + await expect( + database.query( + `SELECT activity_kind, cell_id FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).resolves.toContainEqual({ activity_kind: 'migration', cell_id: target.id }) + }) +}) diff --git a/cloud/apps/relay/src/postgres-idle-client-error.test.ts b/cloud/apps/relay/src/postgres-idle-client-error.test.ts new file mode 100644 index 00000000000..50a7bb48ab5 --- /dev/null +++ b/cloud/apps/relay/src/postgres-idle-client-error.test.ts @@ -0,0 +1,22 @@ +import { describe, expect, it, vi } from 'vitest' +import { absorbPostgresIdleClientErrors } from './database.js' + +describe('PostgreSQL idle-client failure handling', () => { + it('absorbs the pool error without logging connection details', () => { + let listener: ((error: Error) => void) | undefined + const pool = { + on: vi.fn((_event: string, value: (error: Error) => void) => { + listener = value + return pool + }) + } + const warning = vi.spyOn(console, 'warn').mockImplementation(() => {}) + + absorbPostgresIdleClientErrors(pool as never) + expect(() => listener?.(new Error('postgres://user:secret@database'))).not.toThrow() + expect(warning).toHaveBeenCalledWith('[orca-relay] idle PostgreSQL client failed') + expect(JSON.stringify(warning.mock.calls)).not.toContain('secret') + + warning.mockRestore() + }) +}) diff --git a/cloud/apps/relay/src/postgres-maintenance-sweep-plans.test.ts b/cloud/apps/relay/src/postgres-maintenance-sweep-plans.test.ts new file mode 100644 index 00000000000..a8d0e1e3bf9 --- /dev/null +++ b/cloud/apps/relay/src/postgres-maintenance-sweep-plans.test.ts @@ -0,0 +1,61 @@ +import pg from 'pg' +import { afterAll, beforeAll, describe, expect, it } from 'vitest' +import { openRelayDatabase, type RelayDatabase } from './database.js' + +const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL +const describePostgres = databaseUrl ? describe : describe.skip + +// Inactive bases outlive their sweep and are never pruned, so the table only +// grows; production reached ~2.24M rows of which 3 were active. Enough rows +// here that a sequential scan is the cheaper plan without the index. +const INACTIVE_ROWS = 20_000 +const OWNED_PREFIX = 'sweep-plan-' + +describePostgres('PostgreSQL maintenance sweep plans', () => { + let database: RelayDatabase + let client: pg.Client + + beforeAll(async () => { + database = await openRelayDatabase({ databaseUrl, dataDir: '' }) + // Only ever touch this suite's own rows: the database is shared with the + // other PostgreSQL suites running in parallel. + await database.query( + `DELETE FROM relay_connection_bases WHERE basis_conn_id LIKE ?`, + [`${OWNED_PREFIX}%`] + ) + await database.query( + `INSERT INTO relay_connection_bases + (basis_conn_id, user_id, relay_host_id, relay_device_id, + owning_control_generation, credential_kind, deadline, active, created_at) + SELECT '${OWNED_PREFIX}' || generation, 'sweep-plan-user', 'sweepplan01', + 'sweep-plan-device', 1, 'invite', 1000, 0, 1000 + FROM generate_series(1, ${INACTIVE_ROWS}) AS generation` + ) + await database.query(`ANALYZE relay_connection_bases`) + client = new pg.Client({ connectionString: databaseUrl }) + await client.connect() + }) + + afterAll(async () => { + await client.end() + await database.query( + `DELETE FROM relay_connection_bases WHERE basis_conn_id LIKE ?`, + [`${OWNED_PREFIX}%`] + ) + await database.close() + }) + + // Why: a seq scan here held the maintenance transaction open long enough to + // time out assignment lock waits fleet-wide (2026-08-05 incident). + it('matches expired active bases by index instead of scanning the table', async () => { + const result = await client.query( + `EXPLAIN UPDATE relay_connection_bases SET active = $1 + WHERE active = $2 AND deadline <= $3`, + [0, 1, 2000] + ) + const plan = result.rows.map((row) => String(row['QUERY PLAN'])).join('\n') + + expect(plan).not.toMatch(/Seq Scan on relay_connection_bases/) + expect(plan).toMatch(/relay_connection_bases_active_deadline/) + }) +}) diff --git a/cloud/apps/relay/src/postgres-pool-pressure.test.ts b/cloud/apps/relay/src/postgres-pool-pressure.test.ts new file mode 100644 index 00000000000..2e020bf43fb --- /dev/null +++ b/cloud/apps/relay/src/postgres-pool-pressure.test.ts @@ -0,0 +1,48 @@ +import { describe, expect, it, vi } from 'vitest' +import { PostgresPoolPressure } from './postgres-pool-pressure.js' + +describe('PostgreSQL pool pressure', () => { + it('reports current waiters and interval high-water marks', async () => { + let now = 1_000 + let resolveConnection!: (client: unknown) => void + const connection = new Promise((resolve) => { + resolveConnection = resolve + }) + const pool = { + totalCount: 3, + idleCount: 0, + waitingCount: 0, + connect: vi.fn(() => { + pool.waitingCount++ + return connection + }) + } + const pressure = new PostgresPoolPressure(pool as never, () => now) + const pending = pressure.connect() + now = 1_750 + + expect(pressure.consumeCounts()).toMatchObject({ + databasePoolTotal: 3, + databasePoolIdle: 0, + databasePoolWaiting: 1, + databasePoolWaitersMax: 1, + databasePoolOldestWaitMs: 750, + databasePoolWaitMsMax: 750 + }) + + now = 2_250 + pool.waitingCount-- + resolveConnection({ query: vi.fn(), release: vi.fn() }) + await pending + expect(pressure.consumeCounts()).toMatchObject({ + databasePoolWaiting: 0, + databasePoolWaitersMax: 1, + databasePoolOldestWaitMs: 0, + databasePoolWaitMsMax: 1_250 + }) + expect(pressure.consumeCounts()).toMatchObject({ + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + }) + }) +}) diff --git a/cloud/apps/relay/src/postgres-pool-pressure.ts b/cloud/apps/relay/src/postgres-pool-pressure.ts new file mode 100644 index 00000000000..e18bf2bdd42 --- /dev/null +++ b/cloud/apps/relay/src/postgres-pool-pressure.ts @@ -0,0 +1,93 @@ +import type pg from 'pg' + +export type PostgresPoolPressureCounts = { + databasePoolTotal: number + databasePoolIdle: number + databasePoolWaiting: number + databasePoolWaitersMax: number + databasePoolOldestWaitMs: number + databasePoolWaitMsMax: number +} + +const emptyCounts = (): PostgresPoolPressureCounts => ({ + databasePoolTotal: 0, + databasePoolIdle: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolOldestWaitMs: 0, + databasePoolWaitMsMax: 0 +}) + +export class PostgresPoolPressure { + private readonly waiters = new Map() + private waitersMax = 0 + private waitMsMax = 0 + private lastConsumed = emptyCounts() + + constructor( + private readonly pool: pg.Pool, + private readonly now: () => number = Date.now + ) {} + + async connect(): Promise { + const waitingBefore = this.pool.waitingCount + const connection = this.pool.connect() + if (this.pool.waitingCount <= waitingBefore) return await connection + + const waiter = Symbol() + const startedAt = this.now() + this.waiters.set(waiter, startedAt) + this.waitersMax = Math.max(this.waitersMax, this.waiters.size) + try { + return await connection + } finally { + this.waitMsMax = Math.max(this.waitMsMax, this.now() - startedAt) + this.waiters.delete(waiter) + } + } + + consumeCounts(): PostgresPoolPressureCounts { + const counts = this.readCounts() + this.lastConsumed = counts + this.waitersMax = this.waiters.size + this.waitMsMax = counts.databasePoolOldestWaitMs + return counts + } + + peekCounts(): PostgresPoolPressureCounts { + const current = this.readCounts() + return { + ...current, + databasePoolWaitersMax: Math.max( + current.databasePoolWaitersMax, + this.lastConsumed.databasePoolWaitersMax + ), + databasePoolOldestWaitMs: Math.max( + current.databasePoolOldestWaitMs, + this.lastConsumed.databasePoolOldestWaitMs + ), + databasePoolWaitMsMax: Math.max( + current.databasePoolWaitMsMax, + this.lastConsumed.databasePoolWaitMsMax + ) + } + } + + private readCounts(): PostgresPoolPressureCounts { + const now = this.now() + const oldestWaitMs = + this.waiters.size === 0 ? 0 : Math.max(0, now - Math.min(...this.waiters.values())) + return { + databasePoolTotal: this.pool.totalCount, + databasePoolIdle: this.pool.idleCount, + databasePoolWaiting: this.waiters.size, + databasePoolWaitersMax: Math.max(this.waitersMax, this.waiters.size), + databasePoolOldestWaitMs: oldestWaitMs, + databasePoolWaitMsMax: Math.max(this.waitMsMax, oldestWaitMs) + } + } +} + +export function emptyPostgresPoolPressureCounts(): PostgresPoolPressureCounts { + return emptyCounts() +} diff --git a/cloud/apps/relay/src/postgres-schema-concurrency-postgres.test.ts b/cloud/apps/relay/src/postgres-schema-concurrency-postgres.test.ts new file mode 100644 index 00000000000..5208f137ae8 --- /dev/null +++ b/cloud/apps/relay/src/postgres-schema-concurrency-postgres.test.ts @@ -0,0 +1,55 @@ +import pg from 'pg' +import { afterAll, beforeAll, describe, expect, it } from 'vitest' +import { openRelayDatabase, type RelayDatabase } from './database.js' + +const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL +const describePostgres = databaseUrl ? describe : describe.skip +const schema = 'relay_schema_concurrency_test' + +describePostgres('PostgreSQL schema concurrency', () => { + let scopedUrl = '' + + beforeAll(async () => { + const client = new pg.Client({ connectionString: databaseUrl }) + await client.connect() + try { + await client.query(`DROP SCHEMA IF EXISTS ${schema} CASCADE`) + await client.query(`CREATE SCHEMA ${schema}`) + } finally { + await client.end() + } + const url = new URL(databaseUrl!) + url.searchParams.set('options', `-c search_path=${schema}`) + scopedUrl = url.toString() + }) + + afterAll(async () => { + const client = new pg.Client({ connectionString: databaseUrl }) + await client.connect() + try { + await client.query(`DROP SCHEMA IF EXISTS ${schema} CASCADE`) + } finally { + await client.end() + } + }) + + it('opens five directors when one new table is absent', async () => { + const initial = await openRelayDatabase({ databaseUrl: scopedUrl, dataDir: '' }) + await initial.query(`DROP TABLE relay_cell_legacy_fence_adoptions`) + await initial.close() + + const results = await Promise.allSettled( + Array.from({ length: 5 }, async (): Promise => + await openRelayDatabase({ databaseUrl: scopedUrl, dataDir: '' }) + ) + ) + const databases = results.flatMap((result) => + result.status === 'fulfilled' ? [result.value] : [] + ) + try { + expect(results.every((result) => result.status === 'fulfilled')).toBe(true) + } finally { + await Promise.all(databases.map(async (database) => await database.close())) + } + }) +}) diff --git a/cloud/apps/relay/src/postgres-schema-startup.ts b/cloud/apps/relay/src/postgres-schema-startup.ts new file mode 100644 index 00000000000..5e6260ad2fb --- /dev/null +++ b/cloud/apps/relay/src/postgres-schema-startup.ts @@ -0,0 +1,83 @@ +const RETRYABLE_SCHEMA_CODES = new Set(['55P03', '57014']) +const DEFAULT_RETRY_DEADLINE_MS = 30_000 +const RETRY_BASE_DELAY_MS = 250 +const RETRY_MAX_DELAY_MS = 2_000 + +type SchemaStartupOptions = { + now?: () => number + random?: () => number + retryDeadlineMs?: number + wait?: (delayMs: number) => Promise +} + +function retryDelayMs(attempt: number, random: () => number): number { + const ceiling = Math.min( + RETRY_BASE_DELAY_MS * 2 ** (attempt - 1), + RETRY_MAX_DELAY_MS + ) + return Math.ceil(ceiling * (0.5 + random() * 0.5)) +} + +function wait(delayMs: number): Promise { + return new Promise((resolve) => setTimeout(resolve, delayMs)) +} + +function retryableSchemaError(error: unknown, statement: string): boolean { + const value = error as { code?: unknown; constraint?: unknown } + return ( + RETRYABLE_SCHEMA_CODES.has(String(value.code)) || + (value.code === '23505' && + ((value.constraint === 'pg_type_typname_nsp_index' && + /^\s*CREATE\s+TABLE\s+IF\s+NOT\s+EXISTS\b/i.test(statement)) || + (value.constraint === 'pg_class_relname_nsp_index' && + /^\s*CREATE\s+(?:UNIQUE\s+)?INDEX\s+IF\s+NOT\s+EXISTS\b/i.test(statement)))) + ) +} + +export async function applyPostgresSchema( + statements: string[], + query: (statement: string) => Promise, + options: SchemaStartupOptions = {} +): Promise { + const now = options.now ?? Date.now + const random = options.random ?? Math.random + const pause = options.wait ?? wait + const deadlineAt = now() + (options.retryDeadlineMs ?? DEFAULT_RETRY_DEADLINE_MS) + + for (const statement of statements) { + let attempt = 1 + while (true) { + try { + await query(statement) + break + } catch (error) { + const code = String((error as { code?: unknown }).code) + const remainingMs = deadlineAt - now() + const retryable = retryableSchemaError(error, statement) + if (!retryable || remainingMs <= 0) { + if (retryable) { + console.warn( + JSON.stringify({ + event: 'orca_relay_postgres_schema_retry_exhausted', + code, + attempts: attempt + }) + ) + } + throw error + } + const delayMs = Math.min(remainingMs, retryDelayMs(attempt, random)) + console.warn( + JSON.stringify({ + event: 'orca_relay_postgres_schema_retry', + code, + attempt, + delayMs + }) + ) + await pause(delayMs) + attempt += 1 + } + } + } +} diff --git a/cloud/apps/relay/src/postgres-transaction-recovery.test.ts b/cloud/apps/relay/src/postgres-transaction-recovery.test.ts new file mode 100644 index 00000000000..a49c07d2e7d --- /dev/null +++ b/cloud/apps/relay/src/postgres-transaction-recovery.test.ts @@ -0,0 +1,1446 @@ +import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.js' +import { + openRelayDatabase, + type RelayDatabase, + type RelayLockOptions, + type SqlRow +} from './database.js' + +function postgresTestDatabaseUrl(value: string | undefined): string | undefined { + if (!value) return undefined + const url = new URL(value) + // Detect the intended deadlock before the one-second runtime lock deadline. + url.searchParams.set('options', '-c deadlock_timeout=100ms') + return url.toString() +} + +const databaseUrl = postgresTestDatabaseUrl(process.env.ORCA_RELAY_TEST_POSTGRES_URL) +const describePostgres = databaseUrl ? describe : describe.skip + +type QueryLockHook = (phase: 'before' | 'after', sql: string) => Promise + +class TransactionProbeDatabase implements RelayDatabase { + attempts = 0 + + constructor( + private readonly database: RelayDatabase, + private readonly hook: QueryLockHook, + private readonly probeQueries = false + ) {} + + async query(sql: string, params?: unknown[]): Promise { + return await this.database.query(sql, params) + } + + async queryLocked( + sql: string, + params?: unknown[], + options?: RelayLockOptions + ): Promise { + return await this.database.queryLocked(sql, params, options) + } + + async transaction(operation: (transaction: RelayDatabase) => Promise): Promise { + return await this.database.transaction(async (transaction) => { + this.attempts++ + return await operation( + new QueryLockProbeTransaction(transaction, this.hook, this.probeQueries) + ) + }) + } + + async close(): Promise {} +} + +class QueryLockProbeTransaction implements RelayDatabase { + constructor( + private readonly transactionDatabase: RelayDatabase, + private readonly hook: QueryLockHook, + private readonly probeQueries: boolean + ) {} + + async query(sql: string, params?: unknown[]): Promise { + if (this.probeQueries) await this.hook('before', sql) + const rows = await this.transactionDatabase.query(sql, params) + if (this.probeQueries) await this.hook('after', sql) + return rows + } + + async queryLocked( + sql: string, + params?: unknown[], + options?: RelayLockOptions + ): Promise { + await this.hook('before', sql) + const rows = await this.transactionDatabase.queryLocked(sql, params, options) + await this.hook('after', sql) + return rows + } + + async transaction(operation: (transaction: RelayDatabase) => Promise): Promise { + return await operation(this) + } + + async close(): Promise {} +} + +function signal(): { promise: Promise; resolve: () => void } { + let resolve!: () => void + return { promise: new Promise((done) => (resolve = done)), resolve } +} + +describePostgres('PostgreSQL transaction recovery', () => { + let database: RelayDatabase + + beforeAll(async () => { + database = await openRelayDatabase({ databaseUrl, dataDir: '' }) + }) + + afterAll(async () => { + await database.close() + }) + + afterEach(async () => { + for (const identity of [ + { userId: 'released-order-user', relayHostId: 'releasedorder001' }, + { userId: 'normalized-order-user', relayHostId: 'normalizedorder1' }, + { userId: 'atomic-final-user-a', relayHostId: 'atomicfinalhosta' }, + { userId: 'atomic-final-user-b', relayHostId: 'atomicfinalhostb' }, + { userId: 'lease-assignment-contention-user', relayHostId: 'leaseassignment1' } + ]) { + await database.query( + `DELETE FROM relay_assignment_activity_leases WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + await database.query( + `DELETE FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + } + await database.query(`DELETE FROM relay_cells WHERE cell_id IN (?, ?)`, [ + 'released-order-cell', + 'normalized-order-cell' + ]) + await database.query(`DELETE FROM relay_cells WHERE cell_id = ?`, ['atomic-final-cell']) + await database.query(`DELETE FROM relay_cells WHERE cell_id = ?`, [ + 'lease-assignment-contention-cell' + ]) + await database.query( + `DELETE FROM relay_assignment_activity_leases WHERE user_id LIKE 'reconcile-user-%'` + ) + await database.query(`DELETE FROM relay_assignments WHERE user_id LIKE 'reconcile-user-%'`) + await database.query(`DELETE FROM relay_cells WHERE cell_id IN (?, ?)`, [ + 'reconcile-cell-a', + 'reconcile-cell-b' + ]) + await database.query( + `DELETE FROM relay_assignment_activity_leases WHERE user_id = ?`, + ['completion-race-user'] + ) + await database.query(`DELETE FROM relay_assignment_migrations WHERE user_id = ?`, [ + 'completion-race-user' + ]) + await database.query(`DELETE FROM relay_assignments WHERE user_id = ?`, [ + 'completion-race-user' + ]) + await database.query(`DELETE FROM relay_cells WHERE cell_id IN (?, ?)`, [ + 'completion-race-source', + 'completion-race-target' + ]) + await database.query( + `DELETE FROM relay_assignment_activity_leases WHERE user_id = ?`, + ['completion-contention-user'] + ) + await database.query(`DELETE FROM relay_assignment_migrations WHERE user_id = ?`, [ + 'completion-contention-user' + ]) + await database.query(`DELETE FROM relay_assignments WHERE user_id = ?`, [ + 'completion-contention-user' + ]) + await database.query(`DELETE FROM relay_cell_runtime WHERE cell_id IN (?, ?)`, [ + 'completion-contention-source', + 'completion-contention-target' + ]) + await database.query(`DELETE FROM relay_cells WHERE cell_id IN (?, ?)`, [ + 'completion-contention-source', + 'completion-contention-target' + ]) + await database.query(`DELETE FROM relay_cell_fences WHERE cell_id = ?`, [ + 'dead-source-contention-source' + ]) + await database.query( + `DELETE FROM relay_control_connection_reservations WHERE user_id = ?`, + ['dead-source-contention-user'] + ) + await database.query( + `DELETE FROM relay_assignment_activity_leases WHERE user_id = ?`, + ['dead-source-contention-user'] + ) + await database.query( + `DELETE FROM relay_assignment_migration_incarnations WHERE user_id = ?`, + ['dead-source-contention-user'] + ) + await database.query(`DELETE FROM relay_assignment_migrations WHERE user_id = ?`, [ + 'dead-source-contention-user' + ]) + await database.query(`DELETE FROM relay_assignments WHERE user_id = ?`, [ + 'dead-source-contention-user' + ]) + await database.query(`DELETE FROM relay_cell_runtime WHERE cell_id IN (?, ?)`, [ + 'dead-source-contention-source', + 'dead-source-contention-target' + ]) + await database.query(`DELETE FROM relay_cell_admission WHERE cell_id IN (?, ?)`, [ + 'dead-source-contention-source', + 'dead-source-contention-target' + ]) + await database.query(`DELETE FROM relay_cells WHERE cell_id IN (?, ?)`, [ + 'dead-source-contention-source', + 'dead-source-contention-target' + ]) + await database.query( + `DELETE FROM relay_assignment_activity_leases + WHERE user_id IN (?, ?)`, + ['lease-contention-user', 'aggregate-contention-user'] + ) + await database.query( + `DELETE FROM relay_assignments + WHERE user_id IN (?, ?)`, + ['lease-contention-user', 'aggregate-contention-user'] + ) + await database.query(`DELETE FROM relay_cells WHERE cell_id IN (?, ?)`, [ + 'lease-contention-cell', + 'aggregate-contention-cell' + ]) + await database.query( + `DELETE FROM relay_assignment_activity_leases WHERE user_id = ?`, + ['sustained-lock-user'] + ) + await database.query(`DELETE FROM relay_assignments WHERE user_id = ?`, [ + 'sustained-lock-user' + ]) + await database.query(`DELETE FROM relay_cells WHERE cell_id = ?`, [ + 'sustained-lock-cell' + ]) + await database.query( + `DELETE FROM relay_assignment_activity_leases WHERE user_id = ?`, + ['sticky-cell-contention-user'] + ) + await database.query( + `DELETE FROM relay_assignments WHERE user_id = ?`, + ['sticky-cell-contention-user'] + ) + await database.query(`DELETE FROM relay_cells WHERE cell_id = ?`, [ + 'sticky-cell-contention-cell' + ]) + await database.query( + `DELETE FROM relay_assignment_activity_leases WHERE user_id LIKE 'postgres-user-%'` + ) + await database.query(`DELETE FROM relay_assignments WHERE user_id LIKE 'postgres-user-%'`) + await database.query(`DELETE FROM relay_cells WHERE cell_id IN (?, ?, ?)`, [ + 'cell-a', + 'cell-b', + 'cell-c' + ]) + await database.query( + `DELETE FROM relay_assignment_activity_leases WHERE user_id LIKE 'sticky-fast-user-%'` + ) + await database.query( + `DELETE FROM relay_assignments WHERE user_id LIKE 'sticky-fast-user-%'` + ) + await database.query(`DELETE FROM relay_cells WHERE cell_id = ?`, [ + 'sticky-fast-cell' + ]) + }) + + it('retries an entire deadlock victim transaction on a fresh attempt', async () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const keys = ['deadlock-a', 'deadlock-b'] + for (const key of keys) { + await database.query( + `INSERT INTO relay_rate_windows + (scope_key, window_kind, window_started_at, count) VALUES (?, ?, ?, ?) + ON CONFLICT (scope_key, window_kind, window_started_at) DO UPDATE SET count = ?`, + [key, 'transaction-recovery', 1, 0, 0] + ) + } + + let firstAttemptArrivals = 0 + let releaseFirstAttempts!: () => void + const firstAttemptsReady = new Promise((resolve) => { + releaseFirstAttempts = resolve + }) + const attempts = [0, 0] + const mutateWithOppositeLockOrder = async ( + operationIndex: number, + firstKey: string, + secondKey: string + ): Promise => { + await database.transaction(async (transaction) => { + attempts[operationIndex] = (attempts[operationIndex] ?? 0) + 1 + await transaction.queryLocked( + `SELECT * FROM relay_rate_windows + WHERE scope_key = ? AND window_kind = ? AND window_started_at = ?`, + [firstKey, 'transaction-recovery', 1] + ) + if (attempts[operationIndex] === 1) { + firstAttemptArrivals++ + if (firstAttemptArrivals === 2) releaseFirstAttempts() + await firstAttemptsReady + } + await transaction.queryLocked( + `SELECT * FROM relay_rate_windows + WHERE scope_key = ? AND window_kind = ? AND window_started_at = ?`, + [secondKey, 'transaction-recovery', 1] + ) + await transaction.query( + `UPDATE relay_rate_windows SET count = count + 1 + WHERE scope_key = ? AND window_kind = ? AND window_started_at = ?`, + [firstKey, 'transaction-recovery', 1] + ) + }) + } + + await Promise.all([ + mutateWithOppositeLockOrder(0, keys[0]!, keys[1]!), + mutateWithOppositeLockOrder(1, keys[1]!, keys[0]!) + ]) + + expect([...attempts].sort()).toEqual([1, 2]) + const rows = await database.query( + `SELECT scope_key, count FROM relay_rate_windows + WHERE window_kind = ? ORDER BY scope_key`, + ['transaction-recovery'] + ) + expect(rows).toEqual([ + { scope_key: keys[0], count: '1' }, + { scope_key: keys[1], count: '1' } + ]) + expect(warn).toHaveBeenCalledWith( + expect.stringContaining('"event":"orca_relay_postgres_transaction_retry"') + ) + expect(warn).toHaveBeenCalledWith(expect.stringContaining('"phase":"rate-limit"')) + expect(warn).not.toHaveBeenCalledWith( + expect.stringContaining('"event":"orca_relay_postgres_transaction_exhausted"') + ) + warn.mockRestore() + }, 10_000) + + it('defers assignment around a released-cell activity transaction', async () => { + const now = 1_100_000_000_000 + const identity = { userId: 'released-order-user', relayHostId: 'releasedorder001' } + const cellId = 'released-order-cell' + const activityId = 'splice:released-order' + await database.query( + `DELETE FROM relay_assignment_activity_leases WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + await database.query(`DELETE FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, [ + identity.userId, + identity.relayHostId + ]) + const seedStore = new RelayAssignmentStore(database, () => now) + await seedStore.reconcileCells([ + { id: cellId, url: 'https://released-order.example.com', capacityRequests: 100 } + ]) + const assignment = await seedStore.assign(identity) + await seedStore.acquireActivity(identity, { activityId, kind: 'splice', cellId }) + + const assignmentLocked = signal() + const legacyCellLocked = signal() + let gateFirstAssignmentAttempt = true + const directorDatabase = new TransactionProbeDatabase(database, async (phase, sql) => { + if ( + gateFirstAssignmentAttempt && + phase === 'after' && + sql.includes('FROM relay_assignments') + ) { + gateFirstAssignmentAttempt = false + assignmentLocked.resolve() + await legacyCellLocked.promise + } + }) + const directorStore = new RelayAssignmentStore(directorDatabase, () => now) + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + + const directorAssign = directorStore.assign(identity) + await assignmentLocked.promise + let legacyAttempts = 0 + const legacyRelease = database.transaction(async (transaction) => { + legacyAttempts++ + const lease = ( + await transaction.queryLocked( + `SELECT * FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? AND activity_id = ?`, + [identity.userId, identity.relayHostId, activityId] + ) + )[0] + if (!lease) return + await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [cellId]) + legacyCellLocked.resolve() + await transaction.query( + `UPDATE relay_cells SET reserved_requests = reserved_requests - 2 WHERE cell_id = ?`, + [cellId] + ) + await transaction.query( + `DELETE FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? AND activity_id = ?`, + [identity.userId, identity.relayHostId, activityId] + ) + await transaction.query( + `UPDATE relay_assignments SET reserved_splices = reserved_splices - 1 + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + }) + + await Promise.all([directorAssign, legacyRelease]) + + expect(directorDatabase.attempts).toBe(2) + expect(legacyAttempts).toBe(1) + expect(warn).not.toHaveBeenCalledWith( + expect.stringContaining('orca_relay_postgres_transaction_retry') + ) + expect(warn).not.toHaveBeenCalledWith( + expect.stringContaining('orca_relay_postgres_transaction_exhausted') + ) + await expect(seedStore.resolve(identity)).resolves.toMatchObject({ + cellId, + assignmentEpoch: assignment.assignmentEpoch + }) + const state = await database.query( + `SELECT assignment.reserved_controls, assignment.reserved_splices, + cell.reserved_requests + FROM relay_assignments assignment + JOIN relay_cells cell ON cell.cell_id = assignment.cell_id + WHERE assignment.user_id = ? AND assignment.relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + expect(state).toEqual([ + { reserved_controls: '1', reserved_splices: '0', reserved_requests: '1' } + ]) + warn.mockRestore() + }, 15_000) + + it('renews an existing control without waiting on a legacy cell lock', async () => { + const now = 1_150_000_000_000 + const identity = { userId: 'sustained-lock-user', relayHostId: 'sustainedlock01' } + const cell = { + id: 'sustained-lock-cell', + url: 'https://sustained-lock.example.com', + capacityRequests: 100 + } + const seedStore = new RelayAssignmentStore(database, () => now) + await seedStore.reconcileCells([cell]) + await seedStore.assign(identity) + + const legacyCellLocked = signal() + const releaseLegacyCell = signal() + const legacyTransaction = database.transaction(async (transaction) => { + await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [cell.id]) + legacyCellLocked.resolve() + await releaseLegacyCell.promise + }) + await legacyCellLocked.promise + + let inventoryAttempts = 0 + const assignmentDatabase = new TransactionProbeDatabase( + database, + async (phase, sql) => { + if (phase !== 'before' || !sql.includes('FROM relay_cells ORDER BY')) return + inventoryAttempts++ + } + ) + const assignmentStore = new RelayAssignmentStore(assignmentDatabase, () => now) + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + + await expect(assignmentStore.assign(identity)).resolves.toMatchObject({ + cellId: cell.id + }) + releaseLegacyCell.resolve() + await expect(legacyTransaction).resolves.toBeUndefined() + expect(inventoryAttempts).toBe(0) + expect(warn).not.toHaveBeenCalledWith( + expect.stringContaining('orca_relay_postgres_transaction_retry') + ) + warn.mockRestore() + }, 15_000) + + it('retries a new sticky control without forming a legacy cell-first deadlock', async () => { + const now = 1_175_000_000_000 + const identity = { + userId: 'sticky-cell-contention-user', + relayHostId: 'stickycellwait1' + } + const cell = { + id: 'sticky-cell-contention-cell', + url: 'https://sticky-cell-contention.example.com', + capacityRequests: 100 + } + const seedStore = new RelayAssignmentStore(database, () => now) + await seedStore.reconcileCells([cell]) + await seedStore.assign(identity) + await seedStore.changeActivity(identity, 'migration', 1) + // Drops the grant's pending lease too: a sticky control the rows do not + // show is what sends the retry through the NOWAIT cell lock. + await seedStore.releaseActivity(identity, 'control-pending:1') + + const legacyCellLocked = signal() + const directorAssignmentLocked = signal() + const legacyTransaction = database.transaction(async (transaction) => { + await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [cell.id]) + legacyCellLocked.resolve() + await directorAssignmentLocked.promise + await transaction.queryLocked( + `SELECT * FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + }) + await legacyCellLocked.promise + + let signalFirstAttempt = true + const directorLockOrder: string[] = [] + const assignmentDatabase = new TransactionProbeDatabase(database, async (phase, sql) => { + if (phase === 'before') { + if (sql.includes('FROM relay_assignments WHERE user_id = ?')) { + directorLockOrder.push('assignment') + } else if (sql.includes('FROM relay_assignment_activity_leases')) { + directorLockOrder.push('activity') + } else if (sql.includes('FROM relay_cells ORDER BY')) { + directorLockOrder.push('cell-inventory') + } else if (sql.includes('FROM relay_cells WHERE cell_id = ?')) { + directorLockOrder.push('cell') + } + } + if ( + signalFirstAttempt && + phase === 'after' && + sql.includes('FROM relay_assignments WHERE user_id = ?') + ) { + signalFirstAttempt = false + directorAssignmentLocked.resolve() + } + }) + const store = new RelayAssignmentStore(assignmentDatabase, () => now) + + await expect(store.assign(identity)).resolves.toMatchObject({ + cellId: cell.id, + assignmentEpoch: 1 + }) + await expect(legacyTransaction).resolves.toBeUndefined() + expect(assignmentDatabase.attempts).toBe(2) + expect(directorLockOrder).toEqual([ + 'assignment', + 'activity', + 'cell', + 'cell-inventory', + 'assignment', + 'activity' + ]) + const state = await database.query( + `SELECT assignment.reserved_controls, assignment.migration_leases, + cell.reserved_requests + FROM relay_assignments assignment + JOIN relay_cells cell ON cell.cell_id = assignment.cell_id + WHERE assignment.user_id = ? AND assignment.relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + expect(state).toEqual([ + { reserved_controls: '1', migration_leases: '1', reserved_requests: '2' } + ]) + }, 15_000) + + it('serializes normalized assignment and activity release without a retry', async () => { + const now = 1_200_000_000_000 + const identity = { userId: 'normalized-order-user', relayHostId: 'normalizedorder1' } + const cellId = 'normalized-order-cell' + const activityId = 'splice:normalized-order' + await database.query( + `DELETE FROM relay_assignment_activity_leases WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + await database.query(`DELETE FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, [ + identity.userId, + identity.relayHostId + ]) + const seedStore = new RelayAssignmentStore(database, () => now) + await seedStore.reconcileCells([ + { id: cellId, url: 'https://normalized-order.example.com', capacityRequests: 100 } + ]) + await seedStore.assign(identity) + await seedStore.acquireActivity(identity, { activityId, kind: 'splice', cellId }) + + const assignmentLocked = signal() + const releaseReachedAssignment = signal() + const continueAssignment = signal() + let gateFirstAssignmentAttempt = true + const assignDatabase = new TransactionProbeDatabase(database, async (phase, sql) => { + if ( + gateFirstAssignmentAttempt && + phase === 'after' && + sql.includes('FROM relay_assignments') + ) { + gateFirstAssignmentAttempt = false + assignmentLocked.resolve() + await continueAssignment.promise + } + }) + const releaseDatabase = new TransactionProbeDatabase(database, async (phase, sql) => { + if (phase === 'before' && sql.includes('FROM relay_assignments')) { + releaseReachedAssignment.resolve() + } + }) + const assignStore = new RelayAssignmentStore(assignDatabase, () => now) + const releaseStore = new RelayAssignmentStore(releaseDatabase, () => now) + + const assign = assignStore.assign(identity) + await assignmentLocked.promise + const release = releaseStore.releaseActivity(identity, activityId) + await releaseReachedAssignment.promise + continueAssignment.resolve() + await Promise.all([assign, release]) + + expect(assignDatabase.attempts).toBe(1) + expect(releaseDatabase.attempts).toBe(1) + const state = await database.query( + `SELECT assignment.reserved_controls, assignment.reserved_splices, + cell.reserved_requests + FROM relay_assignments assignment + JOIN relay_cells cell ON cell.cell_id = assignment.cell_id + WHERE assignment.user_id = ? AND assignment.relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + expect(state).toEqual([ + { reserved_controls: '1', reserved_splices: '0', reserved_requests: '1' } + ]) + }, 15_000) + + it('takes the shared cell lock only for the final atomic activity write', async () => { + const now = 1_250_000_000_000 + const cell = { + id: 'atomic-final-cell', + url: 'https://atomic-final.example.com', + capacityRequests: 4 + } + const identities = [ + { userId: 'atomic-final-user-a', relayHostId: 'atomicfinalhosta' }, + { userId: 'atomic-final-user-b', relayHostId: 'atomicfinalhostb' } + ] + const seedStore = new RelayAssignmentStore(database, () => now) + await seedStore.reconcileCells([cell]) + await Promise.all(identities.map(async (identity) => await seedStore.assign(identity))) + + const firstAtCellWrite = signal() + const secondAtCellWrite = signal() + const releaseFirst = signal() + const releaseSecond = signal() + const isAtomicCellWrite = (sql: string): boolean => + sql.includes('UPDATE relay_cells SET reserved_requests') && + sql.includes('RETURNING cell_id') + const firstDatabase = new TransactionProbeDatabase( + database, + async (phase, sql) => { + if (phase !== 'before' || !isAtomicCellWrite(sql)) return + firstAtCellWrite.resolve() + await releaseFirst.promise + }, + true + ) + const secondDatabase = new TransactionProbeDatabase( + database, + async (phase, sql) => { + if (phase !== 'before' || !isAtomicCellWrite(sql)) return + secondAtCellWrite.resolve() + await releaseSecond.promise + }, + true + ) + const first = new RelayAssignmentStore(firstDatabase, () => now).acquireActivity( + identities[0]!, + { activityId: 'splice:atomic-final-a', kind: 'splice', cellId: cell.id } + ) + await firstAtCellWrite.promise + const second = new RelayAssignmentStore(secondDatabase, () => now).acquireActivity( + identities[1]!, + { activityId: 'splice:atomic-final-b', kind: 'splice', cellId: cell.id } + ) + let reachTimeout: ReturnType | undefined + const secondReachedCellWrite = await Promise.race([ + secondAtCellWrite.promise.then(() => true), + new Promise((resolve) => { + reachTimeout = setTimeout(() => resolve(false), 2_000) + }) + ]) + if (reachTimeout) clearTimeout(reachTimeout) + if (!secondReachedCellWrite) { + releaseFirst.resolve() + releaseSecond.resolve() + await Promise.allSettled([first, second]) + } + expect(secondReachedCellWrite).toBe(true) + + releaseFirst.resolve() + await expect(first).resolves.toBeUndefined() + releaseSecond.resolve() + await expect(second).rejects.toThrow('relay_capacity_exhausted') + await expect( + database.query( + `SELECT cell.reserved_requests, COUNT(lease.activity_id) AS leases, + COALESCE(SUM(lease.request_units), 0) AS lease_units + FROM relay_cells cell + LEFT JOIN relay_assignment_activity_leases lease ON lease.cell_id = cell.cell_id + WHERE cell.cell_id = ? GROUP BY cell.cell_id, cell.reserved_requests`, + [cell.id] + ) + ).resolves.toEqual([{ reserved_requests: '4', leases: '3', lease_units: '4' }]) + await expect( + database.query( + `SELECT user_id, reserved_splices FROM relay_assignments + WHERE user_id IN (?, ?) ORDER BY user_id`, + [identities[0]!.userId, identities[1]!.userId] + ) + ).resolves.toEqual([ + { user_id: identities[0]!.userId, reserved_splices: '1' }, + { user_id: identities[1]!.userId, reserved_splices: '0' } + ]) + }, 15_000) + + it('reconciles drift under the assignment-first lock order while activity waits', async () => { + const now = 1_300_000_000_000 + const cells = [ + { id: 'reconcile-cell-a', url: 'https://reconcile-a.example.com', capacityRequests: 100 }, + { id: 'reconcile-cell-b', url: 'https://reconcile-b.example.com', capacityRequests: 100 } + ] + const identities = Array.from({ length: 12 }, (_, index) => ({ + userId: `reconcile-user-${index}`, + relayHostId: `reconcilehost${String(index).padStart(4, '0')}` + })) + const seedStore = new RelayAssignmentStore(database, () => now) + await seedStore.reconcileCells(cells) + for (const identity of identities) await seedStore.assign(identity) + await database.query( + `UPDATE relay_assignments SET reserved_controls = 7, reserved_splices = 5 + WHERE user_id LIKE 'reconcile-user-%'` + ) + await database.query( + `UPDATE relay_cells SET reserved_requests = 42 + WHERE cell_id IN (?, ?)`, + ['reconcile-cell-a', 'reconcile-cell-b'] + ) + + const assignmentsLocked = signal() + const activityReachedAssignment = signal() + const continueReconciliation = signal() + let gateReconciliation = true + const reconcileDatabase = new TransactionProbeDatabase(database, async (phase, sql) => { + if ( + gateReconciliation && + phase === 'after' && + sql.includes('SELECT assignment.* FROM relay_assignments assignment') + ) { + gateReconciliation = false + assignmentsLocked.resolve() + await continueReconciliation.promise + } + }) + const activityDatabase = new TransactionProbeDatabase(database, async (phase, sql) => { + if (phase === 'before' && sql.includes('FROM relay_assignments WHERE user_id')) { + activityReachedAssignment.resolve() + } + }) + const reconcileStore = new RelayAssignmentStore(reconcileDatabase, () => now) + const activityStore = new RelayAssignmentStore(activityDatabase, () => now) + + const reconciliation = reconcileStore.cellEvacuationStatus( + 'reconcile-cell-a', + 'reconcile-cell-b', + true + ) + await assignmentsLocked.promise + const activity = activityStore.acquireActivity(identities[0]!, { + activityId: 'splice:reconcile', + kind: 'splice', + cellId: 'reconcile-cell-a' + }) + await activityReachedAssignment.promise + continueReconciliation.resolve() + await Promise.all([reconciliation, activity]) + + expect(reconcileDatabase.attempts).toBe(1) + expect(activityDatabase.attempts).toBe(1) + const reservations = await database.query( + `SELECT cell.cell_id, cell.reserved_requests, + COALESCE(SUM(lease.request_units), 0) AS lease_units + FROM relay_cells cell + LEFT JOIN relay_assignment_activity_leases lease ON lease.cell_id = cell.cell_id + WHERE cell.cell_id IN (?, ?) + GROUP BY cell.cell_id, cell.reserved_requests ORDER BY cell.cell_id`, + ['reconcile-cell-a', 'reconcile-cell-b'] + ) + expect(reservations).toEqual([ + { cell_id: 'reconcile-cell-a', reserved_requests: '8', lease_units: '8' }, + { cell_id: 'reconcile-cell-b', reserved_requests: '6', lease_units: '6' } + ]) + }, 15_000) + + it('fences a source activity queued behind evacuation completion', async () => { + const now = 1_400_000_000_000 + const identity = { + userId: 'completion-race-user', + relayHostId: 'completionrace01' + } + const sourceCellId = 'completion-race-source' + const targetCellId = 'completion-race-target' + const seedStore = new RelayAssignmentStore(database, () => now) + await seedStore.reconcileCells([ + { + id: sourceCellId, + url: 'https://completion-source.example.com', + capacityRequests: 100 + }, + { + id: targetCellId, + url: 'https://completion-target.example.com', + capacityRequests: 100 + } + ]) + const assignment = await seedStore.assign(identity) + const sourceControl = await seedStore.activateControl(identity, { + cellId: sourceCellId, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + const migration = await seedStore.startEvacuation(identity, targetCellId) + await seedStore.activateControl(identity, { + cellId: targetCellId, + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await seedStore.markMigrationTargetRegistered(identity, { + cellId: targetCellId, + assignmentEpoch: migration.assignmentEpoch + }) + await seedStore.releaseActivity(identity, sourceControl) + + const assignmentLocked = signal() + const activityReachedAssignment = signal() + const continueCompletion = signal() + let gateCompletion = true + const completionDatabase = new TransactionProbeDatabase(database, async (phase, sql) => { + if ( + gateCompletion && + phase === 'after' && + sql.includes('FROM relay_assignments WHERE user_id') + ) { + gateCompletion = false + assignmentLocked.resolve() + await continueCompletion.promise + } + }) + const activityDatabase = new TransactionProbeDatabase(database, async (phase, sql) => { + if (phase === 'before' && sql.includes('FROM relay_assignments WHERE user_id')) { + activityReachedAssignment.resolve() + } + }) + const completionStore = new RelayAssignmentStore(completionDatabase, () => now) + const activityStore = new RelayAssignmentStore(activityDatabase, () => now) + + const completion = completionStore.completeReadyEvacuations() + await assignmentLocked.promise + const activity = activityStore.acquireActivity(identity, { + activityId: 'install:queued-source-work', + kind: 'install', + cellId: sourceCellId + }) + const activityRejected = expect(activity).rejects.toThrow( + 'activity_cell_not_authoritative' + ) + await activityReachedAssignment.promise + continueCompletion.resolve() + + await expect(completion).resolves.toBe(1) + await activityRejected + await expect( + seedStore.cellEvacuationStatus(sourceCellId, targetCellId, false) + ).resolves.toMatchObject({ inProgress: 0 }) + }, 15_000) + + it('defers completion instead of deadlocking with a legacy cell-first lock', async () => { + const now = 1_500_000_000_000 + const identity = { + userId: 'completion-contention-user', + relayHostId: 'completionwait01' + } + const sourceCell = { + id: 'completion-contention-source', + url: 'https://completion-contention-source.example.com', + capacityRequests: 100 + } + const targetCell = { + id: 'completion-contention-target', + url: 'https://completion-contention-target.example.com', + capacityRequests: 100 + } + const store = new RelayAssignmentStore(database, () => now, { + requireLiveCells: true + }) + await store.reconcileCells([sourceCell, targetCell]) + await store.recordCellHeartbeat({ + cellId: sourceCell.id, + cellUrl: sourceCell.url, + cellIncarnation: '11111111-1111-4111-8111-111111111111', + startedAt: now - 100, + ready: true, + observedRequests: 0 + }) + await store.recordCellHeartbeat({ + cellId: targetCell.id, + cellUrl: targetCell.url, + cellIncarnation: '22222222-2222-4222-8222-222222222222', + startedAt: now - 100, + ready: true, + observedRequests: 0 + }) + const assignment = await store.assign(identity) + const sourceControl = await store.activateControl(identity, { + cellId: sourceCell.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + const migration = await store.startEvacuation(identity, targetCell.id) + await store.activateControl(identity, { + cellId: targetCell.id, + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await store.markMigrationTargetRegistered(identity, { + cellId: targetCell.id, + assignmentEpoch: migration.assignmentEpoch + }) + await store.releaseActivity(identity, sourceControl) + await store.setCellEnabled(sourceCell.id, false) + + const legacyCellLocked = signal() + const completionAssignmentLocked = signal() + const legacyTransaction = database.transaction(async (transaction) => { + await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [ + sourceCell.id + ]) + legacyCellLocked.resolve() + await completionAssignmentLocked.promise + await transaction.queryLocked( + `SELECT * FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + }) + await legacyCellLocked.promise + + let signalCompletion = true + const completionDatabase = new TransactionProbeDatabase( + database, + async (phase, sql) => { + if ( + signalCompletion && + phase === 'after' && + sql.includes('FROM relay_assignments WHERE user_id') + ) { + signalCompletion = false + completionAssignmentLocked.resolve() + } + } + ) + const completionStore = new RelayAssignmentStore(completionDatabase, () => now, { + requireLiveCells: true + }) + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + + await expect(completionStore.completeReadyEvacuations()).resolves.toBe(0) + await expect(legacyTransaction).resolves.toBeUndefined() + expect(warn).not.toHaveBeenCalledWith( + expect.stringContaining('orca_relay_postgres_transaction_retry') + ) + await expect(store.completeReadyEvacuations()).resolves.toBe(1) + }, 15_000) + + it('normalizes persistent dead-source inventory contention', async () => { + let now = 1_550_000_000_000 + const identity = { + userId: 'dead-source-contention-user', + relayHostId: 'deadsourcewait01' + } + const secondIdentity = { + userId: identity.userId, + relayHostId: 'deadsourcewait02' + } + const sourceCell = { + id: 'dead-source-contention-source', + url: 'https://dead-source-contention-source.example.com', + capacityRequests: 100 + } + const targetCell = { + id: 'dead-source-contention-target', + url: 'https://dead-source-contention-target.example.com', + capacityRequests: 100 + } + const sourceIncarnation = '11111111-1111-4111-8111-111111111111' + const targetIncarnation = '22222222-2222-4222-8222-222222222222' + const store = new RelayAssignmentStore(database, () => now, { + requireLiveCells: true + }) + await store.reconcileCells([sourceCell, targetCell]) + await store.setCellEnabled(targetCell.id, false) + await store.recordCellHeartbeat({ + cellId: sourceCell.id, + cellUrl: sourceCell.url, + cellIncarnation: sourceIncarnation, + startedAt: now - 100, + ready: true, + observedRequests: 0 + }) + await store.recordCellHeartbeat({ + cellId: targetCell.id, + cellUrl: targetCell.url, + cellIncarnation: targetIncarnation, + startedAt: now - 100, + ready: true, + observedRequests: 0 + }) + const assignment = await store.assign(identity) + const sourceControl = await store.activateControl(identity, { + cellId: sourceCell.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + await store.setCellEnabled(targetCell.id, true) + const migration = await store.startEvacuation(identity, targetCell.id) + await store.activateControl(identity, { + cellId: targetCell.id, + assignmentEpoch: migration.assignmentEpoch, + generation: 1 + }) + await store.markMigrationTargetRegistered(identity, { + cellId: targetCell.id, + assignmentEpoch: migration.assignmentEpoch + }) + await store.releaseActivity(identity, sourceControl) + const secondAssignment = await store.assign(secondIdentity) + expect(secondAssignment.cellId).toBe(sourceCell.id) + const secondSourceControl = await store.activateControl(secondIdentity, { + cellId: sourceCell.id, + assignmentEpoch: secondAssignment.assignmentEpoch, + generation: 1 + }) + const secondMigration = await store.startEvacuation(secondIdentity, targetCell.id) + await store.activateControl(secondIdentity, { + cellId: targetCell.id, + assignmentEpoch: secondMigration.assignmentEpoch, + generation: 1 + }) + await store.markMigrationTargetRegistered(secondIdentity, { + cellId: targetCell.id, + assignmentEpoch: secondMigration.assignmentEpoch + }) + await store.releaseActivity(secondIdentity, secondSourceControl) + await store.setCellEnabled(sourceCell.id, false) + now += 45_001 + await store.recordCellHeartbeat({ + cellId: targetCell.id, + cellUrl: targetCell.url, + cellIncarnation: targetIncarnation, + startedAt: now - 45_101, + ready: true, + observedRequests: 0 + }) + await store.attestCellFence(sourceCell.id, sourceIncarnation) + + const legacyCellLocked = signal() + const completionAssignmentLocked = signal() + const legacyTransaction = database.transaction(async (transaction) => { + await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [ + sourceCell.id + ]) + legacyCellLocked.resolve() + await completionAssignmentLocked.promise + await transaction.queryLocked( + `SELECT * FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + }) + await legacyCellLocked.promise + + let signalCompletion = true + const fallbackInventoryLocked = signal() + const freshAssignmentLocked = signal() + const persistentInventoryLocked = signal() + let releasePersistentInventory = false + const completionDatabase = new TransactionProbeDatabase( + database, + async (phase, sql) => { + if ( + signalCompletion && + phase === 'after' && + sql.includes('FROM relay_assignments WHERE user_id') + ) { + signalCompletion = false + completionAssignmentLocked.resolve() + } else if ( + phase === 'after' && + sql.trim() === 'SELECT * FROM relay_cells ORDER BY cell_id ASC' + ) { + fallbackInventoryLocked.resolve() + await freshAssignmentLocked.promise + } + } + ) + const completionStore = new RelayAssignmentStore(completionDatabase, () => now, { + requireLiveCells: true + }) + const freshAssignmentTransaction = database.transaction(async (transaction) => { + await fallbackInventoryLocked.promise + await transaction.queryLocked( + `SELECT * FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + freshAssignmentLocked.resolve() + await transaction.queryLocked(`SELECT * FROM relay_cells ORDER BY cell_id ASC`) + persistentInventoryLocked.resolve() + while (!releasePersistentInventory) { + await new Promise((resolve) => setTimeout(resolve, 250)) + await transaction.query(`SELECT 1`) + } + }) + + const completion = completionStore.cellEvacuationStatus( + sourceCell.id, + targetCell.id, + true + ) + await persistentInventoryLocked.promise + await expect( + completion + ).resolves.toMatchObject({ inProgress: 2, completed: 0, blocked: 2 }) + await expect(legacyTransaction).resolves.toBeUndefined() + releasePersistentInventory = true + await expect(freshAssignmentTransaction).resolves.toBeUndefined() + await expect( + store.cellEvacuationStatus(sourceCell.id, targetCell.id, true) + ).resolves.toMatchObject({ inProgress: 0, completed: 2, blocked: 0 }) + }, 25_000) + + it('defers expired lease cleanup around a legacy cell-first lock', async () => { + const now = 1_600_000_000_000 + const identity = { + userId: 'lease-contention-user', + relayHostId: 'leasecontention1' + } + const cell = { + id: 'lease-contention-cell', + url: 'https://lease-contention-cell.example.com', + capacityRequests: 100 + } + const store = new RelayAssignmentStore(database, () => now) + await store.reconcileCells([cell]) + await store.assign(identity) + await database.query( + `UPDATE relay_assignment_activity_leases SET expires_at = ? + WHERE user_id = ? AND relay_host_id = ?`, + [now - 1, identity.userId, identity.relayHostId] + ) + + const legacyCellLocked = signal() + const cleanupAssignmentLocked = signal() + const legacyTransaction = database.transaction(async (transaction) => { + await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [cell.id]) + legacyCellLocked.resolve() + await cleanupAssignmentLocked.promise + await transaction.queryLocked( + `SELECT * FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + }) + await legacyCellLocked.promise + + let signalCleanup = true + const cleanupDatabase = new TransactionProbeDatabase(database, async (phase, sql) => { + if ( + signalCleanup && + phase === 'after' && + sql.includes('FROM relay_assignments WHERE user_id') + ) { + signalCleanup = false + cleanupAssignmentLocked.resolve() + } + }) + const cleanupStore = new RelayAssignmentStore(cleanupDatabase, () => now) + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + + await expect(cleanupStore.releaseExpiredActivityLeases()).resolves.toBe(0) + await expect(legacyTransaction).resolves.toBeUndefined() + expect(cleanupDatabase.attempts).toBe(1) + expect(warn).not.toHaveBeenCalledWith( + expect.stringContaining('orca_relay_postgres_transaction_retry') + ) + await expect(store.releaseExpiredActivityLeases()).resolves.toBe(1) + }, 15_000) + + it('skips an expired lease when another director owns its assignment lock', async () => { + const now = 1_650_000_000_000 + const identity = { + userId: 'lease-assignment-contention-user', + relayHostId: 'leaseassignment1' + } + const cell = { + id: 'lease-assignment-contention-cell', + url: 'https://lease-assignment-contention-cell.example.com', + capacityRequests: 100 + } + const store = new RelayAssignmentStore(database, () => now) + await store.reconcileCells([cell]) + await store.assign(identity) + await database.query( + `UPDATE relay_assignment_activity_leases SET expires_at = ? + WHERE user_id = ? AND relay_host_id = ?`, + [now - 1, identity.userId, identity.relayHostId] + ) + + const assignmentLocked = signal() + const releaseAssignment = signal() + const legacyTransaction = database.transaction(async (transaction) => { + await transaction.queryLocked( + `SELECT * FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + assignmentLocked.resolve() + await releaseAssignment.promise + }) + await assignmentLocked.promise + + const cleanupDatabase = new TransactionProbeDatabase(database, async () => undefined) + const cleanupStore = new RelayAssignmentStore(cleanupDatabase, () => now) + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + + await expect(cleanupStore.releaseExpiredActivityLeases()).resolves.toBe(0) + expect(cleanupDatabase.attempts).toBe(1) + expect(warn).not.toHaveBeenCalledWith( + expect.stringContaining('orca_relay_postgres_transaction_retry') + ) + warn.mockRestore() + + releaseAssignment.resolve() + await expect(legacyTransaction).resolves.toBeUndefined() + await expect(store.releaseExpiredActivityLeases()).resolves.toBe(1) + }, 15_000) + + it('defers aggregate expiry cleanup around a legacy cell-first lock', async () => { + const now = 1_700_000_000_000 + const identity = { + userId: 'aggregate-contention-user', + relayHostId: 'aggregatewait01' + } + const cell = { + id: 'aggregate-contention-cell', + url: 'https://aggregate-contention-cell.example.com', + capacityRequests: 100 + } + const store = new RelayAssignmentStore(database, () => now) + await store.reconcileCells([cell]) + await store.assign(identity) + await database.query( + `DELETE FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + await database.query( + `UPDATE relay_assignments SET lease_expires_at = ? + WHERE user_id = ? AND relay_host_id = ?`, + [now - 1, identity.userId, identity.relayHostId] + ) + + const legacyCellLocked = signal() + const cleanupAssignmentsLocked = signal() + const legacyTransaction = database.transaction(async (transaction) => { + await transaction.queryLocked(`SELECT * FROM relay_cells WHERE cell_id = ?`, [cell.id]) + legacyCellLocked.resolve() + await cleanupAssignmentsLocked.promise + await transaction.queryLocked( + `SELECT * FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + }) + await legacyCellLocked.promise + + let signalCleanup = true + const cleanupDatabase = new TransactionProbeDatabase(database, async (phase, sql) => { + if ( + signalCleanup && + phase === 'after' && + sql.includes('FROM relay_assignments WHERE lease_expires_at') + ) { + signalCleanup = false + cleanupAssignmentsLocked.resolve() + } + }) + const cleanupStore = new RelayAssignmentStore(cleanupDatabase, () => now) + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + + await expect(cleanupStore.releaseExpiredActivity()).resolves.toBe(0) + await expect(legacyTransaction).resolves.toBeUndefined() + expect(cleanupDatabase.attempts).toBe(1) + expect(warn).not.toHaveBeenCalledWith( + expect.stringContaining('orca_relay_postgres_transaction_retry') + ) + await expect(store.releaseExpiredActivity()).resolves.toBe(1) + }, 15_000) + + it('defers aggregate expiry before waiting on a legacy assignment row', async () => { + const now = 1_700_000_000_000 + const identity = { + userId: 'aggregate-contention-user', + relayHostId: 'aggregatewait01' + } + const cell = { + id: 'aggregate-contention-cell', + url: 'https://aggregate-contention-cell.example.com', + capacityRequests: 100 + } + const store = new RelayAssignmentStore(database, () => now) + await store.reconcileCells([cell]) + await store.assign(identity) + await database.query( + `DELETE FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + await database.query( + `UPDATE relay_assignments SET lease_expires_at = ? + WHERE user_id = ? AND relay_host_id = ?`, + [now - 1, identity.userId, identity.relayHostId] + ) + + const legacyAssignmentLocked = signal() + const releaseLegacyAssignment = signal() + const legacyTransaction = database.transaction(async (transaction) => { + await transaction.queryLocked( + `SELECT * FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + legacyAssignmentLocked.resolve() + await releaseLegacyAssignment.promise + }) + await legacyAssignmentLocked.promise + + const timedOut = Symbol('timed-out') + const cleanup = store.releaseExpiredActivity() + const result = await Promise.race([ + cleanup, + new Promise((resolve) => { + setTimeout(() => resolve(timedOut), 750) + }) + ]) + releaseLegacyAssignment.resolve() + await legacyTransaction + if (result === timedOut) await cleanup + + expect(result).toBe(0) + await expect(store.releaseExpiredActivity()).resolves.toBe(1) + }, 15_000) + + it('keeps concurrent dormant reassignments capacity-consistent', async () => { + const now = 1_000_000_000_000 + const serializedDatabase = new TransactionProbeDatabase(database, async () => undefined) + const store = new RelayAssignmentStore(serializedDatabase, () => now) + await database.query( + `DELETE FROM relay_assignment_activity_leases WHERE user_id LIKE 'postgres-user-%'` + ) + await database.query(`DELETE FROM relay_assignments WHERE user_id LIKE 'postgres-user-%'`) + await database.query(`DELETE FROM relay_cells WHERE cell_id IN (?, ?, ?)`, [ + 'cell-a', + 'cell-b', + 'cell-c' + ]) + await store.reconcileCells([ + { id: 'cell-a', url: 'https://cell-a.example.com', capacityRequests: 100 }, + { id: 'cell-b', url: 'https://cell-b.example.com', capacityRequests: 100 }, + { id: 'cell-c', url: 'https://cell-c.example.com', capacityRequests: 100 } + ]) + const identities = Array.from({ length: 30 }, (_, index) => ({ + userId: `postgres-user-${index}`, + relayHostId: `postgreshost${String(index).padStart(5, '0')}` + })) + await Promise.all(identities.map(async (identity) => await store.assign(identity))) + + await database.query(`DELETE FROM relay_assignment_activity_leases`) + await database.query( + `UPDATE relay_assignments SET lease_expires_at = ?, last_activity_at = ?, + reserved_controls = 0, reserved_splices = 0, reserved_invites = 0, + pending_installs = 0, pending_confirmations = 0, migration_leases = 0`, + [0, 0] + ) + await database.query(`UPDATE relay_cells SET reserved_requests = 0`) + + await Promise.all(identities.map(async (identity) => await store.assign(identity))) + + const reservations = await database.query( + `SELECT cell_id, reserved_requests FROM relay_cells ORDER BY cell_id` + ) + expect(reservations).toEqual([ + { cell_id: 'cell-a', reserved_requests: '10' }, + { cell_id: 'cell-b', reserved_requests: '10' }, + { cell_id: 'cell-c', reserved_requests: '10' } + ]) + expect(serializedDatabase.attempts).toBe(121) + }, 10_000) + + it('renews independent sticky assignments concurrently without the placement queue', async () => { + const now = 1_800_000_000_000 + const identities = Array.from({ length: 6 }, (_, index) => ({ + userId: `sticky-fast-user-${index}`, + relayHostId: `stickyfast${String(index).padStart(6, '0')}` + })) + const cell = { + id: 'sticky-fast-cell', + url: 'https://sticky-fast.example.com', + capacityRequests: 100 + } + const seedStore = new RelayAssignmentStore(database, () => now) + await seedStore.reconcileCells([cell]) + for (const identity of identities) await seedStore.assign(identity) + + const allRenewalsReachedDatabase = signal() + const releaseRenewals = signal() + let renewalArrivals = 0 + const probedDatabase = new TransactionProbeDatabase(database, async (phase, sql) => { + if ( + phase !== 'after' || + !sql.includes('FROM relay_assignments WHERE user_id = ?') + ) { + return + } + renewalArrivals++ + if (renewalArrivals === identities.length) allRenewalsReachedDatabase.resolve() + await releaseRenewals.promise + }) + const store = new RelayAssignmentStore(probedDatabase, () => now) + const renewals = identities.map(async (identity) => await store.assign(identity)) + let reachedBeforeTimeout = false + try { + reachedBeforeTimeout = await Promise.race([ + allRenewalsReachedDatabase.promise.then(() => true), + new Promise((resolve) => setTimeout(() => resolve(false), 1_000)) + ]) + } finally { + releaseRenewals.resolve() + await Promise.all(renewals) + } + + expect(reachedBeforeTimeout).toBe(true) + expect(renewalArrivals).toBe(identities.length) + expect(probedDatabase.attempts).toBe(identities.length) + const state = await database.query( + `SELECT COUNT(*) AS assignments, SUM(reserved_controls) AS controls + FROM relay_assignments WHERE user_id LIKE 'sticky-fast-user-%'` + ) + expect(state).toEqual([{ assignments: '6', controls: '6' }]) + }, 10_000) +}) diff --git a/cloud/apps/relay/src/public-assignment-admission.test.ts b/cloud/apps/relay/src/public-assignment-admission.test.ts new file mode 100644 index 00000000000..dbc9a7278be --- /dev/null +++ b/cloud/apps/relay/src/public-assignment-admission.test.ts @@ -0,0 +1,334 @@ +import { describe, expect, it, vi } from 'vitest' +import { RelayPublicAssignmentAdmission } from './public-assignment-admission.js' + +describe('public assignment admission', () => { + it('bounds global concurrency and repeated work for one relay host', async () => { + let now = 0 + const admission = new RelayPublicAssignmentAdmission({ + maxConcurrent: 2, + minIntervalMs: 5_000, + now: () => now + }) + const first = await admission.acquire('host-a') + expect(first).not.toBeNull() + await expect(admission.acquire('host-a')).resolves.toBeNull() + const second = await admission.acquire('host-b') + expect(second).not.toBeNull() + await expect(admission.acquire('host-c')).resolves.toBeNull() + + first?.release() + await expect(admission.acquire('host-a')).resolves.toBeNull() + now = 5_000 + const recovered = await admission.acquire('host-a') + expect(recovered).not.toBeNull() + recovered?.release() + second?.release() + }) + + it('releases capacity once when callers settle more than once', async () => { + const admission = new RelayPublicAssignmentAdmission({ + maxConcurrent: 1, + minIntervalMs: 0, + now: vi.fn(() => 0) + }) + const lease = await admission.acquire('host-a') + lease?.release() + lease?.release() + await expect(admission.acquire('host-b')).resolves.not.toBeNull() + }) + + it('grants ordinary waiters in FIFO order without raising concurrency', async () => { + const admission = new RelayPublicAssignmentAdmission({ + maxConcurrent: 1, + maxQueued: 2, + waitMs: 1_000, + minIntervalMs: 0, + schedule: () => () => undefined + }) + const first = await admission.acquire('host-a') + const secondPromise = admission.acquire('host-b') + const thirdPromise = admission.acquire('host-c') + + await expect(admission.acquire('host-d')).resolves.toBeNull() + first?.release() + const second = await secondPromise + expect(second).not.toBeNull() + let thirdSettled = false + void thirdPromise.then(() => { thirdSettled = true }) + await Promise.resolve() + expect(thirdSettled).toBe(false) + second?.release() + const third = await thirdPromise + expect(third).not.toBeNull() + third?.release() + }) + + it('fails an ordinary wait closed when its bounded wait expires', async () => { + let expire!: () => void + const admission = new RelayPublicAssignmentAdmission({ + maxConcurrent: 1, + maxQueued: 1, + waitMs: 1_000, + minIntervalMs: 0, + schedule: (callback) => { + expire = callback + return () => undefined + } + }) + const first = await admission.acquire('host-a') + const waiting = admission.acquire('host-b') + + expire() + await expect(waiting).resolves.toBeNull() + first?.release() + }) + + it('gives a bounded reserved waiter the next public slot', async () => { + let expire!: () => void + let cancelled = false + const admission = new RelayPublicAssignmentAdmission({ + maxConcurrent: 2, + maxReservedConcurrent: 1, + reservedWaitMs: 1_000, + minIntervalMs: 0, + schedule: (callback) => { + expire = callback + return () => { + cancelled = true + } + } + }) + const first = await admission.acquire('host-a') + const second = await admission.acquire('host-b') + const reservedPromise = admission.acquireReserved('host-c') + + await expect(admission.acquire('host-d')).resolves.toBeNull() + await expect(admission.acquireReserved('host-e')).resolves.toBeNull() + first?.release() + const reserved = await reservedPromise + + expect(reserved).not.toBeNull() + expect(cancelled).toBe(true) + await expect(admission.acquire('host-d')).resolves.toBeNull() + reserved?.release() + await expect(admission.acquire('host-d')).resolves.not.toBeNull() + second?.release() + expire() + }) + + it('lets reserved recovery displace the same host from the ordinary queue', async () => { + const admission = new RelayPublicAssignmentAdmission({ + maxConcurrent: 1, + maxQueued: 1, + waitMs: 1_000, + maxReservedConcurrent: 1, + reservedWaitMs: 1_000, + minIntervalMs: 0, + schedule: () => () => undefined + }) + const active = await admission.acquire('host-a') + const ordinary = admission.acquire('host-b') + const reserved = admission.acquireReserved('host-b') + + await expect(ordinary).resolves.toBeNull() + active?.release() + const recovered = await reserved + expect(recovered).not.toBeNull() + recovered?.release() + }) + + it('fails a reserved wait closed when its bounded wait expires', async () => { + let expire!: () => void + const admission = new RelayPublicAssignmentAdmission({ + maxConcurrent: 2, + maxReservedConcurrent: 1, + reservedWaitMs: 1_000, + minIntervalMs: 0, + schedule: (callback) => { + expire = callback + return () => undefined + } + }) + const first = await admission.acquire('host-a') + const second = await admission.acquire('host-b') + const reserved = admission.acquireReserved('host-c') + + expire() + await expect(reserved).resolves.toBeNull() + first?.release() + await expect(admission.acquire('host-d')).resolves.not.toBeNull() + second?.release() + }) + + it('serializes same-host assignment and reserved recovery without losing priority', async () => { + const admission = new RelayPublicAssignmentAdmission({ + maxConcurrent: 2, + maxReservedConcurrent: 1, + reservedWaitMs: 1_000, + minIntervalMs: 0, + schedule: () => () => undefined + }) + const assignment = await admission.acquire('host-a') + const reservedPromise = admission.acquireReserved('host-a') + + await Promise.resolve() + await expect(admission.acquire('host-b')).resolves.toBeNull() + assignment?.release() + const reserved = await reservedPromise + + expect(reserved).not.toBeNull() + await expect(admission.acquire('host-a')).resolves.toBeNull() + reserved?.release() + await expect(admission.acquire('host-b')).resolves.not.toBeNull() + }) + + it('separates a self-throttled host from a saturated director', async () => { + let now = 0 + const reasons: string[] = [] + const admission = new RelayPublicAssignmentAdmission({ + maxConcurrent: 1, + maxQueued: 0, + minIntervalMs: 5_000, + now: () => now, + onRejected: (reason) => reasons.push(reason) + }) + const held = await admission.acquire('host-a') + + // Same host while its own attempt is still running. + await expect(admission.acquire('host-a')).resolves.toBeNull() + // A different host with the single slot taken and no queue. + await expect(admission.acquire('host-b')).resolves.toBeNull() + held?.release() + // Same host again, now inside its retry penalty rather than in flight. + await expect(admission.acquire('host-a')).resolves.toBeNull() + + expect(reasons).toEqual(['host-in-flight', 'queue-full', 'host-rate-limited']) + + now = 5_000 + const recovered = await admission.acquire('host-a') + expect(recovered).not.toBeNull() + expect(reasons).toHaveLength(3) + recovered?.release() + }) + + it('reports a timed-out wait separately from a full queue', async () => { + let expire!: () => void + const reasons: string[] = [] + const admission = new RelayPublicAssignmentAdmission({ + maxConcurrent: 1, + maxQueued: 1, + waitMs: 1_000, + minIntervalMs: 0, + schedule: (callback) => { + expire = callback + return () => undefined + }, + onRejected: (reason) => reasons.push(reason) + }) + const first = await admission.acquire('host-a') + const waiting = admission.acquire('host-b') + + expire() + await expect(waiting).resolves.toBeNull() + expect(reasons).toEqual(['wait-timeout']) + first?.release() + }) + + it('reports a displaced queue entry as superseded, not as a timeout', async () => { + const reasons: string[] = [] + const admission = new RelayPublicAssignmentAdmission({ + maxConcurrent: 1, + maxQueued: 1, + waitMs: 1_000, + maxReservedConcurrent: 1, + reservedWaitMs: 1_000, + minIntervalMs: 0, + schedule: () => () => undefined, + onRejected: (reason) => reasons.push(reason) + }) + const active = await admission.acquire('host-a') + const ordinary = admission.acquire('host-b') + const reserved = admission.acquireReserved('host-b') + + await expect(ordinary).resolves.toBeNull() + expect(reasons).toEqual(['superseded']) + active?.release() + const recovered = await reserved + expect(recovered).not.toBeNull() + recovered?.release() + }) + + it('distinguishes a busy reserved lane from a reserved host already in flight', async () => { + const reasons: string[] = [] + const admission = new RelayPublicAssignmentAdmission({ + maxConcurrent: 4, + // Two slots so the lane still has room when the same host asks twice. + maxReservedConcurrent: 2, + reservedWaitMs: 1_000, + minIntervalMs: 0, + schedule: () => () => undefined, + onRejected: (reason) => reasons.push(reason) + }) + const reserved = await admission.acquireReserved('host-a') + expect(reserved).not.toBeNull() + + await expect(admission.acquireReserved('host-a')).resolves.toBeNull() + expect(reasons).toEqual(['host-in-flight']) + + const second = await admission.acquireReserved('host-b') + expect(second).not.toBeNull() + await expect(admission.acquireReserved('host-c')).resolves.toBeNull() + + expect(reasons).toEqual(['host-in-flight', 'reserved-unavailable']) + reserved?.release() + second?.release() + }) + + it('tells each caller its own rejection reason without crossing requests', async () => { + const admission = new RelayPublicAssignmentAdmission({ + maxConcurrent: 1, + maxQueued: 1, + waitMs: 1_000, + maxReservedConcurrent: 1, + reservedWaitMs: 1_000, + minIntervalMs: 5_000, + now: () => 0, + schedule: () => () => undefined + }) + const queued: string[] = [] + const displaced: string[] = [] + const throttled: string[] = [] + + const active = await admission.acquire('host-a') + // 'host-b' waits, then reserved recovery for the same host displaces it. + const ordinary = admission.acquire('host-b', (reason) => displaced.push(reason)) + const reserved = admission.acquireReserved('host-b', (reason) => queued.push(reason)) + await expect(ordinary).resolves.toBeNull() + active?.release() + const recovered = await reserved + recovered?.release() + await admission.acquire('host-a', (reason) => throttled.push(reason)) + + expect(displaced).toEqual(['superseded']) + expect(queued).toEqual([]) + expect(throttled).toEqual(['host-rate-limited']) + }) + + it('stays silent on every granted acquisition', async () => { + const reasons: string[] = [] + const admission = new RelayPublicAssignmentAdmission({ + maxConcurrent: 2, + maxReservedConcurrent: 1, + minIntervalMs: 0, + onRejected: (reason) => reasons.push(reason) + }) + const first = await admission.acquire('host-a') + const second = await admission.acquireReserved('host-b') + + expect(first).not.toBeNull() + expect(second).not.toBeNull() + expect(reasons).toEqual([]) + first?.release() + second?.release() + }) +}) diff --git a/cloud/apps/relay/src/public-assignment-admission.ts b/cloud/apps/relay/src/public-assignment-admission.ts new file mode 100644 index 00000000000..bd6d7f1a488 --- /dev/null +++ b/cloud/apps/relay/src/public-assignment-admission.ts @@ -0,0 +1,239 @@ +type AssignmentAdmissionLease = { release(): void } +type CancelWait = () => void +// Per-call sink: acquire() keeps returning null so callers stay unchanged, and the +// reason rides out of band to whoever made this particular request. +type RejectionSink = (reason: AssignmentAdmissionRejection) => void +type PendingAssignment = { + relayHostId: string + resolve: (lease: AssignmentAdmissionLease | null) => void + cancelWait: CancelWait + notifyRejected?: RejectionSink +} + +// Every rejection here becomes an identical 503, so without the reason a busy +// director and a self-throttling host are indistinguishable in production. +export type AssignmentAdmissionRejection = + | 'host-in-flight' + | 'host-rate-limited' + | 'queue-full' + | 'wait-timeout' + | 'superseded' + | 'reserved-unavailable' + +const MAX_TRACKED_HOSTS = 4_096 + +export class RelayPublicAssignmentAdmission { + private active = 0 + private activeReserved = 0 + private readonly activeAssignmentHosts = new Set() + private readonly activeReservedHosts = new Set() + private readonly queuedAssignmentHosts = new Set() + private readonly lastAttemptByHost = new Map() + private readonly lastReservedAttemptByHost = new Map() + private readonly pendingAssignments: PendingAssignment[] = [] + private pendingReserved: PendingAssignment | undefined + + constructor( + private readonly options: { + maxConcurrent: number + maxQueued?: number + waitMs?: number + maxReservedConcurrent?: number + reservedWaitMs?: number + minIntervalMs: number + now?: () => number + schedule?: (callback: () => void, delayMs: number) => CancelWait + onRejected?: (reason: AssignmentAdmissionRejection) => void + } + ) {} + + async acquire( + relayHostId: string, + notifyRejected?: RejectionSink + ): Promise { + const now = (this.options.now ?? Date.now)() + const lastAttempt = this.lastAttemptByHost.get(relayHostId) + if ( + this.activeAssignmentHosts.has(relayHostId) || + this.activeReservedHosts.has(relayHostId) || + this.queuedAssignmentHosts.has(relayHostId) || + this.pendingReserved?.relayHostId === relayHostId + ) { + return this.reject('host-in-flight', notifyRejected) + } + if (lastAttempt !== undefined && now - lastAttempt < this.options.minIntervalMs) { + return this.reject('host-rate-limited', notifyRejected) + } + if ( + this.active < this.options.maxConcurrent && + this.pendingReserved === undefined && + this.pendingAssignments.length === 0 + ) { + this.recordAttempt(this.lastAttemptByHost, relayHostId, now) + return this.createLease(relayHostId, false) + } + if (this.pendingAssignments.length >= (this.options.maxQueued ?? 0)) { + return this.reject('queue-full', notifyRejected) + } + + return await new Promise((resolve) => { + const schedule = this.options.schedule ?? defaultSchedule + let cancelWait: CancelWait = () => undefined + const pending: PendingAssignment = { + relayHostId, + resolve, + cancelWait: () => cancelWait(), + notifyRejected + } + this.pendingAssignments.push(pending) + this.queuedAssignmentHosts.add(relayHostId) + cancelWait = schedule(() => { + const index = this.pendingAssignments.indexOf(pending) + if (index === -1) return + this.pendingAssignments.splice(index, 1) + this.queuedAssignmentHosts.delete(relayHostId) + resolve(this.reject('wait-timeout', notifyRejected)) + }, this.options.waitMs ?? 1_000) + }) + } + + async acquireReserved( + relayHostId: string, + notifyRejected?: RejectionSink + ): Promise { + const maxReservedConcurrent = this.options.maxReservedConcurrent ?? 0 + const now = (this.options.now ?? Date.now)() + const lastAttempt = this.lastReservedAttemptByHost.get(relayHostId) + if (maxReservedConcurrent === 0 || this.activeReserved >= maxReservedConcurrent) { + return this.reject('reserved-unavailable', notifyRejected) + } + if (this.activeReservedHosts.has(relayHostId)) { + return this.reject('host-in-flight', notifyRejected) + } + if (this.pendingReserved !== undefined) { + return this.reject('reserved-unavailable', notifyRejected) + } + if (lastAttempt !== undefined && now - lastAttempt < this.options.minIntervalMs) { + return this.reject('host-rate-limited', notifyRejected) + } + this.cancelQueuedAssignment(relayHostId) + if ( + this.active < this.options.maxConcurrent && + !this.activeAssignmentHosts.has(relayHostId) + ) { + this.recordAttempt(this.lastReservedAttemptByHost, relayHostId, now) + return this.createLease(relayHostId, true) + } + + return await new Promise((resolve) => { + const schedule = this.options.schedule ?? defaultSchedule + let cancelWait: CancelWait = () => undefined + this.pendingReserved = { + relayHostId, + resolve, + cancelWait: () => cancelWait(), + notifyRejected + } + cancelWait = schedule(() => { + if (this.pendingReserved?.relayHostId !== relayHostId) return + this.pendingReserved = undefined + resolve(this.reject('wait-timeout', notifyRejected)) + this.grantPendingAssignments() + }, this.options.reservedWaitMs ?? 1_000) + }) + } + + private createLease(relayHostId: string, reserved: boolean): AssignmentAdmissionLease { + this.active++ + if (reserved) { + this.activeReserved++ + this.activeReservedHosts.add(relayHostId) + } else { + this.activeAssignmentHosts.add(relayHostId) + } + let released = false + return { + release: () => { + if (released) return + released = true + this.active = Math.max(0, this.active - 1) + if (reserved) { + this.activeReserved = Math.max(0, this.activeReserved - 1) + this.activeReservedHosts.delete(relayHostId) + } else { + this.activeAssignmentHosts.delete(relayHostId) + } + this.grantPendingReserved() + this.grantPendingAssignments() + } + } + } + + private grantPendingReserved(): void { + const pending = this.pendingReserved + if ( + !pending || + this.active >= this.options.maxConcurrent || + this.activeAssignmentHosts.has(pending.relayHostId) + ) { + return + } + this.pendingReserved = undefined + pending.cancelWait() + this.recordAttempt( + this.lastReservedAttemptByHost, + pending.relayHostId, + (this.options.now ?? Date.now)() + ) + pending.resolve(this.createLease(pending.relayHostId, true)) + } + + private grantPendingAssignments(): void { + while ( + this.pendingReserved === undefined && + this.active < this.options.maxConcurrent && + this.pendingAssignments.length > 0 + ) { + const pending = this.pendingAssignments.shift()! + this.queuedAssignmentHosts.delete(pending.relayHostId) + pending.cancelWait() + this.recordAttempt( + this.lastAttemptByHost, + pending.relayHostId, + (this.options.now ?? Date.now)() + ) + pending.resolve(this.createLease(pending.relayHostId, false)) + } + } + + private cancelQueuedAssignment(relayHostId: string): void { + const index = this.pendingAssignments.findIndex( + (pending) => pending.relayHostId === relayHostId + ) + if (index === -1) return + const [pending] = this.pendingAssignments.splice(index, 1) + this.queuedAssignmentHosts.delete(relayHostId) + pending?.cancelWait() + // The sink rides on the pending record: this rejects a different caller's request. + pending?.resolve(this.reject('superseded', pending.notifyRejected)) + } + + private reject(reason: AssignmentAdmissionRejection, notifyRejected?: RejectionSink): null { + this.options.onRejected?.(reason) + notifyRejected?.(reason) + return null + } + + private recordAttempt(attempts: Map, relayHostId: string, now: number): void { + attempts.delete(relayHostId) + attempts.set(relayHostId, now) + if (attempts.size > MAX_TRACKED_HOSTS) { + attempts.delete(attempts.keys().next().value!) + } + } +} + +function defaultSchedule(callback: () => void, delayMs: number): CancelWait { + const timer = setTimeout(callback, delayMs) + return () => clearTimeout(timer) +} diff --git a/cloud/apps/relay/src/public-assignment-circuit-breaker.test.ts b/cloud/apps/relay/src/public-assignment-circuit-breaker.test.ts new file mode 100644 index 00000000000..94be0536497 --- /dev/null +++ b/cloud/apps/relay/src/public-assignment-circuit-breaker.test.ts @@ -0,0 +1,55 @@ +import { describe, expect, it, vi } from 'vitest' +import { createRelayApp } from './app.js' +import type { RelayConfig } from './config.js' + +describe('public assignment circuit breaker', () => { + it('rejects assign and resolve without invoking relay state', async () => { + const assign = vi.fn() + const resolveResume = vi.fn() + const app = createRelayApp(config(), { + store: { resolveResume } as never, + assignments: { assign } as never, + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + + for (const path of ['/v1/assign', '/v1/resolve']) { + const response = await app.request(path, { method: 'POST' }) + expect(response.status).toBe(503) + expect(response.headers.get('retry-after')).toBe('5') + expect(await response.json()).toEqual({ error: 'assignments_temporarily_unavailable' }) + } + expect(assign).not.toHaveBeenCalled() + expect(resolveResume).not.toHaveBeenCalled() + expect((await app.request('/health')).status).toBe(200) + expect((await app.request('/v1/admin/drain', { method: 'POST' })).status).toBe(401) + }) +}) + +function config(): RelayConfig { + return { + port: 8080, + publicUrl: 'https://relay.example.test', + cellUrl: 'https://relay.example.test', + authIssuer: 'https://auth.example.test', + authAudience: 'orca-relay', + jwksUrl: 'https://auth.example.test/jwks', + assignmentSigningKey: new TextEncoder().encode('assignment-key-with-at-least-32-bytes'), + role: 'director', + cellId: 'director', + cells: [], + adminAudience: 'https://relay.example.test/v1/admin/drain', + deployServiceAccount: 'deploy@example.test', + runtimeServiceAccount: 'runtime@example.test', + adminJwksUrl: 'https://auth.example.test/jwks', + databasePoolMax: 3, + publicAssignmentsEnabled: false, + publicAssignmentConcurrency: 2, + publicAssignmentQueueMax: 128, + publicAssignmentWaitMs: 4_000, + publicResolveConcurrency: 1, + publicResolveWaitMs: 5_000, + publicAssignmentRetryAfterSeconds: 5, + dataDir: './data' + } +} diff --git a/cloud/apps/relay/src/public-assignment-overload.test.ts b/cloud/apps/relay/src/public-assignment-overload.test.ts new file mode 100644 index 00000000000..2d0426893c4 --- /dev/null +++ b/cloud/apps/relay/src/public-assignment-overload.test.ts @@ -0,0 +1,236 @@ +import { describe, expect, it, vi } from 'vitest' +import type { RelayAssignment } from './assignment-store.js' +import type { RelayConfig } from './config.js' + +const fakes = vi.hoisted(() => ({ + verifyRelayToken: vi.fn(async (token: string) => ({ + sub: 'user-1', + relayHostId: token + })) +})) + +vi.mock('./relay-token-verifier.js', () => ({ + createRelayTokenVerifier: () => fakes.verifyRelayToken, + readBearer: (value: string | undefined) => value?.replace(/^Bearer /, '') ?? null +})) + +import { createRelayApp } from './app.js' + +describe('public assignment overload', () => { + it('rejects duplicate and excess work before it reaches assignment state', async () => { + const hostA = 'aaaaaaaaaaaaaaaa' + const hostB = 'bbbbbbbbbbbbbbbb' + const hostC = 'cccccccccccccccc' + const pending = new Map>>() + const assign = vi.fn(async ({ relayHostId }: { relayHostId: string }) => { + const operation = deferred() + pending.set(relayHostId, operation) + return await operation.promise + }) + const app = createRelayApp(config({ publicAssignmentQueueMax: 0 }), { + store: {} as never, + assignments: { assign } as never, + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + + const first = app.request('/v1/assign', assignmentRequest(hostA)) + await vi.waitFor(() => expect(pending.has(hostA)).toBe(true)) + const duplicate = await app.request('/v1/assign', assignmentRequest(hostA)) + expect(duplicate.status).toBe(503) + + const second = app.request('/v1/assign', assignmentRequest(hostB)) + await vi.waitFor(() => expect(pending.has(hostB)).toBe(true)) + const excess = await app.request('/v1/assign', assignmentRequest(hostC)) + expect(excess.status).toBe(503) + expect(excess.headers.get('retry-after')).toBe('5') + expect(assign).toHaveBeenCalledTimes(2) + + pending.get(hostA)?.resolve(assignment('cell-a', hostA)) + pending.get(hostB)?.resolve(assignment('cell-b', hostB)) + expect((await first).status).toBe(200) + expect((await second).status).toBe(200) + expect((await app.request('/v1/assign', assignmentRequest(hostA))).status).toBe(503) + expect(assign).toHaveBeenCalledTimes(2) + }) + + it('fairly drains a bounded incident-scale burst without raising database concurrency', async () => { + const pending = new Map>>() + const admittedHosts: string[] = [] + let active = 0 + let highWater = 0 + const assign = vi.fn(async ({ relayHostId }: { relayHostId: string }) => { + admittedHosts.push(relayHostId) + active++ + highWater = Math.max(highWater, active) + const operation = deferred() + pending.set(relayHostId, operation) + try { + return await operation.promise + } finally { + active-- + pending.delete(relayHostId) + } + }) + const app = createRelayApp(config(), { + store: {} as never, + assignments: { assign } as never, + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + const hosts = Array.from( + { length: 430 }, + (_, index) => `host-${String(index).padStart(11, '0')}` + ) + + const firstWave = hosts.map((host) => + app.request('/v1/assign', assignmentRequest(host)) + ) + await vi.waitFor(() => expect(assign).toHaveBeenCalledTimes(2)) + const retryWave = await Promise.all( + hosts.map((host) => app.request('/v1/assign', assignmentRequest(host))) + ) + + expect(retryWave.every((response) => response.status === 503)).toBe(true) + expect(assign).toHaveBeenCalledTimes(2) + while (assign.mock.calls.length < 130) { + const activeOperations = [...pending] + const expectedCalls = assign.mock.calls.length + activeOperations.length + for (const [host, operation] of activeOperations) { + operation.resolve(assignment(`cell-${host}`, host)) + } + await vi.waitFor(() => expect(assign).toHaveBeenCalledTimes(expectedCalls)) + } + for (const [host, operation] of pending) { + operation.resolve(assignment(`cell-${host}`, host)) + } + const firstResponses = await Promise.all(firstWave) + expect(firstResponses.filter((response) => response.status === 200)).toHaveLength(130) + expect(firstResponses.filter((response) => response.status === 503)).toHaveLength(300) + expect(admittedHosts).toEqual(hosts.slice(0, 130)) + expect(highWater).toBe(2) + }) + + it('reserves one bounded resolve lane while assignment work is saturated', async () => { + const pendingAssignments = new Map>>() + const pendingResolve = deferred<{ userId: string; relayDeviceId: string } | null>() + const assign = vi.fn(async ({ relayHostId }: { relayHostId: string }) => { + const operation = deferred() + pendingAssignments.set(relayHostId, operation) + return await operation.promise + }) + const resolveResume = vi.fn(async () => await pendingResolve.promise) + const resolve = vi.fn(async ({ relayHostId }: { relayHostId: string }) => + assignment('target-cell', relayHostId) + ) + const app = createRelayApp(config({ publicAssignmentQueueMax: 0 }), { + store: { resolveResume } as never, + assignments: { assign, resolve } as never, + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + + const assignmentA = app.request('/v1/assign', assignmentRequest('aaaaaaaaaaaaaaaa')) + const assignmentB = app.request('/v1/assign', assignmentRequest('bbbbbbbbbbbbbbbb')) + await vi.waitFor(() => expect(assign).toHaveBeenCalledTimes(2)) + + const recovery = app.request('/v1/resolve', resolveRequest('cccccccccccccccc', 1)) + await Promise.resolve() + expect(resolveResume).not.toHaveBeenCalled() + const excessResolve = await app.request( + '/v1/resolve', + resolveRequest('dddddddddddddddd', 2) + ) + const excessAssign = await app.request('/v1/assign', assignmentRequest('eeeeeeeeeeeeeeee')) + + expect(excessResolve.status).toBe(503) + expect(excessAssign.status).toBe(503) + expect(assign).toHaveBeenCalledTimes(2) + expect(resolveResume).not.toHaveBeenCalled() + + pendingAssignments + .get('aaaaaaaaaaaaaaaa') + ?.resolve(assignment('cell-a', 'aaaaaaaaaaaaaaaa')) + expect((await assignmentA).status).toBe(200) + await vi.waitFor(() => expect(resolveResume).toHaveBeenCalledTimes(1)) + pendingResolve.resolve({ userId: 'user-1', relayDeviceId: 'device-1' }) + expect((await recovery).status).toBe(200) + expect(resolve).toHaveBeenCalledTimes(1) + + for (const [host, operation] of pendingAssignments) { + operation.resolve(assignment(`cell-${host}`, host)) + } + expect((await assignmentB).status).toBe(200) + }) +}) + +function assignmentRequest(relayHostId: string): RequestInit { + return { + method: 'POST', + headers: { + authorization: `Bearer ${relayHostId}`, + 'content-type': 'application/json' + }, + body: JSON.stringify({ v: 1, relayHostId }) + } +} + +function resolveRequest(relayHostId: string, fill: number): RequestInit { + return { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + v: 1, + relayHostId, + resumeToken: Buffer.alloc(32, fill).toString('base64url') + }) + } +} + +function assignment(cellId: string, relayHostId: string): RelayAssignment { + return { + userId: 'user-1', + relayHostId, + cellId, + cellUrl: `https://${cellId}.relay.example.test`, + assignmentEpoch: 1, + leaseExpiresAt: Date.now() + 300_000 + } +} + +function deferred() { + let resolve!: (value: T) => void + const promise = new Promise((resolvePromise) => { + resolve = resolvePromise + }) + return { promise, resolve } +} + +function config(overrides: Partial = {}): RelayConfig { + return { + port: 8080, + publicUrl: 'https://relay.example.test', + cellUrl: 'https://relay.example.test', + authIssuer: 'https://auth.example.test', + authAudience: 'orca-relay', + jwksUrl: 'https://auth.example.test/jwks', + assignmentSigningKey: new TextEncoder().encode('assignment-key-with-at-least-32-bytes'), + role: 'director', + cellId: 'director', + cells: [], + adminAudience: 'https://relay.example.test/v1/admin/drain', + deployServiceAccount: 'deploy@example.test', + runtimeServiceAccount: 'runtime@example.test', + adminJwksUrl: 'https://auth.example.test/jwks', + databasePoolMax: 3, + publicAssignmentsEnabled: true, + publicAssignmentConcurrency: 2, + publicAssignmentQueueMax: 128, + publicAssignmentWaitMs: 4_000, + publicResolveConcurrency: 1, + publicResolveWaitMs: 5_000, + publicAssignmentRetryAfterSeconds: 5, + dataDir: './data', + ...overrides + } +} diff --git a/cloud/apps/relay/src/public-assignment-reconnect-lane.test.ts b/cloud/apps/relay/src/public-assignment-reconnect-lane.test.ts new file mode 100644 index 00000000000..145d2c737cb --- /dev/null +++ b/cloud/apps/relay/src/public-assignment-reconnect-lane.test.ts @@ -0,0 +1,202 @@ +import { describe, expect, it, vi } from 'vitest' +import type { RelayAssignment } from './assignment-store.js' +import type { RelayConfig } from './config.js' + +const fakes = vi.hoisted(() => ({ + verifyRelayToken: vi.fn(async (token: string) => ({ + sub: 'user-1', + relayHostId: token + })) +})) + +vi.mock('./relay-token-verifier.js', () => ({ + createRelayTokenVerifier: () => fakes.verifyRelayToken, + readBearer: (value: string | undefined) => value?.replace(/^Bearer /, '') ?? null +})) + +import { createRelayApp } from './app.js' + +describe('public assignment reconnect lane', () => { + it('admits a verified reconnect past a saturated placement queue', async () => { + const reconnecting = 'rrrrrrrrrrrrrrrr' + const newcomerA = 'aaaaaaaaaaaaaaaa' + const newcomerB = 'bbbbbbbbbbbbbbbb' + const blocked = 'cccccccccccccccc' + const pending = new Map>>() + const assign = vi.fn(async ({ relayHostId }: { relayHostId: string }) => { + if (relayHostId === reconnecting) return assignment('cell-r', reconnecting) + const operation = deferred() + pending.set(relayHostId, operation) + return await operation.promise + }) + const resolve = vi.fn(async ({ relayHostId }: { relayHostId: string }) => + relayHostId === reconnecting ? assignment('cell-r', reconnecting) : null + ) + const outcomes: string[] = [] + const app = createRelayApp(config({ publicAssignmentQueueMax: 0 }), { + store: {} as never, + assignments: { assign, resolve } as never, + drain: vi.fn(), + ready: vi.fn(async () => true), + recordAssignmentAdmission: (outcome) => outcomes.push(outcome) + }) + + // Saturate the placement lane with two in-flight newcomers, zero queue. + const first = app.request('/v1/assign', assignmentRequest(newcomerA)) + await vi.waitFor(() => expect(pending.has(newcomerA)).toBe(true)) + const second = app.request('/v1/assign', assignmentRequest(newcomerB)) + await vi.waitFor(() => expect(pending.has(newcomerB)).toBe(true)) + expect((await app.request('/v1/assign', assignmentRequest(blocked))).status).toBe(503) + + const reconnected = await app.request( + '/v1/assign', + assignmentRequest(reconnecting, { reconnect: true }) + ) + + expect(reconnected.status).toBe(200) + expect(resolve).toHaveBeenCalledWith({ userId: 'user-1', relayHostId: reconnecting }) + expect(assign).toHaveBeenCalledWith({ userId: 'user-1', relayHostId: reconnecting }) + expect(outcomes).toContain('sticky') + + pending.get(newcomerA)?.resolve(assignment('cell-a', newcomerA)) + pending.get(newcomerB)?.resolve(assignment('cell-b', newcomerB)) + expect((await first).status).toBe(200) + expect((await second).status).toBe(200) + }) + + it('rate-limits repeat fast-lane attempts with the short retry-after', async () => { + const host = 'rrrrrrrrrrrrrrrr' + const assign = vi.fn(async () => assignment('cell-r', host)) + const resolve = vi.fn(async () => assignment('cell-r', host)) + const app = createRelayApp(config(), { + store: {} as never, + assignments: { assign, resolve } as never, + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + + expect( + (await app.request('/v1/assign', assignmentRequest(host, { reconnect: true }))).status + ).toBe(200) + const repeat = await app.request('/v1/assign', assignmentRequest(host, { reconnect: true })) + + expect(repeat.status).toBe(503) + expect(repeat.headers.get('retry-after')).toBe('2') + expect(resolve).toHaveBeenCalledTimes(1) + expect(assign).toHaveBeenCalledTimes(1) + }) + + it('sends an unverified reconnect hint through the placement lane unchanged', async () => { + const host = 'nnnnnnnnnnnnnnnn' + const assign = vi.fn(async () => assignment('cell-n', host)) + const resolve = vi.fn(async () => null) + const outcomes: string[] = [] + const app = createRelayApp(config(), { + store: {} as never, + assignments: { assign, resolve } as never, + drain: vi.fn(), + ready: vi.fn(async () => true), + recordAssignmentAdmission: (outcome) => outcomes.push(outcome) + }) + + const response = await app.request('/v1/assign', assignmentRequest(host, { reconnect: true })) + + expect(response.status).toBe(200) + expect(assign).toHaveBeenCalledTimes(1) + expect(outcomes).toEqual(['placement']) + }) + + it('rejects on the fast lane when the verification probe fails transiently', async () => { + const host = 'tttttttttttttttt' + const assign = vi.fn() + const resolve = vi.fn(async () => { + throw Object.assign(new Error('deadlock detected'), { code: '40P01' }) + }) + const app = createRelayApp(config(), { + store: {} as never, + assignments: { assign, resolve } as never, + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + + const response = await app.request('/v1/assign', assignmentRequest(host, { reconnect: true })) + + expect(response.status).toBe(503) + expect(response.headers.get('retry-after')).toBe('2') + expect(assign).not.toHaveBeenCalled() + }) + + it('never probes for unhinted requests', async () => { + const host = 'uuuuuuuuuuuuuuuu' + const assign = vi.fn(async () => assignment('cell-u', host)) + const resolve = vi.fn() + const app = createRelayApp(config(), { + store: {} as never, + assignments: { assign, resolve } as never, + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + + expect((await app.request('/v1/assign', assignmentRequest(host))).status).toBe(200) + expect(resolve).not.toHaveBeenCalled() + }) +}) + +function assignmentRequest(relayHostId: string, extra: { reconnect?: boolean } = {}): RequestInit { + return { + method: 'POST', + headers: { + authorization: `Bearer ${relayHostId}`, + 'content-type': 'application/json' + }, + body: JSON.stringify({ v: 1, relayHostId, ...extra }) + } +} + +function assignment(cellId: string, relayHostId: string): RelayAssignment { + return { + userId: 'user-1', + relayHostId, + cellId, + cellUrl: `https://${cellId}.relay.example.test`, + assignmentEpoch: 1, + leaseExpiresAt: Date.now() + 300_000 + } +} + +function deferred() { + let resolve!: (value: T) => void + const promise = new Promise((resolvePromise) => { + resolve = resolvePromise + }) + return { promise, resolve } +} + +function config(overrides: Partial = {}): RelayConfig { + return { + port: 8080, + publicUrl: 'https://relay.example.test', + cellUrl: 'https://relay.example.test', + authIssuer: 'https://auth.example.test', + authAudience: 'orca-relay', + jwksUrl: 'https://auth.example.test/jwks', + assignmentSigningKey: new TextEncoder().encode('assignment-key-with-at-least-32-bytes'), + role: 'director', + cellId: 'director', + cells: [], + adminAudience: 'https://relay.example.test/v1/admin/drain', + deployServiceAccount: 'deploy@example.test', + runtimeServiceAccount: 'runtime@example.test', + adminJwksUrl: 'https://auth.example.test/jwks', + databasePoolMax: 3, + publicAssignmentsEnabled: true, + publicAssignmentConcurrency: 2, + publicAssignmentQueueMax: 128, + publicAssignmentWaitMs: 4_000, + publicResolveConcurrency: 1, + publicResolveWaitMs: 5_000, + publicAssignmentRetryAfterSeconds: 5, + dataDir: './data', + ...overrides + } +} diff --git a/cloud/apps/relay/src/regional-host-drain-app.test.ts b/cloud/apps/relay/src/regional-host-drain-app.test.ts new file mode 100644 index 00000000000..1cd34902520 --- /dev/null +++ b/cloud/apps/relay/src/regional-host-drain-app.test.ts @@ -0,0 +1,536 @@ +import { describe, expect, it, vi } from 'vitest' +import type { RelayConfig } from './config.js' + +vi.mock('./admin-token-verifier.js', () => ({ + createAdminTokenVerifier: () => async (token: string, route?: string) => + token === 'deploy-token' || + (token === 'monitor-token' && + (!route || route === '/v1/admin/regional-rehome-control')), + createReadOnlyAdminTokenVerifier: () => async () => false, + createRegionalRehomeControlApplyTokenVerifier: () => async (token: string) => + token === 'deploy-token', + createRegionalRehomeRuntimeTokenVerifier: () => async (token: string) => + token === 'runtime-token', + createRegionalRehomeTokenVerifier: () => async (token: string) => token === 'rehome-token', + createRuntimeTokenVerifier: () => async (token: string) => token === 'runtime-token' +})) + +vi.mock('./relay-token-verifier.js', () => ({ + createRelayTokenVerifier: () => async () => null, + readBearer: (value: string | undefined) => value?.replace(/^Bearer /, '') ?? null +})) + +import { createRelayApp } from './app.js' +import { RelayObservability } from './relay-observability.js' +import { emptyPostgresPoolPressureCounts } from './postgres-pool-pressure.js' +import { + REGIONAL_REHOME_TRUST_PROBE_ATTEMPT_ID, + REGIONAL_REHOME_TRUST_PROBE_HOST_ID, + REGIONAL_REHOME_TRUST_PROBE_USER_ID +} from './regional-rehome-trust-probe.js' + +const cellIncarnation = '11111111-1111-4111-8111-111111111111' +const request = { + v: 1, + attemptId: '22222222-2222-4222-8222-222222222222', + userId: 'user-1', + relayHostId: 'abcdefghijklmnop', + sourceCellId: 'production-gce-c7', + sourceCellIncarnation: cellIncarnation, + sourceAssignmentEpoch: 7, + graceMs: 60_000 +} + +describe('regional host drain endpoint', () => { + it('accepts only the dedicated identity and exact cell generation', async () => { + const drainHost = vi.fn(() => 'accepted' as const) + const app = createRelayApp(config(), { + store: {} as never, + assignments: {} as never, + drain: vi.fn(), + drainHost, + cellIncarnation, + ready: vi.fn(async () => true) + }) + + const accepted = await post(app, 'rehome-token', request) + expect(accepted.status).toBe(200) + expect(await accepted.json()).toEqual({ v: 1, outcome: 'accepted' }) + expect(drainHost).toHaveBeenCalledWith(request) + + expect((await post(app, 'deploy-token', request)).status).toBe(401) + expect( + (await post(app, 'rehome-token', { + ...request, + sourceCellIncarnation: '33333333-3333-4333-8333-333333333333' + })).status + ).toBe(409) + expect(drainHost).toHaveBeenCalledOnce() + }) + + it('rejects malformed identities before touching the session registry', async () => { + const drainHost = vi.fn(() => 'accepted' as const) + const app = createRelayApp(config(), { + store: {} as never, + assignments: {} as never, + drain: vi.fn(), + drainHost, + cellIncarnation, + ready: vi.fn(async () => true) + }) + + expect( + (await post(app, 'rehome-token', { ...request, relayHostId: 'raw-host-id' })).status + ).toBe(400) + expect(drainHost).not.toHaveBeenCalled() + }) + + it('is unavailable on the director', async () => { + const app = createRelayApp(config({ role: 'director', cellId: 'director' }), { + store: {} as never, + assignments: {} as never, + drain: vi.fn(), + drainHost: vi.fn(() => 'accepted' as const), + cellIncarnation, + ready: vi.fn(async () => true) + }) + expect((await post(app, 'rehome-token', request)).status).toBe(404) + }) + + it('proves the shared runtime identity is rejected without touching a session', async () => { + const drainHost = vi.fn(() => 'host-not-connected' as const) + const app = createRelayApp(config(), { + store: {} as never, + assignments: {} as never, + drain: vi.fn(), + drainHost, + regionalRehomeTrustProbeHostExists: vi.fn(() => false), + regionalRehomeIdentityToken: vi.fn(async () => 'runtime-token'), + cellIncarnation, + ready: vi.fn(async () => true) + }) + const probe = { + ...request, + attemptId: REGIONAL_REHOME_TRUST_PROBE_ATTEMPT_ID, + userId: REGIONAL_REHOME_TRUST_PROBE_USER_ID, + relayHostId: REGIONAL_REHOME_TRUST_PROBE_HOST_ID, + sourceAssignmentEpoch: 1, + graceMs: 0 + } + + expect(REGIONAL_REHOME_TRUST_PROBE_HOST_ID).toHaveLength(16) + for (let call = 0; call < 2; call++) { + const response = await post(app, 'rehome-token', probe) + expect(response.status).toBe(200) + expect(await response.json()).toEqual({ + v: 1, + outcome: 'host-not-connected', + sharedRuntimeIdentityRejected: true + }) + } + expect(drainHost).toHaveBeenCalledTimes(2) + }) + + it('fails closed if the synthetic host is not provably absent', async () => { + const drainHost = vi.fn(() => 'host-not-connected' as const) + const app = createRelayApp(config(), { + store: {} as never, + assignments: {} as never, + drain: vi.fn(), + drainHost, + regionalRehomeTrustProbeHostExists: vi.fn(() => true), + regionalRehomeIdentityToken: vi.fn(async () => 'runtime-token'), + cellIncarnation, + ready: vi.fn(async () => true) + }) + const response = await post(app, 'rehome-token', { + ...request, + attemptId: REGIONAL_REHOME_TRUST_PROBE_ATTEMPT_ID, + userId: REGIONAL_REHOME_TRUST_PROBE_USER_ID, + relayHostId: REGIONAL_REHOME_TRUST_PROBE_HOST_ID, + sourceAssignmentEpoch: 1, + graceMs: 0 + }) + + expect(response.status).toBe(409) + expect(drainHost).not.toHaveBeenCalled() + }) + + it('rechecks synthetic host absence after asynchronous identity proof', async () => { + let hostExists = false + const drainHost = vi.fn(() => 'host-not-connected' as const) + const app = createRelayApp(config(), { + store: {} as never, + assignments: {} as never, + drain: vi.fn(), + drainHost, + regionalRehomeTrustProbeHostExists: vi.fn(() => hostExists), + regionalRehomeIdentityToken: vi.fn(async () => { + hostExists = true + return 'runtime-token' + }), + cellIncarnation, + ready: vi.fn(async () => true) + }) + const response = await post(app, 'rehome-token', { + ...request, + attemptId: REGIONAL_REHOME_TRUST_PROBE_ATTEMPT_ID, + userId: REGIONAL_REHOME_TRUST_PROBE_USER_ID, + relayHostId: REGIONAL_REHOME_TRUST_PROBE_HOST_ID, + sourceAssignmentEpoch: 1, + graceMs: 0 + }) + + expect(response.status).toBe(409) + expect(drainHost).not.toHaveBeenCalled() + }) + + it('fails closed when the cell cannot prove its runtime identity rejection', async () => { + const drainHost = vi.fn(() => 'host-not-connected' as const) + const app = createRelayApp(config(), { + store: {} as never, + assignments: {} as never, + drain: vi.fn(), + drainHost, + regionalRehomeTrustProbeHostExists: vi.fn(() => false), + regionalRehomeIdentityToken: vi.fn(async () => 'rehome-token'), + cellIncarnation, + ready: vi.fn(async () => true) + }) + const response = await post(app, 'rehome-token', { + ...request, + attemptId: REGIONAL_REHOME_TRUST_PROBE_ATTEMPT_ID, + userId: REGIONAL_REHOME_TRUST_PROBE_USER_ID, + relayHostId: REGIONAL_REHOME_TRUST_PROBE_HOST_ID, + sourceAssignmentEpoch: 1, + graceMs: 0 + }) + + expect(response.status).toBe(409) + expect(drainHost).not.toHaveBeenCalled() + }) +}) + +describe('regional rehome director controls', () => { + it('records cell capability separately from the legacy heartbeat contract', async () => { + const recordCellRegionalRehomeStatus = vi.fn().mockResolvedValue(undefined) + const app = createRelayApp(config({ role: 'director', cellId: 'director' }), { + store: {} as never, + assignments: { recordCellRegionalRehomeStatus } as never, + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + const body = { + v: 1, + cellId: 'production-gce-c7', + cellIncarnation, + regionalRehomeProtocol: 1, + safety: { + observedAt: 100, + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + } + } + const response = await postPath( + app, + '/v1/admin/cell-rehome-status', + 'runtime-token', + body + ) + + expect(response.status).toBe(200) + expect(recordCellRegionalRehomeStatus).toHaveBeenCalledWith(body) + }) + + it('accepts the exact safety payload the cell heartbeat composes', async () => { + // Why: index.ts spreads the FULL pool-pressure counts into safety; a + // strict schema missing any produced field 400s every heartbeat (the + // 2026-08-15..26 outage that left all cells at rehome protocol 0). + const recordCellRegionalRehomeStatus = vi.fn().mockResolvedValue(undefined) + const app = createRelayApp(config({ role: 'director', cellId: 'director' }), { + store: {} as never, + assignments: { recordCellRegionalRehomeStatus } as never, + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + const observability = new RelayObservability( + { role: 'cell', cellId: 'production-gce-c7', region: 'us-central1' }, + () => {} + ) + observability.stop() + const body = { + v: 1, + cellId: 'production-gce-c7', + cellIncarnation, + regionalRehomeProtocol: 1, + safety: { + ...observability.regionalRehomeRuntimeSafety(), + ...emptyPostgresPoolPressureCounts() + } + } + const response = await postPath( + app, + '/v1/admin/cell-rehome-status', + 'runtime-token', + body + ) + + expect(response.status).toBe(200) + expect(recordCellRegionalRehomeStatus).toHaveBeenCalledWith(body) + }) + + it('applies the durable switch only with matching confirmation', async () => { + const inspectRegionalRehomeControl = vi.fn().mockResolvedValue({ + generation: 0, + enabled: false + }) + const applyRegionalRehomeControl = vi.fn(async (input) => ({ + ...input, + generation: input.expectedGeneration + 1 + })) + const app = createRelayApp(config({ role: 'director', cellId: 'director' }), { + store: {} as never, + assignments: { + inspectRegionalRehomeControl, + applyRegionalRehomeControl + } as never, + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + expect((await postPath( + app, + '/v1/admin/regional-rehome-control', + 'deploy-token', + { v: 1, action: 'inspect' } + )).status).toBe(200) + const apply = { + v: 1, + action: 'apply', + expectedGeneration: 0, + enabled: true, + notBefore: 100, + ratePerMinute: 10, + preferenceMaxAgeMs: 24 * 60 * 60_000, + drainGraceMs: 60_000, + confirmation: 'ENABLE_REGIONAL_REHOMING' + } + expect((await postPath( + app, + '/v1/admin/regional-rehome-control', + 'deploy-token', + apply + )).status).toBe(200) + expect(applyRegionalRehomeControl).toHaveBeenCalledOnce() + expect((await postPath( + app, + '/v1/admin/regional-rehome-control', + 'monitor-token', + { v: 1, action: 'inspect' } + )).status).toBe(200) + expect((await postPath( + app, + '/v1/admin/regional-rehome-control', + 'monitor-token', + apply + )).status).toBe(403) + expect((await postPath( + app, + '/v1/admin/regional-rehome-control', + 'deploy-token', + { ...apply, confirmation: 'DISABLE_REGIONAL_REHOMING' } + )).status).toBe(400) + }) + + it('probes dedicated trust twice and returns only aggregate proof', async () => { + const requests: Array<{ url: string; init?: RequestInit }> = [] + const cellDeploymentStatus = vi.fn().mockResolvedValue({ + cellId: 'production-gce-c7', + cellUrl: 'https://c7.relay.example.test', + region: 'us-central1', + runtime: { + cellIncarnation, + ready: true, + heartbeatFresh: true, + regionalRehomeProtocol: 1 + } + }) + const app = createRelayApp(config({ role: 'director', cellId: 'director' }), { + store: {} as never, + assignments: { cellDeploymentStatus } as never, + drain: vi.fn(), + regionalRehomeIdentityToken: vi.fn(async () => 'rehome-token'), + regionalRehomeFetch: (async (url, init) => { + requests.push({ url: String(url), init }) + return Response.json({ + v: 1, + outcome: 'host-not-connected', + sharedRuntimeIdentityRejected: true + }) + }) as typeof fetch, + ready: vi.fn(async () => true) + }) + const response = await postPath( + app, + '/v1/admin/regional-rehome-trust-probe', + 'deploy-token', + { v: 1, sourceCellId: 'production-gce-c7', sourceCellIncarnation: cellIncarnation } + ) + + expect(response.status).toBe(200) + const responseBody = await response.json() + expect(responseBody).toEqual({ + v: 1, + dedicatedIdentity: { + accepted: true, + firstOutcome: 'host-not-connected', + secondOutcome: 'host-not-connected', + idempotent: true + }, + sharedRuntimeIdentityRejected: true, + proven: true + }) + expect(requests).toHaveLength(2) + expect(requests.map(({ url }) => url)).toEqual([ + 'https://c7.relay.example.test/v1/admin/host-drain', + 'https://c7.relay.example.test/v1/admin/host-drain' + ]) + const bodies = requests.map(({ init }) => JSON.parse(String(init?.body))) + expect(bodies[0]).toEqual(bodies[1]) + expect(bodies[0]).toMatchObject({ + attemptId: REGIONAL_REHOME_TRUST_PROBE_ATTEMPT_ID, + userId: REGIONAL_REHOME_TRUST_PROBE_USER_ID, + relayHostId: REGIONAL_REHOME_TRUST_PROBE_HOST_ID, + sourceAssignmentEpoch: 1, + graceMs: 0 + }) + expect(requests.every(({ init }) => init?.signal instanceof AbortSignal)).toBe(true) + expect(JSON.stringify(responseBody)).not.toContain('rehome-token') + }) + + it('restricts trust probes to deploy authorization and strict input', async () => { + const app = createRelayApp(config({ role: 'director', cellId: 'director' }), { + store: {} as never, + assignments: {} as never, + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + const body = { + v: 1, + sourceCellId: 'production-gce-c7', + sourceCellIncarnation: cellIncarnation + } + expect((await postPath( + app, + '/v1/admin/regional-rehome-trust-probe', + 'monitor-token', + body + )).status).toBe(401) + expect((await postPath( + app, + '/v1/admin/regional-rehome-trust-probe', + 'deploy-token', + { ...body, unexpected: true } + )).status).toBe(400) + }) + + it('fails closed when the source rejects the dedicated identity', async () => { + const cellDeploymentStatus = vi.fn().mockResolvedValue({ + cellUrl: 'https://c7.relay.example.test', + region: 'us-central1', + runtime: { + cellIncarnation, + ready: true, + heartbeatFresh: true, + regionalRehomeProtocol: 1 + } + }) + const sourceFetch = vi.fn().mockResolvedValue( + Response.json({ error: 'invalid_token' }, { status: 401 }) + ) + const app = createRelayApp(config({ role: 'director', cellId: 'director' }), { + store: {} as never, + assignments: { cellDeploymentStatus } as never, + drain: vi.fn(), + regionalRehomeIdentityToken: vi.fn(async () => 'rehome-token'), + regionalRehomeFetch: sourceFetch, + ready: vi.fn(async () => true) + }) + const response = await postPath( + app, + '/v1/admin/regional-rehome-trust-probe', + 'deploy-token', + { v: 1, sourceCellId: 'production-gce-c7', sourceCellIncarnation: cellIncarnation } + ) + + expect(response.status).toBe(409) + expect(sourceFetch).toHaveBeenCalledOnce() + expect(JSON.stringify(await response.json())).not.toContain('rehome-token') + }) +}) + +async function post( + app: ReturnType, + token: string, + body: unknown +): Promise { + return await app.request('/v1/admin/host-drain', { + method: 'POST', + headers: { + authorization: `Bearer ${token}`, + 'content-type': 'application/json' + }, + body: JSON.stringify(body) + }) +} + +async function postPath( + app: ReturnType, + path: string, + token: string, + body: unknown +): Promise { + return await app.request(path, { + method: 'POST', + headers: { + authorization: `Bearer ${token}`, + 'content-type': 'application/json' + }, + body: JSON.stringify(body) + }) +} + +function config(overrides: Partial = {}): RelayConfig { + return { + port: 8080, + publicUrl: 'https://c7.relay.example.test', + cellUrl: 'https://c7.relay.example.test', + region: 'us-central1', + authIssuer: 'https://auth.example.test', + authAudience: 'orca-relay', + jwksUrl: 'https://auth.example.test/jwks', + assignmentSigningKey: new Uint8Array(32), + role: 'cell', + cellId: 'production-gce-c7', + cells: [], + adminAudience: 'https://relay.example.test/v1/admin/drain', + deployServiceAccount: 'deploy@example.test', + rehomeDirectorServiceAccount: 'relay-director@example.test', + rehomeAudience: 'https://relay.example.test/v1/admin/host-drain', + runtimeServiceAccount: 'relay-cell@example.test', + adminJwksUrl: 'https://auth.example.test/jwks', + databasePoolMax: 10, + publicAssignmentsEnabled: true, + publicAssignmentConcurrency: 2, + publicAssignmentQueueMax: 128, + publicAssignmentWaitMs: 4_000, + publicResolveConcurrency: 1, + publicResolveWaitMs: 5_000, + publicAssignmentRetryAfterSeconds: 5, + dataDir: './data', + ...overrides + } +} diff --git a/cloud/apps/relay/src/regional-rehome-postgres.test.ts b/cloud/apps/relay/src/regional-rehome-postgres.test.ts new file mode 100644 index 00000000000..d36e26ecd68 --- /dev/null +++ b/cloud/apps/relay/src/regional-rehome-postgres.test.ts @@ -0,0 +1,552 @@ +import { afterAll, beforeAll, beforeEach, describe, expect, it } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.js' +import { openRelayDatabase, type RelayDatabase } from './database.js' +import { + REGIONAL_REHOME_RECONNECTS_PER_CELL_LIMIT, + REGIONAL_REHOME_SQL_FAILURES_LIMIT, + REGIONAL_REHOME_SQL_FAILURES_PER_CELL_LIMIT +} from './regional-rehome-safety.js' + +const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL +const describePostgres = databaseUrl ? describe : describe.skip + +describePostgres('PostgreSQL regional rehoming', () => { + let primary: RelayDatabase + let secondary: RelayDatabase + let sequence = 0 + + beforeAll(async () => { + primary = await openRelayDatabase({ databaseUrl, dataDir: '' }) + secondary = await openRelayDatabase({ databaseUrl, dataDir: '' }) + }) + + beforeEach(async () => await cleanup()) + + afterAll(async () => { + await cleanup() + await secondary.close() + await primary.close() + }) + + async function cleanup(): Promise { + await primary.query( + `DELETE FROM relay_region_rehome_attempts WHERE user_id LIKE 'pg-rehome-user-%'` + ) + await primary.query(`DELETE FROM relay_region_rehome_worker_state`) + await primary.query(`DELETE FROM relay_region_rehome_control`) + await primary.query( + `DELETE FROM relay_control_connection_reservations + WHERE user_id LIKE 'pg-rehome-user-%'` + ) + await primary.query( + `DELETE FROM relay_assignment_migration_incarnations + WHERE user_id LIKE 'pg-rehome-user-%'` + ) + await primary.query( + `DELETE FROM relay_assignment_activity_leases + WHERE user_id LIKE 'pg-rehome-user-%'` + ) + await primary.query( + `DELETE FROM relay_assignment_migrations WHERE user_id LIKE 'pg-rehome-user-%'` + ) + await primary.query( + `DELETE FROM relay_assignment_region_preferences + WHERE user_id LIKE 'pg-rehome-user-%'` + ) + await primary.query(`DELETE FROM relay_assignments WHERE user_id LIKE 'pg-rehome-user-%'`) + for (const table of [ + 'relay_cell_rehome_safety', + 'relay_cell_capabilities', + 'relay_cell_connection_snapshots', + 'relay_cell_connection_runtime', + 'relay_cell_runtime', + 'relay_cell_connection_limits', + 'relay_cell_admission', + 'relay_cell_regions', + 'relay_cells' + ]) { + await primary.query(`DELETE FROM ${table} WHERE cell_id LIKE 'pg-rehome-cell-%'`) + } + } + + it('claims through ambient per-cell sql retry noise', async () => { + const context = await fixture() + await primary.query( + `UPDATE relay_cell_rehome_safety + SET sql_failures = ${REGIONAL_REHOME_SQL_FAILURES_PER_CELL_LIMIT} + WHERE cell_id IN (?, ?)`, + [context.source.id, context.target.id] + ) + + expect(await context.store.claimRegionalRehome()).not.toBeNull() + }) + + it('skips an unclean cell without latching the control off', async () => { + const context = await fixture() + await primary.query( + `UPDATE relay_cell_rehome_safety + SET sql_failures = ${REGIONAL_REHOME_SQL_FAILURES_PER_CELL_LIMIT + 1} + WHERE cell_id = ?`, + [context.target.id] + ) + + expect(await context.store.claimRegionalRehome()).toBeNull() + expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ + generation: 1, + enabled: true + }) + expect(await primary.query( + `SELECT next_dispatch_at FROM relay_region_rehome_worker_state` + )).toEqual([{ next_dispatch_at: String(context.now() + 6_000) }]) + }) + + it('lets only one director claim a host', async () => { + const context = await fixture() + const claims = await Promise.all([ + context.store.claimRegionalRehome(), + context.competingStore.claimRegionalRehome() + ]) + + expect(claims.filter(Boolean)).toHaveLength(1) + expect(await primary.query( + `SELECT COUNT(*) AS count FROM relay_region_rehome_attempts + WHERE user_id = ?`, + [context.identity.userId] + )).toEqual([{ count: '1' }]) + expect(await primary.query( + `SELECT COUNT(*) AS count FROM relay_assignment_migrations + WHERE user_id = ? AND completed_at IS NULL AND aborted_at IS NULL`, + [context.identity.userId] + )).toEqual([{ count: '1' }]) + }) + + it('increments the disable generation once across competing directors', async () => { + const context = await fixture() + const disabled = await Promise.all([ + context.store.disableRegionalRehomeControl(), + context.competingStore.disableRegionalRehomeControl() + ]) + + expect(disabled.sort()).toEqual([false, true]) + expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ + generation: 2, + enabled: false + }) + }) + + it('records one receipt across competing directors', async () => { + const context = await fixture() + const attempt = await context.store.claimRegionalRehome() + const receipts = await Promise.all([ + context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted'), + context.competingStore.recordRegionalRehomeDrainReceipt( + attempt!.attemptId, + 'accepted' + ) + ]) + + expect(receipts.sort()).toEqual([false, true]) + expect(await primary.query( + `SELECT drain_outcome FROM relay_region_rehome_attempts WHERE attempt_id = ?`, + [attempt!.attemptId] + )).toEqual([{ drain_outcome: 'accepted' }]) + }) + + it('rechecks a preference changed while the assignment row is locked', async () => { + const context = await fixture() + let unlock!: () => void + let locked!: () => void + const lockedPromise = new Promise((resolve) => (locked = resolve)) + const unlockPromise = new Promise((resolve) => (unlock = resolve)) + const held = secondary.transaction(async (transaction) => { + await transaction.queryLocked( + `SELECT * FROM relay_assignments WHERE user_id = ? AND relay_host_id = ?`, + [context.identity.userId, context.identity.relayHostId] + ) + locked() + await unlockPromise + }) + await lockedPromise + const claim = context.store.claimRegionalRehome() + await primary.query( + `UPDATE relay_assignment_region_preferences SET preferred_region = 'us-central1', + observed_at = ? WHERE user_id = ? AND relay_host_id = ?`, + [context.now(), context.identity.userId, context.identity.relayHostId] + ) + unlock() + await held + + await expect(claim).resolves.toBeNull() + expect(await primary.query( + `SELECT COUNT(*) AS count FROM relay_assignment_migrations WHERE user_id = ?`, + [context.identity.userId] + )).toEqual([{ count: '0' }]) + }) + + it('rechecks fleet safety under locks before mutating a candidate', async () => { + const context = await fixture() + let unlock!: () => void + let locked!: () => void + const lockedPromise = new Promise((resolve) => (locked = resolve)) + const unlockPromise = new Promise((resolve) => (unlock = resolve)) + const held = secondary.transaction(async (transaction) => { + await transaction.queryLocked( + `SELECT * FROM relay_cell_rehome_safety WHERE cell_id = ?`, + [context.target.id] + ) + await transaction.query( + `UPDATE relay_cell_rehome_safety SET sql_failures = ${REGIONAL_REHOME_SQL_FAILURES_LIMIT + 1} WHERE cell_id = ?`, + [context.target.id] + ) + locked() + await unlockPromise + }) + await lockedPromise + const claim = context.store.claimRegionalRehome() + unlock() + await held + + await expect(claim).resolves.toBeNull() + expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ + generation: 2, + enabled: false + }) + expect(await primary.query( + `SELECT COUNT(*) AS count FROM relay_assignment_migrations WHERE user_id = ?`, + [context.identity.userId] + )).toEqual([{ count: '0' }]) + }) + + it('pauses when one required cell exceeds the reconnect limit', async () => { + const context = await fixture() + await primary.query( + `UPDATE relay_cell_rehome_safety SET reconnects = ? WHERE cell_id = ?`, + [REGIONAL_REHOME_RECONNECTS_PER_CELL_LIMIT + 1, context.source.id] + ) + + await expect(context.store.claimRegionalRehome()).resolves.toBeNull() + await expect(context.store.inspectRegionalRehomeControl()).resolves.toMatchObject({ + generation: 2, + enabled: false + }) + expect(await primary.query( + `SELECT COUNT(*) AS count FROM relay_assignment_migrations WHERE user_id = ?`, + [context.identity.userId] + )).toEqual([{ count: '0' }]) + }) + + it('does not retry a drain against a replacement source incarnation', async () => { + const context = await fixture() + const attempt = await context.store.claimRegionalRehome() + context.advance(31_000) + await heartbeat( + context.store, + context.source, + '33333333-3333-4333-8333-333333333333', + 1, + context.now() + ) + + await expect(context.competingStore.claimRegionalRehome()).resolves.toBeNull() + expect(await primary.query( + `SELECT send_attempts FROM relay_region_rehome_attempts WHERE attempt_id = ?`, + [attempt!.attemptId] + )).toEqual([{ send_attempts: '1' }]) + }) + + it('makes concurrent completion and expiry cleanup idempotent', async () => { + const context = await fixture() + const attempt = await context.store.claimRegionalRehome() + await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + const targetControl = await context.store.activateControl(context.identity, { + cellId: context.target.id, + assignmentEpoch: attempt!.assignmentEpoch, + generation: 1 + }) + await context.store.markMigrationTargetRegistered(context.identity, { + cellId: context.target.id, + assignmentEpoch: attempt!.assignmentEpoch + }) + await context.store.releaseActivity(context.identity, context.sourceControl) + context.advance(24 * 60 * 60_000) + await heartbeat( + context.store, + context.source, + '11111111-1111-4111-8111-111111111111', + 1, + 900_000, + 2 + ) + await heartbeat( + context.store, + context.target, + '22222222-2222-4222-8222-222222222222', + 0, + 900_000, + 2 + ) + await context.store.renewControlActivity(context.identity, { + activityId: targetControl, + cellId: context.target.id, + expiresAt: context.now() + 90_000 + }) + + const outcomes = await Promise.all([ + context.store.completeReadyRegionalRehomes(), + context.competingStore.abortExpiredRegionalRehomes() + ]) + expect(outcomes).toEqual(expect.arrayContaining([0, 1])) + expect(await primary.query( + `SELECT completed_at IS NOT NULL AS completed, aborted_at IS NOT NULL AS aborted + FROM relay_assignment_migrations WHERE user_id = ?`, + [context.identity.userId] + )).toEqual([{ completed: true, aborted: false }]) + }) + + it('will not complete against a replacement target incarnation', async () => { + const context = await fixture() + const attempt = await context.store.claimRegionalRehome() + await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + await context.store.activateControl(context.identity, { + cellId: context.target.id, + assignmentEpoch: attempt!.assignmentEpoch, + generation: 1 + }) + await context.store.markMigrationTargetRegistered(context.identity, { + cellId: context.target.id, + assignmentEpoch: attempt!.assignmentEpoch + }) + await context.store.releaseActivity(context.identity, context.sourceControl) + context.advance(1) + await heartbeat( + context.store, + context.target, + '44444444-4444-4444-8444-444444444444', + 0, + context.now() + ) + + await expect(context.store.completeReadyRegionalRehomes()).resolves.toBe(0) + expect(await primary.query( + `SELECT completed_at, aborted_at FROM relay_assignment_migrations WHERE user_id = ?`, + [context.identity.userId] + )).toEqual([{ completed_at: null, aborted_at: null }]) + }) + + it('does not roll an unregistered target back to a stale regional source', async () => { + const context = await fixture() + await context.store.claimRegionalRehome() + context.advance(6 * 60_000) + await heartbeat( + context.store, + context.target, + '22222222-2222-4222-8222-222222222222', + 0, + 900_000, + 2 + ) + + await expect(context.store.refreshRegionalRehomeLeases()).resolves.toBe(0) + await expect(context.store.abortExpiredEvacuations()).resolves.toBe(0) + expect(await primary.query( + `SELECT cell_id, assignment_epoch FROM relay_assignments WHERE user_id = ?`, + [context.identity.userId] + )).toEqual([{ cell_id: context.target.id, assignment_epoch: '2' }]) + }) + + it('completes after the drained host re-resolves through the director', async () => { + const context = await fixture() + const attempt = await context.store.claimRegionalRehome() + await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + // The drain recovery lands while both controls are still live. + await context.store.assign(context.identity, 'asia-east2') + expect(await controlAccounting(context.identity)).toEqual({ + reservedControls: 2, + controlLeases: 2 + }) + + await context.store.activateControl(context.identity, { + cellId: context.target.id, + assignmentEpoch: attempt!.assignmentEpoch, + generation: 1 + }) + await context.store.markMigrationTargetRegistered(context.identity, { + cellId: context.target.id, + assignmentEpoch: attempt!.assignmentEpoch + }) + await context.store.releaseActivity(context.identity, context.sourceControl) + + await expect(context.store.completeReadyRegionalRehomes()).resolves.toBe(1) + expect(await primary.query( + `SELECT completed_at IS NOT NULL AS completed FROM relay_assignment_migrations + WHERE user_id = ?`, + [context.identity.userId] + )).toEqual([{ completed: true }]) + expect(await controlAccounting(context.identity)).toEqual({ + reservedControls: 1, + controlLeases: 1 + }) + }) + + it('repairs a skewed control counter before completing the rehome', async () => { + const context = await fixture() + const attempt = await context.store.claimRegionalRehome() + await context.store.activateControl(context.identity, { + cellId: context.target.id, + assignmentEpoch: attempt!.assignmentEpoch, + generation: 1 + }) + await context.store.markMigrationTargetRegistered(context.identity, { + cellId: context.target.id, + assignmentEpoch: attempt!.assignmentEpoch + }) + await context.store.releaseActivity(context.identity, context.sourceControl) + // Damage already written by a pre-fix sticky grant. + await primary.query( + `UPDATE relay_assignments SET reserved_controls = 0 WHERE user_id = ?`, + [context.identity.userId] + ) + + await expect(context.store.completeReadyRegionalRehomes()).resolves.toBe(1) + expect(await controlAccounting(context.identity)).toEqual({ + reservedControls: 1, + controlLeases: 1 + }) + }) + + async function controlAccounting(identity: { + userId: string + relayHostId: string + }): Promise<{ reservedControls: number; controlLeases: number }> { + const assignment = ( + await primary.query( + `SELECT reserved_controls FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + )[0]! + const leases = await primary.query( + `SELECT COUNT(*) AS controls FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? AND activity_kind = 'control'`, + [identity.userId, identity.relayHostId] + ) + return { + reservedControls: Number(assignment.reserved_controls), + controlLeases: Number(leases[0]!.controls) + } + } + + async function fixture() { + sequence++ + let now = 1_000_000 + const suffix = String(sequence) + const source = cell(suffix, 'source', 'us-central1') + const target = cell(suffix, 'target', 'asia-east2') + const store = new RelayAssignmentStore(primary, () => now, storeOptions) + const competingStore = new RelayAssignmentStore(secondary, () => now, storeOptions) + await store.inspectRegionalRehomeControl() + now += 24 * 60 * 60_000 + await store.applyRegionalRehomeControl({ + expectedGeneration: 0, + enabled: true, + notBefore: now, + ratePerMinute: 10, + preferenceMaxAgeMs: 24 * 60 * 60_000, + drainGraceMs: 60_000 + }) + await store.reconcileCells([source, target]) + await heartbeat( + store, + source, + '11111111-1111-4111-8111-111111111111', + 1, + 900_000 + ) + await heartbeat( + store, + target, + '22222222-2222-4222-8222-222222222222', + 0, + 900_000 + ) + const identity = { + userId: `pg-rehome-user-${suffix}`, + relayHostId: `rehomehost${suffix.padStart(6, '0')}` + } + const assignment = await store.assign(identity, undefined, 'us-central1') + const sourceControl = await store.activateControl(identity, { + cellId: source.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + await store.assign(identity, 'asia-east2') + return { + store, + competingStore, + identity, + source, + target, + sourceControl, + now: () => now, + advance: (milliseconds: number) => { + now += milliseconds + } + } + } +}) + +const storeOptions = { + requireLiveCells: true, + heartbeatTtlMs: 45_000 +} + +function cell(suffix: string, role: string, region: 'us-central1' | 'asia-east2') { + return { + id: `pg-rehome-cell-${suffix}-${role}`, + url: `https://pg-rehome-${suffix}-${role}.example.test`, + region, + capacityRequests: 100, + connectionHardCap: 1_000 as const, + connectionUnobservedBound: 60 + } +} + +async function heartbeat( + store: RelayAssignmentStore, + cellConfig: ReturnType, + cellIncarnation: string, + regionalRehomeProtocol: number, + startedAt: number, + connectionInclusionWatermark = 1 +): Promise { + await store.recordCellHeartbeat({ + cellId: cellConfig.id, + cellUrl: cellConfig.url, + region: cellConfig.region, + cellIncarnation, + startedAt, + ready: true, + observedRequests: 0, + totalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0, + connectionInclusionWatermark, + connectionHardCap: 1_000, + connectionUnobservedBound: 60 + }) + await store.recordCellRegionalRehomeStatus({ + cellId: cellConfig.id, + cellIncarnation, + regionalRehomeProtocol, + safety: { + observedAt: 1_000_000 + 24 * 60 * 60_000, + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + } + }) +} diff --git a/cloud/apps/relay/src/regional-rehome-safety.test.ts b/cloud/apps/relay/src/regional-rehome-safety.test.ts new file mode 100644 index 00000000000..5c7d10d37a0 --- /dev/null +++ b/cloud/apps/relay/src/regional-rehome-safety.test.ts @@ -0,0 +1,109 @@ +import { describe, expect, it } from 'vitest' +import { + REGIONAL_REHOME_POOL_WAIT_MS_MAX_LIMIT, + REGIONAL_REHOME_POOL_WAITERS_MAX_LIMIT, + REGIONAL_REHOME_RECONNECTS_PER_CELL_LIMIT, + REGIONAL_REHOME_SQL_FAILURES_LIMIT, + regionalRehomeSafetyFailure +} from './regional-rehome-safety.js' + +const NOW = 1_787_900_000_000 + +function safety(overrides: Partial[0]> = {}) { + return { + observedAt: NOW - 1_000, + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0, + ...overrides + } +} + +describe('regionalRehomeSafetyFailure', () => { + it('passes the measured healthy-fleet baseline of pool micro-waits', () => { + // Every production cell idles at 1-2 peak waiters resolved in ~1ms; a + // zero-tolerance bar here disables the worker on its first tick. + expect( + regionalRehomeSafetyFailure( + safety({ databasePoolWaitersMax: 2, databasePoolWaitMsMax: 1 }), + NOW, + 19 + ) + ).toBeNull() + }) + + it('passes routine client reconnect churn', () => { + expect( + regionalRehomeSafetyFailure(safety({ reconnects: 19 * 80 }), NOW, 19) + ).toBeNull() + }) + + it('still fails closed on each pool pressure bound', () => { + for (const overrides of [ + { databasePoolWaitersMax: REGIONAL_REHOME_POOL_WAITERS_MAX_LIMIT + 1 }, + { databasePoolWaitMsMax: REGIONAL_REHOME_POOL_WAIT_MS_MAX_LIMIT + 1 } + ]) { + expect(regionalRehomeSafetyFailure(safety(overrides), NOW, 19)).toBe( + 'database_pool_pressure' + ) + } + }) + + it('still fails closed on a reconnect storm', () => { + expect( + regionalRehomeSafetyFailure( + safety({ reconnects: 19 * REGIONAL_REHOME_RECONNECTS_PER_CELL_LIMIT + 1 }), + NOW, + 19 + ) + ).toBe('elevated_reconnects') + }) + + it('passes ambient 55P03 retry noise', () => { + // Fleet-wide retried lock timeouts peaked at 82 counted failures per + // minute over Aug 25-28; the combined snapshot can span two windows. + expect(regionalRehomeSafetyFailure(safety({ sqlFailures: 164 }), NOW, 19)).toBeNull() + }) + + it('still fails closed on a sql failure storm', () => { + expect( + regionalRehomeSafetyFailure( + safety({ sqlFailures: REGIONAL_REHOME_SQL_FAILURES_LIMIT + 1 }), + NOW, + 19 + ) + ).toBe('sql_failures') + // Literal storm magnitude (measured 2026-08-28) pins the bar itself: the + // limit must sit below real storm scale, not merely exist. + expect(regionalRehomeSafetyFailure(safety({ sqlFailures: 395 }), NOW, 19)).toBe( + 'sql_failures' + ) + // The bar is exclusive: exactly at the limit still passes. + expect( + regionalRehomeSafetyFailure( + safety({ sqlFailures: REGIONAL_REHOME_SQL_FAILURES_LIMIT }), + NOW, + 19 + ) + ).toBeNull() + }) + + it('keeps zero-tolerance for real failure signals', () => { + expect( + regionalRehomeSafetyFailure( + safety({ controlActivityRecoveryFailures: 1 }), + NOW, + 19 + ) + ).toBe('control_recovery_failures') + expect(regionalRehomeSafetyFailure(safety({ observedAt: 0 }), NOW, 19)).toBe( + 'monitoring_stale' + ) + expect( + regionalRehomeSafetyFailure(safety({ observedAt: NOW - 61_000 }), NOW, 19) + ).toBe('monitoring_stale') + }) +}) diff --git a/cloud/apps/relay/src/regional-rehome-safety.ts b/cloud/apps/relay/src/regional-rehome-safety.ts new file mode 100644 index 00000000000..4da51f71b0d --- /dev/null +++ b/cloud/apps/relay/src/regional-rehome-safety.ts @@ -0,0 +1,86 @@ +import type { RegionalRehomeSafetySnapshot } from './relay-observability.js' + +// Limits sit well above the healthy-fleet baseline measured in production on +// 2026-08-28 (peak 2 waiters / 1ms pool waits on every cell; up to ~80 +// reconnects per published two-window row on the busiest cell). Sustained +// pool saturation still trips: waiters-max 16 is 8x baseline yet far under a +// backed-up pool, and 250ms peak wait is 1/10 of the incident-monitor alert. +// Instantaneous databasePoolWaiting is not checked separately: it is bounded +// by databasePoolWaitersMax within every published window. +export const REGIONAL_REHOME_RECONNECTS_PER_CELL_LIMIT = 250 +export const REGIONAL_REHOME_POOL_WAITERS_MAX_LIMIT = 16 +export const REGIONAL_REHOME_POOL_WAIT_MS_MAX_LIMIT = 250 + +// Counted SQL failures are dominated by relay_cells 55P03 lock-timeout +// retries the transaction wrapper heals in place (Aug 25-28: ~1000 retried +// attempts per day vs ~30 exhausted, those clustered in one storm), so a +// zero bar disabled the worker on ambient noise just like the original pool +// bars. Fleet-wide ambient noise peaked at 82 counted failures per minute +// over four days; genuine database distress produced 395-457. The combined +// snapshot spans up to two 30s windows per process (pathological ambient +// alignment ~164), so 250 stays clear of noise while storms still trip. +// Terminal outages also trip the pool bars and the worker's own +// dispatch-failure budget; the sql bar only needs to catch storms. +export const REGIONAL_REHOME_SQL_FAILURES_LIMIT = 250 +// Per-cell candidate cleanliness is a soft skip, not a durable latch; the +// worst ambient per-cell publish carried ~24 counted failures (two windows +// of 12). The fleet bar deliberately dominates: cells that are individually +// clean can sum past 250, and a fleet-wide sum at that scale is a storm. +export const REGIONAL_REHOME_SQL_FAILURES_PER_CELL_LIMIT = 40 + +export function regionalRehomePoolPressure(safety: { + databasePoolWaitersMax: number + databasePoolWaitMsMax: number +}): boolean { + return ( + safety.databasePoolWaitersMax > REGIONAL_REHOME_POOL_WAITERS_MAX_LIMIT || + safety.databasePoolWaitMsMax > REGIONAL_REHOME_POOL_WAIT_MS_MAX_LIMIT + ) +} + +export function regionalRehomeSafetyFailure( + safety: RegionalRehomeSafetySnapshot, + now: number, + requiredCells: number +): string | null { + if (safety.observedAt === 0 || now - safety.observedAt > 60_000) { + return 'monitoring_stale' + } + if (safety.sqlFailures > REGIONAL_REHOME_SQL_FAILURES_LIMIT) return 'sql_failures' + if (regionalRehomePoolPressure(safety)) { + return 'database_pool_pressure' + } + if (safety.controlActivityRecoveryFailures > 0) { + return 'control_recovery_failures' + } + const reconnectLimit = + Math.max(1, requiredCells) * REGIONAL_REHOME_RECONNECTS_PER_CELL_LIMIT + if (safety.reconnects > reconnectLimit) return 'elevated_reconnects' + return null +} + +export function combineRegionalRehomeSafety( + processSafety: RegionalRehomeSafetySnapshot, + fleetSafety: RegionalRehomeSafetySnapshot +): RegionalRehomeSafetySnapshot { + return { + observedAt: Math.min(processSafety.observedAt, fleetSafety.observedAt), + sqlFailures: processSafety.sqlFailures + fleetSafety.sqlFailures, + reconnects: fleetSafety.reconnects, + controlActivityRecoveryFailures: + processSafety.controlActivityRecoveryFailures + + fleetSafety.controlActivityRecoveryFailures, + databasePoolWaiting: Math.max( + processSafety.databasePoolWaiting, + fleetSafety.databasePoolWaiting + ), + databasePoolWaitersMax: Math.max( + processSafety.databasePoolWaitersMax, + fleetSafety.databasePoolWaitersMax + ), + databasePoolWaitMsMax: Math.max( + processSafety.databasePoolWaitMsMax, + fleetSafety.databasePoolWaitMsMax + ) + } +} diff --git a/cloud/apps/relay/src/regional-rehome-store.test.ts b/cloud/apps/relay/src/regional-rehome-store.test.ts new file mode 100644 index 00000000000..43f293b1131 --- /dev/null +++ b/cloud/apps/relay/src/regional-rehome-store.test.ts @@ -0,0 +1,1399 @@ +import { describe, expect, it } from 'vitest' +import { + RelayAssignmentStore, + REGIONAL_REHOME_QUARANTINE_FAILURES, + REGIONAL_REHOME_QUARANTINE_MS, + REGIONAL_REHOME_REDRAIN_SEND_LIMIT +} from './assignment-store.js' +import { openInMemoryRelayDatabase, type RelayDatabase, type SqlRow } from './database.js' +import { + REGIONAL_REHOME_SQL_FAILURES_LIMIT, + REGIONAL_REHOME_SQL_FAILURES_PER_CELL_LIMIT +} from './regional-rehome-safety.js' +import type { RegionalRehomeSafetySnapshot } from './relay-observability.js' + +const source = { + id: 'us-c1', + url: 'https://us-c1.relay.example.test', + region: 'us-central1' as const, + capacityRequests: 100, + connectionHardCap: 1_000 as const, + connectionUnobservedBound: 60 +} +const target = { + id: 'asia-c1', + url: 'https://asia-c1.relay.example.test', + region: 'asia-east2' as const, + capacityRequests: 100, + connectionHardCap: 1_000 as const, + connectionUnobservedBound: 60 +} +const sourceIncarnation = '11111111-1111-4111-8111-111111111111' +const targetIncarnation = '22222222-2222-4222-8222-222222222222' + +describe('regional rehome assignment state', () => { + it('does not open a transaction while the worker is disabled', async () => { + const delegate = await openInMemoryRelayDatabase() + const database = new TransactionCountingDatabase(delegate) + const store = new RelayAssignmentStore(database, () => 1_000_000) + await store.inspectRegionalRehomeControl() + database.transactionCalls = 0 + + await expect(store.claimRegionalRehome()).resolves.toBeNull() + expect(database.transactionCalls).toBe(0) + await database.close() + }) + + it('initializes a missing control row without opening a transaction', async () => { + const delegate = await openInMemoryRelayDatabase() + const database = new TransactionCountingDatabase(delegate) + const store = new RelayAssignmentStore(database, () => 1_000_000) + + await expect(store.claimRegionalRehome()).resolves.toBeNull() + expect(database.transactionCalls).toBe(0) + await expect(store.inspectRegionalRehomeControl()).resolves.toMatchObject({ + generation: 0, + enabled: false + }) + await database.close() + }) + + it('uses a generation-bound durable kill switch', async () => { + const context = await setup() + expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ + generation: 1, + enabled: true, + ratePerMinute: 10 + }) + expect(await context.store.disableRegionalRehomeControl()).toBe(true) + expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ + generation: 2, + enabled: false + }) + await expect(context.store.applyRegionalRehomeControl({ + expectedGeneration: 1, + enabled: true, + notBefore: context.now(), + ratePerMinute: 10, + preferenceMaxAgeMs: 24 * 60 * 60_000, + drainGraceMs: 60_000 + })).rejects.toThrow('regional_rehome_generation_mismatch') + await expect(context.store.applyRegionalRehomeControl({ + expectedGeneration: 2, + enabled: true, + notBefore: context.now(), + ratePerMinute: 10, + preferenceMaxAgeMs: 24 * 60 * 60_000, + drainGraceMs: 60_000 + })).resolves.toMatchObject({ generation: 3, enabled: true }) + await context.database.close() + }) + + it('moves one live preferred host and completes without disabling its source cell', async () => { + const context = await setup() + const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + const neighbor = { userId: 'user-2', relayHostId: 'ponmlkjihgfedcba' } + const sourceControl = await activatePreferredSource(context, identity) + await activateSource(context, neighbor) + + const attempt = await context.store.claimRegionalRehome() + expect(attempt).toMatchObject({ + userId: identity.userId, + relayHostId: identity.relayHostId, + sourceCellId: source.id, + sourceCellIncarnation: sourceIncarnation, + targetCellId: target.id, + targetCellIncarnation: targetIncarnation, + previousEpoch: 1, + assignmentEpoch: 2, + sendAttempts: 1 + }) + expect(await context.store.resolve(neighbor)).toMatchObject({ cellId: source.id }) + expect(await context.store.completeReadyRegionalRehomes()).toBe(0) + expect( + await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + ).toBe(true) + expect( + await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + ).toBe(false) + + const targetControl = await context.store.activateControl(identity, { + cellId: target.id, + assignmentEpoch: 2, + generation: 1 + }) + await context.store.markMigrationTargetRegistered(identity, { + cellId: target.id, + assignmentEpoch: 2 + }) + expect(await context.store.completeReadyRegionalRehomes()).toBe(0) + await context.store.releaseActivity(identity, sourceControl) + expect(await context.store.completeReadyRegionalRehomes()).toBe(1) + + expect(await context.store.resolve(identity)).toMatchObject({ + cellId: target.id, + assignmentEpoch: 2 + }) + expect(await context.store.resolve(neighbor)).toMatchObject({ cellId: source.id }) + expect(await context.database.query( + `SELECT completed_at, aborted_at FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + )).toEqual([{ completed_at: context.now(), aborted_at: null }]) + expect(await context.database.query( + `SELECT completed_at, aborted_at FROM relay_region_rehome_attempts` + )).toEqual([{ completed_at: context.now(), aborted_at: null }]) + expect(targetControl).toMatch(/^control:/) + await context.database.close() + }) + + it('completes from durable activity when the drain response was lost', async () => { + const context = await setup() + const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + const sourceControl = await activatePreferredSource(context, identity) + const attempt = await context.store.claimRegionalRehome() + await context.store.activateControl(identity, { + cellId: target.id, + assignmentEpoch: attempt!.assignmentEpoch, + generation: 1 + }) + await context.store.markMigrationTargetRegistered(identity, { + cellId: target.id, + assignmentEpoch: attempt!.assignmentEpoch + }) + await context.store.releaseActivity(identity, sourceControl) + + expect(await context.store.completeReadyRegionalRehomes()).toBe(1) + expect(await context.database.query( + `SELECT drain_receipt_at, completed_at, aborted_at + FROM relay_region_rehome_attempts` + )).toEqual([{ drain_receipt_at: null, completed_at: context.now(), aborted_at: null }]) + await context.database.close() + }) + + it('requires a fresh preference and an advertised source capability', async () => { + const context = await setup({ sourceProtocol: 0 }) + await activatePreferredSource(context, { + userId: 'user-1', + relayHostId: 'abcdefghijklmnop' + }) + expect(await context.store.claimRegionalRehome()).toBeNull() + expect(await context.database.query(`SELECT * FROM relay_assignment_migrations`)).toEqual([]) + await context.database.close() + }) + + it('fails fleet safety closed until source and target telemetry is fresh', async () => { + const context = await setup() + await context.database.query( + `DELETE FROM relay_cell_rehome_safety WHERE cell_id = ?`, + [target.id] + ) + expect(await context.store.regionalRehomeFleetSafety()).toMatchObject({ + requiredCells: 2, + missingCells: 1, + observedAt: 0 + }) + await heartbeat(context.store, source, sourceIncarnation, 1, 2, { + observedAt: context.now(), + sqlFailures: 0, + reconnects: 2, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + }) + await heartbeat(context.store, target, targetIncarnation, 0, 2, { + observedAt: context.now(), + sqlFailures: 1, + reconnects: 3, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + }) + expect(await context.store.regionalRehomeFleetSafety()).toMatchObject({ + requiredCells: 2, + missingCells: 0, + observedAt: context.now(), + sqlFailures: 1, + reconnects: 5 + }) + await context.database.close() + }) + + it('claims through the measured healthy baseline of pool micro-waits and churn', async () => { + const context = await setup() + const baseline = { + observedAt: context.now(), + sqlFailures: 0, + reconnects: 42, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 2, + databasePoolWaitersMax: 2, + databasePoolWaitMsMax: 1 + } + await heartbeat(context.store, source, sourceIncarnation, 1, 2, baseline) + await heartbeat(context.store, target, targetIncarnation, 0, 2, baseline) + await activatePreferredSource(context, { + userId: 'user-1', + relayHostId: 'abcdefghijklmnop' + }) + + expect(await context.store.claimRegionalRehome()).toMatchObject({ + sourceCellId: source.id, + targetCellId: target.id + }) + expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ + enabled: true + }) + await context.database.close() + }) + + it('latches off on a per-cell reconnect storm even when the fleet sum is low', async () => { + const context = await setup() + await activatePreferredSource(context, { + userId: 'user-1', + relayHostId: 'abcdefghijklmnop' + }) + await context.database.query( + `UPDATE relay_cell_rehome_safety SET reconnects = 251 WHERE cell_id = ?`, + [source.id] + ) + + const warnings = collectDisableWarnings() + try { + expect(await context.store.claimRegionalRehome()).toBeNull() + } finally { + warnings.restore() + } + expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ + generation: 2, + enabled: false + }) + expect(warnings.entries).toMatchObject([ + { reason: 'elevated_reconnects', maxReconnects: 251, controlGeneration: 2 } + ]) + await context.database.close() + }) + + it('latches off on sustained pool pressure and logs the disable exactly once', async () => { + const context = await setup() + await activatePreferredSource(context, { + userId: 'user-1', + relayHostId: 'abcdefghijklmnop' + }) + await context.database.query( + `UPDATE relay_cell_rehome_safety SET database_pool_waiters_max = 17 WHERE cell_id = ?`, + [target.id] + ) + + const warnings = collectDisableWarnings() + try { + expect(await context.store.claimRegionalRehome()).toBeNull() + // Already disabled: the next tick returns before the gate and stays silent. + expect(await context.store.claimRegionalRehome()).toBeNull() + } finally { + warnings.restore() + } + expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ + generation: 2, + enabled: false + }) + expect(warnings.entries).toMatchObject([ + { reason: 'database_pool_pressure', databasePoolWaitersMax: 17 } + ]) + await context.database.close() + }) + + it('claims through ambient per-cell sql retry noise', async () => { + const context = await setup() + await activatePreferredSource(context, { + userId: 'user-1', + relayHostId: 'abcdefghijklmnop' + }) + await context.database.query( + `UPDATE relay_cell_rehome_safety SET sql_failures = ${REGIONAL_REHOME_SQL_FAILURES_PER_CELL_LIMIT}` + ) + + expect(await context.store.claimRegionalRehome()).not.toBeNull() + await context.database.close() + }) + + it('skips an unclean cell without latching the control off', async () => { + const context = await setup() + await activatePreferredSource(context, { + userId: 'user-1', + relayHostId: 'abcdefghijklmnop' + }) + await activatePreferredSource(context, { + userId: 'user-2', + relayHostId: 'ponmlkjihgfedcba' + }) + // Above the per-cell cleanliness bar but below the fleet storm bar: the + // candidate is skipped this tick while the worker stays enabled. + await context.database.query( + `UPDATE relay_cell_rehome_safety SET sql_failures = ${REGIONAL_REHOME_SQL_FAILURES_PER_CELL_LIMIT + 1} WHERE cell_id = ?`, + [target.id] + ) + + const warnings = collectEventWarnings('orca_relay_regional_rehome_candidates_skipped') + try { + expect(await context.store.claimRegionalRehome()).toBeNull() + } finally { + warnings.restore() + } + expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ + generation: 1, + enabled: true + }) + // The skip is visible and named, and both candidates blocked by the one + // unclean cell accumulate into a single entry. + expect(warnings.entries).toMatchObject([ + { + skips: [ + { + reason: 'target_unclean', + cellId: target.id, + sqlFailures: REGIONAL_REHOME_SQL_FAILURES_PER_CELL_LIMIT + 1, + candidates: 2 + } + ] + } + ]) + // A skipped tick is charged the dispatch interval: candidate scans stay + // rate-limited even when nothing claims. + expect( + await context.database.query( + `SELECT next_dispatch_at FROM relay_region_rehome_worker_state` + ) + ).toEqual([{ next_dispatch_at: context.now() + 6_000 }]) + expect(await context.database.query(`SELECT * FROM relay_assignment_migrations`)).toEqual([]) + await context.database.close() + }) + + it('does not throttle or log an idle tick with no candidates', async () => { + const context = await setup() + + const warnings = collectEventWarnings('orca_relay_regional_rehome_candidates_skipped') + try { + expect(await context.store.claimRegionalRehome()).toBeNull() + } finally { + warnings.restore() + } + expect(warnings.entries).toEqual([]) + expect( + await context.database.query( + `SELECT next_dispatch_at FROM relay_region_rehome_worker_state` + ) + ).toEqual([{ next_dispatch_at: 0 }]) + await context.database.close() + }) + + it('atomically latches durable control off when candidate safety changes', async () => { + const context = await setup() + await activatePreferredSource(context, { + userId: 'user-1', + relayHostId: 'abcdefghijklmnop' + }) + await context.database.query( + `UPDATE relay_cell_rehome_safety SET sql_failures = ${REGIONAL_REHOME_SQL_FAILURES_LIMIT + 1} WHERE cell_id = ?`, + [target.id] + ) + + expect(await context.store.claimRegionalRehome()).toBeNull() + expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ + generation: 2, + enabled: false + }) + expect(await context.database.query(`SELECT * FROM relay_assignment_migrations`)).toEqual([]) + await context.database.close() + }) + + it('rechecks locked fleet safety before retrying a drain dispatch', async () => { + const context = await setup() + await activatePreferredSource(context, { + userId: 'user-1', + relayHostId: 'abcdefghijklmnop' + }) + expect(await context.store.claimRegionalRehome()).not.toBeNull() + context.advance(31_000) + await context.database.query( + `UPDATE relay_cell_rehome_safety SET sql_failures = ${REGIONAL_REHOME_SQL_FAILURES_LIMIT + 1} WHERE cell_id = ?`, + [target.id] + ) + + expect(await context.store.claimRegionalRehome()).toBeNull() + expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ + generation: 2, + enabled: false + }) + await context.database.close() + }) + + it('latches off after three dispatch failures and resumes only through CAS', async () => { + const context = await setup() + await activatePreferredSource(context, { + userId: 'user-1', + relayHostId: 'abcdefghijklmnop' + }) + const first = await context.store.claimRegionalRehome() + for (let index = 0; index < 3; index++) { + await context.store.recordRegionalRehomeDispatchFailure(first!.attemptId) + } + context.advance(5 * 60_000 - 1) + expect(await context.store.claimRegionalRehome()).toBeNull() + context.advance(1) + await heartbeat(context.store, source, sourceIncarnation, 1, 2, { + observedAt: context.now(), + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + }) + await heartbeat(context.store, target, targetIncarnation, 0, 2, { + observedAt: context.now(), + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + }) + expect(await context.store.claimRegionalRehome()).toBeNull() + expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ + generation: 2, + enabled: false + }) + await context.store.applyRegionalRehomeControl({ + expectedGeneration: 2, + enabled: true, + notBefore: context.now(), + ratePerMinute: 10, + preferenceMaxAgeMs: 24 * 60 * 60_000, + drainGraceMs: 60_000 + }) + const retry = await context.store.claimRegionalRehome() + expect(retry).toMatchObject({ attemptId: first!.attemptId, sendAttempts: 2 }) + expect(await context.database.query( + `SELECT COUNT(*) AS count FROM relay_assignment_migrations` + )).toEqual([{ count: 1 }]) + await context.database.close() + }) + + it('refreshes only the migration leases while source splices drain', async () => { + const context = await setup() + const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + await activatePreferredSource(context, identity) + await context.store.acquireActivity(identity, { + activityId: 'splice:source', + kind: 'splice', + cellId: source.id + }) + await context.store.claimRegionalRehome() + const before = await context.database.query( + `SELECT activity_id, expires_at FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? ORDER BY activity_id`, + [identity.userId, identity.relayHostId] + ) + context.advance(60_000) + expect(await context.store.refreshRegionalRehomeLeases()).toBe(1) + const after = await context.database.query( + `SELECT activity_id, expires_at FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? ORDER BY activity_id`, + [identity.userId, identity.relayHostId] + ) + const beforeById = new Map(before.map((row) => [row.activity_id, row.expires_at])) + const afterById = new Map(after.map((row) => [row.activity_id, row.expires_at])) + expect(Number(afterById.get('migration:2'))).toBeGreaterThan( + Number(beforeById.get('migration:2')) + ) + expect(Number(afterById.get('control-pending:2'))).toBeGreaterThan( + Number(beforeById.get('control-pending:2')) + ) + expect(afterById.get('splice:source')).toBe(beforeById.get('splice:source')) + await context.database.close() + }) + + it('stops refreshing an unregistered target and lets normal rollback retire it', async () => { + const context = await setup() + const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + await activatePreferredSource(context, identity) + await context.store.claimRegionalRehome() + context.advance(6 * 60_000) + expect(await context.store.refreshRegionalRehomeLeases()).toBe(0) + await heartbeat(context.store, source, sourceIncarnation, 1, 2) + expect(await context.store.abortExpiredEvacuations()).toBe(1) + expect(await context.store.reapRegionalRehomeAttempts()).toBe(1) + expect(await context.store.resolve(identity)).toMatchObject({ + cellId: source.id, + assignmentEpoch: 3 + }) + expect(await context.database.query( + `SELECT completed_at, aborted_at FROM relay_region_rehome_attempts` + )).toEqual([{ completed_at: null, aborted_at: context.now() }]) + await context.database.close() + }) + + it('does not roll an unregistered target back to a stale regional source', async () => { + const context = await setup() + const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + await activatePreferredSource(context, identity) + await context.store.claimRegionalRehome() + context.advance(6 * 60_000) + await heartbeat(context.store, target, targetIncarnation, 0, 2) + + expect(await context.store.refreshRegionalRehomeLeases()).toBe(0) + expect(await context.store.abortExpiredEvacuations()).toBe(0) + expect( + await context.database.query( + `SELECT cell_id, assignment_epoch FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ cell_id: target.id, assignment_epoch: 2 }]) + await context.database.close() + }) + + it('rolls back an inactive registered target only after the 24-hour bound', async () => { + const context = await setup() + const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + const sourceControl = await activatePreferredSource(context, identity) + const attempt = await context.store.claimRegionalRehome() + await context.store.recordRegionalRehomeDrainReceipt( + attempt!.attemptId, + 'accepted' + ) + const targetControl = await context.store.activateControl(identity, { + cellId: target.id, + assignmentEpoch: 2, + generation: 1 + }) + await context.store.markMigrationTargetRegistered(identity, { + cellId: target.id, + assignmentEpoch: 2 + }) + await context.store.releaseActivity(identity, sourceControl) + await context.store.releaseActivity(identity, targetControl) + context.advance(24 * 60 * 60_000) + await heartbeat(context.store, source, sourceIncarnation, 1, 2) + expect(await context.store.abortExpiredRegionalRehomes()).toBe(1) + expect(await context.store.resolve(identity)).toMatchObject({ + cellId: source.id, + assignmentEpoch: 3 + }) + await context.database.close() + }) + + it('redrains a receipted dual-homed attempt once its grace elapses', async () => { + const context = await setup() + const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + const sourceControl = await activatePreferredSource(context, identity) + const attempt = await context.store.claimRegionalRehome() + await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + await context.store.activateControl(identity, { + cellId: target.id, + assignmentEpoch: 2, + generation: 1 + }) + await context.store.markMigrationTargetRegistered(identity, { + cellId: target.id, + assignmentEpoch: 2 + }) + + // Before grace elapses a receipted attempt is not re-dispatched. + context.advance(30 * 60_000) + await freshHeartbeats(context) + expect(await context.store.claimRegionalRehome()).toBeNull() + + context.advance(30 * 60_000 + 1) + await freshHeartbeats(context) + const redrain = await context.store.claimRegionalRehome() + expect(redrain).toMatchObject({ + attemptId: attempt!.attemptId, + drainGraceMs: 0, + sendAttempts: 2 + }) + // The per-dispatch receipt replaces the original without a mismatch. + await expect( + context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'host-not-connected') + ).resolves.toBe(true) + + // Redrains are spaced: nothing new inside the redrain interval. + context.advance(30_000) + await freshHeartbeats(context) + expect(await context.store.claimRegionalRehome()).toBeNull() + context.advance(30_001) + await freshHeartbeats(context) + expect(await context.store.claimRegionalRehome()).toMatchObject({ + attemptId: attempt!.attemptId, + drainGraceMs: 0, + sendAttempts: 3 + }) + + // Once the host actually leaves the source, completion wins over redrain. + await context.store.releaseActivity(identity, sourceControl) + context.advance(60_001) + await freshHeartbeats(context) + await context.store.activateControl(identity, { + cellId: target.id, + assignmentEpoch: 2, + generation: 2 + }) + expect(await context.store.claimRegionalRehome()).toBeNull() + expect(await context.store.completeReadyRegionalRehomes()).toBe(1) + await context.database.close() + }) + + it('resets the failure budget on a repeated redrain receipt outcome', async () => { + const context = await setup() + const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + await activatePreferredSource(context, identity) + const attempt = await context.store.claimRegionalRehome() + await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + await context.store.activateControl(identity, { + cellId: target.id, + assignmentEpoch: 2, + generation: 1 + }) + await context.store.markMigrationTargetRegistered(identity, { + cellId: target.id, + assignmentEpoch: 2 + }) + await context.store.recordRegionalRehomeDispatchFailure(attempt!.attemptId) + await context.store.recordRegionalRehomeDispatchFailure(attempt!.attemptId) + + context.advance(60 * 60_000 + 1) + await freshHeartbeats(context) + expect(await context.store.claimRegionalRehome()).toMatchObject({ + attemptId: attempt!.attemptId, + drainGraceMs: 0 + }) + // The repeated outcome still proves the source answered. + expect( + await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + ).toBe(false) + await context.store.recordRegionalRehomeDispatchFailure(attempt!.attemptId) + expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ + generation: 1, + enabled: true + }) + await context.database.close() + }) + + it('does not redrain before the target registers or when the fleet is unsafe', async () => { + const context = await setup() + const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + await activatePreferredSource(context, identity) + const attempt = await context.store.claimRegionalRehome() + await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + await context.store.activateControl(identity, { + cellId: target.id, + assignmentEpoch: 2, + generation: 1 + }) + + // Past grace but the target never registered: force-closing the source + // would disconnect the host with nowhere proven to land. + context.advance(60 * 60_000 + 1) + await freshHeartbeats(context) + expect(await context.store.claimRegionalRehome()).toBeNull() + + await context.store.markMigrationTargetRegistered(identity, { + cellId: target.id, + assignmentEpoch: 2 + }) + await context.database.query( + `UPDATE relay_cell_rehome_safety SET sql_failures = ${REGIONAL_REHOME_SQL_FAILURES_LIMIT + 1} WHERE cell_id = ?`, + [target.id] + ) + expect(await context.store.claimRegionalRehome()).toBeNull() + expect(await context.store.inspectRegionalRehomeControl()).toMatchObject({ + enabled: false + }) + await context.database.close() + }) + + it('completes healthy candidates past a poisoned attempt and logs it', async () => { + const context = await setup() + const poisoned = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + const healthy = { userId: 'user-2', relayHostId: 'ponmlkjihgfedcba' } + const poisonedSource = await activatePreferredSource(context, poisoned) + const healthySource = await activatePreferredSource(context, healthy) + const first = await context.store.claimRegionalRehome() + context.advance(6_000) + const second = await context.store.claimRegionalRehome() + expect(first!.userId).toBe(poisoned.userId) + expect(second!.userId).toBe(healthy.userId) + for (const [identity, attempt, sourceControl] of [ + [poisoned, first, poisonedSource], + [healthy, second, healthySource] + ] as const) { + await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + await context.store.activateControl(identity, { + cellId: target.id, + assignmentEpoch: attempt!.assignmentEpoch, + generation: 1 + }) + await context.store.markMigrationTargetRegistered(identity, { + cellId: target.id, + assignmentEpoch: attempt!.assignmentEpoch + }) + await context.store.releaseActivity(identity, sourceControl) + } + // The production poison shape: the assignment moved past the attempt. + await context.database.query( + `UPDATE relay_assignments SET assignment_epoch = assignment_epoch + 5 + WHERE user_id = ?`, + [poisoned.userId] + ) + const warnings = collectCandidateFailureWarnings() + try { + expect(await context.store.completeReadyRegionalRehomes()).toBe(1) + } finally { + warnings.restore() + } + expect(warnings.entries).toEqual([ + { + event: 'orca_relay_regional_rehome_candidate_failed', + operation: 'complete', + attemptId: first!.attemptId, + reason: 'regional_rehome_assignment_mismatch' + } + ]) + expect(await context.database.query( + `SELECT completed_at FROM relay_region_rehome_attempts WHERE attempt_id = ?`, + [second!.attemptId] + )).toEqual([{ completed_at: context.now() }]) + await context.database.close() + }) + + it('quarantines a repeatedly failing candidate and redacts free-form errors', async () => { + const context = await setup() + const poisoned = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + const source1 = await activatePreferredSource(context, poisoned) + const attempt = await context.store.claimRegionalRehome() + await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + await context.store.activateControl(poisoned, { + cellId: target.id, + assignmentEpoch: attempt!.assignmentEpoch, + generation: 1 + }) + await context.store.markMigrationTargetRegistered(poisoned, { + cellId: target.id, + assignmentEpoch: attempt!.assignmentEpoch + }) + await context.store.releaseActivity(poisoned, source1) + await context.database.query( + `UPDATE relay_assignments SET assignment_epoch = assignment_epoch + 5 + WHERE user_id = ?`, + [poisoned.userId] + ) + const warnings = collectCandidateFailureWarnings() + try { + for (let round = 0; round < REGIONAL_REHOME_QUARANTINE_FAILURES; round++) { + expect(await context.store.completeReadyRegionalRehomes()).toBe(0) + } + expect(warnings.entries).toHaveLength(REGIONAL_REHOME_QUARANTINE_FAILURES) + // Quarantined: the poisoned row leaves the candidate page entirely. + expect(await context.store.completeReadyRegionalRehomes()).toBe(0) + expect(warnings.entries).toHaveLength(REGIONAL_REHOME_QUARANTINE_FAILURES) + // After the quarantine window it is retried (and fails) once more. + context.advance(REGIONAL_REHOME_QUARANTINE_MS + 1) + await context.store.completeReadyRegionalRehomes() + expect(warnings.entries).toHaveLength(REGIONAL_REHOME_QUARANTINE_FAILURES + 1) + // A free-form error (never a slug) reaches the log only as 'redacted'. + expect( + warnings.entries.every( + (entry) => entry.reason === 'regional_rehome_assignment_mismatch' + ) + ).toBe(true) + context.advance(REGIONAL_REHOME_QUARANTINE_MS + 1) + const database = context.database + const original = database.transaction.bind(database) + database.transaction = () => { + throw new Error('postgresql://secret@database.invalid/relay') + } + try { + await context.store.completeReadyRegionalRehomes() + } finally { + database.transaction = original + } + const last = warnings.entries.at(-1)! + expect(last.reason).toBe('redacted') + expect(JSON.stringify(last)).not.toContain('secret') + } finally { + warnings.restore() + } + await context.database.close() + }) + + it('aborts healthy expired candidates past a poisoned attempt', async () => { + const context = await setup() + const poisoned = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + const healthy = { userId: 'user-2', relayHostId: 'ponmlkjihgfedcba' } + const poisonedSource = await activatePreferredSource(context, poisoned) + const healthySource = await activatePreferredSource(context, healthy) + const first = await context.store.claimRegionalRehome() + context.advance(6_000) + const second = await context.store.claimRegionalRehome() + for (const [identity, attempt, sourceControl] of [ + [poisoned, first, poisonedSource], + [healthy, second, healthySource] + ] as const) { + await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + const targetControl = await context.store.activateControl(identity, { + cellId: target.id, + assignmentEpoch: attempt!.assignmentEpoch, + generation: 1 + }) + await context.store.markMigrationTargetRegistered(identity, { + cellId: target.id, + assignmentEpoch: attempt!.assignmentEpoch + }) + await context.store.releaseActivity(identity, sourceControl) + await context.store.releaseActivity(identity, targetControl) + } + await context.database.query( + `UPDATE relay_assignments SET assignment_epoch = assignment_epoch + 5 + WHERE user_id = ?`, + [poisoned.userId] + ) + context.advance(24 * 60 * 60_000) + await freshHeartbeats(context) + const warnings = collectCandidateFailureWarnings() + try { + expect(await context.store.abortExpiredRegionalRehomes()).toBe(1) + } finally { + warnings.restore() + } + expect(warnings.entries).toEqual([ + { + event: 'orca_relay_regional_rehome_candidate_failed', + operation: 'abort', + attemptId: first!.attemptId, + reason: 'regional_rehome_assignment_mismatch' + } + ]) + expect(await context.database.query( + `SELECT aborted_at FROM relay_region_rehome_attempts WHERE attempt_id = ?`, + [second!.attemptId] + )).toEqual([{ aborted_at: context.now() }]) + await context.database.close() + }) + + it('keeps dual-control accounting when the drained host re-resolves mid-rehome', async () => { + const context = await setup() + const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + const sourceControl = await activatePreferredSource(context, identity) + const attempt = await context.store.claimRegionalRehome() + expect(await controlAccounting(context, identity)).toEqual({ + reservedControls: 2, + controlLeases: 2 + }) + + // The drained host re-resolves through the director while both the source + // control and the target's pending control are still live. + await context.store.assign(identity, 'asia-east2') + expect(await controlAccounting(context, identity)).toEqual({ + reservedControls: 2, + controlLeases: 2 + }) + + await context.store.activateControl(identity, { + cellId: target.id, + assignmentEpoch: attempt!.assignmentEpoch, + generation: 1 + }) + await context.store.markMigrationTargetRegistered(identity, { + cellId: target.id, + assignmentEpoch: attempt!.assignmentEpoch + }) + await context.store.releaseActivity(identity, sourceControl) + expect(await controlAccounting(context, identity)).toEqual({ + reservedControls: 1, + controlLeases: 1 + }) + + expect(await context.store.completeReadyRegionalRehomes()).toBe(1) + expect(await context.database.query( + `SELECT completed_at FROM relay_assignment_migrations + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + )).toEqual([{ completed_at: context.now() }]) + expect(await cellReservations(context)).toEqual({ [source.id]: 0, [target.id]: 1 }) + await context.database.close() + }) + + it('grants a host whose counter was skewed without duplicating its control', async () => { + const context = await setup() + const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + const sourceControl = await activatePreferredSource(context, identity) + const attempt = await context.store.claimRegionalRehome() + await context.store.activateControl(identity, { + cellId: target.id, + assignmentEpoch: attempt!.assignmentEpoch, + generation: 1 + }) + await context.store.releaseActivity(identity, sourceControl) + // Damage already written by a pre-fix sticky grant. + await context.database.query( + `UPDATE relay_assignments SET reserved_controls = 0 + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + + await context.store.assign(identity, 'asia-east2') + expect(await context.database.query( + `SELECT activity_id FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? AND activity_kind = 'control'`, + [identity.userId, identity.relayHostId] + )).toEqual([{ activity_id: `control:${target.id}:1` }]) + expect(await cellReservations(context)).toEqual({ [source.id]: 0, [target.id]: 2 }) + await context.database.close() + }) + + it('repairs a skewed control counter before completing the rehome', async () => { + const context = await setup() + const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + const sourceControl = await activatePreferredSource(context, identity) + const attempt = await context.store.claimRegionalRehome() + await context.store.activateControl(identity, { + cellId: target.id, + assignmentEpoch: attempt!.assignmentEpoch, + generation: 1 + }) + await context.store.markMigrationTargetRegistered(identity, { + cellId: target.id, + assignmentEpoch: attempt!.assignmentEpoch + }) + await context.store.releaseActivity(identity, sourceControl) + // Damage already written by a pre-fix sticky grant. + await context.database.query( + `UPDATE relay_assignments SET reserved_controls = 0 + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + + expect(await context.store.completeReadyRegionalRehomes()).toBe(1) + expect(await controlAccounting(context, identity)).toEqual({ + reservedControls: 1, + controlLeases: 1 + }) + expect(await context.database.query( + `SELECT migration_leases FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + )).toEqual([{ migration_leases: 0 }]) + await context.database.close() + }) + + it('reports a repaired counter only for the candidate it repaired', async () => { + const context = await setup() + const skewed = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + const clean = { userId: 'user-2', relayHostId: 'ponmlkjihgfedcba' } + const skewedSource = await activatePreferredSource(context, skewed) + const cleanSource = await activatePreferredSource(context, clean) + const first = await context.store.claimRegionalRehome() + context.advance(6_000) + const second = await context.store.claimRegionalRehome() + for (const [identity, attempt, sourceControl] of [ + [skewed, first, skewedSource], + [clean, second, cleanSource] + ] as const) { + await context.store.activateControl(identity, { + cellId: target.id, + assignmentEpoch: attempt!.assignmentEpoch, + generation: 1 + }) + await context.store.markMigrationTargetRegistered(identity, { + cellId: target.id, + assignmentEpoch: attempt!.assignmentEpoch + }) + await context.store.releaseActivity(identity, sourceControl) + } + // Damage already written by a pre-fix sticky grant, on one host only. + await context.database.query( + `UPDATE relay_assignments SET reserved_controls = 0 + WHERE user_id = ? AND relay_host_id = ?`, + [skewed.userId, skewed.relayHostId] + ) + + const warnings = collectCandidateFailureWarnings([ + 'orca_relay_regional_rehome_activity_counts_repaired' + ]) + try { + expect(await context.store.completeReadyRegionalRehomes()).toBe(2) + } finally { + warnings.restore() + } + expect(warnings.entries).toEqual([ + { + event: 'orca_relay_regional_rehome_activity_counts_repaired', + attemptId: first!.attemptId + } + ]) + await context.database.close() + }) + + it('never repairs past a migration lease whose shape is wrong', async () => { + const context = await setup() + const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + const sourceControl = await activatePreferredSource(context, identity) + const attempt = await context.store.claimRegionalRehome() + await context.store.activateControl(identity, { + cellId: target.id, + assignmentEpoch: attempt!.assignmentEpoch, + generation: 1 + }) + await context.store.markMigrationTargetRegistered(identity, { + cellId: target.id, + assignmentEpoch: attempt!.assignmentEpoch + }) + await context.store.releaseActivity(identity, sourceControl) + await context.database.query( + `UPDATE relay_assignments SET reserved_controls = 0 + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + await context.database.query( + `UPDATE relay_assignment_migrations SET target_reserved_units = 9 + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + + const warnings = collectCandidateFailureWarnings() + try { + expect(await context.store.completeReadyRegionalRehomes()).toBe(0) + } finally { + warnings.restore() + } + expect(warnings.entries).toEqual([ + { + event: 'orca_relay_regional_rehome_candidate_failed', + operation: 'complete', + attemptId: attempt!.attemptId, + reason: 'migration_activity_lease_shape_mismatch' + } + ]) + expect(await controlAccounting(context, identity)).toEqual({ + reservedControls: 0, + controlLeases: 1 + }) + await context.database.close() + }) + + it('stays silent when a repair cannot make the counts whole', async () => { + const context = await setup() + const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + const sourceControl = await activatePreferredSource(context, identity) + const attempt = await context.store.claimRegionalRehome() + await context.store.activateControl(identity, { + cellId: target.id, + assignmentEpoch: attempt!.assignmentEpoch, + generation: 1 + }) + await context.store.markMigrationTargetRegistered(identity, { + cellId: target.id, + assignmentEpoch: attempt!.assignmentEpoch + }) + await context.store.releaseActivity(identity, sourceControl) + // A vanished migration lease is repairable arithmetic on the first assert + // but still wrong on the re-assert: no repaired event may leak out. + await context.database.query( + `DELETE FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? AND activity_kind = 'migration'`, + [identity.userId, identity.relayHostId] + ) + + const warnings = collectCandidateFailureWarnings([ + 'orca_relay_regional_rehome_candidate_failed', + 'orca_relay_regional_rehome_activity_counts_repaired' + ]) + try { + expect(await context.store.completeReadyRegionalRehomes()).toBe(0) + } finally { + warnings.restore() + } + expect(warnings.entries).toEqual([ + { + event: 'orca_relay_regional_rehome_candidate_failed', + operation: 'complete', + attemptId: attempt!.attemptId, + reason: 'migration_activity_accounting_mismatch' + } + ]) + await context.database.close() + }) + + it('caps redrain dispatches at the send limit', async () => { + const context = await setup() + const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + await activatePreferredSource(context, identity) + const attempt = await context.store.claimRegionalRehome() + await context.store.recordRegionalRehomeDrainReceipt(attempt!.attemptId, 'accepted') + await context.store.activateControl(identity, { + cellId: target.id, + assignmentEpoch: 2, + generation: 1 + }) + await context.store.markMigrationTargetRegistered(identity, { + cellId: target.id, + assignmentEpoch: 2 + }) + await context.database.query( + `UPDATE relay_region_rehome_attempts SET send_attempts = ? WHERE attempt_id = ?`, + [REGIONAL_REHOME_REDRAIN_SEND_LIMIT, attempt!.attemptId] + ) + context.advance(60 * 60_000 + 1) + await freshHeartbeats(context) + expect(await context.store.claimRegionalRehome()).toBeNull() + await context.database.close() + }) +}) + +class TransactionCountingDatabase implements RelayDatabase { + transactionCalls = 0 + + constructor(private readonly delegate: RelayDatabase) {} + + query(sql: string, params?: unknown[]): Promise { + return this.delegate.query(sql, params) + } + + queryLocked( + sql: string, + params?: unknown[], + options?: { failIfUnavailable?: boolean } + ): Promise { + return this.delegate.queryLocked(sql, params, options) + } + + transaction( + operation: (transaction: RelayDatabase) => Promise, + options?: { reportRetries?: boolean } + ): Promise { + this.transactionCalls += 1 + return this.delegate.transaction(operation, options) + } + + close(): Promise { + return this.delegate.close() + } +} + +type Context = Awaited> + +function collectDisableWarnings() { + const entries: Record[] = [] + const original = console.warn + console.warn = (line: unknown, ...rest: unknown[]) => { + try { + const parsed = JSON.parse(line as string) as Record + if (parsed.event === 'orca_relay_regional_rehome_safety_disabled') { + entries.push(parsed) + return + } + } catch { + // fall through to the real console for non-JSON lines + } + original(line, ...rest) + } + return { + entries, + restore: () => { + console.warn = original + } + } +} + +async function setup(options: { sourceProtocol?: number } = {}) { + let clock = 1_000_000 + const database = await openInMemoryRelayDatabase() + const store = new RelayAssignmentStore(database, () => clock, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + await store.inspectRegionalRehomeControl() + clock += 24 * 60 * 60_000 + await store.applyRegionalRehomeControl({ + expectedGeneration: 0, + enabled: true, + notBefore: clock, + ratePerMinute: 10, + preferenceMaxAgeMs: 24 * 60 * 60_000, + drainGraceMs: 60 * 60_000 + }) + await store.reconcileCells([source, target]) + await heartbeat(store, source, sourceIncarnation, options.sourceProtocol ?? 1) + await heartbeat(store, target, targetIncarnation, 0) + return { + database, + store, + now: () => clock, + advance: (milliseconds: number) => { + clock += milliseconds + } + } +} + +function collectEventWarnings(event: string) { + const entries: Record[] = [] + const original = console.warn + console.warn = (line: unknown, ...rest: unknown[]) => { + try { + const parsed = JSON.parse(line as string) as Record + if (parsed.event === event) { + entries.push(parsed) + return + } + } catch { + // fall through to the real console for non-JSON lines + } + original(line, ...rest) + } + return { + entries, + restore: () => { + console.warn = original + } + } +} + +function collectCandidateFailureWarnings( + events: string[] = ['orca_relay_regional_rehome_candidate_failed'] +) { + const entries: Record[] = [] + const original = console.warn + console.warn = (line: unknown, ...rest: unknown[]) => { + try { + const parsed = JSON.parse(line as string) as Record + if (events.includes(parsed.event as string)) { + entries.push(parsed) + return + } + } catch { + // fall through to the real console for non-JSON lines + } + original(line, ...rest) + } + return { + entries, + restore: () => { + console.warn = original + } + } +} + +async function controlAccounting( + context: Context, + identity: { userId: string; relayHostId: string } +): Promise<{ reservedControls: number; controlLeases: number }> { + const assignment = ( + await context.database.query( + `SELECT reserved_controls FROM relay_assignments + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + )[0]! + const leases = await context.database.query( + `SELECT COUNT(*) AS controls FROM relay_assignment_activity_leases + WHERE user_id = ? AND relay_host_id = ? AND activity_kind = 'control'`, + [identity.userId, identity.relayHostId] + ) + return { + reservedControls: Number(assignment.reserved_controls), + controlLeases: Number(leases[0]!.controls) + } +} + +async function cellReservations(context: Context): Promise> { + const rows = await context.database.query( + `SELECT cell_id, reserved_requests FROM relay_cells ORDER BY cell_id` + ) + return Object.fromEntries( + rows.map((row) => [String(row.cell_id), Number(row.reserved_requests)]) + ) +} + +async function freshHeartbeats(context: Context): Promise { + const safety = { + observedAt: context.now(), + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + } + // The clock doubles as a strictly-increasing connection inclusion watermark. + await heartbeat(context.store, source, sourceIncarnation, 1, context.now(), safety) + await heartbeat(context.store, target, targetIncarnation, 0, context.now(), safety) +} + +async function activatePreferredSource( + context: Context, + identity: { userId: string; relayHostId: string } +): Promise { + const assignment = await context.store.assign(identity, undefined, 'us-central1') + const control = await context.store.activateControl(identity, { + cellId: source.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + await context.store.assign(identity, 'asia-east2') + return control +} + +async function activateSource( + context: Context, + identity: { userId: string; relayHostId: string } +): Promise { + const assignment = await context.store.assign(identity, undefined, 'us-central1') + return await context.store.activateControl(identity, { + cellId: source.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) +} + +async function heartbeat( + store: RelayAssignmentStore, + cell: typeof source | typeof target, + cellIncarnation: string, + regionalRehomeProtocol: number, + connectionInclusionWatermark = 1, + regionalRehomeSafety?: RegionalRehomeSafetySnapshot +): Promise { + await store.recordCellHeartbeat({ + cellId: cell.id, + cellUrl: cell.url, + region: cell.region, + cellIncarnation, + startedAt: 900_000, + ready: true, + observedRequests: 0, + totalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0, + connectionInclusionWatermark, + connectionHardCap: 1_000, + connectionUnobservedBound: 60 + }) + await store.recordCellRegionalRehomeStatus({ + cellId: cell.id, + cellIncarnation, + regionalRehomeProtocol, + safety: regionalRehomeSafety ?? { + observedAt: 1_000_000 + 24 * 60 * 60_000, + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + } + }) +} diff --git a/cloud/apps/relay/src/regional-rehome-target-selection.test.ts b/cloud/apps/relay/src/regional-rehome-target-selection.test.ts new file mode 100644 index 00000000000..e2190168735 --- /dev/null +++ b/cloud/apps/relay/src/regional-rehome-target-selection.test.ts @@ -0,0 +1,163 @@ +import { describe, expect, it } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.js' +import { openInMemoryRelayDatabase } from './database.js' +import { REGIONAL_REHOME_SQL_FAILURES_PER_CELL_LIMIT } from './regional-rehome-safety.js' + +const cell = (id: string, region: 'us-central1' | 'asia-east2') => ({ + id, + url: `https://${id}.relay.example.test`, + region, + capacityRequests: 100, + connectionHardCap: 1_000 as const, + connectionUnobservedBound: 60 +}) +const source = cell('us-c1', 'us-central1') +const noHeadroom = cell('asia-a', 'asia-east2') +const unclean = cell('asia-b', 'asia-east2') +const highLoad = cell('asia-c', 'asia-east2') +const lowLoad = cell('asia-d', 'asia-east2') + +const incarnation = (n: number) => + `${String(n).repeat(8)}-${String(n).repeat(4)}-4${String(n).repeat(3)}` + + `-8${String(n).repeat(3)}-${String(n).repeat(12)}` + +async function setup() { + let clock = 1_000_000 + const database = await openInMemoryRelayDatabase() + const store = new RelayAssignmentStore(database, () => clock, { + requireLiveCells: true, + heartbeatTtlMs: 45_000 + }) + await store.inspectRegionalRehomeControl() + clock += 24 * 60 * 60_000 + await store.applyRegionalRehomeControl({ + expectedGeneration: 0, + enabled: true, + notBefore: clock, + ratePerMinute: 10, + preferenceMaxAgeMs: 24 * 60 * 60_000, + drainGraceMs: 60 * 60_000 + }) + await store.reconcileCells([source, noHeadroom, unclean, highLoad, lowLoad]) + const beat = async ( + config: typeof source, + n: number, + protocol: number, + state: { observedRequests: number; enforcedConnections: number; sqlFailures: number } + ) => { + await store.recordCellHeartbeat({ + cellId: config.id, + cellUrl: config.url, + region: config.region, + cellIncarnation: incarnation(n), + startedAt: 900_000, + ready: true, + observedRequests: state.observedRequests, + totalConnections: state.enforcedConnections, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: state.enforcedConnections, + connectionInclusionWatermark: clock, + connectionHardCap: 1_000, + connectionUnobservedBound: 60 + }) + await store.recordCellRegionalRehomeStatus({ + cellId: config.id, + cellIncarnation: incarnation(n), + regionalRehomeProtocol: protocol, + safety: { + observedAt: clock, + sqlFailures: state.sqlFailures, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolWaitMsMax: 0 + } + }) + } + const activatePreferredSource = async () => { + const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' } + const assignment = await store.assign(identity, undefined, 'us-central1') + await store.activateControl(identity, { + cellId: source.id, + assignmentEpoch: assignment.assignmentEpoch, + generation: 1 + }) + await store.assign(identity, 'asia-east2') + } + return { database, store, beat, activatePreferredSource } +} + +const UNCLEAN = REGIONAL_REHOME_SQL_FAILURES_PER_CELL_LIMIT + 1 + +describe('regional rehome target selection', () => { + it('never selects a target without connection headroom, even at lowest load', async () => { + const context = await setup() + await context.beat(source, 1, 1, { + observedRequests: 0, + enforcedConnections: 0, + sqlFailures: 0 + }) + // Lowest load but the connection hard cap is exhausted. + await context.beat(noHeadroom, 2, 0, { + observedRequests: 0, + enforcedConnections: 999, + sqlFailures: 0 + }) + await context.beat(unclean, 3, 0, { + observedRequests: 0, + enforcedConnections: 0, + sqlFailures: UNCLEAN + }) + await context.beat(highLoad, 4, 0, { + observedRequests: 50, + enforcedConnections: 0, + sqlFailures: 0 + }) + await context.beat(lowLoad, 5, 0, { + observedRequests: 10, + enforcedConnections: 0, + sqlFailures: 0 + }) + await context.activatePreferredSource() + + const attempt = await context.store.claimRegionalRehome() + expect(attempt?.targetCellId).toBe(lowLoad.id) + await context.database.close() + }) + + it('falls to the next clean target when the load winner goes unclean', async () => { + const context = await setup() + await context.beat(source, 1, 1, { + observedRequests: 0, + enforcedConnections: 0, + sqlFailures: 0 + }) + await context.beat(noHeadroom, 2, 0, { + observedRequests: 0, + enforcedConnections: 999, + sqlFailures: 0 + }) + await context.beat(unclean, 3, 0, { + observedRequests: 0, + enforcedConnections: 0, + sqlFailures: UNCLEAN + }) + await context.beat(highLoad, 4, 0, { + observedRequests: 50, + enforcedConnections: 0, + sqlFailures: 0 + }) + await context.beat(lowLoad, 5, 0, { + observedRequests: 10, + enforcedConnections: 0, + sqlFailures: UNCLEAN + }) + await context.activatePreferredSource() + + const attempt = await context.store.claimRegionalRehome() + expect(attempt?.targetCellId).toBe(highLoad.id) + await context.database.close() + }) +}) diff --git a/cloud/apps/relay/src/regional-rehome-trust-probe.ts b/cloud/apps/relay/src/regional-rehome-trust-probe.ts new file mode 100644 index 00000000000..57ae2468886 --- /dev/null +++ b/cloud/apps/relay/src/regional-rehome-trust-probe.ts @@ -0,0 +1,107 @@ +import { z } from 'zod' + +export const REGIONAL_REHOME_TRUST_PROBE_ATTEMPT_ID = + '00000000-0000-4000-8000-000000000001' +export const REGIONAL_REHOME_TRUST_PROBE_USER_ID = 'regional-rehome-trust-probe' +export const REGIONAL_REHOME_TRUST_PROBE_HOST_ID = 'trustprobe000001' + +const SOURCE_PROBE_TIMEOUT_MS = 10_000 + +const SourceProbeResponseSchema = z + .object({ + v: z.literal(1), + outcome: z.enum(['accepted', 'already-accepted', 'host-not-connected']), + sharedRuntimeIdentityRejected: z.literal(true) + }) + .strict() + +type SourceProbeOutcome = z.infer['outcome'] + +export type RegionalRehomeTrustProbeResult = { + v: 1 + dedicatedIdentity: { + accepted: boolean + firstOutcome: SourceProbeOutcome + secondOutcome: SourceProbeOutcome + idempotent: boolean + } + sharedRuntimeIdentityRejected: boolean + proven: boolean +} + +export async function probeRegionalRehomeTrust(input: { + sourceCellUrl: string + sourceCellId: string + sourceCellIncarnation: string + audience: string + identityToken: (audience: string) => Promise + fetch: typeof fetch +}): Promise { + const token = await input.identityToken(input.audience) + const body = JSON.stringify({ + v: 1, + attemptId: REGIONAL_REHOME_TRUST_PROBE_ATTEMPT_ID, + userId: REGIONAL_REHOME_TRUST_PROBE_USER_ID, + relayHostId: REGIONAL_REHOME_TRUST_PROBE_HOST_ID, + sourceCellId: input.sourceCellId, + sourceCellIncarnation: input.sourceCellIncarnation, + sourceAssignmentEpoch: 1, + graceMs: 0 + }) + const call = async (): Promise> => { + const response = await input.fetch( + new URL('/v1/admin/host-drain', input.sourceCellUrl), + { + method: 'POST', + headers: { + authorization: `Bearer ${token}`, + 'content-type': 'application/json' + }, + body, + signal: AbortSignal.timeout(SOURCE_PROBE_TIMEOUT_MS) + } + ) + if (!response.ok) throw new Error(`regional_rehome_trust_probe_source_${response.status}`) + const parsed = SourceProbeResponseSchema.safeParse(await response.json()) + if (!parsed.success) throw new Error('regional_rehome_trust_probe_source_invalid_response') + return parsed.data + } + const first = await call() + const second = await call() + const idempotent = first.outcome === second.outcome + const sharedRuntimeIdentityRejected = + first.sharedRuntimeIdentityRejected && second.sharedRuntimeIdentityRejected + const proven = + first.outcome === 'host-not-connected' && + second.outcome === 'host-not-connected' && + idempotent && + sharedRuntimeIdentityRejected + if (!proven) throw new Error('regional_rehome_trust_probe_not_proven') + return { + v: 1, + dedicatedIdentity: { + accepted: true, + firstOutcome: first.outcome, + secondOutcome: second.outcome, + idempotent + }, + sharedRuntimeIdentityRejected, + proven + } +} + +export function isRegionalRehomeTrustProbe(input: { + attemptId: string + userId: string + relayHostId: string + sourceAssignmentEpoch: number + graceMs: number +}): boolean { + return ( + input.attemptId === REGIONAL_REHOME_TRUST_PROBE_ATTEMPT_ID && + input.userId === REGIONAL_REHOME_TRUST_PROBE_USER_ID && + input.relayHostId === REGIONAL_REHOME_TRUST_PROBE_HOST_ID && + input.sourceAssignmentEpoch === 1 && + input.graceMs === 0 + ) +} diff --git a/cloud/apps/relay/src/regional-rehome-worker.test.ts b/cloud/apps/relay/src/regional-rehome-worker.test.ts new file mode 100644 index 00000000000..af905bb9ab2 --- /dev/null +++ b/cloud/apps/relay/src/regional-rehome-worker.test.ts @@ -0,0 +1,227 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import type { RelayAssignmentStore } from './assignment-store.js' +import type { RelayConfig } from './config.js' +import { + combineRegionalRehomeSafety, + REGIONAL_REHOME_RECONNECTS_PER_CELL_LIMIT, + regionalRehomeSafetyFailure +} from './regional-rehome-safety.js' +import { startRegionalRehomeWorker } from './regional-rehome-worker.js' + +describe('regional rehome worker', () => { + afterEach(() => vi.restoreAllMocks()) + + it('sends an incarnation- and source-epoch-bound drain without exposing identity', async () => { + let now = 0 + const attempt = { + attemptId: '11111111-1111-4111-8111-111111111111', + userId: 'private-user', + relayHostId: 'abcdefghijklmnop', + preferredRegion: 'asia-east2', + sourceCellId: 'production-gce-c7', + sourceCellUrl: 'https://c7.relay.example.test', + sourceCellIncarnation: '22222222-2222-4222-8222-222222222222', + targetCellId: 'production-gce-c27', + targetCellIncarnation: '33333333-3333-4333-8333-333333333333', + previousEpoch: 7, + assignmentEpoch: 8, + drainGraceMs: 60_000, + sendAttempts: 1 + } + const claimRegionalRehome = vi.fn().mockResolvedValueOnce(null).mockResolvedValue(attempt) + const recordRegionalRehomeDrainReceipt = vi.fn().mockResolvedValue(true) + const recordRegionalRehomeWorkerFailure = vi.fn().mockResolvedValue(undefined) + const assignments = { + claimRegionalRehome, + recordRegionalRehomeDrainReceipt, + recordRegionalRehomeWorkerFailure + } as unknown as RelayAssignmentStore + const requests: Array<{ url: string; init?: RequestInit }> = [] + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + const worker = startRegionalRehomeWorker(config(), assignments, { + now: () => now, + safetySnapshot: () => safety(now), + intervalMs: 60_000, + identityToken: async (audience) => { + expect(audience).toBe('https://relay.example.test/v1/admin/host-drain') + return 'secret-token' + }, + fetch: (async (url, init) => { + requests.push({ url: String(url), init }) + return Response.json({ v: 1, outcome: 'accepted' }) + }) as typeof fetch + })! + await settleWorker() + now = 1_000 + await worker.run() + worker.stop() + + expect(requests).toHaveLength(1) + expect(requests[0]!.url).toBe('https://c7.relay.example.test/v1/admin/host-drain') + expect(requests[0]!.url).not.toContain('secret-token') + expect(requests[0]!.init?.headers).toMatchObject({ + authorization: 'Bearer secret-token' + }) + expect(JSON.parse(String(requests[0]!.init?.body))).toEqual({ + v: 1, + attemptId: '11111111-1111-4111-8111-111111111111', + userId: 'private-user', + relayHostId: 'abcdefghijklmnop', + sourceCellId: 'production-gce-c7', + sourceCellIncarnation: '22222222-2222-4222-8222-222222222222', + sourceAssignmentEpoch: 7, + graceMs: 60_000 + }) + expect(recordRegionalRehomeDrainReceipt).toHaveBeenCalledWith( + '11111111-1111-4111-8111-111111111111', + 'accepted' + ) + const logs = warn.mock.calls.map((call) => String(call[0])).join('\n') + expect(logs).not.toContain('private-user') + expect(logs).not.toContain('abcdefghijklmnop') + }) + + it('fails closed before the observation gate and records bounded dispatch failures', async () => { + let now = 0 + const attempt = { + attemptId: '11111111-1111-4111-8111-111111111111', + userId: 'private-user', + relayHostId: 'abcdefghijklmnop', + sourceCellId: 'source', + sourceCellUrl: 'https://source.example.test', + sourceCellIncarnation: '22222222-2222-4222-8222-222222222222', + targetCellId: 'target', + previousEpoch: 1, + assignmentEpoch: 2, + drainGraceMs: 60_000, + sendAttempts: 1 + } + const claimRegionalRehome = vi.fn().mockResolvedValueOnce(null).mockResolvedValue(attempt) + const assignments = { + claimRegionalRehome, + recordRegionalRehomeDispatchFailure: vi.fn().mockResolvedValue(undefined), + recordRegionalRehomeWorkerFailure: vi.fn().mockResolvedValue(undefined) + } as unknown as RelayAssignmentStore + const worker = startRegionalRehomeWorker(config(), assignments, { + now: () => now, + safetySnapshot: () => safety(now), + intervalMs: 60_000, + identityToken: async () => { + throw new Error('token unavailable') + } + })! + await settleWorker() + claimRegionalRehome.mockClear() + now = 100 + await worker.run() + worker.stop() + expect(assignments.recordRegionalRehomeDispatchFailure).toHaveBeenCalledWith( + '11111111-1111-4111-8111-111111111111' + ) + expect(assignments.recordRegionalRehomeWorkerFailure).not.toHaveBeenCalled() + }) + + it('passes unsafe process telemetry to the durable claim gate', async () => { + let now = 0 + let sqlFailures = 0 + const claimRegionalRehome = vi.fn().mockResolvedValue(null) + const assignments = { + claimRegionalRehome + } as unknown as RelayAssignmentStore + const worker = startRegionalRehomeWorker(config(), assignments, { + now: () => now, + safetySnapshot: () => ({ ...safety(now), sqlFailures }), + intervalMs: 60_000 + })! + await settleWorker() + claimRegionalRehome.mockClear() + now = 100 + sqlFailures = 1 + await worker.run() + worker.stop() + + expect(claimRegionalRehome).toHaveBeenCalledWith( + expect.objectContaining({ observedAt: 100, sqlFailures: 1 }) + ) + }) + + it('starts inert on directors so durable control can enable without a restart', async () => { + let now = 0 + const claimRegionalRehome = vi.fn().mockResolvedValue(null) + const assignments = { + claimRegionalRehome + } as unknown as RelayAssignmentStore + const worker = startRegionalRehomeWorker( + config(), + assignments, + { + now: () => now, + safetySnapshot: () => safety(now), + intervalMs: 60_000 + } + ) + expect(worker).not.toBeNull() + await settleWorker() + claimRegionalRehome.mockClear() + now = 100 + await worker!.run() + worker!.stop() + expect(claimRegionalRehome).toHaveBeenCalledOnce() + + expect( + startRegionalRehomeWorker( + config({ role: 'cell' }), + {} as RelayAssignmentStore, + { safetySnapshot: () => safety(1) } + ) + ).toBeNull() + }) + + it('treats the reconnect threshold as per-cell and excludes the director', () => { + const cells = 2 + const limit = cells * REGIONAL_REHOME_RECONNECTS_PER_CELL_LIMIT + const processSafety = { ...safety(100), reconnects: limit * 10 } + const fleetSafety = { ...safety(100), reconnects: limit } + expect(regionalRehomeSafetyFailure( + combineRegionalRehomeSafety(processSafety, fleetSafety), + 100, + cells + )).toBeNull() + expect(regionalRehomeSafetyFailure( + combineRegionalRehomeSafety(processSafety, { ...fleetSafety, reconnects: limit + 1 }), + 100, + cells + )).toBe('elevated_reconnects') + }) +}) + +function config(overrides: Partial = {}): RelayConfig { + return { + role: 'director', + rehomeAudience: 'https://relay.example.test/v1/admin/host-drain', + rehomeDirectorServiceAccount: 'relay-director@example.test', + ...overrides + } as RelayConfig +} + +function safety(observedAt: number) { + return { + requiredCells: 2, + missingCells: 0, + observedAt, + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0, + databasePoolTotal: 3, + databasePoolIdle: 3, + databasePoolWaiting: 0, + databasePoolWaitersMax: 0, + databasePoolOldestWaitMs: 0, + databasePoolWaitMsMax: 0 + } +} + +async function settleWorker(): Promise { + await Promise.resolve() + await Promise.resolve() +} diff --git a/cloud/apps/relay/src/regional-rehome-worker.ts b/cloud/apps/relay/src/regional-rehome-worker.ts new file mode 100644 index 00000000000..97c63a61025 --- /dev/null +++ b/cloud/apps/relay/src/regional-rehome-worker.ts @@ -0,0 +1,122 @@ +import { z } from 'zod' +import type { RelayAssignmentStore } from './assignment-store.js' +import type { RelayConfig } from './config.js' +import { googleMetadataIdentityToken } from './google-metadata-identity-token.js' +import type { RegionalRehomeSafetySnapshot } from './relay-observability.js' + +type RegionalRehomeWorkerOptions = { + fetch?: typeof fetch + identityToken?: (audience: string) => Promise + now?: () => number + intervalMs?: number + requestTimeoutMs?: number + safetySnapshot?: () => RegionalRehomeSafetySnapshot +} + +export type RegionalRehomeWorker = { + run: () => Promise + stop: () => void +} + +const RegionalHostDrainResponseSchema = z + .object({ + v: z.literal(1), + outcome: z.enum(['accepted', 'already-accepted', 'host-not-connected']) + }) + .strict() + +export function startRegionalRehomeWorker( + config: RelayConfig, + assignments: RelayAssignmentStore, + options: RegionalRehomeWorkerOptions = {} +): RegionalRehomeWorker | null { + if ( + config.role !== 'director' || + !config.rehomeAudience || + !config.rehomeDirectorServiceAccount || + !options.safetySnapshot + ) { + return null + } + const audience = config.rehomeAudience + const safetySnapshot = options.safetySnapshot + const now = options.now ?? Date.now + const fetchImpl = options.fetch ?? fetch + const tokenProvider = + options.identityToken ?? + ((audience: string) => googleMetadataIdentityToken(audience, fetchImpl)) + let stopped = false + let inFlight = false + const run = async (): Promise => { + if (stopped || inFlight) return + inFlight = true + let attemptId: string | null = null + try { + const processSafety = safetySnapshot() + const attempt = await assignments.claimRegionalRehome(processSafety) + if (!attempt) return + attemptId = attempt.attemptId + const token = await tokenProvider(audience) + const response = await fetchImpl( + new URL('/v1/admin/host-drain', attempt.sourceCellUrl), + { + method: 'POST', + headers: { + authorization: `Bearer ${token}`, + 'content-type': 'application/json' + }, + body: JSON.stringify({ + v: 1, + attemptId: attempt.attemptId, + userId: attempt.userId, + relayHostId: attempt.relayHostId, + sourceCellId: attempt.sourceCellId, + sourceCellIncarnation: attempt.sourceCellIncarnation, + sourceAssignmentEpoch: attempt.previousEpoch, + graceMs: attempt.drainGraceMs + }), + signal: AbortSignal.timeout(options.requestTimeoutMs ?? 10_000) + } + ) + if (!response.ok) throw new Error(`regional_rehome_source_${response.status}`) + const body = RegionalHostDrainResponseSchema.safeParse(await response.json()) + if (!body.success) throw new Error('regional_rehome_source_invalid_response') + await assignments.recordRegionalRehomeDrainReceipt( + attempt.attemptId, + body.data.outcome + ) + console.warn( + JSON.stringify({ + event: 'orca_relay_regional_rehome_dispatched', + sourceCellId: attempt.sourceCellId, + targetCellId: attempt.targetCellId, + outcome: body.data.outcome, + sendAttempts: attempt.sendAttempts + }) + ) + } catch (error) { + await (attemptId + ? assignments.recordRegionalRehomeDispatchFailure(attemptId) + : assignments.recordRegionalRehomeWorkerFailure() + ).catch(() => undefined) + console.warn( + JSON.stringify({ + event: 'orca_relay_regional_rehome_dispatch_failed', + reason: error instanceof Error ? error.message : 'unknown' + }) + ) + } finally { + inFlight = false + } + } + const timer = setInterval(() => void run(), options.intervalMs ?? 1_000) + timer.unref() + void run() + return { + run, + stop: () => { + stopped = true + clearInterval(timer) + } + } +} diff --git a/cloud/apps/relay/src/registered-migration-abandonment.ts b/cloud/apps/relay/src/registered-migration-abandonment.ts new file mode 100644 index 00000000000..9127f935f2a --- /dev/null +++ b/cloud/apps/relay/src/registered-migration-abandonment.ts @@ -0,0 +1,81 @@ +export const REGISTERED_MIGRATION_ABANDON_MS = 24 * 60 * 60 * 1_000 + +export const DURABLY_FENCED_MIGRATION_SOURCE = `( + EXISTS ( + SELECT 1 FROM relay_cell_committed_fences committed_source_fence + JOIN relay_cell_fence_attempts source_fence_attempt + ON source_fence_attempt.attempt_id = committed_source_fence.attempt_id + JOIN relay_cell_runtime source_runtime + ON source_runtime.cell_id = committed_source_fence.cell_id + WHERE committed_source_fence.cell_id = migration.source_cell_id + AND source_fence_attempt.cell_id = committed_source_fence.cell_id + AND source_fence_attempt.cell_incarnation = committed_source_fence.cell_incarnation + AND source_fence_attempt.completed_at IS NOT NULL + AND source_fence_attempt.aborted_at IS NULL + AND source_runtime.cell_incarnation = committed_source_fence.cell_incarnation + AND committed_source_fence.attested_at >= source_runtime.last_heartbeat_at + ) + OR EXISTS ( + SELECT 1 FROM relay_cell_legacy_fence_adoptions adopted_source_fence + JOIN relay_cell_runtime source_runtime + ON source_runtime.cell_id = adopted_source_fence.cell_id + WHERE adopted_source_fence.cell_id = migration.source_cell_id + AND source_runtime.cell_incarnation = adopted_source_fence.cell_incarnation + AND adopted_source_fence.attested_at >= source_runtime.last_heartbeat_at + ) +)` + +// Ordinary recovery waits on admission age; a completed fence proves the source cannot return. +export const ABANDONED_REGISTERED_MIGRATION = `( + migration.target_registered_at IS NOT NULL + AND + migration.expires_at <= ? + AND + ( + EXISTS ( + SELECT 1 FROM relay_cells target_cell + JOIN relay_cell_admission target_admission + ON target_admission.cell_id = target_cell.cell_id + WHERE target_cell.cell_id = migration.target_cell_id + AND target_cell.enabled = 0 + AND target_admission.admission_state = 'existing-only' + AND target_admission.updated_at <= ? + ) + OR ( + EXISTS ( + SELECT 1 FROM relay_cells source_cell + JOIN relay_cell_admission source_admission + ON source_admission.cell_id = source_cell.cell_id + WHERE source_cell.cell_id = migration.source_cell_id + AND source_cell.enabled = 0 + AND source_admission.admission_state = 'existing-only' + AND ( + source_admission.updated_at <= ? + OR ${DURABLY_FENCED_MIGRATION_SOURCE} + ) + ) + AND EXISTS ( + SELECT 1 FROM relay_cells target_cell + JOIN relay_cell_admission target_admission + ON target_admission.cell_id = target_cell.cell_id + WHERE target_cell.cell_id = migration.target_cell_id + AND target_cell.enabled = 1 + AND target_admission.admission_state IN ('migration-only', 'general') + ) + AND NOT EXISTS ( + SELECT 1 FROM relay_assignment_activity_leases source_activity + WHERE source_activity.user_id = migration.user_id + AND source_activity.relay_host_id = migration.relay_host_id + AND source_activity.cell_id = migration.source_cell_id + ) + ) + ) + AND NOT EXISTS ( + SELECT 1 FROM relay_assignment_activity_leases target_control + WHERE target_control.user_id = migration.user_id + AND target_control.relay_host_id = migration.relay_host_id + AND target_control.cell_id = migration.target_cell_id + AND target_control.activity_kind = 'control' + AND target_control.activity_id NOT LIKE 'control-pending:%' + ) +)` diff --git a/cloud/apps/relay/src/registered-migration-inventory-postgres.test.ts b/cloud/apps/relay/src/registered-migration-inventory-postgres.test.ts new file mode 100644 index 00000000000..ad318b2ad3a --- /dev/null +++ b/cloud/apps/relay/src/registered-migration-inventory-postgres.test.ts @@ -0,0 +1,139 @@ +import pg from 'pg' +import { afterAll, beforeAll, describe, expect, it } from 'vitest' +import { openRelayDatabase, type RelayDatabase } from './database.js' +import { REGISTERED_MIGRATION_ABANDON_MS } from './registered-migration-abandonment.js' +import { readRegisteredMigrationInventory } from './registered-migration-inventory.js' + +const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL +const describePostgres = databaseUrl ? describe : describe.skip +const schema = 'relay_migration_inventory_test' + +describePostgres('PostgreSQL registered migration inventory', () => { + let database: RelayDatabase | undefined + + beforeAll(async () => { + const client = new pg.Client({ connectionString: databaseUrl }) + await client.connect() + try { + await client.query(`DROP SCHEMA IF EXISTS ${schema} CASCADE`) + await client.query(`CREATE SCHEMA ${schema}`) + } finally { + await client.end() + } + const url = new URL(databaseUrl!) + url.searchParams.set('options', `-c search_path=${schema}`) + database = await openRelayDatabase({ databaseUrl: url.toString(), dataDir: '' }) + }) + + afterAll(async () => { + await database?.close() + const client = new pg.Client({ connectionString: databaseUrl }) + await client.connect() + try { + await client.query(`DROP SCHEMA IF EXISTS ${schema} CASCADE`) + } finally { + await client.end() + } + }) + + it('counts a disabled-target migration even while its source is active', async () => { + const now = REGISTERED_MIGRATION_ABANDON_MS + 10_000 + await database!.query( + `INSERT INTO relay_cells + (cell_id, cell_url, enabled, capacity_requests, reserved_requests, + observed_requests, last_heartbeat_at, updated_at) + VALUES ('cell-a', 'https://a.example.test', 1, 4000, 1, 0, 0, 1), + ('cell-b', 'https://b.example.test', 0, 4000, 0, 0, 0, 1)` + ) + await database!.query( + `INSERT INTO relay_cell_admission (cell_id, admission_state, updated_at) + VALUES ('cell-a', 'general', 1), ('cell-b', 'existing-only', 1)` + ) + await database!.query( + `INSERT INTO relay_assignments + (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, + last_activity_at, reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases) + VALUES ('user-1', '111111111111', 'cell-b', 2, ?, ?, 0, 0, 0, 0, 0, 0)`, + [now, now] + ) + await database!.query( + `INSERT INTO relay_assignment_activity_leases + (user_id, relay_host_id, activity_id, activity_kind, cell_id, + request_units, expires_at, updated_at) + VALUES ('user-1', '111111111111', 'control:source:1', 'control', 'cell-a', 1, ?, ?)`, + [now, now] + ) + await database!.query( + `INSERT INTO relay_assignment_migrations + (user_id, relay_host_id, source_cell_id, target_cell_id, previous_epoch, + assignment_epoch, source_request_units, target_reserved_units, expires_at, + target_registered_at, created_at, updated_at) + VALUES ('user-1', '111111111111', 'cell-a', 'cell-b', 1, 2, 1, 2, 2, 2, 1, 2)` + ) + + expect(await readRegisteredMigrationInventory(database!, now)).toEqual({ + open: 1, + inactive: 1, + abandoned: 1, + pairs: [ + { + sourceCellId: 'cell-a', + targetCellId: 'cell-b', + open: 1, + inactive: 1, + abandoned: 1 + } + ] + }) + + await database!.query( + `INSERT INTO relay_cells + (cell_id, cell_url, enabled, capacity_requests, reserved_requests, + observed_requests, last_heartbeat_at, updated_at) + VALUES ('cell-c', 'https://c.example.test', 1, 4000, 0, 0, 0, 1)` + ) + await database!.query( + `INSERT INTO relay_cell_admission (cell_id, admission_state, updated_at) + VALUES ('cell-c', 'migration-only', 1)` + ) + await database!.query( + `INSERT INTO relay_assignments + (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, + last_activity_at, reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases) + VALUES ('user-2', '222222222222', 'cell-c', 2, ?, ?, 0, 0, 0, 0, 0, 0)`, + [now, now] + ) + await database!.query( + `INSERT INTO relay_assignment_migrations + (user_id, relay_host_id, source_cell_id, target_cell_id, previous_epoch, + assignment_epoch, source_request_units, target_reserved_units, expires_at, + target_registered_at, created_at, updated_at) + VALUES ('user-2', '222222222222', 'cell-a', 'cell-c', 1, 2, 1, 2, ?, NULL, 1, 2)`, + [now + 60_000] + ) + + expect(await readRegisteredMigrationInventory(database!, now)).toEqual({ + open: 2, + inactive: 1, + abandoned: 1, + pairs: [ + { + sourceCellId: 'cell-a', + targetCellId: 'cell-b', + open: 1, + inactive: 1, + abandoned: 1 + }, + { + sourceCellId: 'cell-a', + targetCellId: 'cell-c', + open: 1, + inactive: 0, + abandoned: 0 + } + ] + }) + }) +}) diff --git a/cloud/apps/relay/src/registered-migration-inventory.test.ts b/cloud/apps/relay/src/registered-migration-inventory.test.ts new file mode 100644 index 00000000000..41aa26ac732 --- /dev/null +++ b/cloud/apps/relay/src/registered-migration-inventory.test.ts @@ -0,0 +1,117 @@ +import { describe, expect, it } from 'vitest' +import { openInMemoryRelayDatabase, type RelayDatabase } from './database.js' +import { REGISTERED_MIGRATION_ABANDON_MS } from './registered-migration-abandonment.js' +import { + formatRegisteredMigrationInventory, + readRegisteredMigrationInventory +} from './registered-migration-inventory.js' + +describe('registered migration inventory', () => { + it('does not restart abandonment when an old source migration lease is refreshed', async () => { + const database = await openInMemoryRelayDatabase() + const now = REGISTERED_MIGRATION_ABANDON_MS + 10_000 + await insertCell(database, 'cell-a', false, 'existing-only') + await insertCell(database, 'cell-b', true, 'migration-only') + await insertMigration(database, now - 1) + + const fresh = await readRegisteredMigrationInventory(database, now) + + expect(fresh).toEqual({ + open: 1, + inactive: 1, + abandoned: 1, + pairs: [ + { + sourceCellId: 'cell-a', + targetCellId: 'cell-b', + open: 1, + inactive: 1, + abandoned: 1 + } + ] + }) + expect(formatRegisteredMigrationInventory(fresh)).toEqual([ + '[orca-relay] migration inventory open=1 expiredRegisteredInactive=1 abandonedRegistered=1', + '[orca-relay] migration pair sourceCellId=cell-a targetCellId=cell-b open=1 expiredRegisteredInactive=1 abandoned=1' + ]) + expect( + ( + await readRegisteredMigrationInventory( + database, + now + 1 + ) + ).abandoned + ).toBe(1) + }) + + it('includes unregistered open migrations without logging a host identity', async () => { + const database = await openInMemoryRelayDatabase() + const now = 10_000 + await insertCell(database, 'cell-a', false, 'existing-only') + await insertCell(database, 'cell-b', true, 'migration-only') + await insertMigration(database, now + 60_000, null) + + const inventory = await readRegisteredMigrationInventory(database, now) + + expect(inventory).toEqual({ + open: 1, + inactive: 0, + abandoned: 0, + pairs: [ + { + sourceCellId: 'cell-a', + targetCellId: 'cell-b', + open: 1, + inactive: 0, + abandoned: 0 + } + ] + }) + expect(formatRegisteredMigrationInventory(inventory).join('\n')).not.toContain( + '111111111111' + ) + }) +}) + +async function insertCell( + database: RelayDatabase, + cellId: string, + enabled: boolean, + admissionState: string +): Promise { + await database.query( + `INSERT INTO relay_cells + (cell_id, cell_url, enabled, capacity_requests, reserved_requests, + observed_requests, last_heartbeat_at, updated_at) + VALUES (?, ?, ?, 4000, 0, 0, 0, 1)`, + [cellId, `https://${cellId}.example.test`, enabled ? 1 : 0] + ) + await database.query( + `INSERT INTO relay_cell_admission (cell_id, admission_state, updated_at) + VALUES (?, ?, 1)`, + [cellId, admissionState] + ) +} + +async function insertMigration( + database: RelayDatabase, + expiresAt: number, + targetRegisteredAt: number | null = 2 +): Promise { + await database.query( + `INSERT INTO relay_assignments + (user_id, relay_host_id, cell_id, assignment_epoch, lease_expires_at, + last_activity_at, reserved_controls, reserved_splices, reserved_invites, + pending_installs, pending_confirmations, migration_leases) + VALUES ('user-1', '111111111111', 'cell-b', 2, ?, ?, 0, 0, 0, 0, 0, 0)`, + [expiresAt, expiresAt] + ) + await database.query( + `INSERT INTO relay_assignment_migrations + (user_id, relay_host_id, source_cell_id, target_cell_id, previous_epoch, + assignment_epoch, source_request_units, target_reserved_units, expires_at, + target_registered_at, created_at, updated_at) + VALUES ('user-1', '111111111111', 'cell-a', 'cell-b', 1, 2, 1, 2, ?, ?, 1, 2)`, + [expiresAt, targetRegisteredAt] + ) +} diff --git a/cloud/apps/relay/src/registered-migration-inventory.ts b/cloud/apps/relay/src/registered-migration-inventory.ts new file mode 100644 index 00000000000..3d47ad31732 --- /dev/null +++ b/cloud/apps/relay/src/registered-migration-inventory.ts @@ -0,0 +1,104 @@ +import type { RelayDatabase, SqlRow } from './database.js' +import { + ABANDONED_REGISTERED_MIGRATION, + REGISTERED_MIGRATION_ABANDON_MS +} from './registered-migration-abandonment.js' + +export type RegisteredMigrationInventory = { + open: number + inactive: number + abandoned: number + pairs: Array<{ + sourceCellId: string + targetCellId: string + open: number + inactive: number + abandoned: number + }> +} + +export async function readRegisteredMigrationInventory( + database: RelayDatabase, + now: number +): Promise { + const abandonedBefore = now - REGISTERED_MIGRATION_ABANDON_MS + const openRows = await database.query( + `SELECT migration.source_cell_id, migration.target_cell_id, COUNT(*) AS open + FROM relay_assignment_migrations migration + WHERE migration.completed_at IS NULL AND migration.aborted_at IS NULL + GROUP BY migration.source_cell_id, migration.target_cell_id + ORDER BY migration.source_cell_id, migration.target_cell_id` + ) + const inactiveRows = await database.query( + `SELECT migration.source_cell_id, migration.target_cell_id, + COUNT(*) AS inactive, + COALESCE(SUM(CASE WHEN ${ABANDONED_REGISTERED_MIGRATION} + THEN 1 ELSE 0 END), 0) AS abandoned + FROM relay_assignment_migrations migration + WHERE migration.target_registered_at IS NOT NULL + AND migration.completed_at IS NULL AND migration.aborted_at IS NULL + AND migration.expires_at <= ? + AND NOT EXISTS ( + SELECT 1 FROM relay_assignment_activity_leases target_control + WHERE target_control.user_id = migration.user_id + AND target_control.relay_host_id = migration.relay_host_id + AND target_control.cell_id = migration.target_cell_id + AND target_control.activity_kind = 'control' + AND target_control.activity_id NOT LIKE 'control-pending:%' + ) + GROUP BY migration.source_cell_id, migration.target_cell_id + ORDER BY migration.source_cell_id, migration.target_cell_id`, + [now, abandonedBefore, abandonedBefore, now] + ) + const inactiveByPair = new Map( + inactiveRows.map((row) => [ + JSON.stringify([text(row, 'source_cell_id'), text(row, 'target_cell_id')]), + row + ]) + ) + const pairs = openRows.map((row) => { + const sourceCellId = text(row, 'source_cell_id') + const targetCellId = text(row, 'target_cell_id') + const inactive = inactiveByPair.get(JSON.stringify([sourceCellId, targetCellId])) + return { + sourceCellId, + targetCellId, + open: integer(row, 'open'), + inactive: integer(inactive, 'inactive'), + abandoned: integer(inactive, 'abandoned') + } + }) + return { + open: pairs.reduce((total, pair) => total + pair.open, 0), + inactive: pairs.reduce((total, pair) => total + pair.inactive, 0), + abandoned: pairs.reduce((total, pair) => total + pair.abandoned, 0), + pairs + } +} + +export function formatRegisteredMigrationInventory( + inventory: RegisteredMigrationInventory +): string[] { + const lines = [ + `[orca-relay] migration inventory open=${inventory.open}` + + ` expiredRegisteredInactive=${inventory.inactive}` + + ` abandonedRegistered=${inventory.abandoned}` + ] + for (const pair of inventory.pairs) { + lines.push( + `[orca-relay] migration pair sourceCellId=${pair.sourceCellId}` + + ` targetCellId=${pair.targetCellId} open=${pair.open}` + + ` expiredRegisteredInactive=${pair.inactive}` + + ` abandoned=${pair.abandoned}` + ) + } + return lines +} + +function text(row: SqlRow | undefined, column: string): string { + return String(row?.[column] ?? '') +} + +function integer(row: SqlRow | undefined, column: string): number { + return Number(row?.[column] ?? 0) +} diff --git a/cloud/apps/relay/src/relay-background-operation.test.ts b/cloud/apps/relay/src/relay-background-operation.test.ts new file mode 100644 index 00000000000..2c85317ba67 --- /dev/null +++ b/cloud/apps/relay/src/relay-background-operation.test.ts @@ -0,0 +1,44 @@ +import { describe, expect, it, vi } from 'vitest' +import { runRelayBackgroundOperation } from './relay-background-operation.js' + +describe('relay background operations', () => { + it('redacts free-form failure messages that could carry secrets', async () => { + const warn = vi.fn() + + await expect( + runRelayBackgroundOperation( + () => Promise.reject(new Error('postgresql://secret@database.invalid/relay')), + '[orca-relay] credential cleanup failed', + warn + ) + ).resolves.toBeUndefined() + + expect(warn).toHaveBeenCalledWith('[orca-relay] credential cleanup failed: Error: redacted') + expect(String(warn.mock.calls[0])).not.toContain('secret') + }) + + it('logs invariant slugs and SQLSTATE codes verbatim', async () => { + const warn = vi.fn() + const locked = Object.assign(new Error('regional_rehome_assignment_mismatch'), { + code: '55P03' + }) + + await runRelayBackgroundOperation( + () => Promise.reject(locked), + '[orca-relay] assignment cleanup failed', + warn + ) + + expect(warn).toHaveBeenCalledWith( + '[orca-relay] assignment cleanup failed: Error: regional_rehome_assignment_mismatch code=55P03' + ) + }) + + it('does not warn after successful maintenance work', async () => { + const warn = vi.fn() + + await runRelayBackgroundOperation(() => Promise.resolve(), 'unused', warn) + + expect(warn).not.toHaveBeenCalled() + }) +}) diff --git a/cloud/apps/relay/src/relay-background-operation.ts b/cloud/apps/relay/src/relay-background-operation.ts new file mode 100644 index 00000000000..be3943b7b1a --- /dev/null +++ b/cloud/apps/relay/src/relay-background-operation.ts @@ -0,0 +1,28 @@ +type Warn = (message: string) => void + +// A swallowed error hides which sweep step is failing (a silently dead +// cleanup chain stalled production rehoming for hours), but raw messages can +// embed secrets such as connection strings. Only two provably inert shapes +// are logged: this codebase's snake_case invariant slugs, and five-character +// SQLSTATE codes; everything else stays redacted. +function describeFailure(error: unknown): string { + const name = error instanceof Error ? error.name : typeof error + const message = error instanceof Error ? error.message : '' + const slug = /^[a-z0-9_]{1,64}$/.test(message) ? message : 'redacted' + const code = (error as { code?: unknown } | null)?.code + const sqlState = typeof code === 'string' && /^[0-9A-Z]{5}$/.test(code) ? ` code=${code}` : '' + return `${name}: ${slug}${sqlState}` +} + +export async function runRelayBackgroundOperation( + operation: () => Promise, + failureMessage: string, + warn: Warn = console.warn +): Promise { + try { + await operation() + } catch (error) { + // Dependency outages must fail readiness without crashing liveness. + warn(`${failureMessage}: ${describeFailure(error)}`) + } +} diff --git a/cloud/apps/relay/src/relay-connection-hard-cap.blackbox.test.ts b/cloud/apps/relay/src/relay-connection-hard-cap.blackbox.test.ts new file mode 100644 index 00000000000..230711e27d8 --- /dev/null +++ b/cloud/apps/relay/src/relay-connection-hard-cap.blackbox.test.ts @@ -0,0 +1,301 @@ +import { createHash, createHmac } from 'node:crypto' +import { generateKeyPair, exportJWK, SignJWT } from 'jose' +import { mkdtempSync, rmSync } from 'node:fs' +import { createServer, type Server } from 'node:http' +import { createServer as createNetServer } from 'node:net' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { + buildHostProofMacInput, + HOST_CHALLENGE_PLAINTEXT_DOMAIN +} from '@orca-cloud/relay-contract' +import nacl from 'tweetnacl' +import { afterEach, describe, expect, it, vi } from 'vitest' +import WebSocket from 'ws' +import type { RawData } from 'ws' +import type { RelayConfig } from './config.js' +import { openRelayDatabase, type RelayDatabase } from './database.js' +import { createRelayServer } from './relay-server.js' + +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 openSocket(url: string, headers: Record = {}): Promise { + return new Promise((resolve, reject) => { + const socket = new WebSocket(url, { headers }) + socket.once('open', () => resolve(socket)) + socket.once('error', reject) + }) +} + +function rejectedUpgrade( + url: string, + headers: Record = {} +): Promise { + return new Promise((resolve, reject) => { + const socket = new WebSocket(url, { headers }) + socket.once('open', () => reject(new Error(`upgrade unexpectedly opened: ${url}`))) + socket.once('unexpected-response', (_request, response) => { + response.resume() + resolve(response.statusCode) + }) + socket.once('error', () => undefined) + }) +} + +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) + }) +} + +async function proveControl( + socket: WebSocket, + hostId: string, + keyPair: nacl.BoxKeyPair, + rebind?: { secret: string; generation: number } +): Promise> { + socket.send( + JSON.stringify({ + type: 'host-hello', + v: 1, + relayHostId: hostId, + assignmentEpoch: 1, + hostPublicKeyB64: Buffer.from(keyPair.publicKey).toString('base64'), + appVersion: 'test', + ...(rebind + ? { + controlResumeSecret: rebind.secret, + previousGeneration: rebind.generation + } + : {}) + }) + ) + const challenge = await nextMessage(socket) + const plaintext = nacl.box.open( + Buffer.from(String(challenge.ciphertextB64), 'base64'), + Buffer.from(String(challenge.nonceB64), 'base64'), + Buffer.from(String(challenge.relayEphemeralPublicKeyB64), 'base64'), + keyPair.secretKey + ) + if (!plaintext) throw new Error('host challenge did not decrypt') + 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 await nextMessage(socket) +} + +describe('relay connection hard cap', () => { + const cleanup: Array<() => Promise | void> = [] + + afterEach(async () => { + for (const close of cleanup.splice(0).reverse()) await close() + }) + + it('preserves control headroom and never disturbs established sockets', async () => { + const keys = await generateKeyPair('ES256') + const publicJwk = await exportJWK(keys.publicKey) + const adminKeys = await generateKeyPair('RS256') + const adminPublicJwk = await exportJWK(adminKeys.publicKey) + const jwks = createServer((_request, response) => { + response.setHeader('content-type', 'application/json') + response.end( + JSON.stringify({ + keys: [ + { ...publicJwk, kid: 'test-key', alg: 'ES256', use: 'sig' }, + { ...adminPublicJwk, kid: 'admin-key', alg: 'RS256', use: 'sig' } + ] + }) + ) + }) + await new Promise((resolve) => jwks.listen(0, '127.0.0.1', resolve)) + cleanup.push(() => new Promise((resolve) => jwks.close(() => resolve()))) + const jwksAddress = jwks.address() + if (!jwksAddress || typeof jwksAddress === 'string') throw new Error('missing JWKS address') + const issuer = `http://127.0.0.1:${jwksAddress.port}` + const port = await unusedPort() + const relayUrl = `http://127.0.0.1:${port}` + const dataDir = mkdtempSync(join(tmpdir(), 'orca-relay-hard-cap-')) + cleanup.push(() => rmSync(dataDir, { recursive: true, force: true })) + const database: RelayDatabase = await openRelayDatabase({ dataDir }) + cleanup.push(() => database.close()) + const config = { + port, + publicUrl: relayUrl, + cellUrl: relayUrl, + authIssuer: issuer, + authAudience: 'orca-relay', + jwksUrl: `${issuer}/jwks`, + assignmentSigningKey: new TextEncoder().encode('test-assignment-key-with-at-least-32-bytes'), + role: 'combined', + cellId: 'combined', + cells: [ + { + id: 'combined', + url: relayUrl, + capacityRequests: 900, + connectionHardCap: 600, + connectionUnobservedBound: 60 + } + ], + adminAudience: `${relayUrl}/admin`, + deployServiceAccount: 'deploy@example.com', + runtimeServiceAccount: 'runtime@example.com', + connectionHardCap: 600, + connectionUnobservedBound: 60, + adminJwksUrl: `${issuer}/jwks`, + databasePoolMax: 10, + publicAssignmentsEnabled: true, + publicAssignmentConcurrency: 2, + publicAssignmentQueueMax: 128, + publicAssignmentWaitMs: 4_000, + publicResolveConcurrency: 1, + publicResolveWaitMs: 5_000, + publicAssignmentRetryAfterSeconds: 5, + dataDir + } satisfies RelayConfig + const relay = createRelayServer(config, database, { + connectionLedgerLimits: { hardCap: 5, controlReserve: 1 } + }) + 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 wsUrl = relayUrl.replace('http:', 'ws:') + const adminToken = await new SignJWT({ + email: 'deploy@example.com', + email_verified: true + }) + .setProtectedHeader({ alg: 'RS256', kid: 'admin-key' }) + .setIssuer('https://accounts.google.com') + .setAudience(`${relayUrl}/admin`) + .setSubject('deploy-subject') + .setIssuedAt() + .setExpirationTime('5m') + .sign(adminKeys.privateKey) + const statusResponse = await fetch(`${relayUrl}/v1/admin/runtime-status`, { + method: 'POST', + headers: { + authorization: `Bearer ${adminToken}`, + 'content-type': 'application/json' + }, + body: JSON.stringify({ v: 1 }) + }) + expect(statusResponse.status).toBe(200) + expect(await statusResponse.json()).toMatchObject({ + connectionCapacity: { + hardCap: 600, + controlRebindReserve: 100, + ordinaryConnectionLimit: 500, + unobservedBound: 60, + normalAdmissionPause: 440 + } + }) + + const unmatchedData = await Promise.all( + ['unmatched-1', 'unmatched-2', 'unmatched-3'].map((connectionId) => + openSocket(`${wsUrl}/v1/host/data/${connectionId}`) + ) + ) + cleanup.push(() => unmatchedData.forEach((socket) => socket.terminate())) + expect(relay.runtimeCounts().enforcedConnectionUnits).toBe(3) + expect(await rejectedUpgrade(`${wsUrl}/v1/connect/abcdefghijklmnop`)).toBe(503) + expect(unmatchedData.every((socket) => socket.readyState === WebSocket.OPEN)).toBe(true) + + const hostKeyPair = nacl.box.keyPair() + const hostId = createHash('sha256') + .update(hostKeyPair.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) + const controlHeaders = { authorization: `Bearer ${token}` } + const control = await openSocket(`${wsUrl}/v1/host/control`, controlHeaders) + cleanup.push(() => control.terminate()) + const controlAck = await proveControl(control, hostId, hostKeyPair) + expect(controlAck).toMatchObject({ type: 'host-hello-ack', generation: 1 }) + expect(relay.runtimeCounts().enforcedConnectionUnits).toBe(4) + expect(await rejectedUpgrade(`${wsUrl}/v1/host/data/another`)).toBe(503) + const unrelatedToken = await new SignJWT({ + prof: 'profile-1', + purpose: 'host-control', + relayHostId: 'ponmlkjihgfedcba' + }) + .setProtectedHeader({ alg: 'ES256', kid: 'test-key' }) + .setIssuer(issuer) + .setAudience('orca-relay') + .setSubject('user-2') + .setIssuedAt() + .setExpirationTime('5m') + .sign(keys.privateKey) + expect( + await rejectedUpgrade(`${wsUrl}/v1/host/control`, { + authorization: `Bearer ${unrelatedToken}` + }) + ).toBe(503) + const failedBorrower = await openSocket(`${wsUrl}/v1/host/control`, controlHeaders) + cleanup.push(() => failedBorrower.terminate()) + expect(relay.runtimeCounts().enforcedConnectionUnits).toBe(5) + const failedBorrowerClosed = new Promise((resolve) => + failedBorrower.once('close', (code) => resolve(code)) + ) + failedBorrower.send(JSON.stringify({ type: 'host-hello' })) + expect(await failedBorrowerClosed).toBe(4401) + await vi.waitFor(() => expect(relay.runtimeCounts().enforcedConnectionUnits).toBe(4)) + expect(control.readyState).toBe(WebSocket.OPEN) + const replacementControl = await openSocket(`${wsUrl}/v1/host/control`, controlHeaders) + cleanup.push(() => replacementControl.terminate()) + expect(relay.runtimeCounts().enforcedConnectionUnits).toBe(5) + const controlClosed = new Promise((resolve) => + control.once('close', (code) => resolve(code)) + ) + const replacementAck = await proveControl(replacementControl, hostId, hostKeyPair, { + secret: String(controlAck.controlResumeSecret), + generation: 1 + }) + expect(replacementAck).toMatchObject({ type: 'host-hello-ack', generation: 1 }) + expect(await controlClosed).toBe(4408) + await vi.waitFor(() => expect(relay.runtimeCounts().enforcedConnectionUnits).toBe(4)) + expect(relay.runtimeCounts().enforcedConnectionUnits).toBe(4) + }) +}) diff --git a/cloud/apps/relay/src/relay-connection-ledger.test.ts b/cloud/apps/relay/src/relay-connection-ledger.test.ts new file mode 100644 index 00000000000..f807c34d20d --- /dev/null +++ b/cloud/apps/relay/src/relay-connection-ledger.test.ts @@ -0,0 +1,172 @@ +import { EventEmitter } from 'node:events' +import { describe, expect, it } from 'vitest' +import type WebSocket from 'ws' +import { RelayConnectionLedger } from './relay-connection-ledger.js' + +function socket(): { webSocket: WebSocket; close: () => void } { + const emitter = new EventEmitter() + return { + webSocket: emitter as WebSocket, + close: () => emitter.emit('close') + } +} + +describe('RelayConnectionLedger', () => { + it('orders connection admissions before covering absolute snapshots', () => { + const ledger = new RelayConnectionLedger(6, 2) + const upgrade = ledger.tryReserveControl(false)! + const beforePromotion = ledger.snapshot() + const controlSocket = socket() + upgrade.promote(controlSocket.webSocket) + const afterPromotion = ledger.snapshot() + + expect(beforePromotion.inclusionWatermark).toBeGreaterThanOrEqual( + upgrade.inclusionWatermark + ) + expect(afterPromotion.inclusionWatermark).toBeGreaterThan( + beforePromotion.inclusionWatermark + ) + expect(afterPromotion).toMatchObject({ + physicalConnections: 1, + inFlightConnections: 0, + enforcedConnectionUnits: 1 + }) + }) + + it.each([600, 1_000, 3_000])( + 'enforces the %i-unit boundary with 100 control units reserved', + (hardCap) => { + const ledger = new RelayConnectionLedger(hardCap, 100) + const sockets: Array<{ webSocket: WebSocket; close: () => void }> = [] + for (let index = 0; index < (hardCap - 100) / 2; index++) { + const admission = ledger.tryReservePhone() + expect(admission).not.toBeNull() + const phoneSocket = socket() + admission!.upgrade.promote(phoneSocket.webSocket) + admission!.hostData.bind(`connection-${index}`) + sockets.push(phoneSocket) + } + expect(ledger.tryReservePhone()).toBeNull() + expect(ledger.tryReserveControl(false)).toBeNull() + for (let index = 0; index < 100; index++) { + const upgrade = ledger.tryReserveControl(true) + expect(upgrade).not.toBeNull() + const controlSocket = socket() + upgrade!.promote(controlSocket.webSocket) + sockets.push(controlSocket) + } + + expect(ledger.counts().enforcedConnectionUnits).toBe(hardCap) + expect(ledger.tryReserveControl(true)).toBeNull() + sockets.at(-1)!.close() + const replacement = ledger.tryReserveControl(true) + expect(replacement).not.toBeNull() + replacement!.release() + } + ) + + it('reserves both phone legs below the control reserve', () => { + const ledger = new RelayConnectionLedger(6, 2) + const first = ledger.tryReservePhone() + const second = ledger.tryReservePhone() + + expect(first).not.toBeNull() + expect(second).not.toBeNull() + expect(ledger.tryReservePhone()).toBeNull() + expect(ledger.counts()).toEqual({ + physicalConnections: 0, + inFlightConnections: 2, + reservedConnectionUnits: 2, + enforcedConnectionUnits: 4 + }) + expect(ledger.tryReserveControl(false)).toBeNull() + expect(ledger.tryReserveControl(true)).not.toBeNull() + expect(ledger.tryReserveControl(true)).not.toBeNull() + expect(ledger.tryReserveControl(true)).toBeNull() + }) + + it('transfers a pending host-data unit without increasing enforced capacity', () => { + const ledger = new RelayConnectionLedger(6, 2) + const phone = ledger.tryReservePhone()! + const phoneSocket = socket() + phone.upgrade.promote(phoneSocket.webSocket) + phone.hostData.bind('connection-1') + + const data = ledger.tryReserveHostData('connection-1')! + expect(ledger.counts().enforcedConnectionUnits).toBe(2) + const dataSocket = socket() + data.promote(dataSocket.webSocket) + expect(ledger.counts().enforcedConnectionUnits).toBe(2) + expect(data.commitHostData()).toBe(true) + + dataSocket.close() + dataSocket.close() + expect(ledger.counts()).toEqual({ + physicalConnections: 1, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 1 + }) + }) + + it('restores a claimed reservation after rejected host-data authentication', () => { + const ledger = new RelayConnectionLedger(6, 2) + const phone = ledger.tryReservePhone()! + const phoneSocket = socket() + phone.upgrade.promote(phoneSocket.webSocket) + phone.hostData.bind('connection-1') + + const rejected = ledger.tryReserveHostData('connection-1')! + const rejectedSocket = socket() + rejected.promote(rejectedSocket.webSocket) + rejectedSocket.close() + + expect(ledger.counts()).toEqual({ + physicalConnections: 1, + inFlightConnections: 0, + reservedConnectionUnits: 1, + enforcedConnectionUnits: 2 + }) + expect(ledger.tryReserveHostData('connection-1')).not.toBeNull() + }) + + it('does not restore a claimed reservation after the phone closes', () => { + const ledger = new RelayConnectionLedger(6, 2) + const phone = ledger.tryReservePhone()! + const phoneSocket = socket() + phone.upgrade.promote(phoneSocket.webSocket) + phone.hostData.bind('connection-1') + const data = ledger.tryReserveHostData('connection-1')! + const dataSocket = socket() + data.promote(dataSocket.webSocket) + + phone.hostData.release() + phoneSocket.close() + dataSocket.close() + + expect(ledger.counts()).toEqual({ + physicalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0 + }) + expect(ledger.tryReserveHostData('connection-1')).not.toBeNull() + }) + + it('releases failed in-flight upgrades exactly once', () => { + const ledger = new RelayConnectionLedger(6, 2) + const phone = ledger.tryReservePhone()! + + phone.upgrade.release() + phone.upgrade.release() + phone.hostData.release() + phone.hostData.release() + + expect(ledger.counts()).toEqual({ + physicalConnections: 0, + inFlightConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0 + }) + }) +}) diff --git a/cloud/apps/relay/src/relay-connection-ledger.ts b/cloud/apps/relay/src/relay-connection-ledger.ts new file mode 100644 index 00000000000..9c678695fa8 --- /dev/null +++ b/cloud/apps/relay/src/relay-connection-ledger.ts @@ -0,0 +1,229 @@ +import type WebSocket from 'ws' + +export type RelayConnectionLedgerCounts = { + physicalConnections: number + inFlightConnections: number + reservedConnectionUnits: number + enforcedConnectionUnits: number +} + +export type RelayConnectionLedgerSnapshot = RelayConnectionLedgerCounts & { + inclusionWatermark: number +} + +export type PendingHostDataReservation = { + bind: (connectionId: string) => void + release: () => void +} + +type ReservationState = 'reserved' | 'claimed' | 'consumed' | 'released' +type UpgradeState = 'in-flight' | 'physical' | 'released' + +class HostDataReservation implements PendingHostDataReservation { + private state: ReservationState = 'reserved' + private connectionId: string | null = null + + constructor(private readonly ledger: RelayConnectionLedger) {} + + bind(connectionId: string): void { + if (this.state !== 'reserved' || this.connectionId !== null) { + throw new Error('host_data_reservation_already_bound') + } + this.connectionId = connectionId + this.ledger.bindHostData(connectionId, this) + } + + release(): void { + if (this.state === 'reserved') { + this.ledger.releaseReserved(this) + } + if (this.state !== 'consumed') { + this.state = 'released' + } + } + + claim(): void { + if (this.state !== 'reserved') throw new Error('host_data_reservation_unavailable') + this.state = 'claimed' + } + + consume(): boolean { + if (this.state !== 'claimed') return false + this.state = 'consumed' + return true + } + + restore(): void { + if (this.state !== 'claimed') return + this.state = this.ledger.restoreReserved(this) ? 'reserved' : 'released' + } + + matches(connectionId: string): boolean { + return this.connectionId === connectionId + } + + get boundConnectionId(): string | null { + return this.connectionId + } +} + +export class RelayConnectionUpgrade { + private state: UpgradeState = 'in-flight' + + constructor( + private readonly ledger: RelayConnectionLedger, + readonly inclusionWatermark: number, + private readonly hostDataReservation: HostDataReservation | null = null + ) {} + + promote(socket: WebSocket): void { + if (this.state !== 'in-flight') throw new Error('connection_upgrade_not_in_flight') + this.state = 'physical' + this.ledger.promoteUpgrade() + socket.once('close', () => this.release()) + } + + commitHostData(): boolean { + return this.hostDataReservation?.consume() ?? false + } + + release(): void { + if (this.state === 'released') return + if (this.state === 'in-flight') this.ledger.releaseUpgrade() + else this.ledger.releasePhysical() + this.state = 'released' + this.hostDataReservation?.restore() + } +} + +export type PhoneConnectionAdmission = { + upgrade: RelayConnectionUpgrade + hostData: PendingHostDataReservation +} + +export class RelayConnectionLedger { + private physicalConnections = 0 + private inFlightConnections = 0 + private reservedConnectionUnits = 0 + private inclusionWatermark = 0 + private readonly hostDataByConnectionId = new Map() + + constructor( + private readonly hardCap: number, + private readonly controlReserve: number + ) { + if (hardCap <= 0 || controlReserve < 0 || controlReserve >= hardCap) { + throw new Error('invalid_connection_ledger_capacity') + } + } + + tryReserveControl(rebind: boolean): RelayConnectionUpgrade | null { + return this.reserveUpgrade(rebind ? this.hardCap : this.normalAdmissionLimit) + } + + tryReservePhone(): PhoneConnectionAdmission | null { + if (this.enforcedConnectionUnits + 2 > this.normalAdmissionLimit) return null + const reservation = new HostDataReservation(this) + const inclusionWatermark = this.advanceWatermark() + this.inFlightConnections++ + this.reservedConnectionUnits++ + return { + upgrade: new RelayConnectionUpgrade(this, inclusionWatermark), + hostData: reservation + } + } + + tryReserveHostData(connectionId: string): RelayConnectionUpgrade | null { + const reservation = this.hostDataByConnectionId.get(connectionId) + if (reservation?.matches(connectionId)) { + this.hostDataByConnectionId.delete(connectionId) + reservation.claim() + const inclusionWatermark = this.advanceWatermark() + this.reservedConnectionUnits-- + this.inFlightConnections++ + return new RelayConnectionUpgrade(this, inclusionWatermark, reservation) + } + return this.reserveUpgrade(this.normalAdmissionLimit) + } + + counts(): RelayConnectionLedgerCounts { + return { + physicalConnections: this.physicalConnections, + inFlightConnections: this.inFlightConnections, + reservedConnectionUnits: this.reservedConnectionUnits, + enforcedConnectionUnits: this.enforcedConnectionUnits + } + } + + snapshot(): RelayConnectionLedgerSnapshot { + return { + ...this.counts(), + inclusionWatermark: this.advanceWatermark() + } + } + + bindHostData(connectionId: string, reservation: HostDataReservation): void { + if (this.hostDataByConnectionId.has(connectionId)) { + throw new Error('host_data_reservation_conflict') + } + this.hostDataByConnectionId.set(connectionId, reservation) + } + + releaseReserved(reservation: HostDataReservation): void { + const connectionId = reservation.boundConnectionId + if (connectionId && this.hostDataByConnectionId.get(connectionId) === reservation) { + this.hostDataByConnectionId.delete(connectionId) + } + this.advanceWatermark() + this.reservedConnectionUnits-- + } + + restoreReserved(reservation: HostDataReservation): boolean { + const connectionId = reservation.boundConnectionId + if (!connectionId || this.hostDataByConnectionId.has(connectionId)) { + return false + } + this.advanceWatermark() + this.reservedConnectionUnits++ + this.hostDataByConnectionId.set(connectionId, reservation) + return true + } + + promoteUpgrade(): void { + this.advanceWatermark() + this.inFlightConnections-- + this.physicalConnections++ + } + + releaseUpgrade(): void { + this.advanceWatermark() + this.inFlightConnections-- + } + + releasePhysical(): void { + this.advanceWatermark() + this.physicalConnections-- + } + + private reserveUpgrade(limit: number): RelayConnectionUpgrade | null { + if (this.enforcedConnectionUnits + 1 > limit) return null + const inclusionWatermark = this.advanceWatermark() + this.inFlightConnections++ + return new RelayConnectionUpgrade(this, inclusionWatermark) + } + + private advanceWatermark(): number { + this.inclusionWatermark++ + return this.inclusionWatermark + } + + private get normalAdmissionLimit(): number { + return this.hardCap - this.controlReserve + } + + private get enforcedConnectionUnits(): number { + return ( + this.physicalConnections + this.inFlightConnections + this.reservedConnectionUnits + ) + } +} diff --git a/cloud/apps/relay/src/relay-first-frame-failure.blackbox.test.ts b/cloud/apps/relay/src/relay-first-frame-failure.blackbox.test.ts new file mode 100644 index 00000000000..57e47a065d3 --- /dev/null +++ b/cloud/apps/relay/src/relay-first-frame-failure.blackbox.test.ts @@ -0,0 +1,124 @@ +import { createServer as createNetServer } from 'node:net' +import { RELAY_CLOSE_CODE } from '@orca-cloud/relay-contract' +import { afterEach, describe, expect, it, vi } from 'vitest' +import WebSocket from 'ws' +import type { RelayConfig } from './config.js' +import type { RelayDatabase } from './database.js' +import { createRelayServer } from './relay-server.js' + +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 openSocket(url: string): Promise { + return new Promise((resolve, reject) => { + const socket = new WebSocket(url) + socket.once('open', () => resolve(socket)) + socket.once('error', reject) + }) +} + +describe('relay first-frame failures', () => { + const cleanup: Array<() => Promise | void> = [] + + afterEach(async () => { + for (const close of cleanup.splice(0).reverse()) await close() + vi.restoreAllMocks() + }) + + it('contains phone authentication failures and releases connection capacity', async () => { + const port = await unusedPort() + const relayUrl = `http://127.0.0.1:${port}` + const poolTimeout = new Error('injected pool connection timeout') + const database: RelayDatabase = { + query: vi.fn(async () => { + throw poolTimeout + }), + queryLocked: vi.fn(async () => { + throw poolTimeout + }), + transaction: vi.fn(async (operation) => await operation(database)), + close: vi.fn(async () => undefined) + } + const config = { + port, + publicUrl: relayUrl, + cellUrl: relayUrl, + authIssuer: 'https://auth.example.com', + authAudience: 'orca-relay', + jwksUrl: 'https://auth.example.com/jwks', + 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: 'https://auth.example.com/admin-jwks', + databasePoolMax: 10, + publicAssignmentsEnabled: true, + publicAssignmentConcurrency: 2, + publicAssignmentQueueMax: 128, + publicAssignmentWaitMs: 4_000, + publicResolveConcurrency: 1, + publicResolveWaitMs: 5_000, + publicAssignmentRetryAfterSeconds: 5, + dataDir: './test-data' + } satisfies RelayConfig + const relay = createRelayServer(config, database, { + connectionLedgerLimits: { hardCap: 5, controlReserve: 1 } + }) + 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()))) + vi.spyOn(console, 'log').mockImplementation(() => undefined) + vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const unhandled: unknown[] = [] + const onUnhandled = (reason: unknown): void => { + unhandled.push(reason) + } + process.on('unhandledRejection', onUnhandled) + cleanup.push(() => { + process.off('unhandledRejection', onUnhandled) + }) + + const socket = await openSocket( + `${relayUrl.replace('http:', 'ws:')}/v1/connect/abcdefghijklmnop` + ) + cleanup.push(() => socket.terminate()) + expect(relay.connectionSnapshot()).toMatchObject({ + physicalConnections: 1, + reservedConnectionUnits: 1, + enforcedConnectionUnits: 2 + }) + const closed = new Promise((resolve) => + socket.once('close', (code) => resolve(code)) + ) + socket.send( + JSON.stringify({ + type: 'relay-auth', + v: 1, + mode: 'connect', + credential: Buffer.alloc(32, 1).toString('base64url') + }) + ) + + expect(await closed).toBe(RELAY_CLOSE_CODE.LIMIT_EXCEEDED) + await vi.waitFor(() => + expect(relay.connectionSnapshot()).toMatchObject({ + physicalConnections: 0, + reservedConnectionUnits: 0, + enforcedConnectionUnits: 0 + }) + ) + await new Promise((resolve) => setImmediate(resolve)) + expect(unhandled).toEqual([]) + }) +}) diff --git a/cloud/apps/relay/src/relay-host-log-digest.ts b/cloud/apps/relay/src/relay-host-log-digest.ts new file mode 100644 index 00000000000..b232170a6d4 --- /dev/null +++ b/cloud/apps/relay/src/relay-host-log-digest.ts @@ -0,0 +1,7 @@ +import { createHash } from 'node:crypto' + +// Store-level 503s were previously silent; the digest keeps per-host log +// correlation possible without emitting the raw relay host id. +export function relayHostLogDigest(relayHostId: string): string { + return createHash('sha256').update(relayHostId).digest('hex').slice(0, 12) +} diff --git a/cloud/apps/relay/src/relay-host-proof-failure.blackbox.test.ts b/cloud/apps/relay/src/relay-host-proof-failure.blackbox.test.ts new file mode 100644 index 00000000000..924523685e3 --- /dev/null +++ b/cloud/apps/relay/src/relay-host-proof-failure.blackbox.test.ts @@ -0,0 +1,153 @@ +import { createHash } from 'node:crypto' +import { createServer, type Server } from 'node:http' +import { createServer as createNetServer } from 'node:net' +import { exportJWK, generateKeyPair, SignJWT } from 'jose' +import { RELAY_CLOSE_CODE } from '@orca-cloud/relay-contract' +import nacl from 'tweetnacl' +import { afterEach, describe, expect, it, vi } from 'vitest' +import WebSocket from 'ws' +import type { RelayConfig } from './config.js' +import type { RelayDatabase } from './database.js' +import { createRelayServer } from './relay-server.js' + +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 +} + +describe('relay host proof failures', () => { + const cleanup: Array<() => Promise | void> = [] + + afterEach(async () => { + for (const close of cleanup.splice(0).reverse()) await close() + vi.restoreAllMocks() + }) + + // The production crash signature: a pg-pool connect timeout rejecting out of + // verifyCellAssignment inside beginProof killed whole cells as an unhandled + // rejection. The guard must contain it to this one handshake. + it('contains a database timeout during host hello to one socket', async () => { + 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 port = await unusedPort() + const relayUrl = `http://127.0.0.1:${port}` + const poolTimeout = new Error('Connection terminated due to connection timeout') + const database: RelayDatabase = { + query: vi.fn(async () => { + throw poolTimeout + }), + queryLocked: vi.fn(async () => { + throw poolTimeout + }), + transaction: vi.fn(async (operation) => await operation(database)), + close: vi.fn(async () => undefined) + } + const config = { + 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: 10, + publicAssignmentsEnabled: true, + publicAssignmentConcurrency: 2, + publicAssignmentQueueMax: 128, + publicAssignmentWaitMs: 4_000, + publicResolveConcurrency: 1, + publicResolveWaitMs: 5_000, + publicAssignmentRetryAfterSeconds: 5, + dataDir: './test-data' + } satisfies RelayConfig + const relay = createRelayServer(config, database) + 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()))) + vi.spyOn(console, 'log').mockImplementation(() => undefined) + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const unhandled: unknown[] = [] + const onUnhandled = (reason: unknown): void => { + unhandled.push(reason) + } + process.on('unhandledRejection', onUnhandled) + cleanup.push(() => { + process.off('unhandledRejection', onUnhandled) + }) + + 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) + const socket = new WebSocket(`${relayUrl.replace('http:', 'ws:')}/v1/host/control`, { + headers: { authorization: `Bearer ${token}` }, + perMessageDeflate: false + }) + cleanup.push(() => socket.terminate()) + await new Promise((resolve, reject) => { + socket.once('open', resolve) + socket.once('error', reject) + }) + const closed = new Promise((resolve) => + socket.once('close', (code) => resolve(code)) + ) + socket.send( + JSON.stringify({ + type: 'host-hello', + v: 1, + relayHostId: hostId, + assignmentEpoch: 1, + hostPublicKeyB64: Buffer.from(keyPair.publicKey).toString('base64'), + appVersion: 'test' + }) + ) + + expect(await closed).toBe(RELAY_CLOSE_CODE.LIMIT_EXCEEDED) + expect( + warn.mock.calls.map((call) => String(call[0])) + ).toContain( + '[orca-relay] host hello proof failed: Connection terminated due to connection timeout' + ) + await new Promise((resolve) => setImmediate(resolve)) + expect(unhandled).toEqual([]) + }) +}) diff --git a/cloud/apps/relay/src/relay-observability.test.ts b/cloud/apps/relay/src/relay-observability.test.ts new file mode 100644 index 00000000000..2b9ceb0b72a --- /dev/null +++ b/cloud/apps/relay/src/relay-observability.test.ts @@ -0,0 +1,246 @@ +import { describe, expect, it, vi } from 'vitest' +import type { RelayDatabase } from './database.js' +import { observeRelayDatabase } from './observed-relay-database.js' +import { + observedRelayRequests, + RelayObservability, + type RelayProcessCounts +} from './relay-observability.js' + +const counts: RelayProcessCounts = { + totalConnections: 9, + preAuthConnections: 1, + controls: 2, + splices: 3, + pendingSplices: 1, + queuedBytes: 4096, + databasePoolTotal: 3, + databasePoolIdle: 0, + databasePoolWaiting: 2, + databasePoolWaitersMax: 3, + databasePoolOldestWaitMs: 750, + databasePoolWaitMsMax: 1_250 +} + +describe('relay observability', () => { + it('emits safe readiness dependency outcomes', () => { + const entries: Array> = [] + const observability = new RelayObservability( + { role: 'cell', cellId: 'production-gce-c28', region: 'asia-east2' }, + (entry) => entries.push(entry) + ) + + observability.recordReadiness({ + ready: false, + failure: 'sql_failed', + jwksLatencyMs: 12, + sqlLatencyMs: 2_001, + totalLatencyMs: 2_013 + }) + + expect(entries).toEqual([ + { + severity: 'WARNING', + message: 'Orca Relay readiness check', + event: 'orca_relay_readiness_check', + metricVersion: 1, + role: 'cell', + cellId: 'production-gce-c28', + region: 'asia-east2', + ready: false, + failure: 'sql_failed', + jwksLatencyMs: 12, + sqlLatencyMs: 2_001, + totalLatencyMs: 2_013 + } + ]) + }) + + it('excludes sockets stuck in closing state from observed relay work', () => { + expect(observedRelayRequests(counts)).toBe(7) + }) + + it('keeps rejection reasons separate per lane and resets them each flush', () => { + const entries: Array> = [] + const observability = new RelayObservability( + { role: 'director', cellId: 'director', region: 'us-central1' }, + (entry) => entries.push(entry) + ) + observability.recordAssignmentAdmission('placement-rejected') + observability.recordAssignmentRejectionReason('placement', 'host-rate-limited') + observability.recordAssignmentRejectionReason('placement', 'host-rate-limited') + observability.recordAssignmentRejectionReason('placement', 'queue-full') + observability.recordAssignmentRejectionReason('sticky', 'wait-timeout') + observability.flush(counts) + observability.flush(counts) + + expect(entries[0]).toMatchObject({ + placementAssignmentRejectionsDelta: 1, + placementRejectionsByReasonDelta: { 'host-rate-limited': 2, 'queue-full': 1 }, + stickyRejectionsByReasonDelta: { 'wait-timeout': 1 } + }) + expect(entries[1]).toMatchObject({ + placementRejectionsByReasonDelta: {}, + stickyRejectionsByReasonDelta: {} + }) + }) + + it('aggregates coarse region requests, selections, fallbacks, and outages', () => { + const entries: Array> = [] + const observability = new RelayObservability( + { role: 'director', cellId: 'director', region: 'us-central1' }, + (entry) => entries.push(entry) + ) + observability.recordRegionRequest('asia-east2') + observability.recordRegionRequest(undefined) + observability.recordRegionSelection({ + targetRegion: 'asia-east2', + selectedRegion: 'us-central1', + fallback: true + }) + observability.recordRegionSelection({ targetRegion: 'asia-east2', fallback: false }) + observability.flush(counts) + observability.flush(counts) + + expect(entries[0]).toMatchObject({ + requestedRegionsDelta: { 'asia-east2': 1, unhinted: 1 }, + selectedRegionsDelta: { 'us-central1': 1 }, + regionFallbacksDelta: { 'asia-east2': 1 }, + unavailableRegionsDelta: { 'asia-east2': 1 } + }) + expect(entries[1]).toMatchObject({ + requestedRegionsDelta: {}, + selectedRegionsDelta: {}, + regionFallbacksDelta: {}, + unavailableRegionsDelta: {} + }) + }) + + it('emits bounded aggregate runtime signals without identities or credentials', () => { + const entries: Array> = [] + const observability = new RelayObservability( + { role: 'cell', cellId: 'staging-c1', region: 'asia-east2' }, + (entry) => entries.push(entry) + ) + observability.recordAuth(true) + observability.recordAuth(false) + observability.recordForwardedBytes(123) + observability.recordHttp(45.6789) + observability.recordReconnect() + observability.recordSql(12.3456, true) + observability.recordSql(4, false) + observability.recordControlRenewal(2, 'renewed') + observability.recordControlRenewal(8, 'control_activity_not_found') + observability.recordControlRenewal(4, 'renewed') + observability.recordControlActivityRecovery(true) + observability.recordControlActivityRecovery(false) + observability.flush(counts) + observability.flush(counts) + + expect(entries[0]).toMatchObject({ + event: 'orca_relay_runtime_metrics', + metricVersion: 2, + role: 'cell', + cellId: 'staging-c1', + region: 'asia-east2', + ...counts, + forwardedBytesDelta: 123, + authSuccessesDelta: 1, + authFailuresDelta: 1, + reconnectsDelta: 1, + sqlQueriesDelta: 2, + sqlFailuresDelta: 1, + sqlLatencyMsMax: 12.346, + controlRenewalsByOutcomeDelta: { renewed: 2, control_activity_not_found: 1 }, + controlRenewalsDelta: 3, + controlRenewalSuccessesDelta: 2, + controlRenewalLeaseMissesDelta: 1, + controlRenewalLatencyMsP50: 4, + controlRenewalLatencyMsP95: 8, + controlRenewalLatencyMsMax: 8, + controlActivityRecoveriesDelta: 1, + controlActivityRecoveryFailuresDelta: 1, + httpLatencyMsMax: 45.679 + }) + expect(entries[1]).toMatchObject({ + forwardedBytesDelta: 0, + authSuccessesDelta: 0, + authFailuresDelta: 0, + reconnectsDelta: 0, + sqlQueriesDelta: 0, + sqlFailuresDelta: 0, + sqlLatencyMsMax: 0, + controlRenewalsByOutcomeDelta: {}, + controlRenewalsDelta: 0, + controlRenewalSuccessesDelta: 0, + controlRenewalLeaseMissesDelta: 0, + controlRenewalLatencyMsP50: 0, + controlRenewalLatencyMsP95: 0, + controlRenewalLatencyMsMax: 0, + controlActivityRecoveriesDelta: 0, + controlActivityRecoveryFailuresDelta: 0, + httpLatencyMsMax: 0 + }) + expect(JSON.stringify(entries)).not.toMatch(/token|credential|userId|relayHostId/) + }) + + it('aggregates control and splice closes as bounded per-reason deltas', () => { + const entries: Array> = [] + const observability = new RelayObservability( + { role: 'cell', cellId: 'staging-c1', region: 'us-central1' }, + (entry) => entries.push(entry) + ) + observability.recordControlClose(1006) + observability.recordControlClose(1006) + observability.recordControlClose(4402) + observability.recordSpliceClose('host-oversize-frame') + observability.recordSpliceClose('queue-limit') + observability.flush(counts) + observability.flush(counts) + + expect(entries[0]).toMatchObject({ + controlClosesByCodeDelta: { 1006: 2, 4402: 1 }, + spliceClosesByTriggerDelta: { 'host-oversize-frame': 1, 'queue-limit': 1 } + }) + expect(entries[1]).toMatchObject({ + controlClosesByCodeDelta: {}, + spliceClosesByTriggerDelta: {} + }) + }) + + it('observes successful and failed database calls including transactions', async () => { + const recordSql = vi.fn() + const underlying: RelayDatabase = { + query: vi.fn(async () => [{ ok: true }]), + queryLocked: vi.fn(async (sql, _params, options) => { + throw new Error( + options?.failIfUnavailable && sql === 'SELECT 3' + ? 'database_lock_unavailable' + : 'database unavailable' + ) + }), + transaction: async (operation) => await operation(underlying), + close: vi.fn(async () => {}) + } + const database = observeRelayDatabase(underlying, { + recordAuth: vi.fn(), + recordForwardedBytes: vi.fn(), + recordHttp: vi.fn(), + recordReconnect: vi.fn(), + recordSql + }) + + await database.query('SELECT 1') + await expect(database.transaction(async (tx) => await tx.queryLocked('SELECT 2'))).rejects.toThrow( + 'database unavailable' + ) + await expect( + database.queryLocked('SELECT 3', [], { failIfUnavailable: true }) + ).rejects.toThrow('database_lock_unavailable') + await expect( + database.queryLocked('SELECT 4', [], { failIfUnavailable: true }) + ).rejects.toThrow('database unavailable') + expect(recordSql).toHaveBeenCalledTimes(4) + expect(recordSql.mock.calls.map((call) => call[1])).toEqual([true, false, true, false]) + }) +}) diff --git a/cloud/apps/relay/src/relay-observability.ts b/cloud/apps/relay/src/relay-observability.ts new file mode 100644 index 00000000000..6125ede8d1a --- /dev/null +++ b/cloud/apps/relay/src/relay-observability.ts @@ -0,0 +1,336 @@ +import { monitorEventLoopDelay, performance } from 'node:perf_hooks' +import type { RelayRegion } from '@orca-cloud/relay-contract' +import type { ControlRenewalOutcome } from './assignment-store.js' +import type { PostgresPoolPressureCounts } from './postgres-pool-pressure.js' +import type { RelayReadinessObservation } from './relay-readiness.js' + +export type RelayRuntimeCounts = { + totalConnections: number + inFlightConnections?: number + reservedConnectionUnits?: number + enforcedConnectionUnits?: number + preAuthConnections: number + controls: number + splices: number + pendingSplices: number + queuedBytes: number +} + +export function observedRelayRequests(counts: RelayRuntimeCounts): number { + return counts.preAuthConnections + counts.controls + counts.splices + counts.pendingSplices +} + +export type RelayProcessCounts = RelayRuntimeCounts & PostgresPoolPressureCounts + +export type RegionalRehomeRuntimeSafety = { + observedAt: number + sqlFailures: number + reconnects: number + controlActivityRecoveryFailures: number +} + +export type RegionalRehomeSafetySnapshot = RegionalRehomeRuntimeSafety & { + databasePoolWaiting: number + databasePoolWaitersMax: number + databasePoolWaitMsMax: number +} + +export type AssignmentAdmissionOutcome = + | 'sticky' + | 'sticky-rejected' + | 'placement' + | 'placement-rejected' + +export type AssignmentAdmissionLane = 'sticky' | 'placement' + +export interface RelayRuntimeObserver { + recordAuth(success: boolean): void + recordForwardedBytes(bytes: number): void + recordHttp(durationMs: number): void + recordReconnect(): void + recordSql(durationMs: number, success: boolean): void + recordControlRenewal?(durationMs: number, outcome: ControlRenewalOutcome): void + recordControlActivityRecovery?(success: boolean): void + recordAssignmentAdmission?(outcome: AssignmentAdmissionOutcome): void + recordAssignmentRejectionReason?(lane: AssignmentAdmissionLane, reason: string): void + recordRegionRequest?(region: RelayRegion | undefined): void + recordRegionSelection?(input: { + targetRegion: RelayRegion + selectedRegion?: RelayRegion + fallback: boolean + }): void + recordControlClose?(code: number): void + recordSpliceClose?(trigger: string): void +} + +type RelayMetricDeltas = { + forwardedBytes: number + authSuccesses: number + authFailures: number + reconnects: number + sqlQueries: number + sqlFailures: number + sqlLatencyMsMax: number + httpLatencyMsMax: number + stickyAssignments: number + stickyAssignmentRejections: number + placementAssignments: number + placementAssignmentRejections: number + stickyRejectionsByReason: Record + placementRejectionsByReason: Record + requestedRegions: Record + selectedRegions: Record + regionFallbacks: Record + unavailableRegions: Record + controlClosesByCode: Record + spliceClosesByTrigger: Record + controlRenewalLatenciesMs: number[] + controlRenewalsByOutcome: Record + controlActivityRecoveries: number + controlActivityRecoveryFailures: number +} + +type MetricWriter = (entry: Record) => void + +const emptyDeltas = (): RelayMetricDeltas => ({ + forwardedBytes: 0, + authSuccesses: 0, + authFailures: 0, + reconnects: 0, + sqlQueries: 0, + sqlFailures: 0, + sqlLatencyMsMax: 0, + httpLatencyMsMax: 0, + stickyAssignments: 0, + stickyAssignmentRejections: 0, + placementAssignments: 0, + placementAssignmentRejections: 0, + stickyRejectionsByReason: {}, + placementRejectionsByReason: {}, + requestedRegions: {}, + selectedRegions: {}, + regionFallbacks: {}, + unavailableRegions: {}, + controlClosesByCode: {}, + spliceClosesByTrigger: {}, + controlRenewalLatenciesMs: [], + controlRenewalsByOutcome: {}, + controlActivityRecoveries: 0, + controlActivityRecoveryFailures: 0 +}) + +function percentile(values: number[], percentileRank: number): number { + if (values.length === 0) return 0 + const sorted = [...values].sort((left, right) => left - right) + return sorted[Math.ceil(percentileRank * sorted.length) - 1] ?? 0 +} + +export class RelayObservability implements RelayRuntimeObserver { + private readonly eventLoop = monitorEventLoopDelay({ resolution: 20 }) + private deltas = emptyDeltas() + private timer: ReturnType | null = null + private lastFlushAt = 0 + private lastFlushedSafety = { + sqlFailures: 0, + reconnects: 0, + controlActivityRecoveryFailures: 0 + } + + constructor( + private readonly identity: { role: string; cellId: string; region: RelayRegion }, + private readonly write: MetricWriter = (entry) => console.log(JSON.stringify(entry)) + ) {} + + recordAuth(success: boolean): void { + if (success) this.deltas.authSuccesses++ + else this.deltas.authFailures++ + } + + recordForwardedBytes(bytes: number): void { + this.deltas.forwardedBytes += bytes + } + + recordHttp(durationMs: number): void { + this.deltas.httpLatencyMsMax = Math.max(this.deltas.httpLatencyMsMax, durationMs) + } + + recordReconnect(): void { + this.deltas.reconnects++ + } + + recordAssignmentAdmission(outcome: AssignmentAdmissionOutcome): void { + if (outcome === 'sticky') this.deltas.stickyAssignments++ + else if (outcome === 'sticky-rejected') this.deltas.stickyAssignmentRejections++ + else if (outcome === 'placement') this.deltas.placementAssignments++ + else this.deltas.placementAssignmentRejections++ + } + + recordAssignmentRejectionReason(lane: AssignmentAdmissionLane, reason: string): void { + const counts = + lane === 'sticky' + ? this.deltas.stickyRejectionsByReason + : this.deltas.placementRejectionsByReason + counts[reason] = (counts[reason] ?? 0) + 1 + } + + recordRegionRequest(region: RelayRegion | undefined): void { + increment(this.deltas.requestedRegions, region ?? 'unhinted') + } + + recordRegionSelection(input: { + targetRegion: RelayRegion + selectedRegion?: RelayRegion + fallback: boolean + }): void { + if (input.selectedRegion) increment(this.deltas.selectedRegions, input.selectedRegion) + else increment(this.deltas.unavailableRegions, input.targetRegion) + if (input.fallback) increment(this.deltas.regionFallbacks, input.targetRegion) + } + + recordSql(durationMs: number, success: boolean): void { + this.deltas.sqlQueries++ + if (!success) this.deltas.sqlFailures++ + this.deltas.sqlLatencyMsMax = Math.max(this.deltas.sqlLatencyMsMax, durationMs) + } + + recordControlRenewal(durationMs: number, outcome: ControlRenewalOutcome): void { + this.deltas.controlRenewalLatenciesMs.push(durationMs) + this.deltas.controlRenewalsByOutcome[outcome] = + (this.deltas.controlRenewalsByOutcome[outcome] ?? 0) + 1 + } + + recordControlActivityRecovery(success: boolean): void { + if (success) this.deltas.controlActivityRecoveries++ + else this.deltas.controlActivityRecoveryFailures++ + } + + recordReadiness(observation: RelayReadinessObservation): void { + this.write({ + severity: observation.ready ? 'INFO' : 'WARNING', + message: 'Orca Relay readiness check', + event: 'orca_relay_readiness_check', + metricVersion: 1, + ...this.identity, + ...observation + }) + } + + recordControlClose(code: number): void { + const key = String(code) + this.deltas.controlClosesByCode[key] = (this.deltas.controlClosesByCode[key] ?? 0) + 1 + } + + recordSpliceClose(trigger: string): void { + this.deltas.spliceClosesByTrigger[trigger] = + (this.deltas.spliceClosesByTrigger[trigger] ?? 0) + 1 + } + + start(readCounts: () => RelayProcessCounts, intervalMs = 30_000): void { + if (this.timer) return + this.eventLoop.enable() + this.timer = setInterval(() => this.flush(readCounts()), intervalMs) + this.timer.unref() + } + + stop(): void { + if (this.timer) clearInterval(this.timer) + this.timer = null + this.eventLoop.disable() + } + + regionalRehomeRuntimeSafety(): RegionalRehomeRuntimeSafety { + return { + observedAt: this.lastFlushAt, + sqlFailures: this.lastFlushedSafety.sqlFailures + this.deltas.sqlFailures, + reconnects: this.lastFlushedSafety.reconnects + this.deltas.reconnects, + controlActivityRecoveryFailures: + this.lastFlushedSafety.controlActivityRecoveryFailures + + this.deltas.controlActivityRecoveryFailures + } + } + + flush(counts: RelayProcessCounts): void { + this.lastFlushAt = Date.now() + const deltas = this.deltas + this.lastFlushedSafety = { + sqlFailures: deltas.sqlFailures, + reconnects: deltas.reconnects, + controlActivityRecoveryFailures: deltas.controlActivityRecoveryFailures + } + this.deltas = emptyDeltas() + const memory = process.memoryUsage() + const p99 = this.eventLoop.count === 0 ? 0 : this.eventLoop.percentile(99) / 1_000_000 + this.eventLoop.reset() + this.write({ + severity: 'INFO', + message: 'Orca Relay runtime metrics', + event: 'orca_relay_runtime_metrics', + metricVersion: 2, + role: this.identity.role, + cellId: this.identity.cellId, + region: this.identity.region, + ...counts, + forwardedBytesDelta: deltas.forwardedBytes, + authSuccessesDelta: deltas.authSuccesses, + authFailuresDelta: deltas.authFailures, + reconnectsDelta: deltas.reconnects, + stickyAssignmentsDelta: deltas.stickyAssignments, + stickyAssignmentRejectionsDelta: deltas.stickyAssignmentRejections, + placementAssignmentsDelta: deltas.placementAssignments, + placementAssignmentRejectionsDelta: deltas.placementAssignmentRejections, + stickyRejectionsByReasonDelta: deltas.stickyRejectionsByReason, + placementRejectionsByReasonDelta: deltas.placementRejectionsByReason, + requestedRegionsDelta: deltas.requestedRegions, + selectedRegionsDelta: deltas.selectedRegions, + regionFallbacksDelta: deltas.regionFallbacks, + unavailableRegionsDelta: deltas.unavailableRegions, + controlClosesByCodeDelta: deltas.controlClosesByCode, + spliceClosesByTriggerDelta: deltas.spliceClosesByTrigger, + sqlQueriesDelta: deltas.sqlQueries, + sqlFailuresDelta: deltas.sqlFailures, + sqlLatencyMsMax: Number(deltas.sqlLatencyMsMax.toFixed(3)), + controlRenewalsByOutcomeDelta: deltas.controlRenewalsByOutcome, + controlRenewalsDelta: deltas.controlRenewalLatenciesMs.length, + controlRenewalSuccessesDelta: deltas.controlRenewalsByOutcome.renewed ?? 0, + controlRenewalLeaseMissesDelta: + deltas.controlRenewalsByOutcome.control_activity_not_found ?? 0, + controlActivityRecoveriesDelta: deltas.controlActivityRecoveries, + controlActivityRecoveryFailuresDelta: deltas.controlActivityRecoveryFailures, + controlRenewalLatencyMsP50: Number( + percentile(deltas.controlRenewalLatenciesMs, 0.5).toFixed(3) + ), + controlRenewalLatencyMsP95: Number( + percentile(deltas.controlRenewalLatenciesMs, 0.95).toFixed(3) + ), + controlRenewalLatencyMsMax: Number( + Math.max(0, ...deltas.controlRenewalLatenciesMs).toFixed(3) + ), + httpLatencyMsMax: Number(deltas.httpLatencyMsMax.toFixed(3)), + heapUsedBytes: memory.heapUsed, + heapTotalBytes: memory.heapTotal, + eventLoopDelayMsP99: Number(p99.toFixed(3)) + }) + } +} + +function increment(counts: Record, key: string): void { + counts[key] = (counts[key] ?? 0) + 1 +} + +export function timedRelayOperation( + operation: () => Promise, + observe: (durationMs: number, success: boolean) => void, + isExpectedError: (error: unknown) => boolean = () => false +): Promise { + const startedAt = performance.now() + return operation().then( + (result) => { + observe(performance.now() - startedAt, true) + return result + }, + (error: unknown) => { + observe(performance.now() - startedAt, isExpectedError(error)) + throw error + } + ) +} diff --git a/cloud/apps/relay/src/relay-readiness.test.ts b/cloud/apps/relay/src/relay-readiness.test.ts new file mode 100644 index 00000000000..a85d9a46aa7 --- /dev/null +++ b/cloud/apps/relay/src/relay-readiness.test.ts @@ -0,0 +1,109 @@ +import { describe, expect, it, vi } from 'vitest' +import type { RelayDatabase } from './database.js' +import { + createRelayReadiness, + type RelayReadinessObservation +} from './relay-readiness.js' + +function database(query: () => Promise[]>): RelayDatabase { + return { + query, + queryLocked: query, + transaction: async (operation) => await operation(database(query)), + close: async () => {} + } +} + +describe('relay readiness', () => { + it('fails readiness while liveness remains independent of SQL and JWKS', async () => { + const jwksFailure = createRelayReadiness(database(async () => [{ ready: 1 }]), 'https://jwks', { + fetch: vi.fn(async () => new Response('', { status: 503 })) as typeof fetch, + cacheMs: 0 + }) + const sqlFailure = createRelayReadiness( + database(async () => { + throw new Error('sql down') + }), + 'https://jwks', + { + fetch: vi.fn(async () => new Response('{}', { status: 200 })) as typeof fetch, + cacheMs: 0 + } + ) + + expect(await jwksFailure()).toBe(false) + expect(await sqlFailure()).toBe(false) + }) + + it.each([ + { + name: 'JWKS HTTP failure', + fetch: vi.fn(async () => new Response('', { status: 503 })) as typeof fetch, + query: vi.fn(async () => [{ ready: 1 }]), + failure: 'jwks_http_failed' + }, + { + name: 'JWKS timeout', + fetch: vi.fn(async () => { + throw new DOMException('redacted', 'TimeoutError') + }) as typeof fetch, + query: vi.fn(async () => [{ ready: 1 }]), + failure: 'jwks_timed_out' + }, + { + name: 'JWKS fetch failure', + fetch: vi.fn(async () => { + throw new Error('redacted') + }) as typeof fetch, + query: vi.fn(async () => [{ ready: 1 }]), + failure: 'jwks_fetch_failed' + }, + { + name: 'SQL failure', + fetch: vi.fn(async () => new Response('{}', { status: 200 })) as typeof fetch, + query: vi.fn(async () => { + throw new Error('redacted') + }), + failure: 'sql_failed' + } + ])('reports a safe reason for $name', async ({ fetch, query, failure }) => { + const observations: RelayReadinessObservation[] = [] + const ready = createRelayReadiness(database(query), 'https://jwks', { + fetch, + cacheMs: 0, + observe: (observation) => observations.push(observation) + }) + + expect(await ready()).toBe(false) + expect(observations).toEqual([ + expect.objectContaining({ ready: false, failure }) + ]) + expect(JSON.stringify(observations)).not.toContain('redacted') + if (failure.startsWith('jwks_')) expect(query).not.toHaveBeenCalled() + }) + + it('reports the initial success but not healthy repeats or cached reads', async () => { + const observations: RelayReadinessObservation[] = [] + let now = 100 + const ready = createRelayReadiness(database(async () => [{ ready: 1 }]), 'https://jwks', { + fetch: vi.fn(async () => new Response('{}', { status: 200 })) as typeof fetch, + cacheMs: 10_000, + now: () => now, + observe: (observation) => observations.push(observation) + }) + + expect(await ready()).toBe(true) + now += 1_000 + expect(await ready()).toBe(true) + now += 10_000 + expect(await ready()).toBe(true) + expect(observations).toEqual([ + { + ready: true, + jwksLatencyMs: 0, + sqlLatencyMs: 0, + totalLatencyMs: 0 + } + ]) + }) +}) diff --git a/cloud/apps/relay/src/relay-readiness.ts b/cloud/apps/relay/src/relay-readiness.ts new file mode 100644 index 00000000000..0b7303367f1 --- /dev/null +++ b/cloud/apps/relay/src/relay-readiness.ts @@ -0,0 +1,89 @@ +import type { RelayDatabase } from './database.js' + +export type RelayReadinessFailure = + | 'jwks_fetch_failed' + | 'jwks_http_failed' + | 'jwks_timed_out' + | 'sql_failed' + +export type RelayReadinessObservation = { + ready: boolean + failure?: RelayReadinessFailure + jwksLatencyMs: number + sqlLatencyMs: number + totalLatencyMs: number +} + +type RelayReadinessOptions = { + fetch?: typeof fetch + timeoutMs?: number + cacheMs?: number + now?: () => number + observe?: (observation: RelayReadinessObservation) => void +} + +function fetchFailure(error: unknown): RelayReadinessFailure { + return error instanceof Error && error.name === 'TimeoutError' + ? 'jwks_timed_out' + : 'jwks_fetch_failed' +} + +export function createRelayReadiness( + database: RelayDatabase, + jwksUrl: string, + options: RelayReadinessOptions = {} +): () => Promise { + const fetchImpl = options.fetch ?? fetch + const timeoutMs = options.timeoutMs ?? 2_000 + const cacheMs = options.cacheMs ?? 10_000 + const now = options.now ?? Date.now + let cachedAt = Number.NEGATIVE_INFINITY + let cached = false + let lastObservedReady: boolean | undefined + + return async () => { + if (now() - cachedAt < cacheMs) return cached + const startedAt = now() + let jwksCompletedAt = startedAt + let sqlStartedAt = startedAt + let failure: RelayReadinessFailure | undefined + try { + let response: Response + try { + response = await fetchImpl(jwksUrl, { signal: AbortSignal.timeout(timeoutMs) }) + } catch (error) { + failure = fetchFailure(error) + throw error + } finally { + jwksCompletedAt = now() + } + if (!response.ok) { + failure = 'jwks_http_failed' + throw new Error(failure) + } + sqlStartedAt = now() + try { + await database.query('SELECT 1 AS ready') + } catch (error) { + failure = 'sql_failed' + throw error + } + } catch { + // The load balancer only needs the boolean; the safe reason is emitted below. + } + const completedAt = now() + cached = failure === undefined + cachedAt = completedAt + if (!cached || cached !== lastObservedReady) { + options.observe?.({ + ready: cached, + ...(failure ? { failure } : {}), + jwksLatencyMs: Math.max(0, jwksCompletedAt - startedAt), + sqlLatencyMs: failure?.startsWith('jwks_') ? 0 : Math.max(0, completedAt - sqlStartedAt), + totalLatencyMs: Math.max(0, completedAt - startedAt) + }) + } + lastObservedReady = cached + return cached + } +} diff --git a/cloud/apps/relay/src/relay-region-app.test.ts b/cloud/apps/relay/src/relay-region-app.test.ts new file mode 100644 index 00000000000..30cf3bf3e26 --- /dev/null +++ b/cloud/apps/relay/src/relay-region-app.test.ts @@ -0,0 +1,155 @@ +import { describe, expect, it, vi } from 'vitest' +import type { RelayConfig } from './config.js' + +const fakes = vi.hoisted(() => ({ + verifyRelayToken: vi.fn(async (token: string) => ({ sub: 'user-1', relayHostId: token })) +})) + +vi.mock('./relay-token-verifier.js', () => ({ + createRelayTokenVerifier: () => fakes.verifyRelayToken, + readBearer: (value: string | undefined) => value?.replace(/^Bearer /, '') ?? null +})) + +import { createRelayApp } from './app.js' + +describe('Relay region API', () => { + it('passes a valid preference to placement and records coarse outcomes', async () => { + const assign = vi.fn(async () => ({ + userId: 'user-1', + relayHostId: 'asiahost00000001', + cellId: 'asia-c1', + cellUrl: 'https://asia-c1.relay.example.test', + region: 'asia-east2' as const, + assignmentEpoch: 1, + leaseExpiresAt: Date.now() + 300_000 + })) + const recordRegionRequest = vi.fn() + const recordRegionSelection = vi.fn() + const app = createRelayApp(config(), { + store: {} as never, + assignments: { assign, resolve: vi.fn(async () => null) } as never, + drain: vi.fn(), + ready: vi.fn(async () => true), + recordRegionRequest, + recordRegionSelection + }) + + const response = await app.request( + '/v1/assign', + assignmentRequest('asiahost00000001', { preferredRegion: 'asia-east2' }) + ) + + expect(response.status).toBe(200) + expect(assign).toHaveBeenCalledWith( + { userId: 'user-1', relayHostId: 'asiahost00000001' }, + 'asia-east2', + 'asia-east2' + ) + expect(recordRegionRequest).toHaveBeenCalledWith('asia-east2') + expect(recordRegionSelection).toHaveBeenCalledWith({ + targetRegion: 'asia-east2', + selectedRegion: 'asia-east2', + fallback: false + }) + }) + + it('keeps the preference for observation while the kill switch places US-first', async () => { + const assign = vi.fn(async () => ({ + userId: 'user-1', + relayHostId: 'killhost00000001', + cellId: 'us-c1', + cellUrl: 'https://us-c1.relay.example.test', + region: 'us-central1' as const, + assignmentEpoch: 1, + leaseExpiresAt: Date.now() + 300_000 + })) + const app = createRelayApp(config({ regionalPlacementEnabled: false }), { + store: {} as never, + assignments: { assign, resolve: vi.fn(async () => null) } as never, + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + + const response = await app.request( + '/v1/assign', + assignmentRequest('killhost00000001', { preferredRegion: 'asia-east2' }) + ) + + expect(response.status).toBe(200) + expect(assign).toHaveBeenCalledWith( + { userId: 'user-1', relayHostId: 'killhost00000001' }, + 'asia-east2', + 'us-central1' + ) + }) + + it('exposes only the store-provided healthy catalog from directors', async () => { + const regionCatalog = vi.fn(async () => [ + { region: 'us-central1' as const, probeOrigins: ['https://us.relay.example.test'] } + ]) + const app = createRelayApp(config(), { + store: {} as never, + assignments: { regionCatalog } as never, + drain: vi.fn(), + ready: vi.fn(async () => true) + }) + + const response = await app.request('/v1/regions') + + expect(response.status).toBe(200) + expect(await response.json()).toEqual({ + v: 1, + regions: [ + { region: 'us-central1', probeOrigins: ['https://us.relay.example.test'] } + ] + }) + + const burst = await Promise.all( + Array.from({ length: 50 }, () => app.request('/v1/regions')) + ) + expect(burst.every(({ status }) => status === 200)).toBe(true) + expect(regionCatalog).toHaveBeenCalledOnce() + }) +}) + +function assignmentRequest(relayHostId: string, extra: Record): RequestInit { + return { + method: 'POST', + headers: { + authorization: `Bearer ${relayHostId}`, + 'content-type': 'application/json' + }, + body: JSON.stringify({ v: 1, relayHostId, ...extra }) + } +} + +function config(overrides: Partial = {}): RelayConfig { + return { + port: 8080, + publicUrl: 'https://relay.example.test', + cellUrl: 'https://relay.example.test', + region: 'us-central1', + authIssuer: 'https://auth.example.test', + authAudience: 'orca-relay', + jwksUrl: 'https://auth.example.test/jwks', + assignmentSigningKey: new TextEncoder().encode('assignment-key-with-at-least-32-bytes'), + role: 'director', + cellId: 'director', + cells: [], + adminAudience: 'https://relay.example.test/v1/admin/drain', + deployServiceAccount: 'deploy@example.test', + runtimeServiceAccount: 'runtime@example.test', + adminJwksUrl: 'https://auth.example.test/jwks', + databasePoolMax: 3, + publicAssignmentsEnabled: true, + regionalPlacementEnabled: true, + publicAssignmentConcurrency: 2, + publicAssignmentQueueMax: 128, + publicAssignmentWaitMs: 4_000, + publicResolveConcurrency: 1, + publicResolveWaitMs: 5_000, + publicAssignmentRetryAfterSeconds: 5, + dataDir: './data', + ...overrides + } +} diff --git a/cloud/apps/relay/src/relay-region-placement.test.ts b/cloud/apps/relay/src/relay-region-placement.test.ts new file mode 100644 index 00000000000..76cf5479a41 --- /dev/null +++ b/cloud/apps/relay/src/relay-region-placement.test.ts @@ -0,0 +1,196 @@ +import { afterEach, beforeEach, describe, expect, it } from 'vitest' +import { RelayAssignmentStore } from './assignment-store.js' +import type { RelayCellConfig } from './config.js' +import { openInMemoryRelayDatabase, type RelayDatabase } from './database.js' + +const CELLS: RelayCellConfig[] = [ + { + id: 'us-c1', + url: 'https://us-c1.relay.example.com', + region: 'us-central1', + capacityRequests: 100 + }, + { + id: 'asia-c1', + url: 'https://asia-c1.relay.example.com', + region: 'asia-east2', + capacityRequests: 100 + }, + { + id: 'asia-c2', + url: 'https://asia-c2.relay.example.com', + region: 'asia-east2', + capacityRequests: 100 + }, + { + id: 'asia-c3', + url: 'https://asia-c3.relay.example.com', + region: 'asia-east2', + capacityRequests: 100 + } +] + +describe('Relay regional placement', () => { + let database: RelayDatabase + let now: number + let store: RelayAssignmentStore + + beforeEach(async () => { + database = await openInMemoryRelayDatabase() + now = 1_000 + store = new RelayAssignmentStore(database, () => now) + await store.reconcileCells(CELLS) + }) + + afterEach(async () => await database.close()) + + it('prefers the requested region and keeps unhinted placement US-first', async () => { + await expect( + store.assign({ userId: 'asia-user', relayHostId: 'asiahost00000001' }, 'asia-east2') + ).resolves.toMatchObject({ cellId: 'asia-c1', region: 'asia-east2' }) + await expect( + store.assign({ userId: 'us-user', relayHostId: 'ushost0000000001' }) + ).resolves.toMatchObject({ cellId: 'us-c1', region: 'us-central1' }) + }) + + it('falls back globally only when the target region has no safe general cell', async () => { + await store.configureCell(CELLS[1]!, 'migration-only') + await store.configureCell(CELLS[2]!, 'migration-only') + await store.configureCell(CELLS[3]!, 'migration-only') + + await expect( + store.assign({ userId: 'fallback-user', relayHostId: 'fallbackhost0001' }, 'asia-east2') + ).resolves.toMatchObject({ cellId: 'us-c1', region: 'us-central1' }) + }) + + it('preserves sticky assignments while updating only explicit preferences', async () => { + const identity = { userId: 'sticky-user', relayHostId: 'stickyhost000001' } + await expect(store.assign(identity)).resolves.toMatchObject({ cellId: 'us-c1' }) + now = 2_000 + await expect(store.assign(identity, 'asia-east2')).resolves.toMatchObject({ cellId: 'us-c1' }) + now = 3_000 + await store.assign(identity) + + expect( + await database.query( + `SELECT preferred_region, observed_at FROM relay_assignment_region_preferences + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).toEqual([{ preferred_region: 'asia-east2', observed_at: 2_000 }]) + }) + + it('uses the explicit placement region as the server-side kill switch', async () => { + await expect( + store.assign( + { userId: 'kill-user', relayHostId: 'killhost00000001' }, + 'asia-east2', + 'us-central1' + ) + ).resolves.toMatchObject({ cellId: 'us-c1', region: 'us-central1' }) + expect(await database.query(`SELECT preferred_region FROM relay_assignment_region_preferences`)) + .toEqual([{ preferred_region: 'asia-east2' }]) + }) + + it('returns at most two deterministic healthy general probe origins per region', async () => { + for (const [index, cell] of CELLS.entries()) { + await store.recordCellHeartbeat({ + cellId: cell.id, + cellUrl: cell.url, + region: cell.region, + cellIncarnation: `00000000-0000-4000-8000-${String(index + 1).padStart(12, '0')}`, + startedAt: 1, + ready: true, + observedRequests: 0 + }) + } + await store.configureCell(CELLS[2]!, 'migration-only') + + await expect(store.regionCatalog()).resolves.toEqual([ + { region: 'us-central1', probeOrigins: ['https://us-c1.relay.example.com'] }, + { + region: 'asia-east2', + probeOrigins: [ + 'https://asia-c1.relay.example.com', + 'https://asia-c3.relay.example.com' + ] + } + ]) + + await database.query(`UPDATE relay_cells SET enabled = 0 WHERE cell_id = ?`, ['asia-c3']) + await expect(store.regionCatalog()).resolves.toEqual([ + { region: 'us-central1', probeOrigins: ['https://us-c1.relay.example.com'] }, + { region: 'asia-east2', probeOrigins: ['https://asia-c1.relay.example.com'] } + ]) + }) + + it('rejects a heartbeat whose explicit region conflicts with registration', async () => { + await expect( + store.recordCellHeartbeat({ + cellId: 'asia-c1', + cellUrl: 'https://asia-c1.relay.example.com', + region: 'us-central1', + cellIncarnation: '00000000-0000-4000-8000-000000000001', + startedAt: 1, + ready: true, + observedRequests: 0 + }) + ).rejects.toThrow('cell_region_mismatch') + }) + + it('defaults cells inserted by an old process after startup to US', async () => { + const oldCell = { + id: 'old-us-cell', + url: 'https://old-us-cell.relay.example.com', + capacityRequests: 100 + } + await database.query( + `INSERT INTO relay_cells + (cell_id, cell_url, enabled, capacity_requests, reserved_requests, + observed_requests, last_heartbeat_at, updated_at) + VALUES (?, ?, 1, 100, 0, 0, ?, ?)`, + [oldCell.id, oldCell.url, now, now] + ) + await database.query( + `INSERT INTO relay_cell_admission (cell_id, admission_state, updated_at) + VALUES (?, 'general', ?)`, + [oldCell.id, now] + ) + await Promise.all(CELLS.map((cell) => store.configureCell(cell, 'migration-only'))) + + const identity = { userId: 'old-user', relayHostId: 'oldhost000000001' } + await expect(store.assign(identity)).resolves.toMatchObject({ + cellId: oldCell.id, + region: 'us-central1' + }) + await expect(store.resolve(identity)).resolves.toMatchObject({ + cellId: oldCell.id, + region: 'us-central1' + }) + await expect( + store.recordCellHeartbeat({ + cellId: oldCell.id, + cellUrl: oldCell.url, + cellIncarnation: '00000000-0000-4000-8000-000000000099', + startedAt: 1, + ready: true, + observedRequests: 0 + }) + ).resolves.toBeUndefined() + }) + + it('expires identity-linked region preferences after 30 days', async () => { + const identity = { userId: 'expiry-user', relayHostId: 'expiryhost000001' } + await store.assign(identity, 'asia-east2') + now += 30 * 24 * 60 * 60_000 + 1 + + await expect(store.releaseExpiredRegionPreferences()).resolves.toBe(1) + await expect( + database.query( + `SELECT preferred_region FROM relay_assignment_region_preferences + WHERE user_id = ? AND relay_host_id = ?`, + [identity.userId, identity.relayHostId] + ) + ).resolves.toEqual([]) + }) +}) diff --git a/cloud/apps/relay/src/relay-server.ts b/cloud/apps/relay/src/relay-server.ts new file mode 100644 index 00000000000..a14240cfa6a --- /dev/null +++ b/cloud/apps/relay/src/relay-server.ts @@ -0,0 +1,533 @@ +import { createAdaptorServer } from '@hono/node-server' +import { + hasAdmissionCapacity, + HostDataAuthSchema, + RELAY_ADMISSION_BUDGETS, + RELAY_CLOSE_CODE, + RELAY_DEFAULT_REGION, + RELAY_PROTOCOL_LIMITS, + RelayAuthSchema +} from '@orca-cloud/relay-contract' +import type { IncomingMessage } from 'node:http' +import { randomUUID } from 'node:crypto' +import { performance } from 'node:perf_hooks' +import { WebSocketServer } from 'ws' +import type WebSocket from 'ws' +import type { RawData } from 'ws' +import { createRelayApp } from './app.js' +import { RelayAssignmentStore } from './assignment-store.js' +import type { RelayConfig } from './config.js' +import { RelayCredentialStore } from './credential-store.js' +import type { RelayDatabase } from './database.js' +import { HostSessionRegistry } from './host-session-registry.js' +import { observeRelayDatabase } from './observed-relay-database.js' +import { RelayObservability } from './relay-observability.js' +import { + RelayConnectionLedger, + type RelayConnectionUpgrade +} from './relay-connection-ledger.js' +import { createRelayReadiness } from './relay-readiness.js' +import { createRelayTokenVerifier, readBearer } from './relay-token-verifier.js' +import { closeRelayWebSocket } from './relay-websocket-close.js' +import { ProcessQueuedByteBudget } from './splice-forwarder.js' + +function rejectUpgrade(socket: NodeJS.WritableStream, status: number, message: string): void { + socket.write(`HTTP/1.1 ${status} ${message}\r\nConnection: close\r\nContent-Length: 0\r\n\r\n`) + if ('destroy' in socket && typeof socket.destroy === 'function') socket.destroy() +} + +function noDelay(socket: WebSocket): void { + const transport = (socket as WebSocket & { _socket?: { setNoDelay: (enabled: boolean) => void } }) + ._socket + transport?.setNoDelay(true) +} + +// Why: a ws receiver error (oversize or malformed frame) with no 'error' +// listener throws process-wide; ws itself already closes the socket after +// emitting it, so logging is all that is left to do. +function guardSocketErrors(socket: WebSocket, kind: string): void { + socket.on('error', (error) => { + console.warn(`[orca-relay] ${kind} socket error: ${error.message}`) + }) +} + +function admissionSource(request: IncomingMessage): string { + const forwarded = request.headers['x-forwarded-for'] + const chain = (Array.isArray(forwarded) ? forwarded.join(',') : forwarded ?? '') + .split(',') + .map((entry) => entry.trim()) + .filter(Boolean) + // Google Front End appends client and load-balancer addresses after any + // caller-supplied values, so only the penultimate hop is trustworthy. + return chain.length >= 2 ? chain.at(-2)! : (request.socket.remoteAddress ?? 'unknown') +} + +function firstPayload(raw: RawData, expectedType: string): unknown { + try { + const parsed = JSON.parse(raw.toString()) as Record + if (parsed.type !== expectedType) return null + const { type: _type, ...rest } = parsed + return rest + } catch { + return null + } +} + +export function createRelayServer( + config: RelayConfig, + database: RelayDatabase, + options: { + now?: () => number + connectionLedgerLimits?: { hardCap: number; controlReserve: number } + cellIncarnation?: string + } = {} +) { + const cellIncarnation = options.cellIncarnation ?? randomUUID() + const observability = new RelayObservability({ + role: config.role, + cellId: config.cellId, + region: config.region ?? RELAY_DEFAULT_REGION + }) + const observedDatabase = observeRelayDatabase(database, observability) + const controls = new WebSocketServer({ + noServer: true, + clientTracking: false, + perMessageDeflate: false, + maxPayload: 1024 * 1024 + }) + const verifyRelayToken = createRelayTokenVerifier(config) + const store = new RelayCredentialStore(observedDatabase, options.now) + const assignments = new RelayAssignmentStore(observedDatabase, options.now, { + requireLiveCells: config.role === 'director', + recordControlRenewal: (durationMs, outcome) => + observability.recordControlRenewal?.(durationMs, outcome) + }) + const ready = createRelayReadiness(observedDatabase, config.jwksUrl, { + observe: (observation) => observability.recordReadiness(observation) + }) + const queuedBytes = new ProcessQueuedByteBudget() + const sessions = new HostSessionRegistry( + config, + verifyRelayToken, + store, + assignments, + queuedBytes, + observability, + options.now + ) + const app = createRelayApp(config, { + store, + assignments, + drain: (graceMs) => sessions.drain(graceMs), + drainHost: (input) => sessions.drainHost(input), + regionalRehomeTrustProbeHostExists: (input) => sessions.get(input) !== null, + cellIncarnation, + isDraining: () => sessions.isDraining(), + runtimeCounts: () => runtimeCounts(), + ready, + recordAssignmentAdmission: (outcome) => observability.recordAssignmentAdmission?.(outcome), + recordAssignmentRejectionReason: (lane, reason) => + observability.recordAssignmentRejectionReason?.(lane, reason), + recordRegionRequest: (region) => observability.recordRegionRequest?.(region), + recordRegionSelection: (input) => observability.recordRegionSelection?.(input) + }) + const observedFetch: typeof app.fetch = async (...args) => { + const startedAt = performance.now() + try { + return await app.fetch(...args) + } finally { + observability.recordHttp(performance.now() - startedAt) + } + } + const server = createAdaptorServer({ fetch: observedFetch }) + const clients = new WebSocketServer({ + noServer: true, + clientTracking: false, + perMessageDeflate: false, + maxPayload: RELAY_PROTOCOL_LIMITS.maxFrameBytes + }) + const dataSockets = new WebSocketServer({ + noServer: true, + clientTracking: false, + perMessageDeflate: false, + maxPayload: RELAY_PROTOCOL_LIMITS.maxFrameBytes + }) + 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() + + const admit = (source: string): boolean => { + const now = Date.now() + const currentWindow = preAuthAttemptsBySource.get(source) + const attempts = + !currentWindow || now - currentWindow.windowStartedAt >= 60_000 + ? { windowStartedAt: now, count: 0 } + : currentWindow + const sourceCount = preAuthBySource.get(source) ?? 0 + if ( + !hasAdmissionCapacity({ + totalRequests: + connectionLedger?.counts().physicalConnections ?? totalConnections, + preAuthConnections, + sourcePreAuthConnections: sourceCount, + totalRequestCeiling: + configuredConnectionLimits === null + ? undefined + : configuredConnectionLimits.hardCap - configuredConnectionLimits.controlReserve + }) || + attempts.count >= RELAY_ADMISSION_BUDGETS.maxPreAuthAttemptsPerSourcePerMinute + ) { + return false + } + attempts.count++ + preAuthAttemptsBySource.delete(source) + preAuthAttemptsBySource.set(source, attempts) + // A bounded LRU keeps source churn from becoming its own memory attack. + if (preAuthAttemptsBySource.size > 4_096) { + preAuthAttemptsBySource.delete(preAuthAttemptsBySource.keys().next().value!) + } + preAuthConnections++ + preAuthBySource.set(source, sourceCount + 1) + return true + } + const authenticated = (source: string): void => { + preAuthConnections = Math.max(0, preAuthConnections - 1) + const next = (preAuthBySource.get(source) ?? 1) - 1 + if (next <= 0) preAuthBySource.delete(source) + else preAuthBySource.set(source, next) + } + + const trackConnection = ( + socket: WebSocket, + upgrade: RelayConnectionUpgrade | null + ): void => { + if (upgrade) { + upgrade.promote(socket) + return + } + totalConnections++ + socket.once('close', () => { + totalConnections = Math.max(0, totalConnections - 1) + }) + } + + const awaitFirstFrame = ( + socket: WebSocket, + source: string, + callback: (raw: RawData) => Promise + ): void => { + let finished = false + const timer = setTimeout(() => { + if (finished) return + finished = true + authenticated(source) + observability.recordAuth(false) + socket.close(RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL, 'first frame timeout') + }, RELAY_PROTOCOL_LIMITS.firstFrameDeadlineMs) + socket.once('message', (raw, binary) => { + if (finished) return + finished = true + clearTimeout(timer) + authenticated(source) + if (binary) { + observability.recordAuth(false) + socket.close(RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL, 'first frame must be text') + return + } + void callback(raw).catch((error: unknown) => { + console.warn( + '[orca-relay] first frame handler failed', + error instanceof Error ? error.message : '' + ) + closeRelayWebSocket( + socket, + RELAY_CLOSE_CODE.LIMIT_EXCEEDED, + 'relay temporarily unavailable' + ) + }) + }) + socket.once('close', () => { + if (!finished) { + finished = true + clearTimeout(timer) + authenticated(source) + } + }) + } + + server.on('upgrade', (request, socket, head) => { + const url = new URL(request.url ?? '/', config.publicUrl) + const source = admissionSource(request) + if (url.search) { + rejectUpgrade(socket, 400, 'Bad Request') + return + } + if (url.pathname.startsWith('/v1/connect/')) { + const hostId = decodeURIComponent(url.pathname.slice('/v1/connect/'.length)) + if (!/^[A-Za-z0-9_-]{16}$/.test(hostId)) { + rejectUpgrade(socket, 429, 'Too Many Requests') + return + } + const phoneAdmission = connectionLedger?.tryReservePhone() ?? null + if (connectionLedger && !phoneAdmission) { + rejectUpgrade(socket, 503, 'Service Unavailable') + return + } + if (!admit(source)) { + phoneAdmission?.upgrade.release() + phoneAdmission?.hostData.release() + rejectUpgrade(socket, 429, 'Too Many Requests') + return + } + const releasePhoneUpgrade = (): void => { + authenticated(source) + phoneAdmission?.upgrade.release() + phoneAdmission?.hostData.release() + } + socket.once('close', releasePhoneUpgrade) + try { + clients.handleUpgrade(request, socket, head, (webSocket) => { + socket.off('close', releasePhoneUpgrade) + trackConnection(webSocket, phoneAdmission?.upgrade ?? null) + guardSocketErrors(webSocket, 'client') + if (phoneAdmission) { + webSocket.once('close', () => phoneAdmission.hostData.release()) + } + noDelay(webSocket) + awaitFirstFrame(webSocket, source, async (raw) => { + const auth = RelayAuthSchema.safeParse(firstPayload(raw, 'relay-auth')) + if (!auth.success) { + phoneAdmission?.hostData.release() + observability.recordAuth(false) + webSocket.send( + JSON.stringify({ type: 'relay-hello', ok: false, code: RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL }) + ) + webSocket.close(RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL, 'invalid relay auth') + return + } + if (config.role === 'director') { + const invite = await store.resolveInviteForMove(hostId, auth.data.credential) + const identity = invite ? { userId: invite.userId, relayHostId: hostId } : null + // Released combined-service invites gain their first durable cell assignment here. + const assignment = identity + ? (await assignments.resolve(identity)) ?? (await assignments.assign(identity)) + : null + if (!invite || !assignment) { + phoneAdmission?.hostData.release() + observability.recordAuth(false) + webSocket.send( + JSON.stringify({ + type: 'relay-hello', + ok: false, + code: RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL + }) + ) + webSocket.close(RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL, 'invalid invite') + return + } + phoneAdmission?.hostData.release() + observability.recordAuth(true) + webSocket.send( + JSON.stringify({ + type: 'relay-moved', + v: 1, + cellUrl: assignment.cellUrl, + assignmentEpoch: assignment.assignmentEpoch + }) + ) + webSocket.close(RELAY_CLOSE_CODE.DRAINING, 'connect to assigned cell') + return + } + await sessions.acceptClient( + webSocket, + hostId, + auth.data.credential, + phoneAdmission?.hostData + ) + }) + }) + } catch { + socket.off('close', releasePhoneUpgrade) + releasePhoneUpgrade() + socket.destroy() + } + return + } + if (url.pathname.startsWith('/v1/host/data/')) { + if (config.role === 'director') { + rejectUpgrade(socket, 404, 'Not Found') + return + } + const connId = decodeURIComponent(url.pathname.slice('/v1/host/data/'.length)) + if (!connId || connId.length > 128) { + rejectUpgrade(socket, 429, 'Too Many Requests') + return + } + const dataUpgrade = connectionLedger?.tryReserveHostData(connId) ?? null + if (connectionLedger && !dataUpgrade) { + rejectUpgrade(socket, 503, 'Service Unavailable') + return + } + if (!admit(source)) { + dataUpgrade?.release() + rejectUpgrade(socket, 429, 'Too Many Requests') + return + } + const releaseDataUpgrade = (): void => { + authenticated(source) + dataUpgrade?.release() + } + socket.once('close', releaseDataUpgrade) + try { + dataSockets.handleUpgrade(request, socket, head, (webSocket) => { + socket.off('close', releaseDataUpgrade) + trackConnection(webSocket, dataUpgrade) + guardSocketErrors(webSocket, 'host-data') + noDelay(webSocket) + awaitFirstFrame(webSocket, source, async (raw) => { + const auth = HostDataAuthSchema.safeParse(firstPayload(raw, 'host-data-auth')) + if (!auth.success) { + observability.recordAuth(false) + webSocket.close(RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL, 'invalid host data auth') + return + } + const accepted = await sessions.acceptHostData( + webSocket, + connId, + auth.data.connTicket, + auth.data.generation + ) + if (accepted) dataUpgrade?.commitHostData() + }) + }) + } catch { + socket.off('close', releaseDataUpgrade) + releaseDataUpgrade() + socket.destroy() + } + return + } + if (url.pathname !== '/v1/host/control') { + rejectUpgrade(socket, 404, 'Not Found') + return + } + if (config.role === 'director') { + rejectUpgrade(socket, 404, 'Not Found') + return + } + const bearer = readBearer(request.headers.authorization) + if (!bearer) { + observability.recordAuth(false) + rejectUpgrade(socket, 401, 'Unauthorized') + return + } + void verifyRelayToken(bearer).then((identity) => { + if (socket.destroyed) return + if (!identity) { + observability.recordAuth(false) + rejectUpgrade(socket, 401, 'Unauthorized') + return + } + const isRebind = sessions.hasActiveControl({ + userId: identity.sub, + relayHostId: identity.relayHostId + }) + const controlUpgrade = connectionLedger?.tryReserveControl(isRebind) ?? null + if ( + (connectionLedger && !controlUpgrade) || + (!connectionLedger && totalConnections >= RELAY_ADMISSION_BUDGETS.cloudRunConcurrency) + ) { + rejectUpgrade( + socket, + connectionLedger ? 503 : 429, + connectionLedger ? 'Service Unavailable' : 'Too Many Requests' + ) + return + } + // Register the release before anything else can throw: the outer catch + // destroys the socket, so a close-registered release cannot leak the + // reserved connection unit. + const releaseControlUpgrade = (): void => controlUpgrade?.release() + socket.once('close', releaseControlUpgrade) + observability.recordAuth(true) + try { + controls.handleUpgrade(request, socket, head, (webSocket) => { + socket.off('close', releaseControlUpgrade) + trackConnection(webSocket, controlUpgrade) + guardSocketErrors(webSocket, 'control') + noDelay(webSocket) + sessions.acceptControl( + webSocket, + identity, + controlUpgrade?.inclusionWatermark + ) + }) + } catch { + socket.off('close', releaseControlUpgrade) + releaseControlUpgrade() + socket.destroy() + } + }).catch((error: unknown) => { + // A throw in the upgrade handling above must cost this socket, not the process. + console.warn( + `[orca-relay] control upgrade failed: ${error instanceof Error ? error.message : 'unknown'}` + ) + socket.destroy() + }) + }) + + server.on('close', () => { + controls.close() + clients.close() + dataSockets.close() + }) + const runtimeCounts = () => { + const ledgerCounts = connectionLedger?.counts() + return { + totalConnections: ledgerCounts?.physicalConnections ?? totalConnections, + preAuthConnections, + ...sessions.runtimeCounts(), + queuedBytes: queuedBytes.current(), + ...(ledgerCounts + ? { + inFlightConnections: ledgerCounts.inFlightConnections, + reservedConnectionUnits: ledgerCounts.reservedConnectionUnits, + enforcedConnectionUnits: ledgerCounts.enforcedConnectionUnits + } + : {}) + } + } + const connectionSnapshot = () => connectionLedger?.snapshot() + return { + server, + sessions, + store, + assignments, + queuedBytes, + observability, + runtimeCounts, + connectionSnapshot, + ready, + cellIncarnation + } +} + +export function closeWithDrain(socket: WebSocket, graceMs: number): void { + socket.send(JSON.stringify({ type: 'drain', graceMs, recovery: 'resolve-director' })) + socket.close(RELAY_CLOSE_CODE.DRAINING, 'resolve configured director') +} diff --git a/cloud/apps/relay/src/relay-token-verifier.ts b/cloud/apps/relay/src/relay-token-verifier.ts new file mode 100644 index 00000000000..d3dc5300a44 --- /dev/null +++ b/cloud/apps/relay/src/relay-token-verifier.ts @@ -0,0 +1,35 @@ +import { createRemoteJWKSet, jwtVerify } from 'jose' +import { z } from 'zod' +import type { RelayConfig } from './config.js' + +const ClaimsSchema = z.object({ + sub: z.string().min(1), + prof: z.string().min(1), + org: z.string().min(1).optional(), + relayHostId: z.string().regex(/^[A-Za-z0-9_-]{16}$/), + purpose: z.literal('host-control'), + exp: z.number().int().positive() +}) + +export type RelayTokenClaims = z.infer + +export function createRelayTokenVerifier(config: RelayConfig): (token: string) => Promise { + const jwks = createRemoteJWKSet(new URL(config.jwksUrl)) + return async (token) => { + try { + const verified = await jwtVerify(token, jwks, { + issuer: config.authIssuer, + audience: config.authAudience, + algorithms: ['ES256'] + }) + return ClaimsSchema.parse(verified.payload) + } catch { + return null + } + } +} + +export function readBearer(value: string | undefined): string | null { + const match = /^Bearer ([^\s]+)$/.exec(value ?? '') + return match?.[1] ?? null +} diff --git a/cloud/apps/relay/src/relay-websocket-close.test.ts b/cloud/apps/relay/src/relay-websocket-close.test.ts new file mode 100644 index 00000000000..222da3f1a7e --- /dev/null +++ b/cloud/apps/relay/src/relay-websocket-close.test.ts @@ -0,0 +1,44 @@ +import { EventEmitter } from 'node:events' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import type WebSocket from 'ws' +import { closeRelayWebSocket } from './relay-websocket-close.js' + +class FakeSocket extends EventEmitter { + readonly OPEN = 1 + readonly CLOSING = 2 + readonly CLOSED = 3 + readyState = this.OPEN + readonly close = vi.fn(() => { + this.readyState = this.CLOSING + }) + readonly terminate = vi.fn(() => { + this.readyState = this.CLOSED + this.emit('close') + }) +} + +describe('relay WebSocket close', () => { + beforeEach(() => vi.useFakeTimers()) + afterEach(() => vi.useRealTimers()) + + it('terminates a peer that does not complete the close handshake', () => { + const socket = new FakeSocket() + closeRelayWebSocket(socket as unknown as WebSocket, 4408, 'relay draining') + + expect(socket.close).toHaveBeenCalledWith(4408, 'relay draining') + vi.advanceTimersByTime(999) + expect(socket.terminate).not.toHaveBeenCalled() + vi.advanceTimersByTime(1) + expect(socket.terminate).toHaveBeenCalledOnce() + }) + + it('cancels forced termination after the peer closes', () => { + const socket = new FakeSocket() + closeRelayWebSocket(socket as unknown as WebSocket, 4408, 'relay draining') + socket.readyState = socket.CLOSED + socket.emit('close') + vi.runAllTimers() + + expect(socket.terminate).not.toHaveBeenCalled() + }) +}) diff --git a/cloud/apps/relay/src/relay-websocket-close.ts b/cloud/apps/relay/src/relay-websocket-close.ts new file mode 100644 index 00000000000..d4a7f61a3ee --- /dev/null +++ b/cloud/apps/relay/src/relay-websocket-close.ts @@ -0,0 +1,22 @@ +import type WebSocket from 'ws' + +const RELAY_WEBSOCKET_FORCE_CLOSE_MS = 1_000 +const forceCloseTimers = new WeakMap>() + +export function closeRelayWebSocket(socket: WebSocket, code: number, reason: string): void { + if (socket.readyState === socket.CLOSED) return + if (!forceCloseTimers.has(socket)) { + const timer = setTimeout(() => { + forceCloseTimers.delete(socket) + if (socket.readyState !== socket.CLOSED) socket.terminate() + }, RELAY_WEBSOCKET_FORCE_CLOSE_MS) + timer.unref() + forceCloseTimers.set(socket, timer) + socket.once('close', () => { + const pending = forceCloseTimers.get(socket) + if (pending) clearTimeout(pending) + forceCloseTimers.delete(socket) + }) + } + if (socket.readyState === socket.OPEN) socket.close(code, reason) +} diff --git a/cloud/apps/relay/src/relay.blackbox.test.ts b/cloud/apps/relay/src/relay.blackbox.test.ts new file mode 100644 index 00000000000..38134213e76 --- /dev/null +++ b/cloud/apps/relay/src/relay.blackbox.test.ts @@ -0,0 +1,2622 @@ +import { spawn, type ChildProcess } from 'node:child_process' +import { createHash, createHmac } from 'node:crypto' +import { mkdtempSync, rmSync, writeFileSync } from 'node:fs' +import { createServer, type Server } from 'node:http' +import { createServer as createNetServer } from 'node:net' +import { dirname, resolve } from 'node:path' +import { tmpdir } from 'node:os' +import { fileURLToPath } from 'node:url' +import { exportJWK, generateKeyPair, jwtVerify, SignJWT } from 'jose' +import { + buildHostProofMacInput, + HOST_CHALLENGE_PLAINTEXT_DOMAIN +} from '@orca-cloud/relay-contract' +import nacl from 'tweetnacl' +import { afterAll, beforeAll, describe, expect, it } from 'vitest' +import WebSocket from 'ws' +import type { RawData } from 'ws' +import { RelayAssignmentStore } from './assignment-store.js' +import { + encodeMembership, + type CellAdmissionMembership +} from './cell-admission-selector.js' +import { openRelayDatabase } from './database.js' + +const appDirectory = resolve(dirname(fileURLToPath(import.meta.url)), '..') +const assignmentKey = 'test-assignment-key-with-at-least-32-bytes' +let relayProcess: ChildProcess +let jwksServer: Server +let relayUrl: string +let issuer: string +let privateKey: Awaited>['privateKey'] +let adminPrivateKey: Awaited>['privateKey'] +let relayDataDirectory: string +let adminAudience: string +let forwardedSourceSequence = 0 + +function forwardedHeaders(): Record { + forwardedSourceSequence++ + return { + 'x-forwarded-for': `spoofed, 192.0.2.${forwardedSourceSequence}, 35.191.0.1` + } +} + +async function unusedPort(): Promise { + const server = createNetServer() + await new Promise((resolveListen) => server.listen(0, '127.0.0.1', resolveListen)) + const address = server.address() + if (!address || typeof address === 'string') throw new Error('missing test port') + await new Promise((resolveClose) => server.close(() => resolveClose())) + return address.port +} + +async function waitForRelay(child: ChildProcess): Promise { + await new Promise((resolveReady, reject) => { + let stderr = '' + const timeout = setTimeout(() => reject(new Error('relay did not start')), 10_000) + child.stdout?.on('data', (chunk: Buffer) => { + if (chunk.toString().includes('[orca-relay] listening')) { + clearTimeout(timeout) + resolveReady() + } + }) + child.stderr?.on('data', (chunk: Buffer) => (stderr += chunk.toString())) + child.once('exit', (code) => + reject(new Error(`relay exited before ready: ${code}\n${stderr}`)) + ) + }) +} + +async function relayToken(audience: string, relayHostId = 'abcdefghijklmnop'): Promise { + return await new SignJWT({ + prof: 'profile-1', + org: 'org-1', + purpose: 'host-control', + relayHostId + }) + .setProtectedHeader({ alg: 'ES256', kid: 'test-key' }) + .setIssuer(issuer) + .setAudience(audience) + .setSubject('user-1') + .setIssuedAt() + .setExpirationTime('5m') + .sign(privateKey) +} + +async function adminToken(): Promise { + return await googleServiceToken(adminAudience) +} + +async function googleServiceToken(audience: string, email = 'deploy@example.com'): Promise { + return await new SignJWT({ email, email_verified: true }) + .setProtectedHeader({ alg: 'RS256', kid: 'admin-key' }) + .setIssuer('https://accounts.google.com') + .setAudience(audience) + .setSubject('deploy-subject') + .setIssuedAt() + .setExpirationTime('5m') + .sign(adminPrivateKey) +} + +async function postCellHeartbeat( + directorUrl: string, + cell: { id: string; url: string }, + overrides: { incarnation?: string; startedAt?: number; ready?: boolean } = {} +): Promise { + const audience = `${directorUrl}/v1/admin/cell-heartbeat` + return await fetch(audience, { + method: 'POST', + headers: { + authorization: `Bearer ${await googleServiceToken(audience)}`, + 'content-type': 'application/json' + }, + body: JSON.stringify({ + v: 1, + cellId: cell.id, + cellUrl: cell.url, + cellIncarnation: overrides.incarnation ?? '11111111-1111-4111-8111-111111111111', + startedAt: overrides.startedAt ?? Date.now(), + ready: overrides.ready ?? true, + observedRequests: 0 + }) + }) +} + +function spawnTopologyRelay(input: { + url: string + dataDirectory: string + role: 'director' | 'cell' + cellId: string + cells?: Array<{ id: string; url: string; capacityRequests: number }> +}): ChildProcess { + return spawn(process.execPath, ['--import', 'tsx', 'src/index.ts'], { + cwd: appDirectory, + env: { + ...process.env, + PORT: new URL(input.url).port, + ORCA_RELAY_PUBLIC_URL: input.url, + ORCA_RELAY_CELL_URL: input.url, + ORCA_RELAY_CELL_ID: input.cellId, + ORCA_RELAY_CELL_CAPACITY: '10', + ORCA_RELAY_CELLS_JSON: JSON.stringify(input.cells ?? []), + ORCA_RELAY_ROLE: input.role, + ORCA_RELAY_AUTH_ISSUER: issuer, + ORCA_RELAY_JWKS_URL: `${issuer}/jwks`, + ORCA_RELAY_ASSIGNMENT_SIGNING_KEY: assignmentKey, + ORCA_RELAY_DATA_DIR: input.dataDirectory, + ORCA_RELAY_ADMIN_AUDIENCE: adminAudience, + ORCA_RELAY_DEPLOY_SERVICE_ACCOUNT: 'deploy@example.com', + ORCA_RELAY_MONITOR_SERVICE_ACCOUNT: 'monitor@example.com', + ORCA_RELAY_FENCE_SERVICE_ACCOUNT: 'fence@example.com', + ORCA_RELAY_FENCE_BROKER_SERVICE_ACCOUNT: 'broker@example.com', + ORCA_RELAY_ADMIN_JWKS_URL: `${issuer}/jwks` + }, + stdio: ['ignore', 'pipe', 'pipe'] + }) +} + +function nextMessage(socket: WebSocket): Promise> { + return new Promise((resolveMessage, reject) => { + const cleanup = (): void => { + socket.off('message', onMessage) + socket.off('error', onError) + socket.off('close', onClose) + } + const onMessage = (data: RawData): void => { + cleanup() + try { + resolveMessage(JSON.parse(data.toString()) as Record) + } catch (error) { + reject(error) + } + } + const onError = (error: Error): void => { + cleanup() + reject(error) + } + const onClose = (code: number, reason: Buffer): void => { + cleanup() + reject(new Error(`socket closed before message: ${code} ${reason.toString()}`)) + } + socket.once('message', onMessage) + socket.once('error', onError) + socket.once('close', onClose) + }) +} + +function nextRawMessage(socket: WebSocket): Promise<{ data: Buffer; binary: boolean }> { + return new Promise((resolveMessage, reject) => { + socket.once('message', (data, binary) => + resolveMessage({ data: Buffer.from(data as ArrayBuffer), binary }) + ) + socket.once('error', reject) + }) +} + +function collectMessages(socket: WebSocket, count: number): Promise[]> { + return new Promise((resolveMessages, reject) => { + const messages: Record[] = [] + const onMessage = (data: RawData): void => { + try { + messages.push(JSON.parse(data.toString()) as Record) + if (messages.length === count) { + socket.off('message', onMessage) + resolveMessages(messages) + } + } catch (error) { + reject(error) + } + } + socket.on('message', onMessage) + socket.once('error', reject) + }) +} + +async function installDirectCredential(input: { + host: WebSocket + relayDeviceId: string + reqId: string + resumeToken: string +}): Promise> { + const response = nextMessage(input.host) + input.host.send( + JSON.stringify({ + type: 'device-credential-install', + v: 1, + reqId: input.reqId, + relayDeviceId: input.relayDeviceId, + newResumeTokenHash: createHash('sha256').update(input.resumeToken).digest('base64url'), + authorization: { mode: 'authenticated-direct', directAuthId: `direct-${input.reqId}` } + }) + ) + return await response +} + +async function attachPhone(input: { + host: WebSocket + hostAck: Record + hostId: string + credential: string +}): Promise<{ + phone: WebSocket + data: WebSocket + connId: string + hello: Record +}> { + const headers = forwardedHeaders() + const phone = new WebSocket(`${relayUrl.replace('http:', 'ws:')}/v1/connect/${input.hostId}`, { + headers + }) + await new Promise((resolveOpen, reject) => { + phone.once('open', resolveOpen) + phone.once('error', reject) + }) + const connOpenPromise = nextMessage(input.host) + phone.send(JSON.stringify({ type: 'relay-auth', v: 1, mode: 'connect', credential: input.credential })) + const connOpen = await connOpenPromise + expect(connOpen.type).toBe('conn-open') + const connId = String(connOpen.connId) + const data = new WebSocket(`${relayUrl.replace('http:', 'ws:')}/v1/host/data/${connId}`, { + headers + }) + await new Promise((resolveOpen, reject) => { + data.once('open', resolveOpen) + data.once('error', reject) + }) + const helloPromise = nextMessage(phone) + data.send( + JSON.stringify({ + type: 'host-data-auth', + v: 1, + connTicket: connOpen.connTicket, + generation: input.hostAck.generation + }) + ) + const hello = await helloPromise + expect(hello).toMatchObject({ type: 'relay-hello', ok: true }) + return { phone, data, connId, hello } +} + +async function openHostControl(input?: { + controlResumeSecret?: string + previousGeneration?: number + keyPair?: nacl.BoxKeyPair + assignmentEpoch?: number +}): Promise<{ socket: WebSocket; ack: Record; keyPair: nacl.BoxKeyPair }> { + const keyPair = input?.keyPair ?? nacl.box.keyPair() + const hostId = createHash('sha256').update(keyPair.publicKey).digest('base64url').slice(0, 16) + const socket = new WebSocket(`${relayUrl.replace('http:', 'ws:')}/v1/host/control`, { + headers: { authorization: `Bearer ${await relayToken('orca-relay', hostId)}` }, + perMessageDeflate: false + }) + await new Promise((resolveOpen, reject) => { + socket.once('open', resolveOpen) + socket.once('error', reject) + }) + socket.send( + JSON.stringify({ + type: 'host-hello', + v: 1, + relayHostId: hostId, + assignmentEpoch: input?.assignmentEpoch ?? 1, + hostPublicKeyB64: Buffer.from(keyPair.publicKey).toString('base64'), + appVersion: 'test', + ...(input?.controlResumeSecret + ? { controlResumeSecret: input.controlResumeSecret } + : {}), + ...(input?.previousGeneration === undefined + ? {} + : { previousGeneration: input.previousGeneration }) + }) + ) + const challenge = await nextMessage(socket) + expect(challenge.type).toBe('host-challenge') + const plaintext = nacl.box.open( + Buffer.from(String(challenge.ciphertextB64), 'base64'), + Buffer.from(String(challenge.nonceB64), 'base64'), + Buffer.from(String(challenge.relayEphemeralPublicKeyB64), 'base64'), + keyPair.secretKey + ) + if (!plaintext) throw new Error('host challenge did not decrypt') + const domain = new TextEncoder().encode(`${HOST_CHALLENGE_PLAINTEXT_DOMAIN}\0`) + expect(plaintext.slice(0, domain.length)).toEqual(domain) + const transcriptLength = new DataView( + plaintext.buffer, + plaintext.byteOffset + domain.length, + 4 + ).getUint32(0, false) + const transcriptStart = domain.length + 4 + const transcript = plaintext.slice(transcriptStart, transcriptStart + transcriptLength) + const secret = plaintext.slice(transcriptStart + transcriptLength) + const proofB64 = createHmac('sha256', secret) + .update(buildHostProofMacInput(transcript)) + .digest('base64') + socket.send( + JSON.stringify({ + type: 'host-challenge-ack', + challengeId: challenge.challengeId, + proofB64 + }) + ) + return { socket, ack: await nextMessage(socket), keyPair } +} + +beforeAll(async () => { + const keys = await generateKeyPair('ES256') + privateKey = keys.privateKey + const publicJwk = await exportJWK(keys.publicKey) + const adminKeys = await generateKeyPair('RS256') + adminPrivateKey = adminKeys.privateKey + const adminPublicJwk = await exportJWK(adminKeys.publicKey) + jwksServer = createServer((_request, response) => { + response.setHeader('content-type', 'application/json') + response.end( + JSON.stringify({ + keys: [ + { ...publicJwk, kid: 'test-key', alg: 'ES256', use: 'sig' }, + { ...adminPublicJwk, kid: 'admin-key', alg: 'RS256', use: 'sig' } + ] + }) + ) + }) + await new Promise((resolveListen) => jwksServer.listen(0, '127.0.0.1', resolveListen)) + const jwksAddress = jwksServer.address() + if (!jwksAddress || typeof jwksAddress === 'string') throw new Error('missing JWKS address') + issuer = `http://127.0.0.1:${jwksAddress.port}` + const relayPort = await unusedPort() + relayUrl = `http://127.0.0.1:${relayPort}` + adminAudience = `${relayUrl}/v1/admin/drain` + relayDataDirectory = mkdtempSync(resolve(tmpdir(), 'orca-relay-blackbox-')) + relayProcess = spawn(process.execPath, ['--import', 'tsx', 'src/index.ts'], { + cwd: appDirectory, + env: { + ...process.env, + PORT: String(relayPort), + ORCA_RELAY_PUBLIC_URL: relayUrl, + ORCA_RELAY_CELL_URL: relayUrl, + ORCA_RELAY_AUTH_ISSUER: issuer, + ORCA_RELAY_JWKS_URL: `${issuer}/jwks`, + ORCA_RELAY_ASSIGNMENT_SIGNING_KEY: assignmentKey, + ORCA_RELAY_DATA_DIR: relayDataDirectory, + ORCA_RELAY_ADMIN_AUDIENCE: adminAudience, + ORCA_RELAY_DEPLOY_SERVICE_ACCOUNT: 'deploy@example.com', + ORCA_RELAY_MONITOR_SERVICE_ACCOUNT: 'monitor@example.com', + ORCA_RELAY_FENCE_SERVICE_ACCOUNT: 'fence@example.com', + ORCA_RELAY_IMAGE_DIGEST: `sha256:${'a'.repeat(64)}`, + ORCA_RELAY_ADMIN_JWKS_URL: `${issuer}/jwks` + }, + stdio: ['ignore', 'pipe', 'pipe'] + }) + await waitForRelay(relayProcess) +}) + +afterAll(async () => { + relayProcess?.kill('SIGTERM') + await new Promise((resolveClose) => jwksServer?.close(() => resolveClose())) + rmSync(relayDataDirectory, { recursive: true, force: true }) +}) + +describe('served relay URL', () => { + it('exposes only /health, never /healthz', async () => { + expect(await (await fetch(`${relayUrl}/health`)).json()).toEqual({ + ok: true, + connectionCapacityProtocol: 2 + }) + expect(await (await fetch(`${relayUrl}/ready`)).json()).toEqual({ ok: true }) + expect((await fetch(`${relayUrl}/healthz`)).status).toBe(404) + }) + + it('keeps liveness healthy when dependency readiness fails', async () => { + const port = await unusedPort() + const url = `http://127.0.0.1:${port}` + const dataDirectory = mkdtempSync(resolve(tmpdir(), 'orca-relay-unready-')) + const processUnderTest = spawn(process.execPath, ['--import', 'tsx', 'src/index.ts'], { + cwd: appDirectory, + env: { + ...process.env, + PORT: String(port), + ORCA_RELAY_PUBLIC_URL: url, + ORCA_RELAY_CELL_URL: url, + ORCA_RELAY_AUTH_ISSUER: issuer, + ORCA_RELAY_JWKS_URL: 'http://127.0.0.1:1/jwks', + ORCA_RELAY_ASSIGNMENT_SIGNING_KEY: assignmentKey, + ORCA_RELAY_DATA_DIR: dataDirectory, + ORCA_RELAY_ADMIN_AUDIENCE: adminAudience, + ORCA_RELAY_DEPLOY_SERVICE_ACCOUNT: 'deploy@example.com', + ORCA_RELAY_IMAGE_DIGEST: `sha256:${'a'.repeat(64)}`, + ORCA_RELAY_ADMIN_JWKS_URL: `${issuer}/jwks` + }, + stdio: ['ignore', 'pipe', 'pipe'] + }) + try { + await waitForRelay(processUnderTest) + expect((await fetch(`${url}/health`)).status).toBe(200) + expect((await fetch(`${url}/ready`)).status).toBe(503) + } finally { + processUnderTest.kill('SIGTERM') + await new Promise((resolveExit) => processUnderTest.once('exit', () => resolveExit())) + rmSync(dataDirectory, { recursive: true, force: true }) + } + }) + + it('rejects missing and broad-audience bearer tokens', async () => { + const body = JSON.stringify({ v: 1, relayHostId: 'abcdefghijklmnop' }) + expect((await fetch(`${relayUrl}/v1/assign`, { method: 'POST', body })).status).toBe(401) + expect( + ( + await fetch(`${relayUrl}/v1/assign`, { + method: 'POST', + headers: { authorization: `Bearer ${await relayToken('orca-cloud')}` }, + body + }) + ).status + ).toBe(401) + }) + + it('bounds slow first-frame admission and rejects URL credentials before upgrade', async () => { + const slow: WebSocket[] = [] + for (let index = 0; index < 4; index++) { + const socket = new WebSocket( + `${relayUrl.replace('http:', 'ws:')}/v1/connect/abcdefghijklmnop` + ) + await new Promise((resolveOpen, reject) => { + socket.once('open', resolveOpen) + socket.once('error', reject) + }) + slow.push(socket) + } + const limited = new WebSocket( + `${relayUrl.replace('http:', 'ws:')}/v1/connect/abcdefghijklmnop` + ) + const limitedStatus = await new Promise((resolveStatus) => { + limited.once('unexpected-response', (_request, response) => + resolveStatus(response.statusCode ?? 0) + ) + limited.once('error', () => {}) + }) + expect(limitedStatus).toBe(429) + const isolated = new WebSocket( + `${relayUrl.replace('http:', 'ws:')}/v1/connect/abcdefghijklmnop`, + { headers: { 'x-forwarded-for': 'spoofed, 198.51.100.9, 35.191.0.1' } } + ) + await new Promise((resolveOpen, reject) => { + isolated.once('open', resolveOpen) + isolated.once('error', reject) + }) + isolated.close() + for (const socket of slow) socket.close() + + const leaked = new WebSocket( + `${relayUrl.replace('http:', 'ws:')}/v1/connect/abcdefghijklmnop?credential=secret` + ) + const leakedStatus = await new Promise((resolveStatus) => { + leaked.once('unexpected-response', (_request, response) => + resolveStatus(response.statusCode ?? 0) + ) + leaked.once('error', () => {}) + }) + expect(leakedStatus).toBe(400) + }) + + it('enforces process-wide slow-auth and per-source rate budgets', async () => { + const slow: WebSocket[] = [] + for (let index = 0; index < 45; index++) { + const socket = new WebSocket( + `${relayUrl.replace('http:', 'ws:')}/v1/connect/abcdefghijklmnop`, + { headers: { 'x-forwarded-for': `spoofed, 198.51.100.${index + 1}, 35.191.0.1` } } + ) + await new Promise((resolveOpen, reject) => { + socket.once('open', resolveOpen) + socket.once('error', reject) + }) + slow.push(socket) + } + const globallyLimited = new WebSocket( + `${relayUrl.replace('http:', 'ws:')}/v1/connect/abcdefghijklmnop`, + { headers: { 'x-forwarded-for': 'spoofed, 203.0.113.1, 35.191.0.1' } } + ) + const globalStatus = await new Promise((resolveStatus) => { + globallyLimited.once('unexpected-response', (_request, response) => + resolveStatus(response.statusCode ?? 0) + ) + globallyLimited.once('error', () => {}) + }) + expect(globalStatus).toBe(429) + await Promise.all( + slow.map( + (socket) => + new Promise((resolveClose) => { + socket.once('close', () => resolveClose()) + socket.close() + }) + ) + ) + + for (let index = 0; index < 30; index++) { + const socket = new WebSocket( + `${relayUrl.replace('http:', 'ws:')}/v1/connect/abcdefghijklmnop`, + { headers: { 'x-forwarded-for': 'spoofed, 203.0.113.2, 35.191.0.1' } } + ) + await new Promise((resolveOpen, reject) => { + socket.once('open', resolveOpen) + socket.once('error', reject) + }) + const closed = new Promise((resolveClose) => socket.once('close', () => resolveClose())) + socket.send('{}') + await closed + } + const rateLimited = new WebSocket( + `${relayUrl.replace('http:', 'ws:')}/v1/connect/abcdefghijklmnop`, + { headers: { 'x-forwarded-for': 'spoofed, 203.0.113.2, 35.191.0.1' } } + ) + const rateStatus = await new Promise((resolveStatus) => { + rateLimited.once('unexpected-response', (_request, response) => + resolveStatus(response.statusCode ?? 0) + ) + rateLimited.once('error', () => {}) + }) + expect(rateStatus).toBe(429) + }) + + it('returns a signed combined-service assignment for the bound host', async () => { + const response = await fetch(`${relayUrl}/v1/assign`, { + method: 'POST', + headers: { + authorization: `Bearer ${await relayToken('orca-relay')}`, + 'content-type': 'application/json' + }, + body: JSON.stringify({ v: 1, relayHostId: 'abcdefghijklmnop' }) + }) + expect(response.status).toBe(200) + const assignment = (await response.json()) as { + v: number + cellUrl: string + assignmentEpoch: number + lease: string + } + expect(assignment).toMatchObject({ v: 1, cellUrl: relayUrl, assignmentEpoch: 1 }) + const verified = await jwtVerify(assignment.lease, new TextEncoder().encode(assignmentKey), { + issuer: relayUrl, + audience: 'orca-relay-cell', + algorithms: ['HS256'] + }) + expect(verified.payload).toMatchObject({ relayHostId: 'abcdefghijklmnop' }) + }) + + it('requires canonical host key possession and supports same-generation rebind', async () => { + const first = await openHostControl() + expect(first.ack).toMatchObject({ type: 'host-hello-ack', v: 1, generation: 1 }) + const firstClose = new Promise((resolveClose) => + first.socket.once('close', (code) => resolveClose(code)) + ) + const rebound = await openHostControl({ + keyPair: first.keyPair, + controlResumeSecret: String(first.ack.controlResumeSecret), + previousGeneration: 1 + }) + expect(rebound.ack).toMatchObject({ type: 'host-hello-ack', v: 1, generation: 1 }) + expect(await firstClose).toBe(4408) + rebound.socket.close() + }) + + it('rejects a scoped token presented by a different host key', async () => { + const claimedKey = nacl.box.keyPair() + const wrongKey = nacl.box.keyPair() + const hostId = createHash('sha256').update(claimedKey.publicKey).digest('base64url').slice(0, 16) + const socket = new WebSocket(`${relayUrl.replace('http:', 'ws:')}/v1/host/control`, { + headers: { authorization: `Bearer ${await relayToken('orca-relay', hostId)}` } + }) + await new Promise((resolveOpen) => socket.once('open', resolveOpen)) + const closed = new Promise((resolveClose) => + socket.once('close', (code) => resolveClose(code)) + ) + socket.send( + JSON.stringify({ + type: 'host-hello', + v: 1, + relayHostId: hostId, + assignmentEpoch: 1, + hostPublicKeyB64: Buffer.from(wrongKey.publicKey).toString('base64'), + appVersion: 'test' + }) + ) + expect(await closed).toBe(4401) + }) + + it('returns typed 4409 without a cell URL for a stale assignment epoch', async () => { + const keyPair = nacl.box.keyPair() + const hostId = createHash('sha256').update(keyPair.publicKey).digest('base64url').slice(0, 16) + const socket = new WebSocket(`${relayUrl.replace('http:', 'ws:')}/v1/host/control`, { + headers: { authorization: `Bearer ${await relayToken('orca-relay', hostId)}` } + }) + await new Promise((resolveOpen, reject) => { + socket.once('open', resolveOpen) + socket.once('error', reject) + }) + const closed = new Promise<{ code: number; reason: string }>((resolveClose) => + socket.once('close', (code, reason) => + resolveClose({ code, reason: reason.toString() }) + ) + ) + socket.send( + JSON.stringify({ + type: 'host-hello', + v: 1, + relayHostId: hostId, + assignmentEpoch: 2, + hostPublicKeyB64: Buffer.from(keyPair.publicKey).toString('base64'), + appVersion: 'test' + }) + ) + const result = await closed + expect(result.code).toBe(4409) + expect(result.reason).not.toContain('http') + }) + + it('keeps a pending attach usable after a bad ticket and rejects ticket replay', async () => { + const host = await openHostControl() + const hostId = createHash('sha256') + .update(host.keyPair.publicKey) + .digest('base64url') + .slice(0, 16) + const inviteResponse = nextMessage(host.socket) + host.socket.send( + JSON.stringify({ type: 'invite-create', reqId: 'ticket-invite', relayDeviceId: 'ticket-device' }) + ) + const invite = await inviteResponse + const phone = new WebSocket(`${relayUrl.replace('http:', 'ws:')}/v1/connect/${hostId}`, { + headers: forwardedHeaders() + }) + await new Promise((resolveOpen, reject) => { + phone.once('open', resolveOpen) + phone.once('error', reject) + }) + const connectionPromise = nextMessage(host.socket) + phone.send( + JSON.stringify({ type: 'relay-auth', v: 1, mode: 'connect', credential: invite.inviteToken }) + ) + const connection = await connectionPromise + const dataUrl = `${relayUrl.replace('http:', 'ws:')}/v1/host/data/${connection.connId}` + + const badData = new WebSocket(dataUrl, { headers: forwardedHeaders() }) + await new Promise((resolveOpen, reject) => { + badData.once('open', resolveOpen) + badData.once('error', reject) + }) + const badClosed = new Promise((resolveClose) => + badData.once('close', (code) => resolveClose(code)) + ) + badData.send( + JSON.stringify({ + type: 'host-data-auth', + v: 1, + connTicket: Buffer.alloc(32, 1).toString('base64url'), + generation: host.ack.generation + }) + ) + expect(await badClosed).toBe(4401) + + const data = new WebSocket(dataUrl, { headers: forwardedHeaders() }) + await new Promise((resolveOpen, reject) => { + data.once('open', resolveOpen) + data.once('error', reject) + }) + const hello = nextMessage(phone) + data.send( + JSON.stringify({ + type: 'host-data-auth', + v: 1, + connTicket: connection.connTicket, + generation: host.ack.generation + }) + ) + expect(await hello).toMatchObject({ type: 'relay-hello', ok: true }) + + const replay = new WebSocket(dataUrl, { headers: forwardedHeaders() }) + await new Promise((resolveOpen, reject) => { + replay.once('open', resolveOpen) + replay.once('error', reject) + }) + const replayClosed = new Promise((resolveClose) => + replay.once('close', (code) => resolveClose(code)) + ) + replay.send( + JSON.stringify({ + type: 'host-data-auth', + v: 1, + connTicket: connection.connTicket, + generation: host.ack.generation + }) + ) + expect(await replayClosed).toBe(4401) + phone.close() + data.close() + host.socket.close() + }) + + it('attaches before success, preserves text/binary opcodes, installs, and confirms resume', async () => { + const host = await openHostControl() + const hostId = createHash('sha256') + .update(host.keyPair.publicKey) + .digest('base64url') + .slice(0, 16) + const invitePromise = nextMessage(host.socket) + host.socket.send( + JSON.stringify({ type: 'invite-create', reqId: 'invite-1', relayDeviceId: 'device-1' }) + ) + const invite = await invitePromise + expect(invite.type).toBe('invite-created') + const first = await attachPhone({ + host: host.socket, + hostAck: host.ack, + hostId, + credential: String(invite.inviteToken) + }) + + const hostText = nextRawMessage(first.data) + first.phone.send('phone-text') + expect(await hostText).toEqual({ data: Buffer.from('phone-text'), binary: false }) + const phoneBinary = nextRawMessage(first.phone) + first.data.send(Buffer.from([1, 2, 3]), { binary: true }) + expect(await phoneBinary).toEqual({ data: Buffer.from([1, 2, 3]), binary: true }) + + const resumeToken = Buffer.alloc(32, 9).toString('base64url') + const installedPromise = nextMessage(host.socket) + host.socket.send( + JSON.stringify({ + type: 'device-credential-install', + v: 1, + reqId: 'install-1', + relayDeviceId: 'device-1', + newResumeTokenHash: createHash('sha256').update(resumeToken).digest('base64url'), + authorization: { mode: 'relay-basis', basisConnId: first.connId } + }) + ) + const installed = await installedPromise + expect(installed).toMatchObject({ + type: 'device-credential-installed', + authorizationMode: 'relay-basis', + currentVersion: 1 + }) + const resolved = await fetch(`${relayUrl}/v1/resolve`, { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ v: 1, relayHostId: hostId, resumeToken }) + }) + expect(resolved.status).toBe(200) + expect(await resolved.json()).toMatchObject({ v: 1, cellUrl: relayUrl, assignmentEpoch: 1 }) + first.phone.close() + + const resumed = await attachPhone({ + host: host.socket, + hostAck: host.ack, + hostId, + credential: resumeToken + }) + const confirmedPromise = nextMessage(host.socket) + host.socket.send( + JSON.stringify({ + type: 'device-resume-confirm', + v: 1, + reqId: 'confirm-1', + basisConnId: resumed.connId + }) + ) + expect(await confirmedPromise).toMatchObject({ + type: 'device-resume-confirmed', + renewed: true, + acceptedAs: 'current' + }) + resumed.phone.close() + resumed.data.close() + host.socket.close() + }) + + it('does not renew outer-only or injected confirmations and reports offline/peer loss', async () => { + const host = await openHostControl() + const hostId = createHash('sha256') + .update(host.keyPair.publicKey) + .digest('base64url') + .slice(0, 16) + const resumeToken = Buffer.alloc(32, 11).toString('base64url') + expect( + await installDirectCredential({ + host: host.socket, + relayDeviceId: 'outer-device', + reqId: 'outer-install', + resumeToken + }) + ).toMatchObject({ type: 'device-credential-installed', currentVersion: 1 }) + + const first = await attachPhone({ host: host.socket, hostAck: host.ack, hostId, credential: resumeToken }) + const originalExpiry = first.hello.resumeExpiresAt + first.phone.close() + first.data.close() + await new Promise((resolveWait) => setTimeout(resolveWait, 25)) + const closedBasis = nextMessage(host.socket) + host.socket.send( + JSON.stringify({ + type: 'device-resume-confirm', + v: 1, + reqId: 'closed-basis-confirm', + basisConnId: first.connId + }) + ) + expect(await closedBasis).toMatchObject({ + type: 'control-error', + reqId: 'closed-basis-confirm', + code: 'confirmation_not_active' + }) + const second = await attachPhone({ host: host.socket, hostAck: host.ack, hostId, credential: resumeToken }) + expect(second.hello.resumeExpiresAt).toBe(originalExpiry) + const injectedResult = nextMessage(host.socket) + host.socket.send( + JSON.stringify({ + type: 'device-resume-confirm', + v: 1, + reqId: 'injected-confirm', + basisConnId: second.connId, + relayDeviceId: 'injected', + acceptedCredentialVersion: 99 + }) + ) + expect(await injectedResult).toMatchObject({ type: 'control-error', reqId: 'injected-confirm' }) + const phoneClosed = new Promise((resolveClose) => + second.phone.once('close', (code) => resolveClose(code)) + ) + second.data.close() + expect(await phoneClosed).toBe(4408) + + await new Promise((resolveClose) => { + host.socket.once('close', () => resolveClose()) + host.socket.close() + }) + const offline = new WebSocket(`${relayUrl.replace('http:', 'ws:')}/v1/connect/${hostId}`, { + headers: forwardedHeaders() + }) + await new Promise((resolveOpen) => offline.once('open', resolveOpen)) + const offlineHello = nextMessage(offline) + offline.send( + JSON.stringify({ type: 'relay-auth', v: 1, mode: 'connect', credential: resumeToken }) + ) + expect(await offlineHello).toEqual({ type: 'relay-hello', ok: false, code: 4404 }) + }) + + it('rejects the first post-E2EE confirmation after its server-owned deadline', async () => { + const originalUrl = relayUrl + const clockPort = await unusedPort() + const clockUrl = `http://127.0.0.1:${clockPort}` + const clockData = mkdtempSync(resolve(tmpdir(), 'orca-relay-clock-')) + const clockFile = resolve(clockData, 'offset-ms') + writeFileSync(clockFile, '0') + const clockProcess = spawn( + process.execPath, + ['--import', 'tsx', 'src/fault-injection-test-entry.ts'], + { + cwd: appDirectory, + env: { + ...process.env, + NODE_ENV: 'test', + PORT: String(clockPort), + ORCA_RELAY_PUBLIC_URL: clockUrl, + ORCA_RELAY_CELL_URL: clockUrl, + ORCA_RELAY_AUTH_ISSUER: issuer, + ORCA_RELAY_JWKS_URL: `${issuer}/jwks`, + ORCA_RELAY_ASSIGNMENT_SIGNING_KEY: assignmentKey, + ORCA_RELAY_DATA_DIR: clockData, + ORCA_RELAY_ADMIN_AUDIENCE: adminAudience, + ORCA_RELAY_DEPLOY_SERVICE_ACCOUNT: 'deploy@example.com', + ORCA_RELAY_ADMIN_JWKS_URL: `${issuer}/jwks`, + ORCA_RELAY_TEST_CLOCK_FILE: clockFile + }, + stdio: ['ignore', 'pipe', 'pipe'] + } + ) + relayUrl = clockUrl + try { + await waitForRelay(clockProcess) + const host = await openHostControl() + const hostId = createHash('sha256') + .update(host.keyPair.publicKey) + .digest('base64url') + .slice(0, 16) + const resumeToken = Buffer.alloc(32, 25).toString('base64url') + await installDirectCredential({ + host: host.socket, + relayDeviceId: 'late-confirm-device', + reqId: 'late-confirm-install', + resumeToken + }) + const splice = await attachPhone({ + host: host.socket, + hostAck: host.ack, + hostId, + credential: resumeToken + }) + writeFileSync(clockFile, '31000') + const rejected = nextMessage(host.socket) + host.socket.send( + JSON.stringify({ + type: 'device-resume-confirm', + v: 1, + reqId: 'late-first-confirm', + basisConnId: splice.connId + }) + ) + expect(await rejected).toMatchObject({ + type: 'control-error', + reqId: 'late-first-confirm', + code: 'confirmation_not_active' + }) + splice.phone.close() + splice.data.close() + host.socket.close() + } finally { + clockProcess.kill('SIGKILL') + await new Promise((resolveExit) => clockProcess.once('exit', () => resolveExit())) + relayUrl = originalUrl + rmSync(clockData, { recursive: true, force: true }) + } + }) + + it('fences a competing generation and rejects its old immutable basis', async () => { + const firstHost = await openHostControl() + const hostId = createHash('sha256') + .update(firstHost.keyPair.publicKey) + .digest('base64url') + .slice(0, 16) + const resumeToken = Buffer.alloc(32, 12).toString('base64url') + await installDirectCredential({ + host: firstHost.socket, + relayDeviceId: 'fenced-device', + reqId: 'fenced-install', + resumeToken + }) + const splice = await attachPhone({ + host: firstHost.socket, + hostAck: firstHost.ack, + hostId, + credential: resumeToken + }) + const replacement = await openHostControl({ keyPair: firstHost.keyPair, previousGeneration: 1 }) + expect(replacement.ack.generation).toBe(2) + const rejected = nextMessage(replacement.socket) + replacement.socket.send( + JSON.stringify({ + type: 'device-resume-confirm', + v: 1, + reqId: 'wrong-generation-confirm', + basisConnId: splice.connId + }) + ) + expect(await rejected).toMatchObject({ + type: 'control-error', + reqId: 'wrong-generation-confirm', + code: 'confirmation_not_active' + }) + replacement.socket.close() + }) + + it('serializes late direct and invite authorization modes into one served result', async () => { + const host = await openHostControl() + const hostId = createHash('sha256') + .update(host.keyPair.publicKey) + .digest('base64url') + .slice(0, 16) + const inviteResponse = nextMessage(host.socket) + host.socket.send( + JSON.stringify({ type: 'invite-create', reqId: 'race-invite', relayDeviceId: 'race-device' }) + ) + const invite = await inviteResponse + const splice = await attachPhone({ + host: host.socket, + hostAck: host.ack, + hostId, + credential: String(invite.inviteToken) + }) + const responses = collectMessages(host.socket, 2) + const base = { + type: 'device-credential-install', + v: 1, + reqId: 'race-install', + relayDeviceId: 'race-device', + newResumeTokenHash: createHash('sha256').update('race-resume').digest('base64url') + } + host.socket.send( + JSON.stringify({ + ...base, + authorization: { mode: 'relay-basis', basisConnId: splice.connId } + }) + ) + host.socket.send( + JSON.stringify({ + ...base, + authorization: { mode: 'authenticated-direct', directAuthId: 'late-direct' } + }) + ) + const installed = await responses + expect(installed).toHaveLength(2) + expect(installed[0]).toEqual(installed[1]) + expect(installed[0]).toMatchObject({ type: 'device-credential-installed', currentVersion: 1 }) + splice.phone.close() + splice.data.close() + host.socket.close() + }) + + it('reconciles a direct commit after its response is ignored and rejects bad authorization', async () => { + const firstHost = await openHostControl() + const hostId = createHash('sha256') + .update(firstHost.keyPair.publicKey) + .digest('base64url') + .slice(0, 16) + const inviteResponse = nextMessage(firstHost.socket) + firstHost.socket.send( + JSON.stringify({ + type: 'invite-create', + reqId: 'lost-response-invite', + relayDeviceId: 'lost-response-device' + }) + ) + const invite = await inviteResponse + const resumeToken = Buffer.alloc(32, 26).toString('base64url') + firstHost.socket.send( + JSON.stringify({ + type: 'device-credential-install', + v: 1, + reqId: 'lost-response-install', + relayDeviceId: 'lost-response-device', + newResumeTokenHash: createHash('sha256').update(resumeToken).digest('base64url'), + authorization: { mode: 'authenticated-direct', directAuthId: 'lost-response-direct' } + }) + ) + // The coordinator deliberately ignores the acknowledgement and recovers + // only from the durable status after its control transport disappears. + await new Promise((resolveWait) => setTimeout(resolveWait, 25)) + firstHost.socket.terminate() + const rebound = await openHostControl({ + keyPair: firstHost.keyPair, + controlResumeSecret: String(firstHost.ack.controlResumeSecret), + previousGeneration: Number(firstHost.ack.generation) + }) + const status = nextMessage(rebound.socket) + rebound.socket.send( + JSON.stringify({ + type: 'device-credential-install-status', + v: 1, + reqId: 'lost-response-install', + relayDeviceId: 'lost-response-device' + }) + ) + expect(await status).toMatchObject({ + type: 'device-credential-install-status-result', + state: 'committed', + result: { authorizationMode: 'authenticated-direct', currentVersion: 1 } + }) + + const invalidatedInvite = new WebSocket( + `${relayUrl.replace('http:', 'ws:')}/v1/connect/${hostId}`, + { headers: forwardedHeaders() } + ) + await new Promise((resolveOpen, reject) => { + invalidatedInvite.once('open', resolveOpen) + invalidatedInvite.once('error', reject) + }) + const rejectedInvite = nextMessage(invalidatedInvite) + invalidatedInvite.send( + JSON.stringify({ type: 'relay-auth', v: 1, mode: 'connect', credential: invite.inviteToken }) + ) + expect(await rejectedInvite).toEqual({ type: 'relay-hello', ok: false, code: 4401 }) + + const unauthorized = nextMessage(rebound.socket) + rebound.socket.send( + JSON.stringify({ + type: 'device-credential-install', + v: 1, + reqId: 'unauthorized-install', + relayDeviceId: 'unauthorized-device', + newResumeTokenHash: createHash('sha256').update('unauthorized').digest('base64url'), + authorization: { mode: 'relay-basis', basisConnId: 'unknown-basis' } + }) + ) + expect(await unauthorized).toMatchObject({ + type: 'control-error', + reqId: 'unauthorized-install', + code: 'invalid_relay_basis' + }) + const missing = nextMessage(rebound.socket) + rebound.socket.send( + JSON.stringify({ + type: 'device-credential-install-status', + v: 1, + reqId: 'unauthorized-install', + relayDeviceId: 'unauthorized-device' + }) + ) + expect(await missing).toMatchObject({ + type: 'device-credential-install-status-result', + state: 'not-found' + }) + rebound.socket.close() + }) + + it('serializes confirmation against direct rotation, relay rotation, and revoke', async () => { + const host = await openHostControl() + const hostId = createHash('sha256') + .update(host.keyPair.publicKey) + .digest('base64url') + .slice(0, 16) + const firstToken = Buffer.alloc(32, 21).toString('base64url') + await installDirectCredential({ + host: host.socket, + relayDeviceId: 'confirm-race-device', + reqId: 'confirm-race-initial', + resumeToken: firstToken + }) + + const firstResume = await attachPhone({ + host: host.socket, + hostAck: host.ack, + hostId, + credential: firstToken + }) + const secondToken = Buffer.alloc(32, 22).toString('base64url') + const directResponses = collectMessages(host.socket, 2) + host.socket.send( + JSON.stringify({ + type: 'device-resume-confirm', + v: 1, + reqId: 'confirm-before-direct', + basisConnId: firstResume.connId + }) + ) + host.socket.send( + JSON.stringify({ + type: 'device-credential-install', + v: 1, + reqId: 'direct-after-confirm', + relayDeviceId: 'confirm-race-device', + newResumeTokenHash: createHash('sha256').update(secondToken).digest('base64url'), + expectedCurrentHash: createHash('sha256').update(firstToken).digest('base64url'), + authorization: { mode: 'authenticated-direct', directAuthId: 'direct-after-confirm' } + }) + ) + const directResults = await directResponses + expect(directResults.find((result) => result.type === 'device-resume-confirmed')).toMatchObject({ + reqId: 'confirm-before-direct', + currentVersion: 1, + renewed: true + }) + expect(directResults.find((result) => result.type === 'device-credential-installed')).toMatchObject({ + reqId: 'direct-after-confirm', + currentVersion: 2 + }) + const replayPromise = nextMessage(host.socket) + host.socket.send( + JSON.stringify({ + type: 'device-resume-confirm', + v: 1, + reqId: 'confirm-before-direct', + basisConnId: firstResume.connId + }) + ) + expect(await replayPromise).toEqual( + directResults.find((result) => result.type === 'device-resume-confirmed') + ) + + const secondResume = await attachPhone({ + host: host.socket, + hostAck: host.ack, + hostId, + credential: secondToken + }) + const inviteResponse = nextMessage(host.socket) + host.socket.send( + JSON.stringify({ + type: 'invite-create', + reqId: 'relay-rotation-invite', + relayDeviceId: 'confirm-race-device' + }) + ) + const invite = await inviteResponse + const inviteSplice = await attachPhone({ + host: host.socket, + hostAck: host.ack, + hostId, + credential: String(invite.inviteToken) + }) + const thirdToken = Buffer.alloc(32, 23).toString('base64url') + const relayResponses = collectMessages(host.socket, 2) + host.socket.send( + JSON.stringify({ + type: 'device-credential-install', + v: 1, + reqId: 'relay-before-confirm', + relayDeviceId: 'confirm-race-device', + newResumeTokenHash: createHash('sha256').update(thirdToken).digest('base64url'), + expectedCurrentHash: createHash('sha256').update(secondToken).digest('base64url'), + authorization: { mode: 'relay-basis', basisConnId: inviteSplice.connId } + }) + ) + host.socket.send( + JSON.stringify({ + type: 'device-resume-confirm', + v: 1, + reqId: 'confirm-after-relay', + basisConnId: secondResume.connId + }) + ) + const relayResults = await relayResponses + expect(relayResults.find((result) => result.type === 'device-credential-installed')).toMatchObject({ + reqId: 'relay-before-confirm', + currentVersion: 3 + }) + expect(relayResults.find((result) => result.type === 'device-resume-confirmed')).toMatchObject({ + reqId: 'confirm-after-relay', + currentVersion: 3, + renewed: false + }) + + const retiredResume = await attachPhone({ + host: host.socket, + hostAck: host.ack, + hostId, + credential: secondToken + }) + const fourthToken = Buffer.alloc(32, 24).toString('base64url') + expect( + await installDirectCredential({ + host: host.socket, + relayDeviceId: 'confirm-race-device', + reqId: 'retire-before-confirm', + resumeToken: fourthToken + }) + ).toMatchObject({ type: 'device-credential-installed', currentVersion: 4 }) + const retiredResult = nextMessage(host.socket) + host.socket.send( + JSON.stringify({ + type: 'device-resume-confirm', + v: 1, + reqId: 'confirm-retired', + basisConnId: retiredResume.connId + }) + ) + expect(await retiredResult).toMatchObject({ + type: 'control-error', + reqId: 'confirm-retired', + code: 'reject-retired' + }) + + const currentResume = await attachPhone({ + host: host.socket, + hostAck: host.ack, + hostId, + credential: fourthToken + }) + const revokeResponses = collectMessages(host.socket, 2) + host.socket.send( + JSON.stringify({ + type: 'device-revoke', + reqId: 'revoke-before-confirm', + relayDeviceId: 'confirm-race-device' + }) + ) + host.socket.send( + JSON.stringify({ + type: 'device-resume-confirm', + v: 1, + reqId: 'confirm-after-revoke', + basisConnId: currentResume.connId + }) + ) + const revokeResults = await revokeResponses + expect(revokeResults.find((result) => result.type === 'device-revoked')).toMatchObject({ + reqId: 'revoke-before-confirm' + }) + expect(revokeResults.find((result) => result.type === 'control-error')).toMatchObject({ + reqId: 'confirm-after-revoke', + code: 'reject-revoked' + }) + + for (const splice of [firstResume, secondResume, inviteSplice, retiredResume, currentResume]) { + splice.phone.close() + splice.data.close() + } + host.socket.close() + }) + + it('enforces the eight-splice host limit with typed 4429 recovery', async () => { + const host = await openHostControl() + const hostId = createHash('sha256') + .update(host.keyPair.publicKey) + .digest('base64url') + .slice(0, 16) + const resumeToken = Buffer.alloc(32, 13).toString('base64url') + await installDirectCredential({ + host: host.socket, + relayDeviceId: 'limit-device', + reqId: 'limit-install', + resumeToken + }) + const splices = [] + for (let index = 0; index < 8; index++) { + splices.push( + await attachPhone({ host: host.socket, hostAck: host.ack, hostId, credential: resumeToken }) + ) + } + const ninth = new WebSocket(`${relayUrl.replace('http:', 'ws:')}/v1/connect/${hostId}`, { + headers: forwardedHeaders() + }) + await new Promise((resolveOpen) => ninth.once('open', resolveOpen)) + const rejected = nextMessage(ninth) + ninth.send( + JSON.stringify({ type: 'relay-auth', v: 1, mode: 'connect', credential: resumeToken }) + ) + expect(await rejected).toEqual({ type: 'relay-hello', ok: false, code: 4429 }) + for (const splice of splices) { + splice.phone.close() + splice.data.close() + } + host.socket.close() + }) + + it('rolls an aborted invite reservation into bounded cooldown before retry', async () => { + const host = await openHostControl() + const hostId = createHash('sha256') + .update(host.keyPair.publicKey) + .digest('base64url') + .slice(0, 16) + const inviteResponse = nextMessage(host.socket) + host.socket.send( + JSON.stringify({ type: 'invite-create', reqId: 'cooldown-invite', relayDeviceId: 'cooldown-device' }) + ) + const invite = await inviteResponse + const first = new WebSocket(`${relayUrl.replace('http:', 'ws:')}/v1/connect/${hostId}`, { + headers: forwardedHeaders() + }) + await new Promise((resolveOpen) => first.once('open', resolveOpen)) + const firstOpen = nextMessage(host.socket) + first.send( + JSON.stringify({ type: 'relay-auth', v: 1, mode: 'connect', credential: invite.inviteToken }) + ) + await firstOpen + first.close() + await new Promise((resolveWait) => setTimeout(resolveWait, 50)) + + const cooldown = new WebSocket(`${relayUrl.replace('http:', 'ws:')}/v1/connect/${hostId}`, { + headers: forwardedHeaders() + }) + await new Promise((resolveOpen) => cooldown.once('open', resolveOpen)) + const cooldownHello = nextMessage(cooldown) + cooldown.send( + JSON.stringify({ type: 'relay-auth', v: 1, mode: 'connect', credential: invite.inviteToken }) + ) + expect(await cooldownHello).toEqual({ type: 'relay-hello', ok: false, code: 4401 }) + + await new Promise((resolveWait) => setTimeout(resolveWait, 2_050)) + const retry = new WebSocket(`${relayUrl.replace('http:', 'ws:')}/v1/connect/${hostId}`, { + headers: forwardedHeaders() + }) + await new Promise((resolveOpen) => retry.once('open', resolveOpen)) + const retriedOpen = nextMessage(host.socket) + retry.send( + JSON.stringify({ type: 'relay-auth', v: 1, mode: 'connect', credential: invite.inviteToken }) + ) + expect(await retriedOpen).toMatchObject({ type: 'conn-open', kind: 'invite' }) + retry.close() + host.socket.close() + }) + + it('keeps install status and every effect not-found after injected SQL failure', async () => { + const originalUrl = relayUrl + const faultPort = await unusedPort() + const faultUrl = `http://127.0.0.1:${faultPort}` + const faultData = mkdtempSync(resolve(tmpdir(), 'orca-relay-fault-')) + const faultProcess = spawn( + process.execPath, + ['--import', 'tsx', 'src/fault-injection-test-entry.ts'], + { + cwd: appDirectory, + env: { + ...process.env, + NODE_ENV: 'test', + PORT: String(faultPort), + ORCA_RELAY_PUBLIC_URL: faultUrl, + ORCA_RELAY_CELL_URL: faultUrl, + ORCA_RELAY_AUTH_ISSUER: issuer, + ORCA_RELAY_JWKS_URL: `${issuer}/jwks`, + ORCA_RELAY_ASSIGNMENT_SIGNING_KEY: assignmentKey, + ORCA_RELAY_DATA_DIR: faultData, + ORCA_RELAY_ADMIN_AUDIENCE: adminAudience, + ORCA_RELAY_DEPLOY_SERVICE_ACCOUNT: 'deploy@example.com', + ORCA_RELAY_ADMIN_JWKS_URL: `${issuer}/jwks`, + ORCA_RELAY_TEST_FAULT_SQL: 'INSERT INTO relay_install_results' + }, + stdio: ['ignore', 'pipe', 'pipe'] + } + ) + relayUrl = faultUrl + try { + await waitForRelay(faultProcess) + const host = await openHostControl() + const hostId = createHash('sha256') + .update(host.keyPair.publicKey) + .digest('base64url') + .slice(0, 16) + const inviteResponse = nextMessage(host.socket) + host.socket.send( + JSON.stringify({ type: 'invite-create', reqId: 'fault-invite', relayDeviceId: 'fault-device' }) + ) + const invite = await inviteResponse + const splice = await attachPhone({ + host: host.socket, + hostAck: host.ack, + hostId, + credential: String(invite.inviteToken) + }) + const install = { + type: 'device-credential-install', + v: 1, + reqId: 'fault-install', + relayDeviceId: 'fault-device', + newResumeTokenHash: createHash('sha256').update('fault-resume').digest('base64url'), + authorization: { mode: 'relay-basis', basisConnId: splice.connId } + } + const failed = nextMessage(host.socket) + host.socket.send(JSON.stringify(install)) + expect(await failed).toMatchObject({ + type: 'control-error', + reqId: 'fault-install', + code: 'injected SQL failure' + }) + const status = nextMessage(host.socket) + host.socket.send( + JSON.stringify({ + type: 'device-credential-install-status', + v: 1, + reqId: 'fault-install', + relayDeviceId: 'fault-device' + }) + ) + expect(await status).toMatchObject({ + type: 'device-credential-install-status-result', + state: 'not-found' + }) + const retried = nextMessage(host.socket) + host.socket.send(JSON.stringify(install)) + expect(await retried).toMatchObject({ + type: 'device-credential-installed', + currentVersion: 1 + }) + splice.phone.close() + splice.data.close() + host.socket.close() + } finally { + faultProcess.kill('SIGKILL') + await new Promise((resolveExit) => faultProcess.once('exit', () => resolveExit())) + relayUrl = originalUrl + rmSync(faultData, { recursive: true, force: true }) + } + }) + + it('falls back to an invite only after a failed direct attempt is authoritatively not-found', async () => { + const originalUrl = relayUrl + const faultPort = await unusedPort() + const faultUrl = `http://127.0.0.1:${faultPort}` + const faultData = mkdtempSync(resolve(tmpdir(), 'orca-relay-direct-fault-')) + const faultProcess = spawn( + process.execPath, + ['--import', 'tsx', 'src/fault-injection-test-entry.ts'], + { + cwd: appDirectory, + env: { + ...process.env, + NODE_ENV: 'test', + PORT: String(faultPort), + ORCA_RELAY_PUBLIC_URL: faultUrl, + ORCA_RELAY_CELL_URL: faultUrl, + ORCA_RELAY_AUTH_ISSUER: issuer, + ORCA_RELAY_JWKS_URL: `${issuer}/jwks`, + ORCA_RELAY_ASSIGNMENT_SIGNING_KEY: assignmentKey, + ORCA_RELAY_DATA_DIR: faultData, + ORCA_RELAY_ADMIN_AUDIENCE: adminAudience, + ORCA_RELAY_DEPLOY_SERVICE_ACCOUNT: 'deploy@example.com', + ORCA_RELAY_ADMIN_JWKS_URL: `${issuer}/jwks`, + ORCA_RELAY_TEST_FAULT_SQL: 'INSERT INTO relay_direct_authorizations' + }, + stdio: ['ignore', 'pipe', 'pipe'] + } + ) + relayUrl = faultUrl + try { + await waitForRelay(faultProcess) + const host = await openHostControl() + const hostId = createHash('sha256') + .update(host.keyPair.publicKey) + .digest('base64url') + .slice(0, 16) + const inviteResponse = nextMessage(host.socket) + host.socket.send( + JSON.stringify({ + type: 'invite-create', + reqId: 'fallback-invite', + relayDeviceId: 'fallback-device' + }) + ) + const invite = await inviteResponse + const splice = await attachPhone({ + host: host.socket, + hostAck: host.ack, + hostId, + credential: String(invite.inviteToken) + }) + const installBase = { + type: 'device-credential-install', + v: 1, + reqId: 'fallback-install', + relayDeviceId: 'fallback-device', + newResumeTokenHash: createHash('sha256').update('fallback-resume').digest('base64url') + } + const directFailure = nextMessage(host.socket) + host.socket.send( + JSON.stringify({ + ...installBase, + authorization: { mode: 'authenticated-direct', directAuthId: 'failed-direct' } + }) + ) + expect(await directFailure).toMatchObject({ + type: 'control-error', + reqId: 'fallback-install', + code: 'injected SQL failure' + }) + const status = nextMessage(host.socket) + host.socket.send( + JSON.stringify({ + type: 'device-credential-install-status', + v: 1, + reqId: 'fallback-install', + relayDeviceId: 'fallback-device' + }) + ) + expect(await status).toMatchObject({ + type: 'device-credential-install-status-result', + state: 'not-found' + }) + const fallback = nextMessage(host.socket) + host.socket.send( + JSON.stringify({ + ...installBase, + authorization: { mode: 'relay-basis', basisConnId: splice.connId } + }) + ) + expect(await fallback).toMatchObject({ + type: 'device-credential-installed', + authorizationMode: 'relay-basis', + currentVersion: 1 + }) + splice.phone.close() + splice.data.close() + host.socket.close() + } finally { + faultProcess.kill('SIGKILL') + await new Promise((resolveExit) => faultProcess.once('exit', () => resolveExit())) + relayUrl = originalUrl + rmSync(faultData, { recursive: true, force: true }) + } + }) + + it('keeps director HTTP routes off cells and enforces the durable cell epoch', async () => { + const originalUrl = relayUrl + const cellPort = await unusedPort() + const cellUrl = `http://127.0.0.1:${cellPort}` + const cellData = mkdtempSync(resolve(tmpdir(), 'orca-relay-cell-')) + const keyPair = nacl.box.keyPair() + const hostId = createHash('sha256') + .update(keyPair.publicKey) + .digest('base64url') + .slice(0, 16) + const database = await openRelayDatabase({ dataDir: cellData }) + const assignments = new RelayAssignmentStore(database, () => 100) + await assignments.reconcileCells([{ id: 'cell-a', url: cellUrl, capacityRequests: 10 }]) + await assignments.assign({ userId: 'user-1', relayHostId: hostId }) + await database.close() + const cellProcess = spawn(process.execPath, ['--import', 'tsx', 'src/index.ts'], { + cwd: appDirectory, + env: { + ...process.env, + PORT: String(cellPort), + ORCA_RELAY_PUBLIC_URL: cellUrl, + ORCA_RELAY_CELL_URL: cellUrl, + ORCA_RELAY_CELL_ID: 'cell-a', + ORCA_RELAY_CELL_CAPACITY: '10', + ORCA_RELAY_ROLE: 'cell', + ORCA_RELAY_AUTH_ISSUER: issuer, + ORCA_RELAY_JWKS_URL: `${issuer}/jwks`, + ORCA_RELAY_ASSIGNMENT_SIGNING_KEY: assignmentKey, + ORCA_RELAY_DATA_DIR: cellData, + ORCA_RELAY_ADMIN_AUDIENCE: adminAudience, + ORCA_RELAY_DEPLOY_SERVICE_ACCOUNT: 'deploy@example.com', + ORCA_RELAY_ADMIN_JWKS_URL: `${issuer}/jwks` + }, + stdio: ['ignore', 'pipe', 'pipe'] + }) + relayUrl = cellUrl + try { + await waitForRelay(cellProcess) + const token = await relayToken('orca-relay', hostId) + const assignmentResponse = await fetch(`${cellUrl}/v1/assign`, { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: JSON.stringify({ v: 1, relayHostId: hostId }) + }) + expect(assignmentResponse.status).toBe(404) + const cellStatusResponse = await fetch(`${cellUrl}/v1/admin/cell-status`, { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: JSON.stringify({ v: 1, cellId: 'cell-a' }) + }) + expect(cellStatusResponse.status).toBe(404) + + const stale = new WebSocket(`${cellUrl.replace('http:', 'ws:')}/v1/host/control`, { + headers: { authorization: `Bearer ${token}` } + }) + await new Promise((resolveOpen) => stale.once('open', resolveOpen)) + const staleClosed = new Promise((resolveClose) => + stale.once('close', (code) => resolveClose(code)) + ) + stale.send( + JSON.stringify({ + type: 'host-hello', + v: 1, + relayHostId: hostId, + assignmentEpoch: 2, + hostPublicKeyB64: Buffer.from(keyPair.publicKey).toString('base64'), + appVersion: 'test' + }) + ) + expect(await staleClosed).toBe(4409) + + const assigned = await openHostControl({ keyPair }) + expect(assigned.ack).toMatchObject({ type: 'host-hello-ack', generation: 1 }) + const invitePromise = nextMessage(assigned.socket) + assigned.socket.send( + JSON.stringify({ type: 'invite-create', reqId: 'cell-invite', relayDeviceId: 'cell-device' }) + ) + const invite = await invitePromise + const splice = await attachPhone({ + host: assigned.socket, + hostAck: assigned.ack, + hostId, + credential: String(invite.inviteToken) + }) + const observedDatabase = await openRelayDatabase({ dataDir: cellData }) + const activities = await observedDatabase.query( + `SELECT activity_kind, request_units FROM relay_assignment_activity_leases + ORDER BY activity_kind, activity_id` + ) + await observedDatabase.close() + expect(activities).toEqual([ + { activity_kind: 'control', request_units: 1 }, + { activity_kind: 'invite', request_units: 1 }, + { activity_kind: 'invite', request_units: 1 }, + { activity_kind: 'splice', request_units: 2 } + ]) + splice.phone.close() + splice.data.close() + assigned.socket.close() + } finally { + cellProcess.kill('SIGTERM') + await new Promise((resolveExit) => cellProcess.once('exit', () => resolveExit())) + relayUrl = originalUrl + rmSync(cellData, { recursive: true, force: true }) + } + }) + + it('runs target-first dual-cell evacuation before releasing the source control', async () => { + const originalUrl = relayUrl + const dataDirectory = mkdtempSync(resolve(tmpdir(), 'orca-relay-topology-')) + const directorUrl = `http://127.0.0.1:${await unusedPort()}` + const cellAUrl = `http://127.0.0.1:${await unusedPort()}` + const cellBUrl = `http://127.0.0.1:${await unusedPort()}` + const cells = [ + { id: 'cell-a', url: cellAUrl, capacityRequests: 10 }, + { id: 'cell-b', url: cellBUrl, capacityRequests: 10 } + ] + const keyPair = nacl.box.keyPair() + const hostId = createHash('sha256') + .update(keyPair.publicKey) + .digest('base64url') + .slice(0, 16) + const seedDatabase = await openRelayDatabase({ dataDir: dataDirectory }) + const seedAssignments = new RelayAssignmentStore(seedDatabase) + await seedAssignments.reconcileCells(cells) + await seedAssignments.assign({ userId: 'user-1', relayHostId: hostId }) + await seedDatabase.close() + + const cellA = spawnTopologyRelay({ + url: cellAUrl, + dataDirectory, + role: 'cell', + cellId: 'cell-a' + }) + let director: ChildProcess | undefined + let cellB: ChildProcess | undefined + try { + await waitForRelay(cellA) + director = spawnTopologyRelay({ + url: directorUrl, + dataDirectory, + role: 'director', + cellId: 'director', + cells + }) + await waitForRelay(director) + expect( + await fetch(`${directorUrl}/v1/admin/cell-heartbeat`, { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: '{}' + }) + ).toMatchObject({ status: 401 }) + const heartbeatAudience = `${directorUrl}/v1/admin/cell-heartbeat` + expect( + await fetch(heartbeatAudience, { + method: 'POST', + headers: { + authorization: `Bearer ${await googleServiceToken(`${directorUrl}/wrong`)}`, + 'content-type': 'application/json' + }, + body: '{}' + }) + ).toMatchObject({ status: 401 }) + expect( + await fetch(heartbeatAudience, { + method: 'POST', + headers: { + authorization: `Bearer ${await googleServiceToken(heartbeatAudience, 'wrong@example.com')}`, + 'content-type': 'application/json' + }, + body: '{}' + }) + ).toMatchObject({ status: 401 }) + expect(await postCellHeartbeat(directorUrl, cells[0]!, { startedAt: 1_000 })).toMatchObject({ + status: 200 + }) + expect( + await postCellHeartbeat(directorUrl, cells[0]!, { + incarnation: '33333333-3333-4333-8333-333333333333', + startedAt: 999 + }) + ).toMatchObject({ status: 409 }) + expect( + await postCellHeartbeat(directorUrl, cells[1]!, { + incarnation: '22222222-2222-4222-8222-222222222222', + startedAt: 2_000 + }) + ).toMatchObject({ status: 200 }) + relayUrl = cellAUrl + const sourceHost = await openHostControl({ keyPair, assignmentEpoch: 1 }) + const token = await adminToken() + const recoveryRequests = [ + { + path: 'cell-fence-attempt-prepare', + body: { + v: 1, + attemptId: '44444444-4444-4444-8444-444444444444', + environment: 'production', + cellId: 'cell-a', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + migName: 'orca-relay-c1', + instanceGroup: 'https://compute.example/instanceGroups/orca-relay-c1', + generationIdentity: + 'https://compute.example/instanceTemplates/orca-relay-c1-abc', + fenceCommit: 'a'.repeat(40), + planSha256: 'b'.repeat(64), + planObjectName: + 'terraform/state/relay-fence-plans/production/44444444-4444-4444-8444-444444444444.tfplan', + varFileSha256: 'c'.repeat(64), + terraformStateLineage: 'c739dab4-e6e1-e627-02a9-504b3dda1a2c', + terraformStateSerial: 7, + terraformStateObjectGeneration: '987654321', + terraformStateObjectSha256: 'd'.repeat(64), + requestReason: + 'orca-relay-fence/44444444-4444-4444-8444-444444444444' + }, + confirmation: 'PREPARE_TERRAFORM_CELL_FENCE', + expected: { status: 409, body: { error: 'cell_fence_admission_enabled' } } + }, + { + path: 'cell-fence-attempt-abort', + body: { + v: 1, + attemptId: '44444444-4444-4444-8444-444444444444', + environment: 'production', + cellId: 'cell-a', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + migName: 'orca-relay-c1', + instanceGroup: 'https://compute.example/instanceGroups/orca-relay-c1', + generationIdentity: + 'https://compute.example/instanceTemplates/orca-relay-c1-abc', + fenceCommit: 'a'.repeat(40), + planSha256: 'b'.repeat(64), + planObjectName: + 'terraform/state/relay-fence-plans/production/44444444-4444-4444-8444-444444444444.tfplan', + varFileSha256: 'c'.repeat(64), + terraformStateLineage: 'c739dab4-e6e1-e627-02a9-504b3dda1a2c', + terraformStateSerial: 7, + terraformStateObjectGeneration: '987654321', + terraformStateObjectSha256: 'd'.repeat(64), + requestReason: + 'orca-relay-fence/44444444-4444-4444-8444-444444444444' + }, + confirmation: 'ABORT_UNSTARTED_TERRAFORM_CELL_FENCE', + expected: { status: 409, body: { error: 'cell_fence_attempt_not_found' } } + }, + { + path: 'migration-supersede-cell', + body: { + v: 1, + sourceCellId: 'cell-a', + currentTargetCellId: 'cell-b', + replacementTargetCellId: 'cell-c', + limit: 100 + }, + confirmation: 'SUPERSEDE_REGISTERED_CELL_MIGRATIONS', + expected: { status: 200, body: { v: 1, superseded: 0 } } + }, + { + path: 'drain-attempt-prepare', + body: { + v: 1, + attemptId: '55555555-5555-4555-8555-555555555555', + cellId: 'cell-a', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + traceValue: '66666666-6666-4666-8666-666666666666', + graceMs: 120_000 + }, + confirmation: 'PREPARE_LEGACY_DRAIN', + expected: { status: 409, body: { error: 'drain_attempt_admission_enabled' } } + }, + { + path: 'drain-attempt-recover-forward', + body: { + v: 1, + cellId: 'cell-a', + cellIncarnation: '11111111-1111-4111-8111-111111111111' + }, + confirmation: 'RECOVER_LEGACY_DRAIN', + expected: { status: 409, body: { error: 'drain_attempt_admission_enabled' } } + } + ] + for (const request of recoveryRequests) { + const url = `${directorUrl}/v1/admin/${request.path}` + expect( + await fetch(url, { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: JSON.stringify(request.body) + }) + ).toMatchObject({ status: 400 }) + const confirmed = await fetch(url, { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: JSON.stringify({ ...request.body, confirmation: request.confirmation }) + }) + expect(confirmed.status).toBe(request.expected.status) + expect(await confirmed.json()).toEqual(request.expected.body) + } + const statusUrl = `${directorUrl}/v1/admin/cell-status` + expect( + await fetch(statusUrl, { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ v: 1, cellId: 'cell-a' }) + }) + ).toMatchObject({ status: 401 }) + expect( + await fetch(statusUrl, { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: JSON.stringify({ v: 1, cellId: 'cell-a', userId: 'must-not-be-accepted' }) + }) + ).toMatchObject({ status: 400 }) + const sourceStatusResponse = await fetch(statusUrl, { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: JSON.stringify({ v: 1, cellId: 'cell-a' }) + }) + expect(sourceStatusResponse.status).toBe(200) + const sourceStatusText = await sourceStatusResponse.text() + expect(sourceStatusText).not.toContain('user-1') + expect(sourceStatusText).not.toContain(hostId) + expect(JSON.parse(sourceStatusText)).toMatchObject({ + v: 1, + status: { + cellId: 'cell-a', + cellUrl: cellAUrl, + enabled: true, + assignments: 1, + activityLeases: 1, + activityRequestUnits: 1, + restartBlockingActivityLeases: 1, + restartBlockingActivityRequestUnits: 1, + restartBlockingReservedRequests: 1, + runtime: { cellUrl: cellAUrl, ready: true, heartbeatFresh: true } + } + }) + const cellState = async (cellId: string, enabled: boolean): Promise => + await fetch(`${directorUrl}/v1/admin/cell-state`, { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: JSON.stringify({ v: 1, cellId, enabled }) + }) + const configuredTarget = await fetch(`${directorUrl}/v1/admin/cell-config`, { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: JSON.stringify({ + v: 1, + cellId: 'cell-b', + cellUrl: cellBUrl, + capacityRequests: 10, + state: 'existing-only' + }) + }) + expect(configuredTarget.status).toBe(200) + expect(await configuredTarget.json()).toEqual({ ok: true }) + const incompleteLimit = await fetch(`${directorUrl}/v1/admin/cell-config`, { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: JSON.stringify({ + v: 1, + cellId: 'cell-b', + cellUrl: cellBUrl, + capacityRequests: 10, + connectionHardCap: 600, + enabled: false + }) + }) + expect(incompleteLimit.status).toBe(400) + const capacity = await fetch(`${directorUrl}/v1/admin/evacuation-capacity`, { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: JSON.stringify({ v: 1, sourceCellId: 'cell-a', targetCellId: 'cell-b' }) + }) + expect(capacity.status).toBe(200) + expect(await capacity.json()).toEqual({ + v: 1, + sourceAssignments: 1, + requiredTargetUnits: 2, + availableTargetUnits: 10 + }) + const disabledTarget = await fetch(`${directorUrl}/v1/admin/evacuate`, { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: JSON.stringify({ + v: 1, + userId: 'user-1', + relayHostId: hostId, + targetCellId: 'cell-b' + }) + }) + expect(disabledTarget.status).toBe(409) + expect(await disabledTarget.json()).toEqual({ error: 'target_cell_unavailable' }) + expect((await cellState('cell-b', true)).status).toBe(200) + const migrationResponse = await fetch(`${directorUrl}/v1/admin/evacuate-cell`, { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: JSON.stringify({ + v: 1, + sourceCellId: 'cell-a', + targetCellId: 'cell-b', + limit: 10 + }) + }) + expect(migrationResponse.status).toBe(200) + expect(await migrationResponse.json()).toEqual({ v: 1, started: 1 }) + const pendingStatus = await fetch(`${directorUrl}/v1/admin/evacuation-status`, { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: JSON.stringify({ + v: 1, + sourceCellId: 'cell-a', + targetCellId: 'cell-b' + }) + }) + expect(pendingStatus.status).toBe(200) + expect(await pendingStatus.json()).toMatchObject({ + v: 1, + inProgress: 1, + targetRegistered: 0, + registeredSourceActive: 0, + registeredCompletable: 0, + registeredTargetInactive: 0, + completed: 0, + blocked: 0, + expiredUnregistered: 0, + repairableExpiredUnregistered: 0, + abortableExpiredUnregistered: 0, + blockedExpiredUnregistered: 0, + blockedExpiredOnNewerTargetAssignment: 0 + }) + // Completion must prove the source has stopped admitting new work, even in + // the served admin workflow that performs its own topology preflight. + expect((await cellState('cell-a', false)).status).toBe(200) + + cellB = spawnTopologyRelay({ + url: cellBUrl, + dataDirectory, + role: 'cell', + cellId: 'cell-b' + }) + await waitForRelay(cellB) + relayUrl = cellBUrl + const targetHost = await openHostControl({ keyPair, assignmentEpoch: 2 }) + const completionBody = JSON.stringify({ + v: 1, + sourceCellId: 'cell-a', + targetCellId: 'cell-b', + completeReady: true + }) + const premature = await fetch(`${directorUrl}/v1/admin/evacuation-status`, { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: completionBody + }) + expect(premature.status).toBe(200) + expect(await premature.json()).toMatchObject({ + v: 1, + inProgress: 1, + targetRegistered: 1, + registeredSourceActive: 1, + registeredCompletable: 0, + registeredTargetInactive: 0, + completed: 0, + blocked: 1, + expiredUnregistered: 0, + repairableExpiredUnregistered: 0, + abortableExpiredUnregistered: 0, + blockedExpiredUnregistered: 0, + blockedExpiredOnNewerTargetAssignment: 0 + }) + + const sourceClosed = new Promise((resolveClose) => + sourceHost.socket.once('close', (code) => resolveClose(code)) + ) + const drain = await fetch(`${cellAUrl}/v1/admin/drain`, { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: JSON.stringify({ v: 1, graceMs: 0 }) + }) + expect(drain.status).toBe(200) + expect(await sourceClosed).toBe(4503) + + let completed: Response | undefined + for (let attempt = 0; attempt < 20; attempt++) { + completed = await fetch(`${directorUrl}/v1/admin/evacuation-status`, { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: completionBody + }) + if ( + completed.status === 200 && + ((await completed.clone().json()) as { inProgress?: number }).inProgress === 0 + ) { + break + } + await new Promise((resolveWait) => setTimeout(resolveWait, 10)) + } + expect(completed?.status).toBe(200) + expect(await completed?.json()).toMatchObject({ + v: 1, + inProgress: 0, + targetRegistered: 0, + registeredSourceActive: 0, + registeredCompletable: 0, + registeredTargetInactive: 0, + completed: 1, + blocked: 0, + expiredUnregistered: 0, + repairableExpiredUnregistered: 0, + abortableExpiredUnregistered: 0, + blockedExpiredUnregistered: 0, + blockedExpiredOnNewerTargetAssignment: 0 + }) + const observedDatabase = await openRelayDatabase({ dataDir: dataDirectory }) + const assignment = await new RelayAssignmentStore(observedDatabase).resolve({ + userId: 'user-1', + relayHostId: hostId + }) + const reservations = await observedDatabase.query( + `SELECT cell_id, reserved_requests FROM relay_cells ORDER BY cell_id` + ) + await observedDatabase.close() + expect(assignment).toMatchObject({ cellId: 'cell-b', assignmentEpoch: 2 }) + expect(reservations).toEqual([ + { cell_id: 'cell-a', reserved_requests: 0 }, + { cell_id: 'cell-b', reserved_requests: 1 } + ]) + const selectorStatusBefore = await fetch( + `${directorUrl}/v1/admin/admission-selector/status`, + { + method: 'POST', + headers: { + authorization: `Bearer ${token}`, + 'content-type': 'application/json' + }, + body: JSON.stringify({ v: 1 }) + } + ) + expect(selectorStatusBefore.status).toBe(200) + const selectorBefore = (await selectorStatusBefore.json()) as { + selector: { membership: CellAdmissionMembership } + } + const selectorInput = { + v: 1, + attemptId: 'blackbox_cutover', + expectedGeneration: 0, + expectedMembershipSha256: createHash('sha256') + .update(encodeMembership(selectorBefore.selector.membership)) + .digest('hex'), + membership: { + existingOnly: ['cell-a'], + migrationOnly: [], + general: ['cell-b'] + } + } + const unsafeSelectorInput = { ...selectorInput, expectedMembershipSha256: undefined } + const unsafeSelectorResponse = await fetch( + `${directorUrl}/v1/admin/admission-selector/apply`, + { + method: 'POST', + headers: { + authorization: `Bearer ${token}`, + 'content-type': 'application/json' + }, + body: JSON.stringify(unsafeSelectorInput) + } + ) + expect(unsafeSelectorResponse.status).toBe(400) + const selectorResponse = await fetch( + `${directorUrl}/v1/admin/admission-selector/apply`, + { + method: 'POST', + headers: { + authorization: `Bearer ${token}`, + 'content-type': 'application/json' + }, + body: JSON.stringify(selectorInput) + } + ) + expect(selectorResponse.status).toBe(200) + expect(await selectorResponse.json()).toMatchObject({ + v: 1, + changed: true, + selector: { generation: 1, membership: selectorInput.membership } + }) + const selectorStatus = await fetch( + `${directorUrl}/v1/admin/admission-selector/status`, + { + method: 'POST', + headers: { + authorization: `Bearer ${token}`, + 'content-type': 'application/json' + }, + body: JSON.stringify({ v: 1, attemptId: selectorInput.attemptId }) + } + ) + expect(selectorStatus.status).toBe(200) + expect(await selectorStatus.json()).toMatchObject({ + v: 1, + selector: { generation: 1, membership: selectorInput.membership }, + intent: { attemptId: selectorInput.attemptId, state: 'committed' } + }) + const addCellsInput = { + v: 1, + attemptId: 'blackbox_add_cells', + expectedGeneration: 1, + cells: [ + { + cellId: 'cell-c', + cellUrl: 'https://relay-c.example.com', + capacityRequests: 20, + connectionHardCap: 1_000, + connectionUnobservedBound: 60 + } + ] + } + const addCellsResponse = await fetch( + `${directorUrl}/v1/admin/admission-selector/add-migration-cells`, + { + method: 'POST', + headers: { + authorization: `Bearer ${token}`, + 'content-type': 'application/json' + }, + body: JSON.stringify(addCellsInput) + } + ) + expect(addCellsResponse.status).toBe(200) + expect(await addCellsResponse.json()).toMatchObject({ + v: 1, + changed: true, + selector: { + generation: 2, + membership: { + existingOnly: ['cell-a'], + migrationOnly: ['cell-c'], + general: ['cell-b'] + } + } + }) + expect((await cellState('cell-a', true)).status).toBe(409) + targetHost.socket.close() + } finally { + for (const child of [cellA, cellB, director]) { + if (child && child.exitCode === null) child.kill('SIGKILL') + } + relayUrl = originalUrl + rmSync(dataDirectory, { recursive: true, force: true }) + } + }) + + it('recovers an unexpired invite through a restarted configured director without reservation', async () => { + const combinedUrl = relayUrl + const host = await openHostControl() + const hostId = createHash('sha256') + .update(host.keyPair.publicKey) + .digest('base64url') + .slice(0, 16) + const invitePromise = nextMessage(host.socket) + host.socket.send( + JSON.stringify({ type: 'invite-create', reqId: 'move-invite', relayDeviceId: 'move-device' }) + ) + const invite = await invitePromise + host.socket.close() + + relayProcess.kill('SIGKILL') + await new Promise((resolveExit) => relayProcess.once('exit', () => resolveExit())) + const directorPort = await unusedPort() + relayUrl = `http://127.0.0.1:${directorPort}` + relayProcess = spawn(process.execPath, ['--import', 'tsx', 'src/index.ts'], { + cwd: appDirectory, + env: { + ...process.env, + PORT: String(directorPort), + ORCA_RELAY_PUBLIC_URL: relayUrl, + ORCA_RELAY_CELL_URL: 'https://relay-c2.onorca.dev', + ORCA_RELAY_AUTH_ISSUER: issuer, + ORCA_RELAY_JWKS_URL: `${issuer}/jwks`, + ORCA_RELAY_ASSIGNMENT_SIGNING_KEY: assignmentKey, + ORCA_RELAY_DATA_DIR: relayDataDirectory, + ORCA_RELAY_ROLE: 'director', + ORCA_RELAY_CELLS_JSON: JSON.stringify([ + { + id: 'cell-c2', + url: 'https://relay-c2.onorca.dev', + capacityRequests: 900 + } + ]), + ORCA_RELAY_ADMIN_AUDIENCE: adminAudience, + ORCA_RELAY_DEPLOY_SERVICE_ACCOUNT: 'deploy@example.com', + ORCA_RELAY_ADMIN_JWKS_URL: `${issuer}/jwks` + }, + stdio: ['ignore', 'pipe', 'pipe'] + }) + await waitForRelay(relayProcess) + expect( + await postCellHeartbeat(relayUrl, { + id: 'cell-c2', + url: 'https://relay-c2.onorca.dev' + }) + ).toMatchObject({ status: 200 }) + + const phone = new WebSocket(`${relayUrl.replace('http:', 'ws:')}/v1/connect/${hostId}`, { + headers: forwardedHeaders() + }) + await new Promise((resolveOpen) => phone.once('open', resolveOpen)) + const movedPromise = nextMessage(phone) + const closed = new Promise((resolveClose) => + phone.once('close', (code) => resolveClose(code)) + ) + phone.send( + JSON.stringify({ + type: 'relay-auth', + v: 1, + mode: 'connect', + credential: invite.inviteToken + }) + ) + expect(await movedPromise).toEqual({ + type: 'relay-moved', + v: 1, + cellUrl: 'https://relay-c2.onorca.dev', + assignmentEpoch: 1 + }) + expect(await closed).toBe(4503) + + relayProcess.kill('SIGTERM') + await new Promise((resolveExit) => relayProcess.once('exit', () => resolveExit())) + relayUrl = combinedUrl + relayProcess = spawn(process.execPath, ['--import', 'tsx', 'src/index.ts'], { + cwd: appDirectory, + env: { + ...process.env, + PORT: new URL(combinedUrl).port, + ORCA_RELAY_PUBLIC_URL: combinedUrl, + ORCA_RELAY_CELL_URL: combinedUrl, + ORCA_RELAY_AUTH_ISSUER: issuer, + ORCA_RELAY_JWKS_URL: `${issuer}/jwks`, + ORCA_RELAY_ASSIGNMENT_SIGNING_KEY: assignmentKey, + ORCA_RELAY_DATA_DIR: relayDataDirectory, + ORCA_RELAY_ADMIN_AUDIENCE: adminAudience, + ORCA_RELAY_DEPLOY_SERVICE_ACCOUNT: 'deploy@example.com', + ORCA_RELAY_CAPACITY_SERVICE_ACCOUNT: 'capacity@example.com', + ORCA_RELAY_MONITOR_SERVICE_ACCOUNT: 'monitor@example.com', + ORCA_RELAY_FENCE_SERVICE_ACCOUNT: 'fence@example.com', + ORCA_RELAY_IMAGE_DIGEST: `sha256:${'a'.repeat(64)}`, + ORCA_RELAY_ADMIN_JWKS_URL: `${issuer}/jwks` + }, + stdio: ['ignore', 'pipe', 'pipe'] + }) + await waitForRelay(relayProcess) + }) + + it('uses configured-director recovery and typed 4503 on unplanned SIGTERM', async () => { + const originalUrl = relayUrl + const signalPort = await unusedPort() + const signalUrl = `http://127.0.0.1:${signalPort}` + const signalData = mkdtempSync(resolve(tmpdir(), 'orca-relay-sigterm-')) + const signalProcess = spawn(process.execPath, ['--import', 'tsx', 'src/index.ts'], { + cwd: appDirectory, + env: { + ...process.env, + PORT: String(signalPort), + ORCA_RELAY_PUBLIC_URL: signalUrl, + ORCA_RELAY_CELL_URL: signalUrl, + ORCA_RELAY_AUTH_ISSUER: issuer, + ORCA_RELAY_JWKS_URL: `${issuer}/jwks`, + ORCA_RELAY_ASSIGNMENT_SIGNING_KEY: assignmentKey, + ORCA_RELAY_DATA_DIR: signalData, + ORCA_RELAY_ADMIN_AUDIENCE: adminAudience, + ORCA_RELAY_DEPLOY_SERVICE_ACCOUNT: 'deploy@example.com', + ORCA_RELAY_ADMIN_JWKS_URL: `${issuer}/jwks` + }, + stdio: ['ignore', 'pipe', 'pipe'] + }) + relayUrl = signalUrl + try { + await waitForRelay(signalProcess) + const host = await openHostControl() + const drain = nextMessage(host.socket) + const closed = new Promise((resolveClose) => + host.socket.once('close', (code) => resolveClose(code)) + ) + const exited = new Promise((resolveExit) => + signalProcess.once('exit', () => resolveExit()) + ) + signalProcess.kill('SIGTERM') + expect(await drain).toEqual({ + type: 'drain', + graceMs: 0, + recovery: 'resolve-director' + }) + expect(await closed).toBe(4503) + await exited + } finally { + if (signalProcess.exitCode === null) signalProcess.kill('SIGKILL') + relayUrl = originalUrl + rmSync(signalData, { recursive: true, force: true }) + } + }) + + it('authenticates admin drain with exact Google audience and deploy identity', async () => { + const host = await openHostControl() + const drainMessage = nextMessage(host.socket) + const closed = new Promise((resolveClose) => + host.socket.once('close', (code) => resolveClose(code)) + ) + const token = await new SignJWT({ email: 'deploy@example.com', email_verified: true }) + .setProtectedHeader({ alg: 'RS256', kid: 'admin-key' }) + .setIssuer('https://accounts.google.com') + .setAudience(adminAudience) + .setSubject('deploy-subject') + .setIssuedAt() + .setExpirationTime('5m') + .sign(adminPrivateKey) + const response = await fetch(`${relayUrl}/v1/admin/drain`, { + method: 'POST', + headers: { + authorization: `Bearer ${token}`, + 'content-type': 'application/json' + }, + body: JSON.stringify({ v: 1, graceMs: 0 }) + }) + expect(response.status).toBe(200) + expect(await drainMessage).toEqual({ + type: 'drain', + graceMs: 0, + recovery: 'resolve-director' + }) + expect(await closed).toBe(4503) + }) + + it('keeps dedicated capacity, monitor, and fence identities on exact routes', async () => { + const capacity = await googleServiceToken(adminAudience, 'capacity@example.com') + const monitor = await googleServiceToken(adminAudience, 'monitor@example.com') + const fence = await googleServiceToken(adminAudience, 'fence@example.com') + const broker = await googleServiceToken(adminAudience, 'broker@example.com') + const post = async (path: string, token: string, body: unknown): Promise => + await fetch(`${relayUrl}${path}`, { + method: 'POST', + headers: { + authorization: `Bearer ${token}`, + 'content-type': 'application/json' + }, + body: JSON.stringify(body) + }) + + expect((await post('/v1/admin/runtime-status', capacity, { v: 1 })).status).toBe(200) + expect( + (await post('/v1/admin/evacuation-status', capacity, { + v: 1, + sourceCellId: 'source', + targetCellId: 'target', + completeReady: false + })).status + ).toBe(401) + expect( + (await post('/v1/admin/cell-fence-attempt-status', capacity, { + v: 1, + cellId: 'production-gce-c1' + })).status + ).toBe(401) + + expect((await post('/v1/admin/runtime-status', monitor, { v: 1 })).status).toBe(200) + expect((await post('/v1/admin/drain', monitor, { v: 1, graceMs: 0 })).status).toBe(401) + expect( + (await post('/v1/admin/cell-fence-attempt-status', monitor, { + v: 1, + cellId: 'production-gce-c1' + })).status + ).toBe(401) + + expect((await post('/v1/admin/runtime-status', fence, { v: 1 })).status).toBe(200) + expect((await post('/v1/admin/drain', fence, { v: 1, graceMs: 0 })).status).toBe(401) + expect( + (await post('/v1/admin/cell-fence-attempt-status', fence, { + v: 1, + cellId: 'production-gce-c1' + })).status + ).toBe(404) + expect( + (await post('/v1/admin/cell-fence-attest', fence, { + v: 1 + })).status + ).toBe(401) + + const directorPort = await unusedPort() + const directorUrl = `http://127.0.0.1:${directorPort}` + const directorData = mkdtempSync(resolve(tmpdir(), 'orca-relay-monitor-auth-')) + const director = spawnTopologyRelay({ + url: directorUrl, + dataDirectory: directorData, + role: 'director', + cellId: 'director', + cells: [{ id: 'cell-a', url: 'https://cell-a.example.com', capacityRequests: 10 }] + }) + try { + await waitForRelay(director) + const adoption = { + v: 1, + cellId: 'cell-a', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + confirmation: 'ADOPT_LEGACY_TERRAFORM_CELL_FENCE' + } + const postDirector = async (token: string): Promise => + await fetch(`${directorUrl}/v1/admin/cell-fence-adopt-legacy`, { + method: 'POST', + headers: { + authorization: `Bearer ${token}`, + 'content-type': 'application/json' + }, + body: JSON.stringify(adoption) + }) + expect( + (await postDirector(fence)).status + ).toBe(401) + expect( + (await postDirector(broker)).status + ).toBe(409) + const commitAdoption = { + v: 1, + cellId: 'cell-a', + cellIncarnation: '11111111-1111-4111-8111-111111111111', + confirmation: 'COMMIT_LEGACY_TERRAFORM_CELL_FENCE_ADOPTION' + } + const postCommit = async (token: string): Promise => + await fetch(`${directorUrl}/v1/admin/cell-fence-commit-legacy-adoption`, { + method: 'POST', + headers: { + authorization: `Bearer ${token}`, + 'content-type': 'application/json' + }, + body: JSON.stringify(commitAdoption) + }) + expect((await postCommit(fence)).status).toBe(401) + expect((await postCommit(broker)).status).toBe(409) + const response = await fetch(`${directorUrl}/v1/admin/evacuation-status`, { + method: 'POST', + headers: { + authorization: `Bearer ${monitor}`, + 'content-type': 'application/json' + }, + body: JSON.stringify({ + v: 1, + sourceCellId: 'cell-a', + targetCellId: 'cell-b', + completeReady: true + }) + }) + expect(response.status).toBe(403) + const fenceResponse = await fetch(`${directorUrl}/v1/admin/evacuation-status`, { + method: 'POST', + headers: { + authorization: `Bearer ${fence}`, + 'content-type': 'application/json' + }, + body: JSON.stringify({ + v: 1, + sourceCellId: 'cell-a', + targetCellId: 'cell-b', + completeReady: true + }) + }) + expect(fenceResponse.status).toBe(403) + } finally { + if (director.exitCode === null) director.kill('SIGKILL') + rmSync(directorData, { recursive: true, force: true }) + } + }) + + it('reports only authenticated aggregate runtime identity and the served digest', async () => { + const token = await adminToken() + expect( + await fetch(`${relayUrl}/v1/admin/runtime-status`, { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ v: 1 }) + }) + ).toMatchObject({ status: 401 }) + expect( + await fetch(`${relayUrl}/v1/admin/runtime-status`, { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: JSON.stringify({ v: 1, hostId: 'must-not-be-accepted' }) + }) + ).toMatchObject({ status: 400 }) + const response = await fetch(`${relayUrl}/v1/admin/runtime-status`, { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: JSON.stringify({ v: 1 }) + }) + expect(response.status).toBe(200) + expect(await response.json()).toEqual({ + v: 1, + role: 'combined', + cellId: 'combined', + cellUrl: relayUrl, + region: 'us-central1', + imageDigest: `sha256:${'a'.repeat(64)}`, + draining: true, + regionalRehomeProtocol: 0, + connectionCapacity: null, + runtime: { + totalConnections: 0, + preAuthConnections: 0, + controls: 0, + splices: 0, + pendingSplices: 0, + queuedBytes: 0 + } + }) + }) +}) diff --git a/cloud/apps/relay/src/splice-forwarder.test.ts b/cloud/apps/relay/src/splice-forwarder.test.ts new file mode 100644 index 00000000000..3aa5e75fc97 --- /dev/null +++ b/cloud/apps/relay/src/splice-forwarder.test.ts @@ -0,0 +1,135 @@ +import { EventEmitter } from 'node:events' +import { describe, expect, it, vi } from 'vitest' +import type WebSocket from 'ws' +import { + ProcessQueuedByteBudget, + wireSplice, + type SpliceCloseInfo +} from './splice-forwarder.js' + +class FakeSocket extends EventEmitter { + readonly OPEN = 1 + readyState = 1 + bufferedAmount = 0 + sent: Array<{ data: unknown; binary: boolean }> = [] + closes: Array<{ code: number; reason: string }> = [] + _socket = { pause: vi.fn(), resume: vi.fn(), setNoDelay: vi.fn() } + + send(data: unknown, options: { binary: boolean }): void { + this.sent.push({ data, binary: options.binary }) + } + + close(code: number, reason: string): void { + this.readyState = 3 + this.closes.push({ code, reason }) + this.emit('close', code, Buffer.from(reason)) + } +} + +describe('splice forwarder', () => { + it('preserves text/binary opcodes and propagates peer close', () => { + const client = new FakeSocket() + const host = new FakeSocket() + const onClose = vi.fn() + const onForwardedBytes = vi.fn() + wireSplice({ + client: client as unknown as WebSocket, + host: host as unknown as WebSocket, + budget: new ProcessQueuedByteBudget(), + onClose, + onForwardedBytes + }) + client.emit('message', Buffer.from('text'), false) + client.emit('message', Buffer.from([1, 2]), true) + expect(host.sent).toEqual([ + { data: Buffer.from('text'), binary: false }, + { data: Buffer.from([1, 2]), binary: true } + ]) + expect(onForwardedBytes.mock.calls).toEqual([[4], [2]]) + client.emit('close', 1000, Buffer.alloc(0)) + expect(host.closes[0]?.code).toBe(4408) + expect(onClose).toHaveBeenCalledOnce() + }) + + it('hard-closes a wedged splice and releases its global queued-byte reservation', () => { + const client = new FakeSocket() + const host = new FakeSocket() + host.bufferedAmount = 1024 * 1024 + const budget = new ProcessQueuedByteBudget() + wireSplice({ + client: client as unknown as WebSocket, + host: host as unknown as WebSocket, + budget, + onClose: vi.fn() + }) + client.emit('message', Buffer.alloc(8 * 1024 * 1024), true) + expect(budget.current()).toBe(8 * 1024 * 1024) + client.emit('message', Buffer.alloc(512 * 1024), true) + expect(client.closes[0]?.code).toBe(4429) + expect(host.closes[0]?.code).toBe(4429) + expect(budget.current()).toBe(0) + }) + + it('reports the close trigger so limit and oversize kills are attributable', () => { + const wire = (): { client: FakeSocket; host: FakeSocket; closes: SpliceCloseInfo[] } => { + const client = new FakeSocket() + const host = new FakeSocket() + const closes: SpliceCloseInfo[] = [] + wireSplice({ + client: client as unknown as WebSocket, + host: host as unknown as WebSocket, + budget: new ProcessQueuedByteBudget(), + onClose: vi.fn(), + onClosed: (closeInfo) => closes.push(closeInfo) + }) + return { client, host, closes } + } + + const queueLimit = wire() + queueLimit.host.bufferedAmount = 1024 * 1024 + queueLimit.client.emit('message', Buffer.alloc(8 * 1024 * 1024), true) + queueLimit.client.emit('message', Buffer.alloc(512 * 1024), true) + expect(queueLimit.closes).toEqual([ + { code: 4429, reason: 'relay queue limit exceeded', trigger: 'queue-limit' } + ]) + + // The ws receiver error for a frame above maxPayload names the payload cap. + const oversize = wire() + oversize.host.emit('error', new RangeError('Max payload size exceeded')) + expect(oversize.closes[0]?.trigger).toBe('host-oversize-frame') + + const peerClose = wire() + peerClose.client.emit('close', 1001, Buffer.alloc(0)) + expect(peerClose.closes[0]?.trigger).toBe('client-closed') + expect(peerClose.closes).toHaveLength(1) + }) + + it('queues one full catalog-sized frame for a backpressured peer without closing', async () => { + // Why: the desktop's worktree catalog response exceeds 1MiB on large + // workspaces; a single maxFrameBytes frame must survive backpressure. + const client = new FakeSocket() + const host = new FakeSocket() + host.bufferedAmount = 1024 * 1024 + const budget = new ProcessQueuedByteBudget() + const onClose = vi.fn() + wireSplice({ + client: client as unknown as WebSocket, + host: host as unknown as WebSocket, + budget, + onClose + }) + client.emit('message', Buffer.alloc(8 * 1024 * 1024), true) + expect(client.closes).toEqual([]) + expect(host.closes).toEqual([]) + expect(budget.current()).toBe(8 * 1024 * 1024) + + // Peer drains; the queued frame flushes and the reservation releases. + host.bufferedAmount = 0 + await new Promise((resolve) => setTimeout(resolve, 60)) + expect(host.sent.some((frame) => (frame.data as Buffer).byteLength === 8 * 1024 * 1024)).toBe( + true + ) + expect(budget.current()).toBe(0) + expect(onClose).not.toHaveBeenCalled() + }) +}) diff --git a/cloud/apps/relay/src/splice-forwarder.ts b/cloud/apps/relay/src/splice-forwarder.ts new file mode 100644 index 00000000000..b0a31b49830 --- /dev/null +++ b/cloud/apps/relay/src/splice-forwarder.ts @@ -0,0 +1,158 @@ +import { RELAY_ADMISSION_BUDGETS, RELAY_CLOSE_CODE } from '@orca-cloud/relay-contract' +import type WebSocket from 'ws' +import type { RawData } from 'ws' +import { closeRelayWebSocket } from './relay-websocket-close.js' + +type QueuedFrame = { data: RawData; binary: boolean; bytes: number } + +export class ProcessQueuedByteBudget { + private queued = 0 + + reserve(bytes: number): boolean { + if (this.queued + bytes > RELAY_ADMISSION_BUDGETS.maxProcessQueuedBytes) return false + this.queued += bytes + return true + } + + release(bytes: number): void { + this.queued = Math.max(0, this.queued - bytes) + } + + current(): number { + return this.queued + } +} + +function frameBytes(data: RawData): number { + if (typeof data === 'string') return Buffer.byteLength(data) + if (Array.isArray(data)) return data.reduce((total, part) => total + part.byteLength, 0) + return data.byteLength +} + +function transport(socket: WebSocket): { pause(): void; resume(): void; setNoDelay(value: boolean): void } | null { + return ( + socket as WebSocket & { + _socket?: { pause(): void; resume(): void; setNoDelay(value: boolean): void } + } + )._socket ?? null +} + +export type SpliceCloseInfo = { code: number; reason: string; trigger: string } + +// The ws receiver kills a connection whose frame exceeds maxPayload with this +// message; surfacing it separately is what makes catalog-growth kills visible. +function errorTrigger(side: 'client' | 'host', error: Error): string { + return /max payload/i.test(error.message) ? `${side}-oversize-frame` : `${side}-error` +} + +export function wireSplice(input: { + client: WebSocket + host: WebSocket + budget: ProcessQueuedByteBudget + onClose: () => void + onForwardedBytes?: (bytes: number) => void + onClosed?: (closeInfo: SpliceCloseInfo) => void +}): (code?: number, reason?: string) => void { + let closed = false + const timers = new Set>() + const cleanups: Array<() => void> = [] + + const close = ( + code: number = RELAY_CLOSE_CODE.PEER_DROPPED, + reason = 'peer connection dropped', + trigger = 'external' + ): void => { + if (closed) return + closed = true + for (const timer of timers) clearTimeout(timer) + timers.clear() + for (const cleanup of cleanups) cleanup() + closeRelayWebSocket(input.client, code, reason) + closeRelayWebSocket(input.host, code, reason) + input.onClosed?.({ code, reason, trigger }) + input.onClose() + } + + const direction = (source: WebSocket, target: WebSocket): void => { + const queue: QueuedFrame[] = [] + let queuedBytes = 0 + let wedgedSince: number | null = null + cleanups.push(() => { + input.budget.release(queuedBytes) + queuedBytes = 0 + queue.length = 0 + transport(source)?.resume() + }) + + const flush = (): void => { + if (closed) return + if (target.readyState !== target.OPEN || source.readyState !== source.OPEN) { + close(undefined, undefined, 'peer-gone') + return + } + while ( + queue.length > 0 && + target.bufferedAmount <= RELAY_ADMISSION_BUDGETS.spliceLowWaterBytes + ) { + const frame = queue.shift()! + queuedBytes -= frame.bytes + input.budget.release(frame.bytes) + target.send(frame.data, { binary: frame.binary }) + } + if (queue.length === 0) { + wedgedSince = null + transport(source)?.resume() + return + } + if ( + wedgedSince !== null && + Date.now() - wedgedSince >= RELAY_ADMISSION_BUDGETS.spliceWedgedTimeoutMs + ) { + close(RELAY_CLOSE_CODE.LIMIT_EXCEEDED, 'wedged relay link', 'wedged') + return + } + const timer = setTimeout(() => { + timers.delete(timer) + flush() + }, 25) + timers.add(timer) + } + + source.on('message', (data, binary) => { + if (closed) return + const bytes = frameBytes(data) + if ( + queue.length === 0 && + target.readyState === target.OPEN && + target.bufferedAmount <= RELAY_ADMISSION_BUDGETS.spliceHighWaterBytes + ) { + target.send(data, { binary }) + input.onForwardedBytes?.(bytes) + return + } + if ( + queuedBytes + bytes > RELAY_ADMISSION_BUDGETS.spliceHardQueuedBytes || + !input.budget.reserve(bytes) + ) { + close(RELAY_CLOSE_CODE.LIMIT_EXCEEDED, 'relay queue limit exceeded', 'queue-limit') + return + } + queue.push({ data, binary, bytes }) + queuedBytes += bytes + input.onForwardedBytes?.(bytes) + wedgedSince ??= Date.now() + transport(source)?.pause() + if (queue.length === 1) flush() + }) + } + + transport(input.client)?.setNoDelay(true) + transport(input.host)?.setNoDelay(true) + direction(input.client, input.host) + direction(input.host, input.client) + input.client.once('close', () => close(undefined, undefined, 'client-closed')) + input.host.once('close', () => close(undefined, undefined, 'host-closed')) + input.client.once('error', (error) => close(undefined, undefined, errorTrigger('client', error))) + input.host.once('error', (error) => close(undefined, undefined, errorTrigger('host', error))) + return close +} diff --git a/cloud/apps/relay/src/staging-asia-proof-admission.test.ts b/cloud/apps/relay/src/staging-asia-proof-admission.test.ts new file mode 100644 index 00000000000..3bbee3ab9a3 --- /dev/null +++ b/cloud/apps/relay/src/staging-asia-proof-admission.test.ts @@ -0,0 +1,37 @@ +import { describe, expect, it } from 'vitest' +import { stagingAsiaProofMembership } from './staging-asia-proof-admission.js' + +describe('staging Asia proof admission', () => { + it('promotes only C4 and preserves every other cell', () => { + expect(stagingAsiaProofMembership({ + existingOnly: ['staging-gce-c1'], + migrationOnly: ['staging-gce-c4'], + general: ['staging-gce-c2'] + }, 'general')).toEqual({ + existingOnly: ['staging-gce-c1'], + migrationOnly: [], + general: ['staging-gce-c2', 'staging-gce-c4'] + }) + }) + + it('rolls back only C4', () => { + expect(stagingAsiaProofMembership({ + existingOnly: ['staging-gce-c1'], + migrationOnly: [], + general: ['staging-gce-c2', 'staging-gce-c4'] + }, 'migration-only')).toEqual({ + existingOnly: ['staging-gce-c1'], + migrationOnly: ['staging-gce-c4'], + general: ['staging-gce-c2'] + }) + }) + + it('rejects absent or irreversible C4 state', () => { + expect(() => stagingAsiaProofMembership({ + existingOnly: [], migrationOnly: [], general: ['staging-gce-c2'] + }, 'general')).toThrow('staging_asia_proof_cell_not_transitionable') + expect(() => stagingAsiaProofMembership({ + existingOnly: ['staging-gce-c4'], migrationOnly: [], general: [] + }, 'migration-only')).toThrow('staging_asia_proof_cell_not_transitionable') + }) +}) diff --git a/cloud/apps/relay/src/staging-asia-proof-admission.ts b/cloud/apps/relay/src/staging-asia-proof-admission.ts new file mode 100644 index 00000000000..a45a338768d --- /dev/null +++ b/cloud/apps/relay/src/staging-asia-proof-admission.ts @@ -0,0 +1,30 @@ +import type { CellAdmissionMembership } from './cell-admission-selector.js' + +export const STAGING_ASIA_PROOF_CELL_ID = 'staging-gce-c4' + +export type StagingAsiaProofAdmissionState = 'general' | 'migration-only' + +export function stagingAsiaProofMembership( + current: CellAdmissionMembership, + state: StagingAsiaProofAdmissionState +): CellAdmissionMembership { + const currentStates = [ + current.existingOnly.includes(STAGING_ASIA_PROOF_CELL_ID), + current.migrationOnly.includes(STAGING_ASIA_PROOF_CELL_ID), + current.general.includes(STAGING_ASIA_PROOF_CELL_ID) + ] + if (currentStates.filter(Boolean).length !== 1 || currentStates[0]) { + throw new Error('staging_asia_proof_cell_not_transitionable') + } + return { + existingOnly: [...current.existingOnly], + migrationOnly: current.migrationOnly + .filter((cellId) => cellId !== STAGING_ASIA_PROOF_CELL_ID) + .concat(state === 'migration-only' ? [STAGING_ASIA_PROOF_CELL_ID] : []) + .sort(), + general: current.general + .filter((cellId) => cellId !== STAGING_ASIA_PROOF_CELL_ID) + .concat(state === 'general' ? [STAGING_ASIA_PROOF_CELL_ID] : []) + .sort() + } +} diff --git a/cloud/apps/relay/tsconfig.build.json b/cloud/apps/relay/tsconfig.build.json new file mode 100644 index 00000000000..489ddfd34d6 --- /dev/null +++ b/cloud/apps/relay/tsconfig.build.json @@ -0,0 +1,10 @@ +{ + "extends": "./tsconfig.json", + "compilerOptions": { + "declaration": true, + "noEmit": false, + "outDir": "dist", + "rootDir": "src" + }, + "exclude": ["src/**/*.test.ts"] +} diff --git a/cloud/apps/relay/tsconfig.json b/cloud/apps/relay/tsconfig.json new file mode 100644 index 00000000000..a552e34dbe9 --- /dev/null +++ b/cloud/apps/relay/tsconfig.json @@ -0,0 +1,5 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { "noEmit": true }, + "include": ["src/**/*.ts"] +} diff --git a/cloud/apps/relay/vitest.config.ts b/cloud/apps/relay/vitest.config.ts new file mode 100644 index 00000000000..f56bbd7be3a --- /dev/null +++ b/cloud/apps/relay/vitest.config.ts @@ -0,0 +1,44 @@ +import { readdirSync, readFileSync } from 'node:fs' +import { fileURLToPath } from 'node:url' +import { defaultExclude, defineConfig } from 'vitest/config' + +// Every test file that opens ORCA_RELAY_TEST_POSTGRES_URL shares one CI +// database, and those tests take session-level locks with a 1s lock_timeout. +// Running them alongside anything else collides into lock timeouts, capacity +// exhaustion, and afterAll hangs. Keep them in their own serialized project so +// only they give up file parallelism; the rest of the suite opens SQLite data +// directories and stays fully parallel. +const sourceDirectory = fileURLToPath(new URL('src', import.meta.url)) +const sharedPostgresTests = readdirSync(sourceDirectory) + .filter((entry) => entry.endsWith('.test.ts')) + .filter((entry) => + readFileSync(`${sourceDirectory}/${entry}`, 'utf8').includes( + 'ORCA_RELAY_TEST_POSTGRES_URL' + ) + ) + .map((entry) => `src/${entry}`) + +const timeouts = { testTimeout: 15_000, hookTimeout: 15_000 } + +export default defineConfig({ + test: { + projects: [ + { + test: { + ...timeouts, + name: 'relay', + include: ['src/**/*.test.ts'], + exclude: [...defaultExclude, ...sharedPostgresTests] + } + }, + { + test: { + ...timeouts, + name: 'relay-postgres', + include: sharedPostgresTests, + fileParallelism: false + } + } + ] + } +}) diff --git a/cloud/package.json b/cloud/package.json new file mode 100644 index 00000000000..30970c3b24a --- /dev/null +++ b/cloud/package.json @@ -0,0 +1,23 @@ +{ + "name": "orca-cloud", + "private": true, + "version": "0.0.0", + "packageManager": "pnpm@10.24.0", + "engines": { + "node": ">=24 <27", + "pnpm": ">=10" + }, + "scripts": { + "build": "pnpm -r build", + "dev": "pnpm --filter @orca-cloud/relay dev", + "lint": "pnpm -r lint", + "test": "pnpm -r test", + "typecheck": "pnpm -r typecheck" + }, + "devDependencies": { + "@types/node": "^24.10.0", + "tsx": "^4.21.0", + "typescript": "^5.9.3", + "vitest": "^4.0.8" + } +} diff --git a/cloud/packages/relay-contract/package.json b/cloud/packages/relay-contract/package.json new file mode 100644 index 00000000000..4c224b51085 --- /dev/null +++ b/cloud/packages/relay-contract/package.json @@ -0,0 +1,23 @@ +{ + "name": "@orca-cloud/relay-contract", + "private": true, + "version": "0.0.0", + "type": "module", + "main": "dist/index.js", + "types": "dist/index.d.ts", + "scripts": { + "build": "pnpm clean && tsc -p tsconfig.build.json", + "clean": "node -e \"require('fs').rmSync('dist', { recursive: true, force: true })\"", + "lint": "tsc -p tsconfig.json --noEmit", + "test": "vitest run", + "typecheck": "tsc -p tsconfig.json --noEmit" + }, + "dependencies": { + "zod": "^3.25.76" + }, + "devDependencies": { + "@types/node": "^24.10.0", + "typescript": "^5.9.3", + "vitest": "^4.0.8" + } +} diff --git a/cloud/packages/relay-contract/src/admission-budgets.ts b/cloud/packages/relay-contract/src/admission-budgets.ts new file mode 100644 index 00000000000..6dcf18ec3af --- /dev/null +++ b/cloud/packages/relay-contract/src/admission-budgets.ts @@ -0,0 +1,76 @@ +export const RELAY_ADMISSION_BUDGETS = { + cloudRunConcurrency: 900, + maxPreAuthConnections: 45, + maxPreAuthPerSource: 4, + maxPreAuthAttemptsPerSourcePerMinute: 30, + reservedHostControls: 100, + reservedHostDataSockets: 150, + maxProcessQueuedBytes: 64 * 1024 * 1024, + spliceLowWaterBytes: 64 * 1024, + spliceHighWaterBytes: 256 * 1024, + // Why: must admit at least one maxFrameBytes frame for a backpressured + // peer, or large catalog responses close the splice with LIMIT_EXCEEDED. + spliceHardQueuedBytes: 8 * 1024 * 1024 + 256 * 1024, + spliceWedgedTimeoutMs: 10 * 1000 +} as const + +export const RELAY_CELL_CONNECTION_HARD_CAPS = [600, 1_000, 3_000] as const +export type RelayCellConnectionHardCap = (typeof RELAY_CELL_CONNECTION_HARD_CAPS)[number] +export const RELAY_CELL_CONNECTION_HARD_CAP: RelayCellConnectionHardCap = 600 + +export function isRelayCellConnectionHardCap( + value: unknown +): value is RelayCellConnectionHardCap { + return RELAY_CELL_CONNECTION_HARD_CAPS.some((hardCap) => hardCap === value) +} + +// Legacy/default bounds remain available for fixed-600 operational gates. +export const RELAY_CELL_ADMISSION_BOUNDS = { + hardCap: RELAY_CELL_CONNECTION_HARD_CAP, + // Ordinary sockets stop here; the remainder is held for same-host control rebinds. + socketAdmissionCeiling: + RELAY_CELL_CONNECTION_HARD_CAP - RELAY_ADMISSION_BUDGETS.reservedHostControls, + // A cell must leave at least one unit admissible after its unobserved allowance. + maxUnobservedBound: + RELAY_CELL_CONNECTION_HARD_CAP - RELAY_ADMISSION_BUDGETS.reservedHostControls - 1 +} as const + +export function relayCellAdmissionBounds(hardCap: RelayCellConnectionHardCap): { + hardCap: RelayCellConnectionHardCap + socketAdmissionCeiling: number + maxUnobservedBound: number +} { + return { + hardCap, + socketAdmissionCeiling: hardCap - RELAY_ADMISSION_BUDGETS.reservedHostControls, + maxUnobservedBound: hardCap - RELAY_ADMISSION_BUDGETS.reservedHostControls - 1 + } +} + +export const RELAY_MAX_CELL_CONNECTION_UNOBSERVED_BOUND = relayCellAdmissionBounds( + RELAY_CELL_CONNECTION_HARD_CAPS.at(-1)! +).maxUnobservedBound + +// Director placement stops short of the socket ceiling by the cell's own unobserved allowance. +export function cellPlacementCeiling( + connectionHardCap: RelayCellConnectionHardCap, + connectionUnobservedBound: number +): number { + return relayCellAdmissionBounds(connectionHardCap).socketAdmissionCeiling - connectionUnobservedBound +} + +export function hasAdmissionCapacity(input: { + totalRequests: number + preAuthConnections: number + sourcePreAuthConnections: number + totalRequestCeiling?: number +}): boolean { + if (input.preAuthConnections >= RELAY_ADMISSION_BUDGETS.maxPreAuthConnections) return false + if (input.sourcePreAuthConnections >= RELAY_ADMISSION_BUDGETS.maxPreAuthPerSource) return false + const totalRequestCeiling = + input.totalRequestCeiling ?? + RELAY_ADMISSION_BUDGETS.cloudRunConcurrency - + RELAY_ADMISSION_BUDGETS.reservedHostControls - + RELAY_ADMISSION_BUDGETS.reservedHostDataSockets + return input.totalRequests < totalRequestCeiling +} diff --git a/cloud/packages/relay-contract/src/assignment-invariants.ts b/cloud/packages/relay-contract/src/assignment-invariants.ts new file mode 100644 index 00000000000..dc60a1ec7b8 --- /dev/null +++ b/cloud/packages/relay-contract/src/assignment-invariants.ts @@ -0,0 +1,56 @@ +import { z } from 'zod' +import { EpochMsSchema, GenerationSchema, OpaqueIdSchema, RelayHostIdSchema } from './wire-scalars.js' + +export const ASSIGNMENT_LIMITS = { + activityLeaseMs: 90 * 1000, + dormantTtlMs: 24 * 60 * 60 * 1000, + migrationLeaseMs: 15 * 60 * 1000 +} as const + +export const AssignmentActivitySchema = z + .object({ + relayHostId: RelayHostIdSchema, + cellId: OpaqueIdSchema, + assignmentEpoch: GenerationSchema, + leaseExpiresAt: EpochMsSchema, + lastActivityAt: EpochMsSchema, + reservedControls: z.number().int().nonnegative(), + reservedSplices: z.number().int().nonnegative(), + reservedInvites: z.number().int().nonnegative(), + pendingInstalls: z.number().int().nonnegative(), + pendingConfirmations: z.number().int().nonnegative(), + migrationLeases: z.number().int().nonnegative() + }) + .strict() + +export function hasAssignmentActivity(record: z.infer): boolean { + return ( + record.reservedControls + + record.reservedSplices + + record.reservedInvites + + record.pendingInstalls + + record.pendingConfirmations + + record.migrationLeases > + 0 + ) +} + +export function mayNormallyReassign( + record: z.infer, + now: number +): boolean { + return !hasAssignmentActivity(record) && now >= record.lastActivityAt + ASSIGNMENT_LIMITS.dormantTtlMs +} + +export const EvacuationCommitSchema = z + .object({ + relayHostId: RelayHostIdSchema, + sourceCellId: OpaqueIdSchema, + targetCellId: OpaqueIdSchema, + previousEpoch: GenerationSchema, + assignmentEpoch: GenerationSchema, + targetCapacityReserved: z.literal(true) + }) + .strict() + .refine((move) => move.sourceCellId !== move.targetCellId, 'target cell must differ') + .refine((move) => move.assignmentEpoch === move.previousEpoch + 1, 'epoch must increment exactly once') diff --git a/cloud/packages/relay-contract/src/close-codes.ts b/cloud/packages/relay-contract/src/close-codes.ts new file mode 100644 index 00000000000..cb2c6307b90 --- /dev/null +++ b/cloud/packages/relay-contract/src/close-codes.ts @@ -0,0 +1,10 @@ +export const RELAY_CLOSE_CODE = { + BAD_OUTER_CREDENTIAL: 4401, + HOST_OFFLINE: 4404, + PEER_DROPPED: 4408, + WRONG_CELL: 4409, + LIMIT_EXCEEDED: 4429, + DRAINING: 4503 +} as const + +export type RelayCloseCode = (typeof RELAY_CLOSE_CODE)[keyof typeof RELAY_CLOSE_CODE] diff --git a/cloud/packages/relay-contract/src/contract.test.ts b/cloud/packages/relay-contract/src/contract.test.ts new file mode 100644 index 00000000000..805cb8ea698 --- /dev/null +++ b/cloud/packages/relay-contract/src/contract.test.ts @@ -0,0 +1,347 @@ +import { describe, expect, it } from 'vitest' +import { RELAY_CLOSE_CODE } from './close-codes.js' +import { + AuthRefreshSchema, + DrainSchema, + HostChallengeSchema, + HostDataAuthSchema, + HostHelloAckSchema, + HostHelloSchema +} from './control-messages.js' +import { + DeviceCredentialInstallSchema, + DeviceResumeConfirmSchema, + RelayAuthSchema, + RelayHelloSchema +} from './credential-messages.js' +import { + AssignmentRequestSchema, + AssignmentResponseSchema, + isTrustedNewerMove, + RelayMovedSchema, + ResolveRequestSchema +} from './director-messages.js' +import { RELAY_PROTOCOL_LIMITS } from './protocol-limits.js' +import { RelayRegionCatalogResponseSchema } from './relay-regions.js' +import { + buildHostChallengePlaintext, + buildHostProofMacInput, + buildHostProofTranscript +} from './host-proof-transcript.js' +import { ConfirmableResumeTupleSchema } from './resume-confirmation-contract.js' +import { + canAdvanceSplice, + mayAcknowledgeClient, + SPLICE_STATE +} from './splice-state-machine.js' + +const TOKEN = 'abcdefghijklmnopqrstuvwxyzABCDEFGH012345678' +const NONCE = 'abcdefghijklmnopqrstuvwxyzABCDEF' +const KEY_B64 = Buffer.alloc(32, 1).toString('base64') + +describe('relay protocol contract', () => { + it('locks close codes and fixed normative limits', () => { + expect(RELAY_CLOSE_CODE).toEqual({ + BAD_OUTER_CREDENTIAL: 4401, + HOST_OFFLINE: 4404, + PEER_DROPPED: 4408, + WRONG_CELL: 4409, + LIMIT_EXCEEDED: 4429, + DRAINING: 4503 + }) + expect(RELAY_PROTOCOL_LIMITS).toMatchObject({ + firstFrameDeadlineMs: 2_000, + maxHttpBodyBytes: 4_096, + maxFrameBytes: 8 * 1024 * 1024, + maxConnectionsPerHost: 8, + idleTimeoutMs: 600_000, + inviteMaxAttempts: 5, + inviteReservationLeaseMs: 15_000, + hostAttachDeadlineMs: 10_000, + resumeConfirmationDeadlineMs: 30_000 + }) + }) + + it('strictly validates host registration and continuity fields', () => { + const hello = { + v: 1, + relayHostId: 'abcdefghijklmnop', + assignmentEpoch: 7, + hostPublicKeyB64: KEY_B64, + appVersion: '1.2.3' + } + expect(HostHelloSchema.safeParse(hello).success).toBe(true) + expect(HostHelloSchema.safeParse({ ...hello, userId: 'injected' }).success).toBe(false) + expect( + HostHelloAckSchema.safeParse({ + v: 1, + generation: 8, + controlResumeSecret: TOKEN, + leaseExpiresAt: 1_800_000_000_000, + activeConnIds: [], + pendingConns: [] + }).success + ).toBe(true) + }) + + it('locks bounded director assignment and resume-only resolve payloads', () => { + const assignment = { v: 1, relayHostId: 'abcdefghijklmnop' } + expect(AssignmentRequestSchema.safeParse(assignment).success).toBe(true) + expect( + AssignmentRequestSchema.safeParse({ ...assignment, preferredRegion: 'asia-east2' }).success + ).toBe(true) + expect( + AssignmentRequestSchema.safeParse({ ...assignment, preferredRegion: 'europe-west1' }).success + ).toBe(false) + expect(AssignmentRequestSchema.safeParse({ ...assignment, relayJwt: 'url-secret' }).success).toBe( + false + ) + expect( + AssignmentResponseSchema.safeParse({ + v: 1, + cellUrl: 'https://relay-c1.onorca.dev', + assignmentEpoch: 3, + lease: 'signed-lease' + }).success + ).toBe(true) + expect( + ResolveRequestSchema.safeParse({ + v: 1, + relayHostId: 'abcdefghijklmnop', + resumeToken: TOKEN + }).success + ).toBe(true) + }) + + it('accepts only unique regions with unique canonical HTTPS probe origins', () => { + const catalog = { + v: 1, + regions: [ + { + region: 'us-central1', + probeOrigins: ['https://relay-c1.example.test', 'https://relay-c2.example.test'] + }, + { region: 'asia-east2', probeOrigins: ['https://relay-c27.example.test'] } + ] + } + expect(RelayRegionCatalogResponseSchema.safeParse(catalog).success).toBe(true) + expect( + RelayRegionCatalogResponseSchema.safeParse({ + ...catalog, + regions: [catalog.regions[0], { ...catalog.regions[0] }] + }).success + ).toBe(false) + expect( + RelayRegionCatalogResponseSchema.safeParse({ + ...catalog, + regions: [ + catalog.regions[0], + { region: 'asia-east2', probeOrigins: ['https://relay-c1.example.test'] } + ] + }).success + ).toBe(false) + for (const origin of [ + 'http://relay-c1.example.test', + 'https://relay-c1.example.test/', + 'https://relay-c1.example.test/health', + 'https://relay-c1.example.test?probe=1' + ]) { + expect( + RelayRegionCatalogResponseSchema.safeParse({ + v: 1, + regions: [{ region: 'us-central1', probeOrigins: [origin] }] + }).success + ).toBe(false) + } + }) + + it('accepts moves only from the configured director at a strictly newer epoch', () => { + const move = RelayMovedSchema.parse({ + v: 1, + cellUrl: 'https://relay-c2.onorca.dev', + assignmentEpoch: 4 + }) + const base = { + configuredDirectorOrigin: 'https://relay.onorca.dev', + currentAssignmentEpoch: 3, + move + } + expect(isTrustedNewerMove({ ...base, sourceOrigin: 'https://relay.onorca.dev' })).toBe(true) + expect(isTrustedNewerMove({ ...base, sourceOrigin: move.cellUrl })).toBe(false) + expect( + isTrustedNewerMove({ + ...base, + sourceOrigin: 'https://relay.onorca.dev', + currentAssignmentEpoch: 4 + }) + ).toBe(false) + }) + + it('bounds challenge material and fixes director-only drain recovery', () => { + expect( + HostChallengeSchema.safeParse({ + challengeId: 'challenge-1', + relayEphemeralPublicKeyB64: KEY_B64, + nonceB64: NONCE, + ciphertextB64: KEY_B64, + expiresAt: 1_800_000_000_000 + }).success + ).toBe(true) + expect(DrainSchema.safeParse({ graceMs: 0, recovery: 'resolve-director' }).success).toBe(true) + expect(DrainSchema.safeParse({ graceMs: 0, recovery: 'follow-server-url' }).success).toBe(false) + expect(AuthRefreshSchema.safeParse({ relayJwt: 'jwt', identity: 'injected' }).success).toBe(false) + }) + + it('strictly binds one-shot host data tickets to a control generation', () => { + expect(HostDataAuthSchema.safeParse({ v: 1, connTicket: TOKEN, generation: 3 }).success).toBe( + true + ) + expect( + HostDataAuthSchema.safeParse({ + v: 1, + connTicket: TOKEN, + generation: 3, + relayDeviceId: 'injected' + }).success + ).toBe(false) + }) + + it('cannot acknowledge a splice before forwarding handlers exist', () => { + expect( + canAdvanceSplice(SPLICE_STATE.PRE_AUTH_ADMITTED, SPLICE_STATE.CREDENTIAL_LEASE_RESERVED) + ).toBe(true) + expect(canAdvanceSplice(SPLICE_STATE.PRE_AUTH_ADMITTED, SPLICE_STATE.SPLICED)).toBe(false) + expect(mayAcknowledgeClient(SPLICE_STATE.HOST_ATTACHED, false)).toBe(false) + expect(mayAcknowledgeClient(SPLICE_STATE.HOST_ATTACHED, true)).toBe(true) + }) + + it('locks the immutable server-owned resume tuple shape', () => { + const tuple = { + basisConnId: 'conn-1', + owningControlGeneration: 4, + relayDeviceId: 'device-1', + acceptedCredentialVersion: 2, + acceptedAs: 'current', + confirmDeadline: 1_800_000_000_000 + } + expect(ConfirmableResumeTupleSchema.safeParse(tuple).success).toBe(true) + expect(ConfirmableResumeTupleSchema.safeParse({ ...tuple, userId: 'caller-value' }).success).toBe( + false + ) + }) + + it('locks the complete host key-possession transcript', () => { + const transcript = buildHostProofTranscript({ + relayOrigin: 'https://relay.onorca.dev', + relayEphemeralPublicKey: new Uint8Array(32).fill(1), + challengeNonce: new Uint8Array(24).fill(4), + challengeId: 'challenge-1', + issuedAt: 1_700_000_000_000, + expiresAt: 1_700_000_010_000, + userId: 'user-1', + profileId: 'profile-1', + organizationId: 'org-1', + relayHostId: 'abcdefghijklmnop', + hostPublicKey: new Uint8Array(32).fill(2), + assignmentEpoch: 7, + previousGeneration: 6, + resumeRequested: true + }) + const challenge = buildHostChallengePlaintext(transcript, new Uint8Array(32).fill(3)) + const proofInput = buildHostProofMacInput(transcript) + + expect(Buffer.from(transcript).toString('base64url')).toBe( + 'AAAACHByb3RvY29sAAAAGG9yY2EtcmVsYXktaG9zdC1wcm9vZi92MQAAAAd2ZXJzaW9uAAAAAQEAAAALcmVsYXlPcmlnaW4AAAAYaHR0cHM6Ly9yZWxheS5vbm9yY2EuZGV2AAAAF3JlbGF5RXBoZW1lcmFsUHVibGljS2V5AAAAIAEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAAAADmNoYWxsZW5nZU5vbmNlAAAAGAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAAAAAtjaGFsbGVuZ2VJZAAAAAtjaGFsbGVuZ2UtMQAAAAhpc3N1ZWRBdAAAAAgAAAGLz-VoAAAAAAlleHBpcmVzQXQAAAAIAAABi8_ljxAAAAAGdXNlcklkAAAABnVzZXItMQAAAAlwcm9maWxlSWQAAAAJcHJvZmlsZS0xAAAADm9yZ2FuaXphdGlvbklkAAAABW9yZy0xAAAAC3JlbGF5SG9zdElkAAAAEGFiY2RlZmdoaWprbG1ub3AAAAANaG9zdFB1YmxpY0tleQAAACACAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgAAAA9hc3NpZ25tZW50RXBvY2gAAAAIAAAAAAAAAAcAAAAScHJldmlvdXNHZW5lcmF0aW9uAAAACAAAAAAAAAAGAAAAD3Jlc3VtZVJlcXVlc3RlZAAAAAEB' + ) + expect(challenge.byteLength).toBe(transcript.byteLength + 65) + expect(proofInput.byteLength).toBe(transcript.byteLength + 29) + expect(() => buildHostChallengePlaintext(transcript, new Uint8Array(31))).toThrow( + 'challengeSecret must be 32 bytes' + ) + expect(() => + buildHostProofTranscript({ + relayOrigin: 'https://relay.onorca.dev', + relayEphemeralPublicKey: new Uint8Array(32), + challengeNonce: new Uint8Array(32), + challengeId: 'challenge-1', + issuedAt: 1, + expiresAt: 2, + userId: 'user-1', + profileId: 'profile-1', + organizationId: 'org-1', + relayHostId: 'abcdefghijklmnop', + hostPublicKey: new Uint8Array(32), + assignmentEpoch: 1, + resumeRequested: false + }) + ).toThrow('challengeNonce must be 24 bytes') + }) + + it('accepts only the bounded first-frame credential shape', () => { + expect(RelayAuthSchema.parse({ v: 1, mode: 'connect', credential: TOKEN })).toEqual({ + v: 1, + mode: 'connect', + credential: TOKEN + }) + expect( + RelayAuthSchema.safeParse({ v: 1, mode: 'connect', credential: TOKEN, extra: true }).success + ).toBe(false) + }) + + it('makes install authorization modes mutually exclusive', () => { + const base = { + v: 1, + reqId: 'req-1', + relayDeviceId: 'device-1', + newResumeTokenHash: TOKEN + } + expect( + DeviceCredentialInstallSchema.safeParse({ + ...base, + authorization: { mode: 'relay-basis', basisConnId: 'conn-1' } + }).success + ).toBe(true) + expect( + DeviceCredentialInstallSchema.safeParse({ + ...base, + authorization: { + mode: 'relay-basis', + basisConnId: 'conn-1', + directAuthId: 'injected' + } + }).success + ).toBe(false) + }) + + it('does not let invite attach observations impersonate resume state', () => { + expect( + RelayHelloSchema.safeParse({ + ok: true, + credentialKind: 'invite', + leaseExpiresAt: 1_800_000_000_000 + }).success + ).toBe(true) + expect( + RelayHelloSchema.safeParse({ + ok: true, + credentialKind: 'invite', + leaseExpiresAt: 1_800_000_000_000, + acceptedCredentialVersion: 7, + acceptedAs: 'current', + resumeExpiresAt: 1_800_000_000_000 + }).success + ).toBe(false) + }) + + it('rejects caller-supplied device and credential metadata on resume confirmation', () => { + const confirmation = { v: 1, reqId: 'req-1', basisConnId: 'conn-1' } + expect(DeviceResumeConfirmSchema.safeParse(confirmation).success).toBe(true) + expect( + DeviceResumeConfirmSchema.safeParse({ + ...confirmation, + relayDeviceId: 'injected', + acceptedCredentialVersion: 99 + }).success + ).toBe(false) + }) +}) diff --git a/cloud/packages/relay-contract/src/control-continuity.ts b/cloud/packages/relay-contract/src/control-continuity.ts new file mode 100644 index 00000000000..11821fc537e --- /dev/null +++ b/cloud/packages/relay-contract/src/control-continuity.ts @@ -0,0 +1,43 @@ +export const CONTROL_GENERATION_STATE = { + ACTIVE: 'active', + ORPHANED: 'orphaned', + DRAIN_ONLY: 'drain-only', + FENCED: 'fenced', + CLOSED: 'closed' +} as const + +export const CONTROL_CONTINUITY_LIMITS = { + orphanGraceMs: 30 * 1000, + authRefreshMinBeforeExpiryMs: 60 * 1000, + authRefreshMaxBeforeExpiryMs: 120 * 1000, + expiredAuthExistingSpliceGraceMs: 60 * 1000 +} as const + +export function controlLossDisposition(input: { + matchingResumeSecret: boolean + competingGeneration: boolean +}): 'rebind' | 'fence-old' | 'orphan-grace' { + if (input.matchingResumeSecret) return 'rebind' + if (input.competingGeneration) return 'fence-old' + return 'orphan-grace' +} + +export function mayStartNewRelayWork(state: string, authExpired: boolean): boolean { + return state === CONTROL_GENERATION_STATE.ACTIVE && !authExpired +} + +export interface RelayAuthIdentity { + sub: string + prof: string + org?: string + relayHostId: string +} + +export function preservesRelayAuthIdentity(previous: RelayAuthIdentity, refreshed: RelayAuthIdentity): boolean { + return ( + previous.sub === refreshed.sub && + previous.prof === refreshed.prof && + previous.org === refreshed.org && + previous.relayHostId === refreshed.relayHostId + ) +} diff --git a/cloud/packages/relay-contract/src/control-messages.ts b/cloud/packages/relay-contract/src/control-messages.ts new file mode 100644 index 00000000000..0d5f8d1b851 --- /dev/null +++ b/cloud/packages/relay-contract/src/control-messages.ts @@ -0,0 +1,119 @@ +import { z } from 'zod' +import { + Base6432ByteSchema, + Base64Raw24ByteSchema, + Base64Url32ByteSchema, + EpochMsSchema, + GenerationSchema, + OpaqueIdSchema, + PositiveDurationMsSchema, + RelayHostIdSchema +} from './wire-scalars.js' + +const AppVersionSchema = z.string().min(1).max(128) +const BoundedCiphertextSchema = z + .string() + .min(1) + .max(16 * 1024) + .regex(/^(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?$/) +const ConnectionKindSchema = z.enum(['invite', 'resume']) + +export const HostHelloSchema = z + .object({ + v: z.literal(1), + relayHostId: RelayHostIdSchema, + assignmentEpoch: GenerationSchema, + hostPublicKeyB64: Base6432ByteSchema, + appVersion: AppVersionSchema, + previousGeneration: GenerationSchema.optional(), + controlResumeSecret: Base64Url32ByteSchema.optional() + }) + .strict() + +export const HostChallengeSchema = z + .object({ + challengeId: OpaqueIdSchema, + relayEphemeralPublicKeyB64: Base6432ByteSchema, + nonceB64: Base64Raw24ByteSchema, + ciphertextB64: BoundedCiphertextSchema, + expiresAt: EpochMsSchema + }) + .strict() + +export const HostChallengeAckSchema = z + .object({ challengeId: OpaqueIdSchema, proofB64: Base6432ByteSchema }) + .strict() + +const PendingConnectionSchema = z + .object({ connId: OpaqueIdSchema, connTicket: Base64Url32ByteSchema }) + .strict() + +export const HostHelloAckSchema = z + .object({ + v: z.literal(1), + generation: GenerationSchema, + controlResumeSecret: Base64Url32ByteSchema, + leaseExpiresAt: EpochMsSchema, + activeConnIds: z.array(OpaqueIdSchema).max(8), + pendingConns: z.array(PendingConnectionSchema).max(8) + }) + .strict() + +export const ConnectionOpenSchema = z + .object({ + connId: OpaqueIdSchema, + connTicket: Base64Url32ByteSchema, + kind: ConnectionKindSchema, + relayDeviceId: OpaqueIdSchema, + attachDeadlineMs: PositiveDurationMsSchema + }) + .strict() + +export const HostDataAuthSchema = z + .object({ + v: z.literal(1), + connTicket: Base64Url32ByteSchema, + generation: GenerationSchema + }) + .strict() + +export const InviteCreateSchema = z + .object({ reqId: OpaqueIdSchema, relayDeviceId: OpaqueIdSchema }) + .strict() + +export const InviteCreatedSchema = z + .object({ + reqId: OpaqueIdSchema, + inviteToken: Base64Url32ByteSchema, + expiresAt: EpochMsSchema, + maxAttempts: z.number().int().positive().max(16) + }) + .strict() + +export const DeviceRevokeSchema = z + .object({ reqId: OpaqueIdSchema, relayDeviceId: OpaqueIdSchema }) + .strict() + +export const AuthRefreshSchema = z.object({ relayJwt: z.string().min(1).max(8 * 1024) }).strict() + +export const DrainSchema = z + .object({ + graceMs: z.number().int().nonnegative().max(60 * 60 * 1000), + recovery: z.literal('resolve-director') + }) + .strict() + +export const HeartbeatSchema = z.object({ t: EpochMsSchema }).strict() + +export type HostHello = z.infer +export type HostChallenge = z.infer +export type HostChallengeAck = z.infer +export type HostHelloAck = z.infer +export type ConnectionOpen = z.infer +export type HostDataAuth = z.infer +export type InviteCreate = z.infer +export type InviteCreated = z.infer +export type DeviceRevoke = z.infer +export type AuthRefresh = z.infer +export type Drain = z.infer +export type Heartbeat = z.infer diff --git a/cloud/packages/relay-contract/src/credential-messages.ts b/cloud/packages/relay-contract/src/credential-messages.ts new file mode 100644 index 00000000000..658d2f25f89 --- /dev/null +++ b/cloud/packages/relay-contract/src/credential-messages.ts @@ -0,0 +1,105 @@ +import { z } from 'zod' +import { Base64Url32ByteSchema, EpochMsSchema, OpaqueIdSchema } from './wire-scalars.js' + +export const RelayAuthSchema = z + .object({ v: z.literal(1), mode: z.literal('connect'), credential: Base64Url32ByteSchema }) + .strict() + +const RelayErrorCodeSchema = z.union([ + z.literal(4401), + z.literal(4404), + z.literal(4408), + z.literal(4409), + z.literal(4429), + z.literal(4503) +]) + +export const RelayHelloSchema = z.union([ + z.object({ ok: z.literal(false), code: RelayErrorCodeSchema }).strict(), + z + .object({ + ok: z.literal(true), + credentialKind: z.literal('invite'), + leaseExpiresAt: EpochMsSchema + }) + .strict(), + z + .object({ + ok: z.literal(true), + credentialKind: z.literal('resume'), + leaseExpiresAt: EpochMsSchema, + acceptedCredentialVersion: z.number().int().positive(), + acceptedAs: z.enum(['current', 'grace']), + resumeExpiresAt: EpochMsSchema, + graceExpiresAt: EpochMsSchema.optional() + }) + .strict() +]) + +export const DeviceCredentialInstallSchema = z + .object({ + v: z.literal(1), + reqId: OpaqueIdSchema, + relayDeviceId: OpaqueIdSchema, + newResumeTokenHash: Base64Url32ByteSchema, + expectedCurrentHash: Base64Url32ByteSchema.optional(), + authorization: z.discriminatedUnion('mode', [ + z.object({ mode: z.literal('relay-basis'), basisConnId: OpaqueIdSchema }).strict(), + z.object({ mode: z.literal('authenticated-direct'), directAuthId: OpaqueIdSchema }).strict() + ]) + }) + .strict() + +export const DeviceCredentialInstallStatusSchema = z + .object({ v: z.literal(1), reqId: OpaqueIdSchema, relayDeviceId: OpaqueIdSchema }) + .strict() + +export const DeviceResumeConfirmSchema = z + .object({ v: z.literal(1), reqId: OpaqueIdSchema, basisConnId: OpaqueIdSchema }) + .strict() + +export const DeviceCredentialInstalledSchema = z + .object({ + v: z.literal(1), + reqId: OpaqueIdSchema, + authorizationMode: z.enum(['relay-basis', 'authenticated-direct']), + currentVersion: z.number().int().positive(), + resumeExpiresAt: EpochMsSchema, + graceExpiresAt: EpochMsSchema.optional() + }) + .strict() + +export const DeviceCredentialInstallStatusResultSchema = z.union([ + z.object({ v: z.literal(1), reqId: OpaqueIdSchema, state: z.literal('not-found') }).strict(), + z + .object({ + v: z.literal(1), + reqId: OpaqueIdSchema, + state: z.literal('committed'), + result: DeviceCredentialInstalledSchema + }) + .strict() +]) + +export const DeviceResumeConfirmedSchema = z + .object({ + v: z.literal(1), + reqId: OpaqueIdSchema, + currentVersion: z.number().int().positive(), + acceptedAs: z.enum(['current', 'grace']), + renewed: z.boolean(), + resumeExpiresAt: EpochMsSchema, + graceExpiresAt: EpochMsSchema.optional() + }) + .strict() + +export type RelayAuth = z.infer +export type RelayHello = z.infer +export type DeviceCredentialInstall = z.infer +export type DeviceCredentialInstalled = z.infer +export type DeviceCredentialInstallStatus = z.infer +export type DeviceCredentialInstallStatusResult = z.infer< + typeof DeviceCredentialInstallStatusResultSchema +> +export type DeviceResumeConfirm = z.infer +export type DeviceResumeConfirmed = z.infer diff --git a/cloud/packages/relay-contract/src/director-messages.ts b/cloud/packages/relay-contract/src/director-messages.ts new file mode 100644 index 00000000000..e697135b68e --- /dev/null +++ b/cloud/packages/relay-contract/src/director-messages.ts @@ -0,0 +1,75 @@ +import { z } from 'zod' +import { + Base64Url32ByteSchema, + CanonicalHttpsOriginSchema, + EpochMsSchema, + GenerationSchema, + RelayHostIdSchema +} from './wire-scalars.js' +import { RelayRegionSchema } from './relay-regions.js' + +const SignedAssignmentLeaseSchema = z.string().min(1).max(8 * 1024) + +export const AssignmentRequestSchema = z + .object({ + v: z.literal(1), + relayHostId: RelayHostIdSchema, + // Client-declared reconnection; the director verifies it against the + // durable assignment before granting fast-lane admission. + reconnect: z.boolean().optional(), + preferredRegion: RelayRegionSchema.optional() + }) + .strict() + +export const AssignmentResponseSchema = z + .object({ + v: z.literal(1), + cellUrl: CanonicalHttpsOriginSchema, + assignmentEpoch: GenerationSchema, + lease: SignedAssignmentLeaseSchema + }) + .strict() + +export const ResolveRequestSchema = z + .object({ + v: z.literal(1), + relayHostId: RelayHostIdSchema, + resumeToken: Base64Url32ByteSchema + }) + .strict() + +export const ResolveResponseSchema = z + .object({ + v: z.literal(1), + cellUrl: CanonicalHttpsOriginSchema, + assignmentEpoch: GenerationSchema, + leaseExpiresAt: EpochMsSchema + }) + .strict() + +export const RelayMovedSchema = z + .object({ + v: z.literal(1), + cellUrl: CanonicalHttpsOriginSchema, + assignmentEpoch: GenerationSchema + }) + .strict() + +export function isTrustedNewerMove(input: { + sourceOrigin: string + configuredDirectorOrigin: string + currentAssignmentEpoch: number + move: z.infer +}): boolean { + // Why: cells and stale director responses must never redirect a credential-bearing client. + return ( + input.sourceOrigin === input.configuredDirectorOrigin && + input.move.assignmentEpoch > input.currentAssignmentEpoch + ) +} + +export type AssignmentRequest = z.infer +export type AssignmentResponse = z.infer +export type ResolveRequest = z.infer +export type ResolveResponse = z.infer +export type RelayMoved = z.infer diff --git a/cloud/packages/relay-contract/src/host-proof-transcript.ts b/cloud/packages/relay-contract/src/host-proof-transcript.ts new file mode 100644 index 00000000000..853986f23f2 --- /dev/null +++ b/cloud/packages/relay-contract/src/host-proof-transcript.ts @@ -0,0 +1,103 @@ +const textEncoder = new TextEncoder() + +export const HOST_PROOF_TRANSCRIPT_DOMAIN = 'orca-relay-host-proof/v1' +export const HOST_CHALLENGE_PLAINTEXT_DOMAIN = 'orca-relay-host-challenge/v1' +export const HOST_CHALLENGE_BOX_ALGORITHM = 'Curve25519-XSalsa20-Poly1305' +export const HOST_PROOF_ALGORITHM = 'HMAC-SHA-256' + +export interface HostProofTranscriptInput { + relayOrigin: string + relayEphemeralPublicKey: Uint8Array + challengeNonce: Uint8Array + challengeId: string + issuedAt: number + expiresAt: number + userId: string + profileId: string + organizationId: string + relayHostId: string + hostPublicKey: Uint8Array + assignmentEpoch: number + previousGeneration?: number + resumeRequested: boolean +} + +function uint32(value: number): Uint8Array { + const bytes = new Uint8Array(4) + new DataView(bytes.buffer).setUint32(0, value, false) + return bytes +} + +function uint64(value: number): Uint8Array { + const bytes = new Uint8Array(8) + new DataView(bytes.buffer).setBigUint64(0, BigInt(value), false) + return bytes +} + +function concat(parts: readonly Uint8Array[]): Uint8Array { + const output = new Uint8Array(parts.reduce((total, part) => total + part.byteLength, 0)) + let offset = 0 + for (const part of parts) { + output.set(part, offset) + offset += part.byteLength + } + return output +} + +function field(name: string, value: Uint8Array): Uint8Array { + const encodedName = textEncoder.encode(name) + return concat([uint32(encodedName.byteLength), encodedName, uint32(value.byteLength), value]) +} + +function text(value: string): Uint8Array { + return textEncoder.encode(value) +} + +function requireByteLength(value: Uint8Array, expected: number, name: string): void { + if (value.byteLength !== expected) throw new Error(`${name} must be ${expected} bytes`) +} + +export function buildHostProofTranscript(input: HostProofTranscriptInput): Uint8Array { + requireByteLength(input.relayEphemeralPublicKey, 32, 'relayEphemeralPublicKey') + requireByteLength(input.challengeNonce, 24, 'challengeNonce') + requireByteLength(input.hostPublicKey, 32, 'hostPublicKey') + return concat([ + field('protocol', text(HOST_PROOF_TRANSCRIPT_DOMAIN)), + field('version', new Uint8Array([1])), + field('relayOrigin', text(input.relayOrigin)), + field('relayEphemeralPublicKey', input.relayEphemeralPublicKey), + field('challengeNonce', input.challengeNonce), + field('challengeId', text(input.challengeId)), + field('issuedAt', uint64(input.issuedAt)), + field('expiresAt', uint64(input.expiresAt)), + field('userId', text(input.userId)), + field('profileId', text(input.profileId)), + field('organizationId', text(input.organizationId)), + field('relayHostId', text(input.relayHostId)), + field('hostPublicKey', input.hostPublicKey), + field('assignmentEpoch', uint64(input.assignmentEpoch)), + field( + 'previousGeneration', + input.previousGeneration === undefined ? new Uint8Array() : uint64(input.previousGeneration) + ), + field('resumeRequested', new Uint8Array([input.resumeRequested ? 1 : 0])) + ]) +} + +export function buildHostChallengePlaintext( + transcript: Uint8Array, + challengeSecret: Uint8Array +): Uint8Array { + if (challengeSecret.byteLength !== 32) throw new Error('challengeSecret must be 32 bytes') + // Why: the encrypted random secret makes the public transcript insufficient to forge the ack. + return concat([ + text(`${HOST_CHALLENGE_PLAINTEXT_DOMAIN}\0`), + uint32(transcript.byteLength), + transcript, + challengeSecret + ]) +} + +export function buildHostProofMacInput(transcript: Uint8Array): Uint8Array { + return concat([text(`${HOST_PROOF_TRANSCRIPT_DOMAIN}\0ack\0`), transcript]) +} diff --git a/cloud/packages/relay-contract/src/index.ts b/cloud/packages/relay-contract/src/index.ts new file mode 100644 index 00000000000..2a7d7d0feda --- /dev/null +++ b/cloud/packages/relay-contract/src/index.ts @@ -0,0 +1,14 @@ +export * from './close-codes.js' +export * from './admission-budgets.js' +export * from './assignment-invariants.js' +export * from './control-messages.js' +export * from './control-continuity.js' +export * from './credential-messages.js' +export * from './director-messages.js' +export * from './host-proof-transcript.js' +export * from './persistence-invariants.js' +export * from './protocol-limits.js' +export * from './resume-confirmation-contract.js' +export * from './relay-regions.js' +export * from './splice-state-machine.js' +export * from './wire-scalars.js' diff --git a/cloud/packages/relay-contract/src/persistence-invariants.test.ts b/cloud/packages/relay-contract/src/persistence-invariants.test.ts new file mode 100644 index 00000000000..291fa7462da --- /dev/null +++ b/cloud/packages/relay-contract/src/persistence-invariants.test.ts @@ -0,0 +1,156 @@ +import { describe, expect, it } from 'vitest' +import { hasAdmissionCapacity, RELAY_ADMISSION_BUDGETS } from './admission-budgets.js' +import { + ASSIGNMENT_LIMITS, + AssignmentActivitySchema, + EvacuationCommitSchema, + hasAssignmentActivity, + mayNormallyReassign +} from './assignment-invariants.js' +import { + CONTROL_GENERATION_STATE, + controlLossDisposition, + mayStartNewRelayWork, + preservesRelayAuthIdentity +} from './control-continuity.js' +import { + CredentialInstallIdentitySchema, + decideResumeCommit, + INSTALL_TRANSACTION_CONTRACT, + INSTALL_STATUS, + InviteRecordSchema, + INVITE_STATE, + nextCredentialVersion, + PAIRING_RECOVERY_MATRIX, + transitionInvite, + TokenFamilySchema +} from './persistence-invariants.js' + +const HASH = 'abcdefghijklmnopqrstuvwxyzABCDEFGH012345678' + +describe('relay persistence invariants', () => { + it('admits pre-auth sockets only outside every reserved budget', () => { + expect(RELAY_ADMISSION_BUDGETS.maxProcessQueuedBytes).toBe(64 * 1024 * 1024) + expect(hasAdmissionCapacity({ totalRequests: 649, preAuthConnections: 44, sourcePreAuthConnections: 3 })).toBe(true) + expect(hasAdmissionCapacity({ totalRequests: 650, preAuthConnections: 0, sourcePreAuthConnections: 0 })).toBe(false) + expect(hasAdmissionCapacity({ totalRequests: 0, preAuthConnections: 45, sourcePreAuthConnections: 0 })).toBe(false) + expect(hasAdmissionCapacity({ totalRequests: 0, preAuthConnections: 0, sourcePreAuthConnections: 4 })).toBe(false) + expect(hasAdmissionCapacity({ + totalRequests: 2_899, + preAuthConnections: 0, + sourcePreAuthConnections: 0, + totalRequestCeiling: 2_900 + })).toBe(true) + expect(hasAdmissionCapacity({ + totalRequests: 2_900, + preAuthConnections: 0, + sourcePreAuthConnections: 0, + totalRequestCeiling: 2_900 + })).toBe(false) + }) + + it('represents persisted invite leases without an intermediate install status', () => { + expect(Object.values(INSTALL_STATUS)).toEqual(['not-found', 'committed']) + expect(INSTALL_TRANSACTION_CONTRACT.atomicEffects).toEqual([ + 'idempotency-result', 'token-family', 'affected-invites' + ]) + expect(nextCredentialVersion(undefined)).toBe(1) + expect(nextCredentialVersion(3)).toBe(4) + expect( + InviteRecordSchema.safeParse({ + relayDeviceId: 'device-1', tokenHash: HASH, state: INVITE_STATE.RESERVED, + attemptCount: 1, maxAttempts: 5, expiresAt: 100, reservationId: 'reservation-1', + reservationExpiresAt: 50 + }).success + ).toBe(true) + expect( + InviteRecordSchema.safeParse({ + relayDeviceId: 'device-1', tokenHash: HASH, state: INVITE_STATE.RESERVED, + attemptCount: 1, maxAttempts: 5, expiresAt: 100 + }).success + ).toBe(false) + expect( + CredentialInstallIdentitySchema.safeParse({ + userId: 'user-1', relayHostId: 'abcdefghijklmnop', relayDeviceId: 'device-1', reqId: 'req-1' + }).success + ).toBe(true) + }) + + it('leases, cools down, rolls back, consumes, invalidates, and cleans invites durably', () => { + const available = InviteRecordSchema.parse({ + relayDeviceId: 'device-1', tokenHash: HASH, state: INVITE_STATE.AVAILABLE, + attemptCount: 0, maxAttempts: 2, expiresAt: 1_000 + }) + const reserved = transitionInvite( + available, + { type: 'reserve', reservationId: 'reservation-1', reservationExpiresAt: 50 }, + 10 + ) + expect(reserved.state).toBe(INVITE_STATE.RESERVED) + expect(transitionInvite(reserved, { type: 'transaction-rollback' }, 20)).toBe(reserved) + const cooldown = transitionInvite(reserved, { type: 'lease-expired', cooldownUntil: 60 }, 50) + expect(cooldown.state).toBe(INVITE_STATE.COOLDOWN) + const second = transitionInvite( + cooldown, + { type: 'reserve', reservationId: 'reservation-2', reservationExpiresAt: 80 }, + 60 + ) + expect(transitionInvite(second, { type: 'attach-failed', cooldownUntil: 90 }, 70).state).toBe( + INVITE_STATE.INVALIDATED + ) + expect(transitionInvite(reserved, { type: 'provision-committed' }, 20).state).toBe( + INVITE_STATE.CONSUMED + ) + expect(transitionInvite(available, { type: 'direct-install-committed' }, 20).state).toBe( + INVITE_STATE.INVALIDATED + ) + expect(transitionInvite(available, { type: 'cleanup' }, 1_000).state).toBe(INVITE_STATE.EXPIRED) + expect(PAIRING_RECOVERY_MATRIX).toHaveLength(4) + }) + + it('makes commit-time current, grace, retired, expired, and revoked outcomes exact', () => { + const family = TokenFamilySchema.parse({ + currentVersion: 3, currentHash: HASH, currentExpiresAt: 200, + graceVersion: 2, graceHash: HASH, graceExpiresAt: 150 + }) + expect(decideResumeCommit(family, 3, 100)).toBe('renew-current') + expect(decideResumeCommit(family, 2, 100)).toBe('return-unchanged-grace') + expect(decideResumeCommit(family, 1, 100)).toBe('reject-retired') + expect(decideResumeCommit(family, 2, 151)).toBe('reject-expired') + expect(decideResumeCommit({ ...family, revokedAt: 90 }, 3, 100)).toBe('reject-revoked') + }) + + it('distinguishes same-process rebind, replacement fencing, and drain-only work', () => { + expect(controlLossDisposition({ matchingResumeSecret: true, competingGeneration: true })).toBe('rebind') + expect(controlLossDisposition({ matchingResumeSecret: false, competingGeneration: true })).toBe('fence-old') + expect(controlLossDisposition({ matchingResumeSecret: false, competingGeneration: false })).toBe('orphan-grace') + expect(mayStartNewRelayWork(CONTROL_GENERATION_STATE.DRAIN_ONLY, false)).toBe(false) + expect(mayStartNewRelayWork(CONTROL_GENERATION_STATE.ACTIVE, true)).toBe(false) + const identity = { sub: 'u', prof: 'p', org: 'o', relayHostId: 'abcdefghijklmnop' } + expect(preservesRelayAuthIdentity(identity, identity)).toBe(true) + expect(preservesRelayAuthIdentity(identity, { ...identity, org: 'other' })).toBe(false) + }) + + it('counts every durable activity class before normal reassignment', () => { + const record = AssignmentActivitySchema.parse({ + relayHostId: 'abcdefghijklmnop', cellId: 'cell-1', assignmentEpoch: 1, + leaseExpiresAt: 100, lastActivityAt: 100, reservedControls: 0, reservedSplices: 0, + reservedInvites: 0, pendingInstalls: 0, pendingConfirmations: 0, migrationLeases: 1 + }) + expect(hasAssignmentActivity(record)).toBe(true) + expect(mayNormallyReassign(record, 100 + ASSIGNMENT_LIMITS.dormantTtlMs)).toBe(false) + expect(mayNormallyReassign({ ...record, migrationLeases: 0 }, 100 + ASSIGNMENT_LIMITS.dormantTtlMs)).toBe(true) + expect( + EvacuationCommitSchema.safeParse({ + relayHostId: 'abcdefghijklmnop', sourceCellId: 'cell-1', targetCellId: 'cell-2', + previousEpoch: 1, assignmentEpoch: 2, targetCapacityReserved: true + }).success + ).toBe(true) + expect( + EvacuationCommitSchema.safeParse({ + relayHostId: 'abcdefghijklmnop', sourceCellId: 'cell-1', targetCellId: 'cell-2', + previousEpoch: 1, assignmentEpoch: 3, targetCapacityReserved: true + }).success + ).toBe(false) + }) +}) diff --git a/cloud/packages/relay-contract/src/persistence-invariants.ts b/cloud/packages/relay-contract/src/persistence-invariants.ts new file mode 100644 index 00000000000..0f07250d008 --- /dev/null +++ b/cloud/packages/relay-contract/src/persistence-invariants.ts @@ -0,0 +1,172 @@ +import { z } from 'zod' +import { Base64Url32ByteSchema, EpochMsSchema, OpaqueIdSchema, RelayHostIdSchema } from './wire-scalars.js' + +export const INVITE_STATE = { + AVAILABLE: 'available', + RESERVED: 'reserved', + COOLDOWN: 'cooldown', + CONSUMED: 'consumed', + EXPIRED: 'expired', + INVALIDATED: 'invalidated' +} as const + +export const InviteRecordSchema = z + .object({ + relayDeviceId: OpaqueIdSchema, + tokenHash: Base64Url32ByteSchema, + state: z.nativeEnum(INVITE_STATE), + attemptCount: z.number().int().nonnegative(), + maxAttempts: z.number().int().positive(), + expiresAt: EpochMsSchema, + reservationId: OpaqueIdSchema.optional(), + reservationExpiresAt: EpochMsSchema.optional(), + cooldownUntil: EpochMsSchema.optional() + }) + .strict() + .superRefine((invite, context) => { + if ( + invite.state === INVITE_STATE.RESERVED && + (!invite.reservationId || !invite.reservationExpiresAt) + ) { + context.addIssue({ code: z.ZodIssueCode.custom, message: 'reserved invite requires a lease' }) + } + if (invite.state === INVITE_STATE.COOLDOWN && invite.cooldownUntil === undefined) { + context.addIssue({ + code: z.ZodIssueCode.custom, + message: 'cooldown invite requires cooldownUntil' + }) + } + }) + +export type InviteRecord = z.infer + +export type InviteEvent = + | { type: 'reserve'; reservationId: string; reservationExpiresAt: number } + | { type: 'attach-failed'; cooldownUntil: number } + | { type: 'lease-expired'; cooldownUntil: number } + | { type: 'provision-committed' } + | { type: 'direct-install-committed' } + | { type: 'transaction-rollback' } + | { type: 'cleanup' } + +function withoutLease(invite: InviteRecord): InviteRecord { + const { reservationId: _reservationId, reservationExpiresAt: _reservationExpiresAt, ...rest } = invite + return rest +} + +export function transitionInvite(invite: InviteRecord, event: InviteEvent, now: number): InviteRecord { + if (event.type === 'transaction-rollback') return invite + if (event.type === 'direct-install-committed') { + return { ...withoutLease(invite), state: INVITE_STATE.INVALIDATED } + } + if (event.type === 'provision-committed') { + return { ...withoutLease(invite), state: INVITE_STATE.CONSUMED } + } + if (event.type === 'cleanup') { + if (now >= invite.expiresAt) return { ...withoutLease(invite), state: INVITE_STATE.EXPIRED } + if ( + invite.state === INVITE_STATE.RESERVED && + invite.reservationExpiresAt !== undefined && + now >= invite.reservationExpiresAt + ) { + const state = + invite.attemptCount >= invite.maxAttempts ? INVITE_STATE.INVALIDATED : INVITE_STATE.COOLDOWN + return { ...withoutLease(invite), state, cooldownUntil: now } + } + return invite + } + if (event.type === 'reserve') { + const available = + invite.state === INVITE_STATE.AVAILABLE || + (invite.state === INVITE_STATE.COOLDOWN && (invite.cooldownUntil ?? 0) <= now) + if (!available || now >= invite.expiresAt || invite.attemptCount >= invite.maxAttempts) return invite + return { + ...invite, + state: INVITE_STATE.RESERVED, + attemptCount: invite.attemptCount + 1, + reservationId: event.reservationId, + reservationExpiresAt: event.reservationExpiresAt, + cooldownUntil: undefined + } + } + if (invite.state !== INVITE_STATE.RESERVED) return invite + const state = + invite.attemptCount >= invite.maxAttempts ? INVITE_STATE.INVALIDATED : INVITE_STATE.COOLDOWN + return { ...withoutLease(invite), state, cooldownUntil: event.cooldownUntil } +} + +export const PAIRING_RECOVERY_MATRIX = [ + 'direct-commit-ack', + 'direct-commit-response-lost', + 'direct-no-commit-invite-fallback', + 'late-direct-versus-invite-global-key' +] as const + +export const CredentialInstallIdentitySchema = z + .object({ + userId: OpaqueIdSchema, + relayHostId: RelayHostIdSchema, + relayDeviceId: OpaqueIdSchema, + reqId: OpaqueIdSchema + }) + .strict() + +export const INSTALL_STATUS = { + NOT_FOUND: 'not-found', + COMMITTED: 'committed' +} as const + +export const INSTALL_TRANSACTION_CONTRACT = { + idempotencyKeyFields: ['userId', 'relayHostId', 'relayDeviceId', 'reqId'], + transactionDeadlineMs: 10 * 1000, + lockAttemptTimeoutMs: 2 * 1000, + maxDeadlockRetries: 3, + atomicEffects: ['idempotency-result', 'token-family', 'affected-invites'] +} as const + +export function nextCredentialVersion(currentVersion: number | undefined): number { + return currentVersion === undefined ? 1 : currentVersion + 1 +} + +export const TokenFamilySchema = z + .object({ + currentVersion: z.number().int().positive(), + currentHash: Base64Url32ByteSchema, + currentExpiresAt: EpochMsSchema, + graceVersion: z.number().int().positive().optional(), + graceHash: Base64Url32ByteSchema.optional(), + graceExpiresAt: EpochMsSchema.optional(), + revokedAt: EpochMsSchema.optional() + }) + .strict() + .superRefine((family, context) => { + const graceFields = [family.graceVersion, family.graceHash, family.graceExpiresAt] + if (graceFields.some((value) => value !== undefined) && graceFields.some((value) => value === undefined)) { + context.addIssue({ code: z.ZodIssueCode.custom, message: 'grace fields must be all present or absent' }) + } + if (family.graceVersion !== undefined && family.graceVersion >= family.currentVersion) { + context.addIssue({ code: z.ZodIssueCode.custom, message: 'grace version must precede current' }) + } + }) + +export type ResumeCommitDecision = + | 'renew-current' + | 'return-unchanged-grace' + | 'reject-retired' + | 'reject-expired' + | 'reject-revoked' + +export function decideResumeCommit( + family: z.infer, + acceptedVersion: number, + now: number +): ResumeCommitDecision { + if (family.revokedAt !== undefined && family.revokedAt <= now) return 'reject-revoked' + if (acceptedVersion === family.currentVersion) { + return family.currentExpiresAt > now ? 'renew-current' : 'reject-expired' + } + if (acceptedVersion === family.graceVersion) { + return (family.graceExpiresAt ?? 0) > now ? 'return-unchanged-grace' : 'reject-expired' + } + return 'reject-retired' +} diff --git a/cloud/packages/relay-contract/src/protocol-limits.ts b/cloud/packages/relay-contract/src/protocol-limits.ts new file mode 100644 index 00000000000..6257f9cecb7 --- /dev/null +++ b/cloud/packages/relay-contract/src/protocol-limits.ts @@ -0,0 +1,23 @@ +export const RELAY_PROTOCOL_LIMITS = { + firstFrameDeadlineMs: 2_000, + maxHttpBodyBytes: 4 * 1024, + // Why: the splice is an opaque E2EE stream; the desktop's worktree catalog + // response already exceeds 1MiB on large workspaces (~775KiB at 415 + // worktrees, growing), and an oversized frame kills the session on every + // reconnect. 8MiB buys years of headroom; catalog pagination is the + // long-term fix on the desktop side. + maxFrameBytes: 8 * 1024 * 1024, + maxConnectionsPerHost: 8, + idleTimeoutMs: 10 * 60 * 1000, + inviteTtlMs: 10 * 60 * 1000, + inviteMaxAttempts: 5, + inviteReservationLeaseMs: 15 * 1000, + inviteAttemptCooldownMs: 2 * 1000, + hostAttachDeadlineMs: 10 * 1000, + resumeConfirmationDeadlineMs: 30 * 1000, + resumeTtlMs: 30 * 24 * 60 * 60 * 1000, + relayTokenTtlMs: 5 * 60 * 1000, + expiredAuthExistingSpliceGraceMs: 60 * 1000, + controlPingIntervalMs: 15 * 1000, + controlSilenceTimeoutMs: 75 * 1000 +} as const diff --git a/cloud/packages/relay-contract/src/relay-regions.ts b/cloud/packages/relay-contract/src/relay-regions.ts new file mode 100644 index 00000000000..38ac36cd738 --- /dev/null +++ b/cloud/packages/relay-contract/src/relay-regions.ts @@ -0,0 +1,62 @@ +import { z } from 'zod' + +export const RELAY_REGIONS = ['us-central1', 'asia-east2'] as const + +export const RelayRegionSchema = z.enum(RELAY_REGIONS) + +export type RelayRegion = z.infer + +export const RELAY_DEFAULT_REGION: RelayRegion = 'us-central1' + +const RelayProbeOriginSchema = z.string().url().max(2_048).refine(isCanonicalHttpsOrigin) + +export const RelayRegionCatalogResponseSchema = z + .object({ + v: z.literal(1), + regions: z + .array( + z + .object({ + region: RelayRegionSchema, + probeOrigins: z.array(RelayProbeOriginSchema).min(1).max(2) + }) + .strict() + ) + .max(RELAY_REGIONS.length) + }) + .strict() + .superRefine((catalog, context) => { + const regions = new Set() + const origins = new Set() + for (const [regionIndex, entry] of catalog.regions.entries()) { + if (regions.has(entry.region)) { + context.addIssue({ + code: 'custom', + message: 'duplicate relay region', + path: ['regions', regionIndex, 'region'] + }) + } + regions.add(entry.region) + for (const [originIndex, origin] of entry.probeOrigins.entries()) { + if (origins.has(origin)) { + context.addIssue({ + code: 'custom', + message: 'duplicate relay probe origin', + path: ['regions', regionIndex, 'probeOrigins', originIndex] + }) + } + origins.add(origin) + } + } + }) + +export type RelayRegionCatalogResponse = z.infer + +function isCanonicalHttpsOrigin(value: string): boolean { + try { + const url = new URL(value) + return url.protocol === 'https:' && url.origin === value + } catch { + return false + } +} diff --git a/cloud/packages/relay-contract/src/resume-confirmation-contract.ts b/cloud/packages/relay-contract/src/resume-confirmation-contract.ts new file mode 100644 index 00000000000..0f429fc9e81 --- /dev/null +++ b/cloud/packages/relay-contract/src/resume-confirmation-contract.ts @@ -0,0 +1,23 @@ +import { z } from 'zod' +import { EpochMsSchema, GenerationSchema, OpaqueIdSchema } from './wire-scalars.js' + +export const ConfirmableResumeTupleSchema = z + .object({ + basisConnId: OpaqueIdSchema, + owningControlGeneration: GenerationSchema, + relayDeviceId: OpaqueIdSchema, + acceptedCredentialVersion: z.number().int().positive(), + acceptedAs: z.enum(['current', 'grace']), + confirmDeadline: EpochMsSchema + }) + .strict() + +export const RESUME_CONFIRMATION_COMMIT_OUTCOME = { + RENEW_CURRENT: 'renew-current', + RETURN_UNCHANGED_GRACE: 'return-unchanged-grace', + REJECT_RETIRED: 'reject-retired', + REJECT_EXPIRED: 'reject-expired', + REJECT_REVOKED: 'reject-revoked' +} as const + +export type ConfirmableResumeTuple = z.infer diff --git a/cloud/packages/relay-contract/src/splice-state-machine.ts b/cloud/packages/relay-contract/src/splice-state-machine.ts new file mode 100644 index 00000000000..3944df74584 --- /dev/null +++ b/cloud/packages/relay-contract/src/splice-state-machine.ts @@ -0,0 +1,34 @@ +export const SPLICE_STATE = { + PRE_AUTH_ADMITTED: 'pre-auth-admitted', + CREDENTIAL_LEASE_RESERVED: 'credential-lease-reserved', + HOST_NOTIFIED: 'host-notified', + ATTACH_PENDING: 'attach-pending', + HOST_ATTACHED: 'host-attached', + CLIENT_ACKNOWLEDGED: 'client-acknowledged', + SPLICED: 'spliced', + E2EE_CONFIRMABLE: 'e2ee-confirmable', + TEARDOWN: 'teardown' +} as const + +export type SpliceState = (typeof SPLICE_STATE)[keyof typeof SPLICE_STATE] + +export const SPLICE_FORWARD_TRANSITIONS: Readonly> = { + [SPLICE_STATE.PRE_AUTH_ADMITTED]: [SPLICE_STATE.CREDENTIAL_LEASE_RESERVED, SPLICE_STATE.TEARDOWN], + [SPLICE_STATE.CREDENTIAL_LEASE_RESERVED]: [SPLICE_STATE.HOST_NOTIFIED, SPLICE_STATE.TEARDOWN], + [SPLICE_STATE.HOST_NOTIFIED]: [SPLICE_STATE.ATTACH_PENDING, SPLICE_STATE.TEARDOWN], + [SPLICE_STATE.ATTACH_PENDING]: [SPLICE_STATE.HOST_ATTACHED, SPLICE_STATE.TEARDOWN], + [SPLICE_STATE.HOST_ATTACHED]: [SPLICE_STATE.CLIENT_ACKNOWLEDGED, SPLICE_STATE.TEARDOWN], + [SPLICE_STATE.CLIENT_ACKNOWLEDGED]: [SPLICE_STATE.SPLICED, SPLICE_STATE.TEARDOWN], + [SPLICE_STATE.SPLICED]: [SPLICE_STATE.E2EE_CONFIRMABLE, SPLICE_STATE.TEARDOWN], + [SPLICE_STATE.E2EE_CONFIRMABLE]: [SPLICE_STATE.TEARDOWN], + [SPLICE_STATE.TEARDOWN]: [] +} + +export function canAdvanceSplice(from: SpliceState, to: SpliceState): boolean { + return SPLICE_FORWARD_TRANSITIONS[from].includes(to) +} + +export function mayAcknowledgeClient(state: SpliceState, forwardingHandlersInstalled: boolean): boolean { + // Why: success before both forwarding handlers exist can strand a client on a fake splice. + return state === SPLICE_STATE.HOST_ATTACHED && forwardingHandlersInstalled +} diff --git a/cloud/packages/relay-contract/src/wire-scalars.ts b/cloud/packages/relay-contract/src/wire-scalars.ts new file mode 100644 index 00000000000..27dd3a8b30f --- /dev/null +++ b/cloud/packages/relay-contract/src/wire-scalars.ts @@ -0,0 +1,20 @@ +import { z } from 'zod' + +export const Base64Url32ByteSchema = z.string().regex(/^[A-Za-z0-9_-]{43}$/) +export const Base64Url24ByteSchema = z.string().regex(/^[A-Za-z0-9_-]{32}$/) +export const Base6432ByteSchema = z.string().regex(/^(?:[A-Za-z0-9+/]{4}){10}[A-Za-z0-9+/]{3}=$/) +export const Base64Raw24ByteSchema = z.string().regex(/^(?:[A-Za-z0-9+/]{4}){8}$/) +export const RelayHostIdSchema = z.string().regex(/^[A-Za-z0-9_-]{16}$/) +export const OpaqueIdSchema = z.string().min(1).max(128) +export const EpochMsSchema = z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER) +export const GenerationSchema = z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER) +export const PositiveDurationMsSchema = z.number().int().positive().max(24 * 60 * 60 * 1000) + +export const CanonicalHttpsOriginSchema = z.string().max(2048).refine((value) => { + try { + const url = new URL(value) + return url.protocol === 'https:' && url.origin === value && url.pathname === '/' + } catch { + return false + } +}, 'must be a canonical HTTPS origin') diff --git a/cloud/packages/relay-contract/tsconfig.build.json b/cloud/packages/relay-contract/tsconfig.build.json new file mode 100644 index 00000000000..94c84b60803 --- /dev/null +++ b/cloud/packages/relay-contract/tsconfig.build.json @@ -0,0 +1,11 @@ +{ + "extends": "./tsconfig.json", + "compilerOptions": { + "declaration": true, + "emitDeclarationOnly": false, + "noEmit": false, + "outDir": "dist", + "rootDir": "src" + }, + "exclude": ["src/**/*.test.ts"] +} diff --git a/cloud/packages/relay-contract/tsconfig.json b/cloud/packages/relay-contract/tsconfig.json new file mode 100644 index 00000000000..a552e34dbe9 --- /dev/null +++ b/cloud/packages/relay-contract/tsconfig.json @@ -0,0 +1,5 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { "noEmit": true }, + "include": ["src/**/*.ts"] +} diff --git a/cloud/pnpm-lock.yaml b/cloud/pnpm-lock.yaml new file mode 100644 index 00000000000..2dc8314c93d --- /dev/null +++ b/cloud/pnpm-lock.yaml @@ -0,0 +1,1286 @@ +lockfileVersion: '9.0' + +settings: + autoInstallPeers: true + excludeLinksFromLockfile: false + +importers: + + .: + devDependencies: + '@types/node': + specifier: ^24.10.0 + version: 24.13.2 + tsx: + specifier: ^4.21.0 + version: 4.22.4 + typescript: + specifier: ^5.9.3 + version: 5.9.3 + vitest: + specifier: ^4.0.8 + version: 4.1.9(@types/node@24.13.2)(vite@8.0.16(@types/node@24.13.2)(esbuild@0.28.1)(tsx@4.22.4)) + + apps/relay: + dependencies: + '@hono/node-server': + specifier: ^1.19.14 + version: 1.19.14(hono@4.12.27) + '@orca-cloud/relay-contract': + specifier: workspace:* + version: link:../../packages/relay-contract + hono: + specifier: ^4.12.27 + version: 4.12.27 + jose: + specifier: ^6.1.3 + version: 6.2.3 + pg: + specifier: ^8.22.0 + version: 8.22.0 + tweetnacl: + specifier: ^1.0.3 + version: 1.0.3 + ws: + specifier: ^8.18.3 + version: 8.21.0 + zod: + specifier: ^3.25.76 + version: 3.25.76 + devDependencies: + '@types/node': + specifier: ^24.10.0 + version: 24.13.2 + '@types/pg': + specifier: ^8.20.0 + version: 8.20.0 + '@types/ws': + specifier: ^8.18.1 + version: 8.18.1 + tsx: + specifier: ^4.21.0 + version: 4.22.4 + typescript: + specifier: ^5.9.3 + version: 5.9.3 + vitest: + specifier: ^4.0.8 + version: 4.1.9(@types/node@24.13.2)(vite@8.0.16(@types/node@24.13.2)(esbuild@0.28.1)(tsx@4.22.4)) + + packages/relay-contract: + dependencies: + zod: + specifier: ^3.25.76 + version: 3.25.76 + devDependencies: + '@types/node': + specifier: ^24.10.0 + version: 24.13.2 + typescript: + specifier: ^5.9.3 + version: 5.9.3 + vitest: + specifier: ^4.0.8 + version: 4.1.9(@types/node@24.13.2)(vite@8.0.16(@types/node@24.13.2)(esbuild@0.28.1)(tsx@4.22.4)) + +packages: + + '@emnapi/core@1.10.0': + resolution: {integrity: sha512-yq6OkJ4p82CAfPl0u9mQebQHKPJkY7WrIuk205cTYnYe+k2Z8YBh11FrbRG/H6ihirqcacOgl2BIO8oyMQLeXw==} + + '@emnapi/runtime@1.10.0': + resolution: {integrity: sha512-ewvYlk86xUoGI0zQRNq/mC+16R1QeDlKQy21Ki3oSYXNgLb45GV1P6A0M+/s6nyCuNDqe5VpaY84BzXGwVbwFA==} + + '@emnapi/wasi-threads@1.2.1': + resolution: {integrity: sha512-uTII7OYF+/Mes/MrcIOYp5yOtSMLBWSIoLPpcgwipoiKbli6k322tcoFsxoIIxPDqW01SQGAgko4EzZi2BNv2w==} + + '@esbuild/aix-ppc64@0.28.1': + resolution: {integrity: sha512-Svl7tq8k/08+p6CXPpRjQ1fKX+1odH/BQbb48fV6fj3CWHhsoIOoY87w1oHXm0qEpkIK3ZfVgp0hed3XBXzXMQ==} + engines: {node: '>=18'} + cpu: [ppc64] + os: [aix] + + '@esbuild/android-arm64@0.28.1': + resolution: {integrity: sha512-34EGEbCIAgosYz6goLcopX6Mo7NyGv9tfwEM2/7Ce2VcVRk568iSvniGWcUXIy7wEDR1wzolcxcriFVrWYcwBg==} + engines: {node: '>=18'} + cpu: [arm64] + os: [android] + + '@esbuild/android-arm@0.28.1': + resolution: {integrity: sha512-0k2F129Xdio1TdJfzJ8sy1Q47vUD2NnwdhiAf7drUN1EBTfPf4hsFCtmMgu/6m8JSzsBrlmVjudMBQqOfG8usQ==} + engines: {node: '>=18'} + cpu: [arm] + os: [android] + + '@esbuild/android-x64@0.28.1': + resolution: {integrity: sha512-dbwY7ltSMDWsRatcRpCnES4F+im88OCUgGZjy52shC7GqHRE/cYlxNbB4Z4UpJswpcc4Qxd2oE/ufM0p61IKng==} + engines: {node: '>=18'} + cpu: [x64] + os: [android] + + '@esbuild/darwin-arm64@0.28.1': + resolution: {integrity: sha512-TZbWkQY7kvTAXbXUT7uVACR5cMHsDiSz9z7ZKAX/RTq/WJEk3QyRr0wZpNhBDX+/0CtdqUIJlOiodQcta6tY3Q==} + engines: {node: '>=18'} + cpu: [arm64] + os: [darwin] + + '@esbuild/darwin-x64@0.28.1': + resolution: {integrity: sha512-zfdzgK9ACBNZLI/CyHTOx81SyNbM6YXn7rxSgX97VjyiPl9W1i4Ka4fgKECEoFCKGpvBj5qArWIGgQjOwkgskQ==} + engines: {node: '>=18'} + cpu: [x64] + os: [darwin] + + '@esbuild/freebsd-arm64@0.28.1': + resolution: {integrity: sha512-wG2EA8ENdEI0qhkSZMjfqrdY+ziCYCPMmtZjjIwOmXFjmyzEHn+UUxk5of+SYsjtfs3VpnlC7QLzSI5hY/rOAw==} + engines: {node: '>=18'} + cpu: [arm64] + os: [freebsd] + + '@esbuild/freebsd-x64@0.28.1': + resolution: {integrity: sha512-i7dZ9vQgnvSCzi/rYCXNgtF/U+eKZNJBzu3eTQbRgHnM7tNSizLOkRFAl3qzVc/Op/u5YkHHa4pf/3DOYHthLQ==} + engines: {node: '>=18'} + cpu: [x64] + os: [freebsd] + + '@esbuild/linux-arm64@0.28.1': + resolution: {integrity: sha512-yHs+0uc8+nvEAfAfxrWQKK5peSNzBc4PegcMO0EJ2hT71uA7vB8Ihg2e77R2P7SG5uYjPbHlLLmve4LLLRCf0g==} + engines: {node: '>=18'} + cpu: [arm64] + os: [linux] + + '@esbuild/linux-arm@0.28.1': + resolution: {integrity: sha512-qVXBOHQS+d5Y722GwJzJUtOLlX7km3CraOaGormF1pDtPd2C/l1SHRPgjLunLGe51Sh5YYWKMFDyV4SxgMQYTQ==} + engines: {node: '>=18'} + cpu: [arm] + os: [linux] + + '@esbuild/linux-ia32@0.28.1': + resolution: {integrity: sha512-d1z4ZuP0ajrfz/FhGT4vv278rX8KnPPJx8i5+AtK7TYbx9Le9F1hyzurZpkEyjkGa9dUGhQow4C1NmeGvqxN2w==} + engines: {node: '>=18'} + cpu: [ia32] + os: [linux] + + '@esbuild/linux-loong64@0.28.1': + resolution: {integrity: sha512-M5sRjUVZrkm1OAPR3dlOYzNmN+loZKGVi1VUQGrwuqLcbR6qeAz+famMhjASeH3YVKvZz+zT1jlh/keC3Rj/lg==} + engines: {node: '>=18'} + cpu: [loong64] + os: [linux] + + '@esbuild/linux-mips64el@0.28.1': + resolution: {integrity: sha512-mRObBZeHh2OxcBFPWE/FjylkRgZdYuiTR3vaTozquCGOH14iP9oN4x4Ge81CoIDYQrXmIxpFumJBu5MtZpnQJQ==} + engines: {node: '>=18'} + cpu: [mips64el] + os: [linux] + + '@esbuild/linux-ppc64@0.28.1': + resolution: {integrity: sha512-slScBsMAb3GFDcdrCgLwZtPYRoH2H/youv10QiZyRjmsP48fznoveWytSgCI/R0ZcUgpc0ZhIUEx6LHts8yrfQ==} + engines: {node: '>=18'} + cpu: [ppc64] + os: [linux] + + '@esbuild/linux-riscv64@0.28.1': + resolution: {integrity: sha512-kw0owk1o0GFETUJyW0jc0G4Yzs0BHZn0JDZ8JRT088vjJYX777BAs1fDGxAC+q831qOs2DTC96mNsG2opdfyyQ==} + engines: {node: '>=18'} + cpu: [riscv64] + os: [linux] + + '@esbuild/linux-s390x@0.28.1': + resolution: {integrity: sha512-/lAIjX8aYFRByhh6L5rYtPEDRqa9de/4V/juOXcta5frjvzXO4/sqEtyytse0g3zZFuWu5cDN0MkLz2qRDD2Ag==} + engines: {node: '>=18'} + cpu: [s390x] + os: [linux] + + '@esbuild/linux-x64@0.28.1': + resolution: {integrity: sha512-u/anNYF2mmVOEDwLtnQ1wOr3EZ9sTNGLWrsYGYwHWzGA3Si84IOkHXlbWTD1NB+9/1lcnweYKO54uhxZydNzfA==} + engines: {node: '>=18'} + cpu: [x64] + os: [linux] + + '@esbuild/netbsd-arm64@0.28.1': + resolution: {integrity: sha512-oks0DYbLwWMmaakTsCb+zL4E+aHRVLom9IJZOAthMQEPiQmydXHkziYEsGYRx0uNV/IjEKGAV941JzH02pflqw==} + engines: {node: '>=18'} + cpu: [arm64] + os: [netbsd] + + '@esbuild/netbsd-x64@0.28.1': + resolution: {integrity: sha512-aeL6lAnN89Hz43Mlh1G8ARasbuoYvSITDEx0tHh5b7jJnHcssqgjy9Yx430GDpmCa6OyrKoS0aNRjKundRizGg==} + engines: {node: '>=18'} + cpu: [x64] + os: [netbsd] + + '@esbuild/openbsd-arm64@0.28.1': + resolution: {integrity: sha512-MEFJe5C3R8pwXdZ5Y21oo6m7ePiS0d9pWucn99O/wvyJZChoIQKrQDxKrGeW8F5+T0okTHesAmDeiHDTIq0V/Q==} + engines: {node: '>=18'} + cpu: [arm64] + os: [openbsd] + + '@esbuild/openbsd-x64@0.28.1': + resolution: {integrity: sha512-i/ZLIOafE0Z8cI/XANJAixoJL/uRAoS2xOA3rb0xN+KK0K177cMAsQYkzHtBrtMXAKuAc7HGgcWiZ/sRC1Nxgw==} + engines: {node: '>=18'} + cpu: [x64] + os: [openbsd] + + '@esbuild/openharmony-arm64@0.28.1': + resolution: {integrity: sha512-ge+Z7EXFNt2BO1oAMsVpiQ8EwndV9i1xXerAeTIK7AtPs3bKFXQM7nlRxDSIUIMeueR1CNXxqztLzdNeReKBJg==} + engines: {node: '>=18'} + cpu: [arm64] + os: [openharmony] + + '@esbuild/sunos-x64@0.28.1': + resolution: {integrity: sha512-BEjgtECkL3vY+SaSQ6nzVfiALUeFxpawyp8Jmf5PtYhf1Ug40N1h/hxlhts+f1FvSvarEigdxS3BlSMI2PJLcQ==} + engines: {node: '>=18'} + cpu: [x64] + os: [sunos] + + '@esbuild/win32-arm64@0.28.1': + resolution: {integrity: sha512-lCv9eK/H6ZJWbE7bh2nw54CZ9M2nupBxJcTsdk/QQnWkdSjKGuxmmH8/GWrlT1eMmZfn4dGcCjRte397WqfQXA==} + engines: {node: '>=18'} + cpu: [arm64] + os: [win32] + + '@esbuild/win32-ia32@0.28.1': + resolution: {integrity: sha512-zvb/mB2bSCoJOpoCBgYKKpX6YM6mJBlBUVUtVj41DlZJVEB6/0CKlRYxP5wWl1C1ILiCoAU5wZZ4q1P3qeS6Eg==} + engines: {node: '>=18'} + cpu: [ia32] + os: [win32] + + '@esbuild/win32-x64@0.28.1': + resolution: {integrity: sha512-bm4Mowrv+GXMlpWX++EcXw/iLyd1o3+bJkC2DkWXYVvgZCqD/bSj9ctZeAMC3cIxgjRVR2Dufaiu4YPxr5gW1A==} + engines: {node: '>=18'} + cpu: [x64] + os: [win32] + + '@hono/node-server@1.19.14': + resolution: {integrity: sha512-GwtvgtXxnWsucXvbQXkRgqksiH2Qed37H9xHZocE5sA3N8O8O8/8FA3uclQXxXVzc9XBZuEOMK7+r02FmSpHtw==} + engines: {node: '>=18.14.1'} + peerDependencies: + hono: ^4 + + '@jridgewell/sourcemap-codec@1.5.5': + resolution: {integrity: sha512-cYQ9310grqxueWbl+WuIUIaiUaDcj7WOq5fVhEljNVgRfOUhY9fy2zTvfoqWsnebh8Sl70VScFbICvJnLKB0Og==} + + '@napi-rs/wasm-runtime@1.1.5': + resolution: {integrity: sha512-AWPoBRJ9tsnVhor4sjO7rkni+7p+2IAEFj6cx06UgP10jkQHqay/36uRV/bFkgrh18D9vb4cr8Q0Pthskgzy+Q==} + peerDependencies: + '@emnapi/core': ^1.7.1 + '@emnapi/runtime': ^1.7.1 + + '@oxc-project/types@0.133.0': + resolution: {integrity: sha512-KzkdCd6Uxqnf6l3HOw1xfatAlUURA0g14cvBYFyJ5SaNOQbOUvBr9PKArcPcrNIeRsBdgcUzOGrhKveVpvOIGA==} + + '@rolldown/binding-android-arm64@1.0.3': + resolution: {integrity: sha512-454rs7jHngixp/NMxd5srYD57OnzSlZ/eFTETjORQHLwJG1lRtmNOJcBerZlfu4GjKqeq8aCCIQrMdHyhI51Hw==} + engines: {node: ^20.19.0 || >=22.12.0} + cpu: [arm64] + os: [android] + + '@rolldown/binding-darwin-arm64@1.0.3': + resolution: {integrity: sha512-PcAhP+ynjURNyy8SKGl5DQP94aGuB/7JrXJb/t7P+hanXvQVMWzUvRRhBAcg/lNRadBhoUPqSoP4xw5tR/KBEA==} + engines: {node: ^20.19.0 || >=22.12.0} + cpu: [arm64] + os: [darwin] + + '@rolldown/binding-darwin-x64@1.0.3': + resolution: {integrity: sha512-9YpfeUvSE2RS7wysJ81uOZkXJz7f7Q55H2Gvp3VEw/EsahqDtrphrZ0EwDLK5vvKOzaCrBsjF8JmnMLcUt78Gg==} + engines: {node: ^20.19.0 || >=22.12.0} + cpu: [x64] + os: [darwin] + + '@rolldown/binding-freebsd-x64@1.0.3': + resolution: {integrity: sha512-yB1IlAsSNHncV6SCTL27/MVGR5htvQsoGxIv5KMGXALp+Ll1wYsn+x98M9MW7qa+NdSbvrrY7ANI4wLJ0n1e6g==} + engines: {node: ^20.19.0 || >=22.12.0} + cpu: [x64] + os: [freebsd] + + '@rolldown/binding-linux-arm-gnueabihf@1.0.3': + resolution: {integrity: sha512-Yi30IVAAfLUCy2MseFjbB1jAMDl1VMCAas5StnYp8da9+CKvMd2H2cbEjWcw5NPaPqzvYkVIaF1nNUG+b7u/sw==} + engines: {node: ^20.19.0 || >=22.12.0} + cpu: [arm] + os: [linux] + + '@rolldown/binding-linux-arm64-gnu@1.0.3': + resolution: {integrity: sha512-jsO7R8To+AdlYgUmN5sHSCZbfhtMBkO0WUx8iORQnPcMMdgr7qM2DQmMwgabs3GhNztdmoKkMKQFHD6DTMCIQw==} + engines: {node: ^20.19.0 || >=22.12.0} + cpu: [arm64] + os: [linux] + + '@rolldown/binding-linux-arm64-musl@1.0.3': + resolution: {integrity: sha512-VWkUHwWriDciit80wleYwKILoR/KMvxh/IdwS/paX+ZgpuRpCrKLUdadJbc0NpBEiyhpYawsJ73j9aCvOH+f7Q==} + engines: {node: ^20.19.0 || >=22.12.0} + cpu: [arm64] + os: [linux] + + '@rolldown/binding-linux-ppc64-gnu@1.0.3': + resolution: {integrity: sha512-5f1laC0SlIR0yDbFCd8acUhvJIag6N3zC5P7oUPN6wX0aOma+uKJ0wBDH5aq7I1PVI2ttTlhJwzwRIBnLiSGEg==} + engines: {node: ^20.19.0 || >=22.12.0} + cpu: [ppc64] + os: [linux] + + '@rolldown/binding-linux-s390x-gnu@1.0.3': + resolution: {integrity: sha512-Iq4ko0r4XsgbrF/LunNgHtAGLRRVE2kXonAXQ/MV0mC6jQpMOhW1SvtZja2EhC/kd05++bP78dsqBeIQyYJ6Yg==} + engines: {node: ^20.19.0 || >=22.12.0} + cpu: [s390x] + os: [linux] + + '@rolldown/binding-linux-x64-gnu@1.0.3': + resolution: {integrity: sha512-B8m6tD5+/N5FeNQFbKlLA/2yVq9ycQP1SeedyEYYKWBNR3ZQbkvIUcNnDNM03lO1l5F2roiiFJGgvoLLyZXtSg==} + engines: {node: ^20.19.0 || >=22.12.0} + cpu: [x64] + os: [linux] + + '@rolldown/binding-linux-x64-musl@1.0.3': + resolution: {integrity: sha512-pSdpdUJHkuCxun9LE7jvgUB9qsRgaiyNNCX7m/AvHTcq67AiT/Yhoxvw5zPfhrM8k/BfP8ce/hMOpthKDpEUow==} + engines: {node: ^20.19.0 || >=22.12.0} + cpu: [x64] + os: [linux] + + '@rolldown/binding-openharmony-arm64@1.0.3': + resolution: {integrity: sha512-OXXS3RKJgX2uLwM+gYyuH5omcH8fL1LJs96pZGgtetVCahON57+d4SJHzTgZiOjxgGkSnpXpOsWuPDGAKAigEg==} + engines: {node: ^20.19.0 || >=22.12.0} + cpu: [arm64] + os: [openharmony] + + '@rolldown/binding-wasm32-wasi@1.0.3': + resolution: {integrity: sha512-JTtb8BWFynicNSoPrehsCzBtOKjZ6jhMiPFEmOiuXg1Fl8dn2KHQob+GuPSGR0dryQa1PQJbzjF3dqO/whhjLg==} + engines: {node: ^20.19.0 || >=22.12.0} + cpu: [wasm32] + + '@rolldown/binding-win32-arm64-msvc@1.0.3': + resolution: {integrity: sha512-gEdFFEN70A/jxb2svrWsN3aDL7OUtmvlOy+6fa2jxG8K0wQ1ZbdeLGnidov6Yu5/733dI5ySfzFlQ/cb0bSz1g==} + engines: {node: ^20.19.0 || >=22.12.0} + cpu: [arm64] + os: [win32] + + '@rolldown/binding-win32-x64-msvc@1.0.3': + resolution: {integrity: sha512-eXB7CHuaQdqmJcc3koCNtNPmT/bj2gc999kUFgBxG8Ac0NdgXc4rkCHhqrgrhN3zddvvvrgzj1e90SuSfmyIXA==} + engines: {node: ^20.19.0 || >=22.12.0} + cpu: [x64] + os: [win32] + + '@rolldown/pluginutils@1.0.1': + resolution: {integrity: sha512-2j9bGt5Jh8hj+vPtgzPtl72j0yRxHAyumoo6TNfAjsLB04UtpSvPbPcDcBMxz7n+9CYB0c1GxQFxYRg2jimqGw==} + + '@standard-schema/spec@1.1.0': + resolution: {integrity: sha512-l2aFy5jALhniG5HgqrD6jXLi/rUWrKvqN/qJx6yoJsgKhblVd+iqqU4RCXavm/jPityDo5TCvKMnpjKnOriy0w==} + + '@tybys/wasm-util@0.10.2': + resolution: {integrity: sha512-RoBvJ2X0wuKlWFIjrwffGw1IqZHKQqzIchKaadZZfnNpsAYp2mM0h36JtPCjNDAHGgYez/15uMBpfGwchhiMgg==} + + '@types/chai@5.2.3': + resolution: {integrity: sha512-Mw558oeA9fFbv65/y4mHtXDs9bPnFMZAL/jxdPFUpOHHIXX91mcgEHbS5Lahr+pwZFR8A7GQleRWeI6cGFC2UA==} + + '@types/deep-eql@4.0.2': + resolution: {integrity: sha512-c9h9dVVMigMPc4bwTvC5dxqtqJZwQPePsWjPlpSOnojbor6pGqdk541lfA7AqFQr5pB1BRdq0juY9db81BwyFw==} + + '@types/estree@1.0.9': + resolution: {integrity: sha512-GhdPgy1el4/ImP05X05Uw4cw2/M93BCUmnEvWZNStlCzEKME4Fkk+YpoA5OiHNQmoS7Cafb8Xa3Pya8m1Qrzeg==} + + '@types/node@24.13.2': + resolution: {integrity: sha512-fRa09kZTgu8o71KFcDjUFuc7F+dEbZYZmkI0mg5YBTRs0yMKjYHsq/c0urDKeDb+D5qVgXOdFcuu+DZPKOITwA==} + + '@types/pg@8.20.0': + resolution: {integrity: sha512-bEPFOaMAHTEP1EzpvHTbmwR8UsFyHSKsRisLIHVMXnpNefSbGA1bD6CVy+qKjGSqmZqNqBDV2azOBo8TgkcVow==} + + '@types/ws@8.18.1': + resolution: {integrity: sha512-ThVF6DCVhA8kUGy+aazFQ4kXQ7E1Ty7A3ypFOe0IcJV8O/M511G99AW24irKrW56Wt44yG9+ij8FaqoBGkuBXg==} + + '@vitest/expect@4.1.9': + resolution: {integrity: sha512-vl/rYsUKcBr3SnQn166+XR5ZQcgMx3DQhFWdfli/cWpLnLUmbxZvyrJZotLFUryib+LtArYMSTJ5RbQ57ZqrlA==} + + '@vitest/mocker@4.1.9': + resolution: {integrity: sha512-EVkXzBjrPGM+cK8/ANWgBrkUCfJfb38/EfTSO8h7pWvKkyPkpWxvR7BkD2MyItMF62C97zAEoqdpUixwR/e+Rw==} + peerDependencies: + msw: ^2.4.9 + vite: ^6.0.0 || ^7.0.0 || ^8.0.0 + peerDependenciesMeta: + msw: + optional: true + vite: + optional: true + + '@vitest/pretty-format@4.1.9': + resolution: {integrity: sha512-s0iufns3iIFitdgm+YR7g1whCAaGtXz459VS9/PqyKDEEFgYIhsHOQmXgIgDuYCt7DeQmiZT0Qe2OA2p4ZPu5A==} + + '@vitest/runner@4.1.9': + resolution: {integrity: sha512-KXLMDtc7oe70+3mJfGrPUWPesswH+3sTxAMAMl8DG7I8IUQT4XW718dY5ID3vPUcmlu27CcKfY4P3h3I29SLJg==} + + '@vitest/snapshot@4.1.9': + resolution: {integrity: sha512-Jc7RKGNBo8Z28WYIm0Niej4xdSPByRf6mU58VpHQkd6Zh05rlnA+twjbK5HyeIGHxrzsc3mJgS43uM0CZKzaIA==} + + '@vitest/spy@4.1.9': + resolution: {integrity: sha512-fHpsS6mIi+PiEW+vcRVOMkX1oSaPKne3VOclSFICPcGOmfKgXPU5iAah+wcNcj2xPrCCmfq99IDGf+EojhhvhA==} + + '@vitest/utils@4.1.9': + resolution: {integrity: sha512-A51o8ymO5PpqlWNnBP9ZHPXDIpuMtTLlGSjN7la4US+LJzoUMyhwjA5QXlm39JexgwHKW4Xjs8Z2d3dLCXOeuA==} + + assertion-error@2.0.1: + resolution: {integrity: sha512-Izi8RQcffqCeNVgFigKli1ssklIbpHnCYc6AknXGYoB6grJqyeby7jv12JUQgmTAnIDnbck1uxksT4dzN3PWBA==} + engines: {node: '>=12'} + + chai@6.2.2: + resolution: {integrity: sha512-NUPRluOfOiTKBKvWPtSD4PhFvWCqOi0BGStNWs57X9js7XGTprSmFoz5F0tWhR4WPjNeR9jXqdC7/UpSJTnlRg==} + engines: {node: '>=18'} + + convert-source-map@2.0.0: + resolution: {integrity: sha512-Kvp459HrV2FEJ1CAsi1Ku+MY3kasH19TFykTz2xWmMeq6bk2NU3XXvfJ+Q61m0xktWwt+1HSYf3JZsTms3aRJg==} + + detect-libc@2.1.2: + resolution: {integrity: sha512-Btj2BOOO83o3WyH59e8MgXsxEQVcarkUOpEYrubB0urwnN10yQ364rsiByU11nZlqWYZm05i/of7io4mzihBtQ==} + engines: {node: '>=8'} + + es-module-lexer@2.1.0: + resolution: {integrity: sha512-n27zTYMjYu1aj4MjCWzSP7G9r75utsaoc8m61weK+W8JMBGGQybd43GstCXZ3WNmSFtGT9wi59qQTW6mhTR5LQ==} + + esbuild@0.28.1: + resolution: {integrity: sha512-HrJrvZv5ayxBzPfwphOoNzkzOIIlifzk0KJrGK2c8R4+LKpMtpYLQeUdjnwjWv/LZlkH2laZk+4w78pi99D4Vw==} + engines: {node: '>=18'} + hasBin: true + + estree-walker@3.0.3: + resolution: {integrity: sha512-7RUKfXgSMMkzt6ZuXmqapOurLGPPfgj6l9uRZ7lRGolvk0y2yocc35LdcxKC5PQZdn2DMqioAQ2NoWcrTKmm6g==} + + expect-type@1.3.0: + resolution: {integrity: sha512-knvyeauYhqjOYvQ66MznSMs83wmHrCycNEN6Ao+2AeYEfxUIkuiVxdEa1qlGEPK+We3n0THiDciYSsCcgW/DoA==} + engines: {node: '>=12.0.0'} + + fdir@6.5.0: + resolution: {integrity: sha512-tIbYtZbucOs0BRGqPJkshJUYdL+SDH7dVM8gjy+ERp3WAUjLEFJE+02kanyHtwjWOnwrKYBiwAmM0p4kLJAnXg==} + engines: {node: '>=12.0.0'} + peerDependencies: + picomatch: ^3 || ^4 + peerDependenciesMeta: + picomatch: + optional: true + + fsevents@2.3.3: + resolution: {integrity: sha512-5xoDfX+fL7faATnagmWPpbFtwh/R77WmMMqqHGS65C3vvB0YHrgF+B1YmZ3441tMj5n63k0212XNoJwzlhffQw==} + engines: {node: ^8.16.0 || ^10.6.0 || >=11.0.0} + os: [darwin] + + hono@4.12.27: + resolution: {integrity: sha512-1yrb/+w6HWQJrUCLkJ2IF5jNIPvvFkblV5RNOYl6bV+OA6p9GLcMpHFFGTosSvHvcAUibuUukRqhlYI4z32C7Q==} + engines: {node: '>=16.9.0'} + + jose@6.2.3: + resolution: {integrity: sha512-YYVDInQKFJfR/xa3ojUTl8c2KoTwiL1R5Wg9YCydwH0x0B9grbzlg5HC7mMjCtUJjbQ/YnGEZIhI5tCgfTb4Hw==} + + lightningcss-android-arm64@1.32.0: + resolution: {integrity: sha512-YK7/ClTt4kAK0vo6w3X+Pnm0D2cf2vPHbhOXdoNti1Ga0al1P4TBZhwjATvjNwLEBCnKvjJc2jQgHXH0NEwlAg==} + engines: {node: '>= 12.0.0'} + cpu: [arm64] + os: [android] + + lightningcss-darwin-arm64@1.32.0: + resolution: {integrity: sha512-RzeG9Ju5bag2Bv1/lwlVJvBE3q6TtXskdZLLCyfg5pt+HLz9BqlICO7LZM7VHNTTn/5PRhHFBSjk5lc4cmscPQ==} + engines: {node: '>= 12.0.0'} + cpu: [arm64] + os: [darwin] + + lightningcss-darwin-x64@1.32.0: + resolution: {integrity: sha512-U+QsBp2m/s2wqpUYT/6wnlagdZbtZdndSmut/NJqlCcMLTWp5muCrID+K5UJ6jqD2BFshejCYXniPDbNh73V8w==} + engines: {node: '>= 12.0.0'} + cpu: [x64] + os: [darwin] + + lightningcss-freebsd-x64@1.32.0: + resolution: {integrity: sha512-JCTigedEksZk3tHTTthnMdVfGf61Fky8Ji2E4YjUTEQX14xiy/lTzXnu1vwiZe3bYe0q+SpsSH/CTeDXK6WHig==} + engines: {node: '>= 12.0.0'} + cpu: [x64] + os: [freebsd] + + lightningcss-linux-arm-gnueabihf@1.32.0: + resolution: {integrity: sha512-x6rnnpRa2GL0zQOkt6rts3YDPzduLpWvwAF6EMhXFVZXD4tPrBkEFqzGowzCsIWsPjqSK+tyNEODUBXeeVHSkw==} + engines: {node: '>= 12.0.0'} + cpu: [arm] + os: [linux] + + lightningcss-linux-arm64-gnu@1.32.0: + resolution: {integrity: sha512-0nnMyoyOLRJXfbMOilaSRcLH3Jw5z9HDNGfT/gwCPgaDjnx0i8w7vBzFLFR1f6CMLKF8gVbebmkUN3fa/kQJpQ==} + engines: {node: '>= 12.0.0'} + cpu: [arm64] + os: [linux] + + lightningcss-linux-arm64-musl@1.32.0: + resolution: {integrity: sha512-UpQkoenr4UJEzgVIYpI80lDFvRmPVg6oqboNHfoH4CQIfNA+HOrZ7Mo7KZP02dC6LjghPQJeBsvXhJod/wnIBg==} + engines: {node: '>= 12.0.0'} + cpu: [arm64] + os: [linux] + + lightningcss-linux-x64-gnu@1.32.0: + resolution: {integrity: sha512-V7Qr52IhZmdKPVr+Vtw8o+WLsQJYCTd8loIfpDaMRWGUZfBOYEJeyJIkqGIDMZPwPx24pUMfwSxxI8phr/MbOA==} + engines: {node: '>= 12.0.0'} + cpu: [x64] + os: [linux] + + lightningcss-linux-x64-musl@1.32.0: + resolution: {integrity: sha512-bYcLp+Vb0awsiXg/80uCRezCYHNg1/l3mt0gzHnWV9XP1W5sKa5/TCdGWaR/zBM2PeF/HbsQv/j2URNOiVuxWg==} + engines: {node: '>= 12.0.0'} + cpu: [x64] + os: [linux] + + lightningcss-win32-arm64-msvc@1.32.0: + resolution: {integrity: sha512-8SbC8BR40pS6baCM8sbtYDSwEVQd4JlFTOlaD3gWGHfThTcABnNDBda6eTZeqbofalIJhFx0qKzgHJmcPTnGdw==} + engines: {node: '>= 12.0.0'} + cpu: [arm64] + os: [win32] + + lightningcss-win32-x64-msvc@1.32.0: + resolution: {integrity: sha512-Amq9B/SoZYdDi1kFrojnoqPLxYhQ4Wo5XiL8EVJrVsB8ARoC1PWW6VGtT0WKCemjy8aC+louJnjS7U18x3b06Q==} + engines: {node: '>= 12.0.0'} + cpu: [x64] + os: [win32] + + lightningcss@1.32.0: + resolution: {integrity: sha512-NXYBzinNrblfraPGyrbPoD19C1h9lfI/1mzgWYvXUTe414Gz/X1FD2XBZSZM7rRTrMA8JL3OtAaGifrIKhQ5yQ==} + engines: {node: '>= 12.0.0'} + + magic-string@0.30.21: + resolution: {integrity: sha512-vd2F4YUyEXKGcLHoq+TEyCjxueSeHnFxyyjNp80yg0XV4vUhnDer/lvvlqM/arB5bXQN5K2/3oinyCRyx8T2CQ==} + + nanoid@3.3.13: + resolution: {integrity: sha512-sPdqC6ByMVVGvF1ynvvMo0/o+oD1VX7DaHhijt1bFgjvBkHBib4t49GoNDhf2NDta4oeUNlaGbSt5K7qjZ955Q==} + engines: {node: ^10 || ^12 || ^13.7 || ^14 || >=15.0.1} + hasBin: true + + obug@2.1.3: + resolution: {integrity: sha512-9miFgM2OFba7hB+pRgvtV84pYTBaoTHohvmIgiRt6dRIzbwEOIaNaP+dIlGs2fNFoB0SeISs0Jz5WFVRid6Xyg==} + engines: {node: '>=12.20.0'} + + pathe@2.0.3: + resolution: {integrity: sha512-WUjGcAqP1gQacoQe+OBJsFA7Ld4DyXuUIjZ5cc75cLHvJ7dtNsTugphxIADwspS+AraAUePCKrSVtPLFj/F88w==} + + pg-cloudflare@1.4.0: + resolution: {integrity: sha512-Vo7z/6rrQYxpNRylp4Tlob2elzbh+N/MOQbxFVWCxS7oEx6jF53GTJFxK2WWpKuBRkmiin4Mt+xofFDjx09R0A==} + + pg-connection-string@2.14.0: + resolution: {integrity: sha512-XwWDGcLRGCXAR8F/AM5bG7Q+A3Wm2s6QeEjlOKZLlH3UYcguiqCWKyWXVag5TLTIjR7oOJUY8kcADaZgWPyLeg==} + + pg-int8@1.0.1: + resolution: {integrity: sha512-WCtabS6t3c8SkpDBUlb1kjOs7l66xsGdKpIPZsg4wR+B3+u9UAum2odSsF9tnvxg80h4ZxLWMy4pRjOsFIqQpw==} + engines: {node: '>=4.0.0'} + + pg-pool@3.14.0: + resolution: {integrity: sha512-gKtPkFdQPU3DksooVLi9LsjZxrsBUZIpa+7aVx+LV5pNh0KzP4Zleud2po+ConrxbuXGBJ6Hfer6hdgpIBpBaw==} + peerDependencies: + pg: '>=8.0' + + pg-protocol@1.15.0: + resolution: {integrity: sha512-cq9sECI5s0+uPUXjbz8ioyPJni6RzsRib0US67i5IoTZKw8fNeYlVE7u8F4dG7vEJJtc5wdD1K189lCCUwqWTQ==} + + pg-types@2.2.0: + resolution: {integrity: sha512-qTAAlrEsl8s4OiEQY69wDvcMIdQN6wdz5ojQiOy6YRMuynxenON0O5oCpJI6lshc6scgAY8qvJ2On/p+CXY0GA==} + engines: {node: '>=4'} + + pg@8.22.0: + resolution: {integrity: sha512-8wih1vVIBMxoUM2oB4soJsD9tDnDpLv4OXBJ+EJzFsvycD+lfyIreC2gGHq78f8jbLLt+bvlPTFdFZfJkOuzAA==} + engines: {node: '>= 16.0.0'} + peerDependencies: + pg-native: '>=3.0.1' + peerDependenciesMeta: + pg-native: + optional: true + + pgpass@1.0.5: + resolution: {integrity: sha512-FdW9r/jQZhSeohs1Z3sI1yxFQNFvMcnmfuj4WBMUTxOrAyLMaTcE1aAMBiTlbMNaXvBCQuVi0R7hd8udDSP7ug==} + + picocolors@1.1.1: + resolution: {integrity: sha512-xceH2snhtb5M9liqDsmEw56le376mTZkEX/jEb/RxNFyegNul7eNslCXP9FDj/Lcu0X8KEyMceP2ntpaHrDEVA==} + + picomatch@4.0.4: + resolution: {integrity: sha512-QP88BAKvMam/3NxH6vj2o21R6MjxZUAd6nlwAS/pnGvN9IVLocLHxGYIzFhg6fUQ+5th6P4dv4eW9jX3DSIj7A==} + engines: {node: '>=12'} + + postcss@8.5.15: + resolution: {integrity: sha512-FfR8sjd4em2T6fb3I2MwAJU7HWVMr9zba+enmQeeWFfCbm+UOC/0X4DS8XtpUTMwWMGbjKYP7xjfNekzyGmB3A==} + engines: {node: ^10 || ^12 || >=14} + + postgres-array@2.0.0: + resolution: {integrity: sha512-VpZrUqU5A69eQyW2c5CA1jtLecCsN2U/bD6VilrFDWq5+5UIEVO7nazS3TEcHf1zuPYO/sqGvUvW62g86RXZuA==} + engines: {node: '>=4'} + + postgres-bytea@1.0.1: + resolution: {integrity: sha512-5+5HqXnsZPE65IJZSMkZtURARZelel2oXUEO8rH83VS/hxH5vv1uHquPg5wZs8yMAfdv971IU+kcPUczi7NVBQ==} + engines: {node: '>=0.10.0'} + + postgres-date@1.0.7: + resolution: {integrity: sha512-suDmjLVQg78nMK2UZ454hAG+OAW+HQPZ6n++TNDUX+L0+uUlLywnoxJKDou51Zm+zTCjrCl0Nq6J9C5hP9vK/Q==} + engines: {node: '>=0.10.0'} + + postgres-interval@1.2.0: + resolution: {integrity: sha512-9ZhXKM/rw350N1ovuWHbGxnGh/SNJ4cnxHiM0rxE4VN41wsg8P8zWn9hv/buK00RP4WvlOyr/RBDiptyxVbkZQ==} + engines: {node: '>=0.10.0'} + + rolldown@1.0.3: + resolution: {integrity: sha512-i00lAJ2ks1BYr7rjNjKC7BcqAS7nVfiT3QX1SI5aY+AFHblCmaUf9OE9dbdzDvW6dJxbi2ZCZiy9v3CcwOiX3g==} + engines: {node: ^20.19.0 || >=22.12.0} + hasBin: true + + siginfo@2.0.0: + resolution: {integrity: sha512-ybx0WO1/8bSBLEWXZvEd7gMW3Sn3JFlW3TvX1nREbDLRNQNaeNN8WK0meBwPdAaOI7TtRRRJn/Es1zhrrCHu7g==} + + source-map-js@1.2.1: + resolution: {integrity: sha512-UXWMKhLOwVKb728IUtQPXxfYU+usdybtUrK/8uGE8CQMvrhOpwvzDBwj0QhSL7MQc7vIsISBG8VQ8+IDQxpfQA==} + engines: {node: '>=0.10.0'} + + split2@4.2.0: + resolution: {integrity: sha512-UcjcJOWknrNkF6PLX83qcHM6KHgVKNkV62Y8a5uYDVv9ydGQVwAHMKqHdJje1VTWpljG0WYpCDhrCdAOYH4TWg==} + engines: {node: '>= 10.x'} + + stackback@0.0.2: + resolution: {integrity: sha512-1XMJE5fQo1jGH6Y/7ebnwPOBEkIEnT4QF32d5R1+VXdXveM0IBMJt8zfaxX1P3QhVwrYe+576+jkANtSS2mBbw==} + + std-env@4.1.0: + resolution: {integrity: sha512-Rq7ybcX2RuC55r9oaPVEW7/xu3tj8u4GeBYHBWCychFtzMIr86A7e3PPEBPT37sHStKX3+TiX/Fr/ACmJLVlLQ==} + + tinybench@2.9.0: + resolution: {integrity: sha512-0+DUvqWMValLmha6lr4kD8iAMK1HzV0/aKnCtWb9v9641TnP/MFb7Pc2bxoxQjTXAErryXVgUOfv2YqNllqGeg==} + + tinyexec@1.2.4: + resolution: {integrity: sha512-SHf/r48b7vOrjve9PxJo3MN5v5yuyjHvdUcrQffT3WXMUfnGmHDVbC4k3sHJaJTgZCwpUplIaAo5ANtMyp3YHg==} + engines: {node: '>=18'} + + tinyglobby@0.2.17: + resolution: {integrity: sha512-wXR/dYpcqKmfWpEdZjiKJOwCNFndD0DMnrW/cYjVGttEkBfVgcLFHoNrlj47mjOVic9yyNu65alsgF4NQyTa2g==} + engines: {node: '>=12.0.0'} + + tinyrainbow@3.1.0: + resolution: {integrity: sha512-Bf+ILmBgretUrdJxzXM0SgXLZ3XfiaUuOj/IKQHuTXip+05Xn+uyEYdVg0kYDipTBcLrCVyUzAPz7QmArb0mmw==} + engines: {node: '>=14.0.0'} + + tslib@2.8.1: + resolution: {integrity: sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w==} + + tsx@4.22.4: + resolution: {integrity: sha512-X8EX+XV4QR5xCsrgxaED954zTDfY8KqlDtskKEL0cHhyS/P8b4IFOvGDQpsC9Q1XnLq915wEfwwY/zzskCtmhg==} + engines: {node: '>=18.0.0'} + hasBin: true + + tweetnacl@1.0.3: + resolution: {integrity: sha512-6rt+RN7aOi1nGMyC4Xa5DdYiukl2UWCbcJft7YhxReBGQD7OAM8Pbxw6YMo4r2diNEA8FEmu32YOn9rhaiE5yw==} + + typescript@5.9.3: + resolution: {integrity: sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==} + engines: {node: '>=14.17'} + hasBin: true + + undici-types@7.18.2: + resolution: {integrity: sha512-AsuCzffGHJybSaRrmr5eHr81mwJU3kjw6M+uprWvCXiNeN9SOGwQ3Jn8jb8m3Z6izVgknn1R0FTCEAP2QrLY/w==} + + vite@8.0.16: + resolution: {integrity: sha512-h9bXPmJichP5fLmVQo3PyaGSDE2n3aPuomeAlVRm0JLmt4rY6zmPKd59HYI4LNW8oTK7tlTsuC7l/m7awx9Jcw==} + engines: {node: ^20.19.0 || >=22.12.0} + hasBin: true + peerDependencies: + '@types/node': ^20.19.0 || >=22.12.0 + '@vitejs/devtools': ^0.1.18 + esbuild: ^0.27.0 || ^0.28.0 + jiti: '>=1.21.0' + less: ^4.0.0 + sass: ^1.70.0 + sass-embedded: ^1.70.0 + stylus: '>=0.54.8' + sugarss: ^5.0.0 + terser: ^5.16.0 + tsx: ^4.8.1 + yaml: ^2.4.2 + peerDependenciesMeta: + '@types/node': + optional: true + '@vitejs/devtools': + optional: true + esbuild: + optional: true + jiti: + optional: true + less: + optional: true + sass: + optional: true + sass-embedded: + optional: true + stylus: + optional: true + sugarss: + optional: true + terser: + optional: true + tsx: + optional: true + yaml: + optional: true + + vitest@4.1.9: + resolution: {integrity: sha512-nE3/LEyc0z87uHYLZebqCUOaJr2hdtuPp7BQ4BosVFnfltxgAvMG08NyrSGlPpOUWvR27c5flSmYFTNr78L9GQ==} + engines: {node: ^20.0.0 || ^22.0.0 || >=24.0.0} + hasBin: true + peerDependencies: + '@edge-runtime/vm': '*' + '@opentelemetry/api': ^1.9.0 + '@types/node': ^20.0.0 || ^22.0.0 || >=24.0.0 + '@vitest/browser-playwright': 4.1.9 + '@vitest/browser-preview': 4.1.9 + '@vitest/browser-webdriverio': 4.1.9 + '@vitest/coverage-istanbul': 4.1.9 + '@vitest/coverage-v8': 4.1.9 + '@vitest/ui': 4.1.9 + happy-dom: '*' + jsdom: '*' + vite: ^6.0.0 || ^7.0.0 || ^8.0.0 + peerDependenciesMeta: + '@edge-runtime/vm': + optional: true + '@opentelemetry/api': + optional: true + '@types/node': + optional: true + '@vitest/browser-playwright': + optional: true + '@vitest/browser-preview': + optional: true + '@vitest/browser-webdriverio': + optional: true + '@vitest/coverage-istanbul': + optional: true + '@vitest/coverage-v8': + optional: true + '@vitest/ui': + optional: true + happy-dom: + optional: true + jsdom: + optional: true + + why-is-node-running@2.3.0: + resolution: {integrity: sha512-hUrmaWBdVDcxvYqnyh09zunKzROWjbZTiNy8dBEjkS7ehEDQibXJ7XvlmtbwuTclUiIyN+CyXQD4Vmko8fNm8w==} + engines: {node: '>=8'} + hasBin: true + + ws@8.21.0: + resolution: {integrity: sha512-Vsp28b7DRcimFQvrqu2Wek3z1iYxDCWqHYB8Qsnk/S4RfaCQzPGPyBNuVjJV3cd6UiKtUtp6sNM77gWvzcCH+g==} + engines: {node: '>=10.0.0'} + peerDependencies: + bufferutil: ^4.0.1 + utf-8-validate: '>=5.0.2' + peerDependenciesMeta: + bufferutil: + optional: true + utf-8-validate: + optional: true + + xtend@4.0.2: + resolution: {integrity: sha512-LKYU1iAXJXUgAXn9URjiu+MWhyUXHsvfp7mcuYm9dSUKK0/CjtrUwFAxD82/mCWbtLsGjFIad0wIsod4zrTAEQ==} + engines: {node: '>=0.4'} + + zod@3.25.76: + resolution: {integrity: sha512-gzUt/qt81nXsFGKIFcC3YnfEAx5NkunCfnDlvuBSSFS02bcXu4Lmea0AFIUwbLWxWPx3d9p8S5QoaujKcNQxcQ==} + +snapshots: + + '@emnapi/core@1.10.0': + dependencies: + '@emnapi/wasi-threads': 1.2.1 + tslib: 2.8.1 + optional: true + + '@emnapi/runtime@1.10.0': + dependencies: + tslib: 2.8.1 + optional: true + + '@emnapi/wasi-threads@1.2.1': + dependencies: + tslib: 2.8.1 + optional: true + + '@esbuild/aix-ppc64@0.28.1': + optional: true + + '@esbuild/android-arm64@0.28.1': + optional: true + + '@esbuild/android-arm@0.28.1': + optional: true + + '@esbuild/android-x64@0.28.1': + optional: true + + '@esbuild/darwin-arm64@0.28.1': + optional: true + + '@esbuild/darwin-x64@0.28.1': + optional: true + + '@esbuild/freebsd-arm64@0.28.1': + optional: true + + '@esbuild/freebsd-x64@0.28.1': + optional: true + + '@esbuild/linux-arm64@0.28.1': + optional: true + + '@esbuild/linux-arm@0.28.1': + optional: true + + '@esbuild/linux-ia32@0.28.1': + optional: true + + '@esbuild/linux-loong64@0.28.1': + optional: true + + '@esbuild/linux-mips64el@0.28.1': + optional: true + + '@esbuild/linux-ppc64@0.28.1': + optional: true + + '@esbuild/linux-riscv64@0.28.1': + optional: true + + '@esbuild/linux-s390x@0.28.1': + optional: true + + '@esbuild/linux-x64@0.28.1': + optional: true + + '@esbuild/netbsd-arm64@0.28.1': + optional: true + + '@esbuild/netbsd-x64@0.28.1': + optional: true + + '@esbuild/openbsd-arm64@0.28.1': + optional: true + + '@esbuild/openbsd-x64@0.28.1': + optional: true + + '@esbuild/openharmony-arm64@0.28.1': + optional: true + + '@esbuild/sunos-x64@0.28.1': + optional: true + + '@esbuild/win32-arm64@0.28.1': + optional: true + + '@esbuild/win32-ia32@0.28.1': + optional: true + + '@esbuild/win32-x64@0.28.1': + optional: true + + '@hono/node-server@1.19.14(hono@4.12.27)': + dependencies: + hono: 4.12.27 + + '@jridgewell/sourcemap-codec@1.5.5': {} + + '@napi-rs/wasm-runtime@1.1.5(@emnapi/core@1.10.0)(@emnapi/runtime@1.10.0)': + dependencies: + '@emnapi/core': 1.10.0 + '@emnapi/runtime': 1.10.0 + '@tybys/wasm-util': 0.10.2 + optional: true + + '@oxc-project/types@0.133.0': {} + + '@rolldown/binding-android-arm64@1.0.3': + optional: true + + '@rolldown/binding-darwin-arm64@1.0.3': + optional: true + + '@rolldown/binding-darwin-x64@1.0.3': + optional: true + + '@rolldown/binding-freebsd-x64@1.0.3': + optional: true + + '@rolldown/binding-linux-arm-gnueabihf@1.0.3': + optional: true + + '@rolldown/binding-linux-arm64-gnu@1.0.3': + optional: true + + '@rolldown/binding-linux-arm64-musl@1.0.3': + optional: true + + '@rolldown/binding-linux-ppc64-gnu@1.0.3': + optional: true + + '@rolldown/binding-linux-s390x-gnu@1.0.3': + optional: true + + '@rolldown/binding-linux-x64-gnu@1.0.3': + optional: true + + '@rolldown/binding-linux-x64-musl@1.0.3': + optional: true + + '@rolldown/binding-openharmony-arm64@1.0.3': + optional: true + + '@rolldown/binding-wasm32-wasi@1.0.3': + dependencies: + '@emnapi/core': 1.10.0 + '@emnapi/runtime': 1.10.0 + '@napi-rs/wasm-runtime': 1.1.5(@emnapi/core@1.10.0)(@emnapi/runtime@1.10.0) + optional: true + + '@rolldown/binding-win32-arm64-msvc@1.0.3': + optional: true + + '@rolldown/binding-win32-x64-msvc@1.0.3': + optional: true + + '@rolldown/pluginutils@1.0.1': {} + + '@standard-schema/spec@1.1.0': {} + + '@tybys/wasm-util@0.10.2': + dependencies: + tslib: 2.8.1 + optional: true + + '@types/chai@5.2.3': + dependencies: + '@types/deep-eql': 4.0.2 + assertion-error: 2.0.1 + + '@types/deep-eql@4.0.2': {} + + '@types/estree@1.0.9': {} + + '@types/node@24.13.2': + dependencies: + undici-types: 7.18.2 + + '@types/pg@8.20.0': + dependencies: + '@types/node': 24.13.2 + pg-protocol: 1.15.0 + pg-types: 2.2.0 + + '@types/ws@8.18.1': + dependencies: + '@types/node': 24.13.2 + + '@vitest/expect@4.1.9': + dependencies: + '@standard-schema/spec': 1.1.0 + '@types/chai': 5.2.3 + '@vitest/spy': 4.1.9 + '@vitest/utils': 4.1.9 + chai: 6.2.2 + tinyrainbow: 3.1.0 + + '@vitest/mocker@4.1.9(vite@8.0.16(@types/node@24.13.2)(esbuild@0.28.1)(tsx@4.22.4))': + dependencies: + '@vitest/spy': 4.1.9 + estree-walker: 3.0.3 + magic-string: 0.30.21 + optionalDependencies: + vite: 8.0.16(@types/node@24.13.2)(esbuild@0.28.1)(tsx@4.22.4) + + '@vitest/pretty-format@4.1.9': + dependencies: + tinyrainbow: 3.1.0 + + '@vitest/runner@4.1.9': + dependencies: + '@vitest/utils': 4.1.9 + pathe: 2.0.3 + + '@vitest/snapshot@4.1.9': + dependencies: + '@vitest/pretty-format': 4.1.9 + '@vitest/utils': 4.1.9 + magic-string: 0.30.21 + pathe: 2.0.3 + + '@vitest/spy@4.1.9': {} + + '@vitest/utils@4.1.9': + dependencies: + '@vitest/pretty-format': 4.1.9 + convert-source-map: 2.0.0 + tinyrainbow: 3.1.0 + + assertion-error@2.0.1: {} + + chai@6.2.2: {} + + convert-source-map@2.0.0: {} + + detect-libc@2.1.2: {} + + es-module-lexer@2.1.0: {} + + esbuild@0.28.1: + optionalDependencies: + '@esbuild/aix-ppc64': 0.28.1 + '@esbuild/android-arm': 0.28.1 + '@esbuild/android-arm64': 0.28.1 + '@esbuild/android-x64': 0.28.1 + '@esbuild/darwin-arm64': 0.28.1 + '@esbuild/darwin-x64': 0.28.1 + '@esbuild/freebsd-arm64': 0.28.1 + '@esbuild/freebsd-x64': 0.28.1 + '@esbuild/linux-arm': 0.28.1 + '@esbuild/linux-arm64': 0.28.1 + '@esbuild/linux-ia32': 0.28.1 + '@esbuild/linux-loong64': 0.28.1 + '@esbuild/linux-mips64el': 0.28.1 + '@esbuild/linux-ppc64': 0.28.1 + '@esbuild/linux-riscv64': 0.28.1 + '@esbuild/linux-s390x': 0.28.1 + '@esbuild/linux-x64': 0.28.1 + '@esbuild/netbsd-arm64': 0.28.1 + '@esbuild/netbsd-x64': 0.28.1 + '@esbuild/openbsd-arm64': 0.28.1 + '@esbuild/openbsd-x64': 0.28.1 + '@esbuild/openharmony-arm64': 0.28.1 + '@esbuild/sunos-x64': 0.28.1 + '@esbuild/win32-arm64': 0.28.1 + '@esbuild/win32-ia32': 0.28.1 + '@esbuild/win32-x64': 0.28.1 + + estree-walker@3.0.3: + dependencies: + '@types/estree': 1.0.9 + + expect-type@1.3.0: {} + + fdir@6.5.0(picomatch@4.0.4): + optionalDependencies: + picomatch: 4.0.4 + + fsevents@2.3.3: + optional: true + + hono@4.12.27: {} + + jose@6.2.3: {} + + lightningcss-android-arm64@1.32.0: + optional: true + + lightningcss-darwin-arm64@1.32.0: + optional: true + + lightningcss-darwin-x64@1.32.0: + optional: true + + lightningcss-freebsd-x64@1.32.0: + optional: true + + lightningcss-linux-arm-gnueabihf@1.32.0: + optional: true + + lightningcss-linux-arm64-gnu@1.32.0: + optional: true + + lightningcss-linux-arm64-musl@1.32.0: + optional: true + + lightningcss-linux-x64-gnu@1.32.0: + optional: true + + lightningcss-linux-x64-musl@1.32.0: + optional: true + + lightningcss-win32-arm64-msvc@1.32.0: + optional: true + + lightningcss-win32-x64-msvc@1.32.0: + optional: true + + lightningcss@1.32.0: + dependencies: + detect-libc: 2.1.2 + optionalDependencies: + lightningcss-android-arm64: 1.32.0 + lightningcss-darwin-arm64: 1.32.0 + lightningcss-darwin-x64: 1.32.0 + lightningcss-freebsd-x64: 1.32.0 + lightningcss-linux-arm-gnueabihf: 1.32.0 + lightningcss-linux-arm64-gnu: 1.32.0 + lightningcss-linux-arm64-musl: 1.32.0 + lightningcss-linux-x64-gnu: 1.32.0 + lightningcss-linux-x64-musl: 1.32.0 + lightningcss-win32-arm64-msvc: 1.32.0 + lightningcss-win32-x64-msvc: 1.32.0 + + magic-string@0.30.21: + dependencies: + '@jridgewell/sourcemap-codec': 1.5.5 + + nanoid@3.3.13: {} + + obug@2.1.3: {} + + pathe@2.0.3: {} + + pg-cloudflare@1.4.0: + optional: true + + pg-connection-string@2.14.0: {} + + pg-int8@1.0.1: {} + + pg-pool@3.14.0(pg@8.22.0): + dependencies: + pg: 8.22.0 + + pg-protocol@1.15.0: {} + + pg-types@2.2.0: + dependencies: + pg-int8: 1.0.1 + postgres-array: 2.0.0 + postgres-bytea: 1.0.1 + postgres-date: 1.0.7 + postgres-interval: 1.2.0 + + pg@8.22.0: + dependencies: + pg-connection-string: 2.14.0 + pg-pool: 3.14.0(pg@8.22.0) + pg-protocol: 1.15.0 + pg-types: 2.2.0 + pgpass: 1.0.5 + optionalDependencies: + pg-cloudflare: 1.4.0 + + pgpass@1.0.5: + dependencies: + split2: 4.2.0 + + picocolors@1.1.1: {} + + picomatch@4.0.4: {} + + postcss@8.5.15: + dependencies: + nanoid: 3.3.13 + picocolors: 1.1.1 + source-map-js: 1.2.1 + + postgres-array@2.0.0: {} + + postgres-bytea@1.0.1: {} + + postgres-date@1.0.7: {} + + postgres-interval@1.2.0: + dependencies: + xtend: 4.0.2 + + rolldown@1.0.3: + dependencies: + '@oxc-project/types': 0.133.0 + '@rolldown/pluginutils': 1.0.1 + optionalDependencies: + '@rolldown/binding-android-arm64': 1.0.3 + '@rolldown/binding-darwin-arm64': 1.0.3 + '@rolldown/binding-darwin-x64': 1.0.3 + '@rolldown/binding-freebsd-x64': 1.0.3 + '@rolldown/binding-linux-arm-gnueabihf': 1.0.3 + '@rolldown/binding-linux-arm64-gnu': 1.0.3 + '@rolldown/binding-linux-arm64-musl': 1.0.3 + '@rolldown/binding-linux-ppc64-gnu': 1.0.3 + '@rolldown/binding-linux-s390x-gnu': 1.0.3 + '@rolldown/binding-linux-x64-gnu': 1.0.3 + '@rolldown/binding-linux-x64-musl': 1.0.3 + '@rolldown/binding-openharmony-arm64': 1.0.3 + '@rolldown/binding-wasm32-wasi': 1.0.3 + '@rolldown/binding-win32-arm64-msvc': 1.0.3 + '@rolldown/binding-win32-x64-msvc': 1.0.3 + + siginfo@2.0.0: {} + + source-map-js@1.2.1: {} + + split2@4.2.0: {} + + stackback@0.0.2: {} + + std-env@4.1.0: {} + + tinybench@2.9.0: {} + + tinyexec@1.2.4: {} + + tinyglobby@0.2.17: + dependencies: + fdir: 6.5.0(picomatch@4.0.4) + picomatch: 4.0.4 + + tinyrainbow@3.1.0: {} + + tslib@2.8.1: + optional: true + + tsx@4.22.4: + dependencies: + esbuild: 0.28.1 + optionalDependencies: + fsevents: 2.3.3 + + tweetnacl@1.0.3: {} + + typescript@5.9.3: {} + + undici-types@7.18.2: {} + + vite@8.0.16(@types/node@24.13.2)(esbuild@0.28.1)(tsx@4.22.4): + dependencies: + lightningcss: 1.32.0 + picomatch: 4.0.4 + postcss: 8.5.15 + rolldown: 1.0.3 + tinyglobby: 0.2.17 + optionalDependencies: + '@types/node': 24.13.2 + esbuild: 0.28.1 + fsevents: 2.3.3 + tsx: 4.22.4 + + vitest@4.1.9(@types/node@24.13.2)(vite@8.0.16(@types/node@24.13.2)(esbuild@0.28.1)(tsx@4.22.4)): + dependencies: + '@vitest/expect': 4.1.9 + '@vitest/mocker': 4.1.9(vite@8.0.16(@types/node@24.13.2)(esbuild@0.28.1)(tsx@4.22.4)) + '@vitest/pretty-format': 4.1.9 + '@vitest/runner': 4.1.9 + '@vitest/snapshot': 4.1.9 + '@vitest/spy': 4.1.9 + '@vitest/utils': 4.1.9 + es-module-lexer: 2.1.0 + expect-type: 1.3.0 + magic-string: 0.30.21 + obug: 2.1.3 + pathe: 2.0.3 + picomatch: 4.0.4 + std-env: 4.1.0 + tinybench: 2.9.0 + tinyexec: 1.2.4 + tinyglobby: 0.2.17 + tinyrainbow: 3.1.0 + vite: 8.0.16(@types/node@24.13.2)(esbuild@0.28.1)(tsx@4.22.4) + why-is-node-running: 2.3.0 + optionalDependencies: + '@types/node': 24.13.2 + transitivePeerDependencies: + - msw + + why-is-node-running@2.3.0: + dependencies: + siginfo: 2.0.0 + stackback: 0.0.2 + + ws@8.21.0: {} + + xtend@4.0.2: {} + + zod@3.25.76: {} diff --git a/cloud/pnpm-workspace.yaml b/cloud/pnpm-workspace.yaml new file mode 100644 index 00000000000..cc61e60464c --- /dev/null +++ b/cloud/pnpm-workspace.yaml @@ -0,0 +1,4 @@ +packages: + - apps/* + - packages/* + diff --git a/cloud/tsconfig.base.json b/cloud/tsconfig.base.json new file mode 100644 index 00000000000..d977ad5c9b9 --- /dev/null +++ b/cloud/tsconfig.base.json @@ -0,0 +1,19 @@ +{ + "compilerOptions": { + "allowSyntheticDefaultImports": true, + "declaration": true, + "esModuleInterop": true, + "forceConsistentCasingInFileNames": true, + "isolatedModules": true, + "lib": ["ES2024"], + "module": "NodeNext", + "moduleResolution": "NodeNext", + "noUncheckedIndexedAccess": true, + "outDir": "dist", + "resolveJsonModule": true, + "skipLibCheck": true, + "strict": true, + "target": "ES2024" + } +} + diff --git a/config/oxlint-react-doctor.json b/config/oxlint-react-doctor.json index 3c91b01d58d..c39f0d51bf2 100644 --- a/config/oxlint-react-doctor.json +++ b/config/oxlint-react-doctor.json @@ -25,5 +25,5 @@ "react-doctor/zustand-no-mutating-state": "warn", "react-doctor/zustand-no-whole-store-destructure": "warn" }, - "ignorePatterns": ["**/node_modules", "**/dist", "**/out"] + "ignorePatterns": ["**/node_modules", "**/dist", "**/out", "cloud/**"] } diff --git a/config/scripts/check-changed-code-quality.mjs b/config/scripts/check-changed-code-quality.mjs index 4427c6b7f38..eedf3dbda78 100644 --- a/config/scripts/check-changed-code-quality.mjs +++ b/config/scripts/check-changed-code-quality.mjs @@ -7,6 +7,7 @@ import { resolvePullRequestDiffBase } from './git-pull-request-diff-base.mjs' import { resolveOxlintInvocation } from './oxlint-cli-invocation.mjs' const SOURCE_FILE_PATTERN = /\.(?:[cm]?[jt]sx?)$/ +const ROOT_CODE_QUALITY_IGNORED_PREFIXES = ['cloud/'] export const OXLINT_SCANS = [ { // Why: no --config, so Oxlint keeps discovering nested configs. Pinning the root @@ -73,6 +74,10 @@ function splitNullDelimited(output) { return output.split('\0').filter(Boolean) } +export function isRootCodeQualityPath(file) { + return !ROOT_CODE_QUALITY_IGNORED_PREFIXES.some((prefix) => file.startsWith(prefix)) +} + function resolveBase(root, requestedBase) { for (const candidate of [ requestedBase, @@ -107,7 +112,11 @@ export function collectAddedLineRanges(root, requestedBase) { const rangesByFile = new Map() for (const file of changedFiles) { - if (!SOURCE_FILE_PATTERN.test(file) || !existsSync(path.join(root, file))) { + if ( + !isRootCodeQualityPath(file) || + !SOURCE_FILE_PATTERN.test(file) || + !existsSync(path.join(root, file)) + ) { continue } const diff = runGit(root, ['diff', '--unified=0', '--no-color', comparisonBase, '--', file]) @@ -119,7 +128,11 @@ export function collectAddedLineRanges(root, requestedBase) { for (const file of untrackedFiles) { const absolutePath = path.join(root, file) - if (!SOURCE_FILE_PATTERN.test(file) || !existsSync(absolutePath)) { + if ( + !isRootCodeQualityPath(file) || + !SOURCE_FILE_PATTERN.test(file) || + !existsSync(absolutePath) + ) { continue } const lineCount = readFileSync(absolutePath, 'utf8').split(/\r?\n/).length diff --git a/config/scripts/check-changed-code-quality.test.mjs b/config/scripts/check-changed-code-quality.test.mjs index 76a25802e5c..3a88cf1b02e 100644 --- a/config/scripts/check-changed-code-quality.test.mjs +++ b/config/scripts/check-changed-code-quality.test.mjs @@ -3,6 +3,7 @@ import { OXLINT_SCANS, diagnosticTouchesAddedLines, isMovedCode, + isRootCodeQualityPath, overlapsAddedLines, parseAddedLineRanges } from './check-changed-code-quality.mjs' @@ -52,6 +53,11 @@ describe('changed-code quality line matching', () => { expect(scan.args).not.toContain('--config') expect(scan.args).not.toContain('--disable-nested-config') }) + + it('leaves Cloud source to the independent Cloud quality checks', () => { + expect(isRootCodeQualityPath('cloud/apps/relay/src/index.ts')).toBe(false) + expect(isRootCodeQualityPath('src/main/index.ts')).toBe(true) + }) }) describe('moved-code exemption', () => { diff --git a/config/scripts/check-root-directory-entries.test.mjs b/config/scripts/check-root-directory-entries.test.mjs index 15276e8aa82..dcc5596771b 100644 --- a/config/scripts/check-root-directory-entries.test.mjs +++ b/config/scripts/check-root-directory-entries.test.mjs @@ -105,6 +105,15 @@ describe('root directory guard', () => { expect(output).toContain('new-root.md') }) + it('allows the reviewed cloud workspace directory', () => { + const fixture = makeFixture() + const head = commitFiles(fixture.root, [['cloud/package.json', '{}\n']]) + + const result = runGuard({ ...fixture, head }) + + expect(result.status).toBe(0) + }) + it('rejects a new top-level directory', () => { const fixture = makeFixture() const head = commitFiles(fixture.root, [['new-folder/file.txt', 'too prominent\n']]) diff --git a/config/scripts/pr-code-change-scope.mjs b/config/scripts/pr-code-change-scope.mjs index 67d4f565afa..8bc10fc5b72 100644 --- a/config/scripts/pr-code-change-scope.mjs +++ b/config/scripts/pr-code-change-scope.mjs @@ -237,6 +237,8 @@ const WINDOWS_PACKAGE_TESTS = [ const DESKTOP_IRRELEVANT_PREFIXES = [ 'mobile/', + 'cloud/', + '.github/workflows/cloud-', '.github/workflows/mobile.yml', '.github/workflows/mobile-ios-release.yml', '.github/workflows/mobile-android-release.yml' diff --git a/config/scripts/pr-code-change-scope.test.mjs b/config/scripts/pr-code-change-scope.test.mjs index 9a1c9e649b6..1fe296af265 100644 --- a/config/scripts/pr-code-change-scope.test.mjs +++ b/config/scripts/pr-code-change-scope.test.mjs @@ -86,6 +86,17 @@ describe('docs-only path classification', () => { it('does not start desktop PR Checks for mobile-only diffs', () => { expect(shouldRunPrChecks(['mobile/src/App.tsx', 'mobile/package.json'])).toBe(false) }) + + it('does not start desktop PR Checks for cloud-only diffs', () => { + expect( + shouldRunPrChecks([ + 'cloud/apps/relay/src/index.ts', + 'cloud/package.json', + 'cloud/.gitleaks.toml', + '.github/workflows/cloud-verify.yml' + ]) + ).toBe(false) + }) }) describe('per-job path classification', () => { diff --git a/package.json b/package.json index 9846604fab6..a7886667d25 100644 --- a/package.json +++ b/package.json @@ -290,12 +290,12 @@ "windows-native-registry": "3.2.2" }, "lint-staged": { - "*.{ts,tsx,js,jsx,mjs,mts,cts}": [ + "{*.{ts,tsx,js,jsx,mjs,mts,cts},!(cloud)/**/*.{ts,tsx,js,jsx,mjs,mts,cts}}": [ "oxlint", "oxlint --config config/oxlint-react-doctor.json", "oxfmt --write" ], - "*.{json,css}": [ + "{*.{json,css},!(cloud)/**/*.{json,css}}": [ "oxfmt --write" ] },