starred/orca
1
0
Fork 0
mirror of https://github.com/stablyai/orca.git synced 2026-11-15 10:27:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge origin/main into brennanb2025/orchestration-mail-holes

Browse Source
This commit is contained in:
Brennan Benson
2026-11-13 15:44:59 -08:00
parent a3169dea93 8268704516
commit f481f53a3f
5241 changed files with 51677 additions and 169752 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/actions/cloud-sql-rollout-lease/storage-lease.test.mjs
+11 -95
View File
@@ -1,28 +1,16 @@
import assert from 'node:assert/strict'
import { test } from 'node:test'
import {
CloudSqlRolloutLease,
LEASE_TTL_MS,
LeaseConflict,
LeaseUnreadable
} from './storage-lease.mjs'
import { CloudSqlRolloutLease, LeaseConflict, LeaseUnreadable } from './storage-lease.mjs'
const BUCKET = 'onorca-cloud-terraform-state'
const OBJECT = 'terraform/state/cloud-sql-rollout/production.lock'
const NOW = 1_756_000_000_000
/** Model real GCS generation increments, ifGenerationMatch and 412 conflicts; faults inject failures. */
function fakeStorage({ object = null, faults = [] } = {}) {
const state = { object, requests: [] }
function fakeStorage({ object = null } = {}) {
const state = { object }
const fetcher = async (rawUrl, init = {}) => {
const url = new URL(rawUrl)
const method = init.method ?? 'GET'
const record = { method, url, path: url.pathname, search: url.searchParams }
state.requests.push(record)
const fault = faults.find((candidate) => candidate.when(record))
if (fault) {
return json(fault.status, fault.body ?? {})
}
if (url.pathname.startsWith('/upload/')) {
const want = url.searchParams.get('ifGenerationMatch')
@@ -63,15 +51,10 @@ function json(status, body) {
}
}
function storedRecord({
holderKey,
expiresAt,
repository = 'stablyai/orca',
workflow = 'Deploy Relay Production'
}) {
function storedRecord({ holderKey, expiresAt }) {
return {
repository,
workflow,
repository: 'stablyai/orca',
workflow: 'Deploy Relay Production',
run_id: '9001',
run_url: 'https://github.com/stablyai/orca/actions/runs/9001',
run_attempt: '1',
@@ -81,14 +64,14 @@ function storedRecord({
}
}
function leaseFor(storage, { warn = () => {} } = {}) {
function leaseFor(storage) {
return new CloudSqlRolloutLease({
bucket: BUCKET,
objectName: OBJECT,
accessToken: 'test-token',
fetcher: storage.fetcher,
now: () => NOW,
warn
warn: () => {}
})
}
@@ -114,77 +97,24 @@ test('refuses a live lease held by another run and never writes', async () => {
.catch((thrown) => thrown)
assert.ok(error instanceof LeaseConflict)
assert.equal(error.holder.repository, 'stablyai/orca')
assert.equal(error.holder.run_url, 'https://github.com/stablyai/orca/actions/runs/9001')
assert.equal(
storage.state.requests.filter((request) => request.method !== 'GET').length,
0,
'a foreign live lease must not be written'
)
assert.equal(storage.state.object.generation, '1500')
})
test('re-enters a live lease this run already holds and extends it', async () => {
const storage = fakeStorage({
object: {
generation: '1500',
body: storedRecord({ holderKey: HOLDER.holderKey, expiresAt: NOW + 60_000 })
}
})
const warnings = []
const claim = await leaseFor(storage, { warn: (message) => warnings.push(message) }).acquire(
HOLDER
)
assert.equal(claim.state, 'reentrant')
assert.deepEqual(warnings, [], 're-entering our own lease is not a takeover')
assert.equal(claim.record.acquired_at, NOW - 60_000, 'original acquisition time is preserved')
assert.equal(claim.record.expires_at, NOW + LEASE_TTL_MS)
const upload = storage.state.requests.find((request) => request.path.startsWith('/upload/'))
assert.equal(upload.search.get('ifGenerationMatch'), '1500')
assert.equal(storage.state.object.generation, '1501')
})
test('takes over an expired lease and warns naming the stale holder', async () => {
test('takes over an expired lease', async () => {
const storage = fakeStorage({
object: {
generation: '1500',
body: storedRecord({
holderKey: 'stablyai/orca/9001',
expiresAt: NOW - 1,
repository: 'stablyai/orca',
workflow: 'Deploy Relay Production Capacity'
expiresAt: NOW - 1
})
}
})
const warnings = []
const claim = await leaseFor(storage, { warn: (message) => warnings.push(message) }).acquire(
HOLDER
)
await leaseFor(storage).acquire(HOLDER)
assert.equal(claim.state, 'takeover')
assert.equal(warnings.length, 1)
assert.match(warnings[0], /stablyai\/orca/)
assert.match(warnings[0], /Deploy Relay Production Capacity/)
assert.match(warnings[0], /actions\/runs\/9001/)
const upload = storage.state.requests.find((request) => request.path.startsWith('/upload/'))
assert.equal(upload.search.get('ifGenerationMatch'), '1500')
assert.equal(storage.state.object.body.holder_key, HOLDER.holderKey)
})
test('releases with a generation match and leaves the object gone', async () => {
const storage = fakeStorage()
const lease = leaseFor(storage)
const claim = await lease.acquire(HOLDER)
const released = await lease.release(HOLDER.holderKey)
assert.deepEqual(released, { released: true, generation: claim.generation })
const remove = storage.state.requests.find((request) => request.method === 'DELETE')
assert.equal(remove.search.get('ifGenerationMatch'), claim.generation)
assert.equal(storage.state.object, null)
})
test('refuses to release a lease another run now holds', async () => {
const storage = fakeStorage({
object: {
@@ -212,17 +142,3 @@ test('treats an unreadable record as held, not free', async () => {
assert.ok(error instanceof LeaseUnreadable)
assert.equal(storage.state.object.generation, '1500')
})
test('renewal rewrites only expires_at on the observed generation', async () => {
const storage = fakeStorage()
const lease = leaseFor(storage)
await lease.acquire(HOLDER)
storage.state.object.body.expires_at = NOW - 1
const renewed = await lease.renew(HOLDER.holderKey)
assert.equal(renewed.renewed, true)
assert.equal(storage.state.object.body.expires_at, NOW + LEASE_TTL_MS)
assert.equal(storage.state.object.body.acquired_at, NOW)
assert.equal(storage.state.object.body.holder_key, HOLDER.holderKey)
})
.github/workflows/computer-e2e.yml
+1 -2
View File
@@ -53,7 +53,6 @@ on:
- 'src/shared/computer-use-*.ts'
- 'tests/e2e/computer-linux.e2e.ts'
- 'tests/e2e/computer-windows.e2e.ts'
- 'tests/e2e/computer-windows-store.e2e.ts'
- 'tests/e2e/computer-windows-background.e2e.ts'
- 'tests/e2e/helpers/windows-semantic-controls-fixture.*'
- 'tests/e2e/helpers/computer-cli-driver.ts'
@@ -153,4 +152,4 @@ jobs:
- run: pnpm verify:computer-native
- run: pnpm build:cli
- run: pnpm run build:electron-vite:parallel
- run: pnpm test:e2e:computer --reporter=verbose tests/e2e/computer-windows.e2e.ts tests/e2e/computer-windows-store.e2e.ts tests/e2e/computer-windows-background.e2e.ts
- run: pnpm test:e2e:computer --reporter=verbose tests/e2e/computer-windows.e2e.ts tests/e2e/computer-windows-background.e2e.ts
.github/workflows/computer-macos-native.yml
+2
View File
@@ -38,6 +38,8 @@ jobs:
mac-native-owner-smoke:
if: github.event_name != 'pull_request' || github.event.pull_request.draft != true
runs-on: macos-15
env:
ORCA_BACKGROUND_LAUNCH: '1'
timeout-minutes: 180
permissions:
contents: read
cloud/apps/push/src/delivery/durable-push-claim.test.ts
+1 -33
View File
@@ -1,6 +1,6 @@
import { afterEach, expect, it } from 'vitest'
import type { CloudDatabase as PushDatabase } from '@orca-cloud/cloud-database/cloud-database'
import { DELIVERY_LEASE_MS, DurablePushStore } from './durable-push-store.js'
import { DurablePushStore } from './durable-push-store.js'
import {
cleanupDurablePushFixtures,
@@ -63,20 +63,6 @@ it.skipIf(!durablePushTestDatabaseUrl)(
}
)
it('leases exactly one delivery per device across concurrent claimers', async () => {
const { db, store, clock, advance } = await fixture()
for (let seq = 1; seq <= 5; seq++) {
await store.accept('host', 'phone', notification(seq))
advance(1)
}
const claims = await Promise.all(
Array.from({ length: 8 }, () => new DurablePushStore(db, clock).claim())
)
const leased = claims.filter((claim) => claim !== null)
expect(leased).toHaveLength(1)
expect(leased[0]?.notification.notificationSeq).toBe(1)
})
it.skipIf(!durablePushTestDatabaseUrl)(
'does not lease the row behind a head another claimer holds',
async () => {
@@ -138,21 +124,3 @@ it.skipIf(!durablePushTestDatabaseUrl)(
expect((await first)?.notification.notificationSeq).toBe(1)
}
)
// Only the TTL term keeps the candidate scan off an unpruned backlog; expires_at is not indexed first.
it('frees a device held by a crashed worker once its lease lapses', async () => {
const { db, store, clock, advance } = await fixture()
await store.accept('host', 'phone', notification(1))
advance(1)
await store.accept('host', 'phone', notification(2))
const crashed = (await new DurablePushStore(db, clock).claim())!
advance(DELIVERY_LEASE_MS - 1)
expect(await store.claim()).toBeNull()
advance(1)
const reclaimed = (await store.claim())!
expect(reclaimed.id).toBe(crashed.id)
expect(reclaimed.attempts).toBe(2)
await store.finish(reclaimed)
expect((await store.claim())?.notification.notificationSeq).toBe(2)
})
cloud/apps/push/src/delivery/durable-push-store.test.ts
+3 -73
View File
@@ -1,9 +1,7 @@
import { afterEach, describe, expect, it } from 'vitest'
import { openPushDatabase } from '../database/push-database.js'
import { DurablePushStore, DELIVERY_LEASE_MS } from './durable-push-store.js'
import {
cleanupDurablePushFixtures,
durablePushTestDatabaseUrl,
fixture,
notification
} from './durable-push-store.test-fixture.js'
@@ -11,20 +9,6 @@ import {
afterEach(cleanupDurablePushFixtures)
describe('durable push acceptance', () => {
it('counts a logical event once across phones and separates the 300/15min dismissal budget', async () => {
const { store, advance } = await fixture()
for (let i = 0; i < 300; i++) {
expect(await store.accept('host', 'phone1', notification(i))).toBe('queued')
expect(await store.accept('host', 'phone2', notification(i))).toBe('queued')
expect(await store.accept('host', 'phone1', notification(i, 'dismiss'))).toBe('queued')
}
expect(await store.accept('host', 'phone1', notification(300))).toBe('rate_limited')
expect(await store.accept('host', 'phone1', notification(300, 'dismiss'))).toBe('rate_limited')
expect(await store.accept('another-host', 'phone3', notification(300))).toBe('queued')
advance(15 * 60_000)
expect(await store.accept('host', 'phone1', notification(301))).toBe('queued')
})
it('queues one delivery per event and recovers work across service instances', async () => {
const { db, store, clock, advance } = await fixture()
await store.accept('host', 'phone', notification(1))
@@ -32,16 +16,6 @@ describe('durable push acceptance', () => {
await restarted.accept('host', 'phone', notification(1))
advance(1)
await restarted.accept('host', 'phone', notification(2))
const rows = await db.query(
"SELECT payload_json, due_at, created_at FROM push_delivery_batches WHERE registration_id = ? AND state = 'pending' ORDER BY created_at, batch_id",
['phone']
)
expect(rows).toHaveLength(2)
expect(rows.map((row) => JSON.parse(String(row.payload_json)))).toEqual([
notification(1),
notification(2)
])
expect(rows.every((row) => Number(row.due_at) >= Number(row.created_at))).toBe(true)
const delivery = await restarted.claim()
expect(delivery?.notification.notificationSeq).toBe(1)
expect(await store.claim()).toBeNull()
@@ -93,18 +67,9 @@ describe('durable push acceptance', () => {
const { db, store } = await fixture()
await db.query('ALTER TABLE push_delivery_batches RENAME TO push_delivery_batches_unavailable')
try {
const databaseUrl = durablePushTestDatabaseUrl
if (databaseUrl) {
const concurrent = await openPushDatabase({ databaseUrl, dataDir: '' })
try {
await expect(
concurrent.query('SELECT COUNT(*) FROM push_delivery_batches')
).resolves.toHaveLength(1)
} finally {
await concurrent.close()
}
}
await expect(store.accept('host', 'phone', notification(1))).rejects.toThrow()
await expect(store.accept('host', 'phone', notification(1))).rejects.toThrow(
/push_delivery_batches/
)
expect(await db.query('SELECT * FROM push_events')).toEqual([])
expect(await db.query('SELECT * FROM push_event_recipients')).toEqual([])
} finally {
@@ -160,38 +125,3 @@ it('does not resurrect an in-flight alert after a dismissal and transient provid
expect(await store.claim()).toBeNull()
expect(await store.pendingCount('phone')).toBe(0)
})
it('fences late renew and finish after an expired claim is dismissed', async () => {
const { db, store, advance } = await fixture()
const alert = notification(1)
await store.accept('host', 'phone', alert)
const stale = (await store.claim())!
advance(DELIVERY_LEASE_MS)
await store.accept('host', 'phone', {
...notification(2, 'dismiss'),
notificationId: alert.notificationId
})
const read = async () =>
(
await db.query(
'SELECT state, payload_json, lease_until FROM push_delivery_batches WHERE batch_id = ?',
[stale.id]
)
)[0]
expect(await read()).toBeUndefined()
await store.renew(stale)
await store.finish(stale, 1000)
await store.finish(stale)
expect(await read()).toBeUndefined()
const dismissal = (await store.claim())!
expect(dismissal.notification.kind).toBe('dismiss')
await store.finish(dismissal)
await store.accept('host', 'phone', notification(3))
const fresh = (await store.claim())!
await store.finish(fresh, 1000)
advance(1000)
const retry = (await store.claim())!
expect(retry.id).toBe(fresh.id)
await store.finish(retry)
expect(await store.claim()).toBeNull()
})
cloud/apps/push/src/delivery/durable-push-worker.test.ts
+1 -5
View File
@@ -13,7 +13,6 @@ import { pauseAfter, within } from './durable-push-store.test-fixture.js'
const cleanups: (() => Promise<void>)[] = []
afterEach(async () => {
for (const cleanup of cleanups.splice(0)) await cleanup()
vi.useRealTimers()
vi.restoreAllMocks()
})
const note = (seq: number, overrides: Partial<PushNotification> = {}): PushNotification => ({
@@ -42,9 +41,8 @@ async function fixture() {
const send = vi.fn(async (_delivery: PushDelivery): Promise<PushProviderOutcome> => ({
status: 'sent'
}))
const onRetry = vi.fn()
const dispatcher = new PushDispatcher({ devices, fcm: { send } as never })
const worker = new DurablePushWorker(store, dispatcher, { now, onRetry })
const worker = new DurablePushWorker(store, dispatcher, { now })
cleanups.push(async () => {
await worker.stop()
await db.close()
@@ -57,7 +55,6 @@ async function fixture() {
worker,
dispatcher,
send,
onRetry,
now,
registrationId: device.registrationId,
accept: (notification: PushNotification) =>
@@ -102,7 +99,6 @@ it('stops queued claim admissions without leasing the remaining device', async (
paused.release()
await within(Promise.all([pending, stopping]), 2_000)
expect(h.send).toHaveBeenCalledOnce()
expect(h.send.mock.calls[0]?.[0].orca.notificationSeq).toBe(1)
const [remaining] = await h.db.query(
'SELECT attempts, lease_until FROM push_delivery_batches WHERE registration_id = ?',
[other.registrationId]
cloud/apps/push/src/host-auth/host-challenge-store.test.ts
+1 -9
View File
@@ -26,18 +26,10 @@ describe('push host challenge store', () => {
await database.close()
})
it('refuses a wrong host: the box will not open and a foreign proof will not match', async () => {
it('rejects a proof issued for another challenge', async () => {
const owner = createPushHostKeypair(6)
const intruder = createPushHostKeypair(7)
const ownerChallenge = await store.issue(hostPublicKeyB64(owner))
expect(
answerPushHostChallenge(ownerChallenge!, {
gatewayOrigin: GATEWAY_ORIGIN,
keypair: intruder,
now: () => clock
})
).toBeNull()
const intruderChallenge = await store.issue(hostPublicKeyB64(intruder))
const intruderProof = answerPushHostChallenge(intruderChallenge!, {
gatewayOrigin: GATEWAY_ORIGIN,
cloud/apps/push/src/host-auth/push-session-concurrency.test.ts
-5
View File
@@ -19,11 +19,6 @@ async function concurrentSessions(db: PushDatabase) {
sessions.map((session) => store.resolve(session.sessionToken))
)
expect(decisions.filter((decision) => decision.ok)).toHaveLength(1)
const [row] = await db.query(
'SELECT COUNT(*) AS count FROM push_sessions WHERE host_fingerprint = ?',
[host]
)
expect(Number(row?.count)).toBe(1)
} finally {
await db.query('DELETE FROM push_sessions WHERE host_fingerprint = ?', [host])
}
cloud/apps/push/src/http/push-send-idempotency.test.ts
-36
View File
@@ -1,36 +0,0 @@
import { afterEach, expect, it } from 'vitest'
import { createPushServerHarness, notification } from './push-server-harness.test-fixture.js'
import { createPushHostKeypair } from '../host-auth/host-challenge-answering.test-fixture.js'
const harnesses: Awaited<ReturnType<typeof createPushServerHarness>>[] = []
afterEach(async () => {
await Promise.all(harnesses.splice(0).map((h) => h.close()))
})
it('returns queued for concurrent retries without double quota or delivery', async () => {
const h = await createPushServerHarness()
harnesses.push(h)
const token = await h.signIn(createPushHostKeypair(2))
const registrationId = await h.registerAndroid(token)
const body = { v: 1, registrationIds: [registrationId], notification: notification() }
const responses = await Promise.all(
Array.from({ length: 10 }, () => h.post('/v1/send', body, token))
)
for (const response of responses)
expect(await response.json()).toEqual({ results: [{ registrationId, status: 'queued' }] })
expect(await h.server.deliveryStore.pendingCount(registrationId)).toBe(1)
await h.flushDeliveries()
await h.post('/v1/send', body, token)
await h.flushDeliveries()
expect(h.fcmRequests).toHaveLength(1)
expect(JSON.parse(h.fcmRequests[0]!.body).message.data.coalescedCount).toBeUndefined()
expect(
Number((await h.database.query('SELECT COUNT(*) AS count FROM push_events'))[0]?.count)
).toBe(1)
await h.post(
'/v1/send',
{ ...body, notification: notification({ notificationEpoch: 'new-epoch' }) },
token
)
await h.flushDeliveries()
expect(h.fcmRequests).toHaveLength(2)
})
cloud/apps/push/src/http/push-server-send.test.ts
-12
View File
@@ -2,7 +2,6 @@ import { afterEach, beforeEach, describe, expect, it } from 'vitest'
import { createPushHostKeypair } from '../host-auth/host-challenge-answering.test-fixture.js'
import {
createPushServerHarness,
FCM_TOKEN,
notification
} from './push-server-harness.test-fixture.js'
@@ -34,12 +33,6 @@ describe('push gateway send route', () => {
})
await harness.flushDeliveries()
expect(harness.fcmRequests).toHaveLength(1)
expect(JSON.parse(harness.fcmRequests[0]!.body)).toMatchObject({
message: {
token: FCM_TOKEN,
data: { title: 'Agent needs input', tag: expect.stringMatching(/^[a-f0-9]{64}$/) }
}
})
const afterDeath = await harness.post(
'/v1/send',
@@ -47,11 +40,6 @@ describe('push gateway send route', () => {
sessionToken
)
expect(await afterDeath.json()).toEqual({ results: [{ registrationId, status: 'dead' }] })
const listed = await harness.authorized('/v1/devices', {}, sessionToken)
expect(await listed.json()).toEqual({
devices: [{ registrationId, deviceId: 'device-1', platform: 'android', dead: true }]
})
})
it('reports an error for a registration the host does not own', async () => {
cloud/apps/push/src/providers/apns/apns-topic-recovery.test.ts
-1
View File
@@ -36,7 +36,6 @@ it('delivers again after a topic error without re-registering the phone', async
expect(await sent.json()).toEqual({ results: [{ registrationId, status: 'queued' }] })
await h.flushDeliveries()
if (seq === 1) {
expect(h.server.observability.consume().delivery_error).toBe(1)
h.setApnsResponse({ status: 200, body: '' })
}
}
cloud/apps/push/src/providers/fcm/fcm-client.test.ts
+2 -3
View File
@@ -67,11 +67,10 @@ describe('fcm client', () => {
error: { status: 'INVALID_ARGUMENT', message: 'Invalid value at message.android.ttl' }
})
})
await expect(unnamed.client.send(delivery(), { token: TOKEN })).resolves.toEqual({
await expect(unnamed.client.send(delivery(), { token: TOKEN })).resolves.toMatchObject({
status: 'error',
reason: 'INVALID_ARGUMENT',
retryable: false,
retryAfterMs: 10000
retryable: false
})
})
})
cloud/apps/relay-edge/test/admin-verify.test.ts
+5 -10
View File
@@ -12,7 +12,7 @@ const ERASER = 'https://erasure.test.invalid'
const BODY = JSON.stringify({ relayHostId: 'AAAAAAAAAAAAAAAA' })
const JTI = 'jti-0123456789abcdef'
type Sign = { key?: string; issuer?: string; op?: AdminOp; jti?: unknown; body?: string; iat?: number; exp?: number }
type Sign = { key?: string; issuer?: string; op?: AdminOp; jti?: unknown; body?: string }
async function sign(opts: Sign = {}): Promise<string> {
const claims: Record<string, unknown> = {
@@ -26,8 +26,8 @@ async function sign(opts: Sign = {}): Promise<string> {
.setProtectedHeader({ alg: 'HS256' })
.setIssuer(opts.issuer ?? ISSUER)
.setAudience('orca-relay-edge-admin')
.setIssuedAt(opts.iat ?? now)
.setExpirationTime(opts.exp ?? now + 60)
.setIssuedAt(now)
.setExpirationTime(now + 60)
.sign(new TextEncoder().encode(opts.key ?? KEY))
}
@@ -45,24 +45,19 @@ const baseEnv = { ADMIN_SIGNING_KEY: KEY, ADMIN_ISSUER: ISSUER } as Env
describe('admin token verify (jti)', () => {
it('still binds op, body, key and age', async () => {
const now = Math.floor(Date.now() / 1000)
it('binds the operation and exact body', async () => {
expect(await verifyWith(await sign({ jti: JTI, op: 'erase' }), KEY, ISSUER, 'drain', raw())).toBeNull()
expect(await verifyWith(await sign({ jti: JTI, body: '{}' }), KEY, ISSUER, 'drain', raw())).toBeNull()
expect(await verifyWith(await sign({ jti: JTI, key: NEXT }), KEY, ISSUER, 'drain', raw())).toBeNull()
expect(await verifyWith(await sign({ jti: JTI, iat: now - 180, exp: now + 60 }), KEY, ISSUER, 'drain', raw())).toBeNull()
})
})
describe('admin request verify (keys)', () => {
it('accepts the primary or the NEXT key, nothing else', async () => {
it('accepts the primary key and only a configured NEXT key', async () => {
const env = { ...baseEnv, ADMIN_SIGNING_KEY_NEXT: NEXT } as Env
expect(await verifyAdminRequest(env, request(await sign({ jti: JTI })), 'drain')).not.toBeNull()
expect(await verifyAdminRequest(env, request(await sign({ jti: JTI, key: NEXT })), 'drain')).not.toBeNull()
expect(await verifyAdminRequest(baseEnv, request(await sign({ jti: JTI, key: NEXT })), 'drain')).toBeNull()
const wrong = randomBytes(36).toString('base64url')
expect(await verifyAdminRequest(env, request(await sign({ jti: JTI, key: wrong })), 'drain')).toBeNull()
})
it('the erasure keys (primary or NEXT) verify only erase, under their own issuer', async () => {
cloud/apps/relay-edge/test/jwt-keys.test.ts
+3 -7
View File
@@ -28,7 +28,7 @@ const mint = (key: CryptoKey = privateKey, kid = 'k1') =>
.sign(key)
// A Cache API stand-in shared by every "isolate" (module instance) in the test.
function stubCache(): Map<string, Response> {
function stubCache(): void {
const store = new Map<string, Response>()
vi.stubGlobal('caches', {
default: {
@@ -38,7 +38,6 @@ function stubCache(): Map<string, Response> {
match: async (request: Request) => store.get(request.url)?.clone()
}
})
return store
}
// Each call is a fresh isolate: a new module instance with an empty in-memory key cache.
@@ -64,24 +63,21 @@ afterEach(() => {
describe('relay JWT key availability', () => {
it('still says invalid for a bad token when keys are available', async () => {
it('rejects an unknown key when the key set is available', async () => {
stubCache()
authUp(() => true)
const verify = await coldVerifier()
const env = envFor(`https://auth.test.invalid/jwks/${++urlSeq}`)
expect(await verify(env, await mint(otherKey))).toEqual({ ok: false, reason: 'invalid' })
expect(await verify(env, 'not-a-jwt')).toEqual({ ok: false, reason: 'invalid' })
expect(await verify(env, await mint(privateKey, 'unknown-kid'))).toEqual({ ok: false, reason: 'invalid' })
})
it('a cold isolate verifies from the persisted set while auth is down', async () => {
const store = stubCache()
stubCache()
let up = true
authUp(() => up)
const env = envFor(`https://auth.test.invalid/jwks/${++urlSeq}`)
const warm = await coldVerifier()
expect((await warm(env, await mint())).ok).toBe(true)
expect(store.size).toBe(1)
up = false
const cold = await coldVerifier()
const result = await cold(env, await mint())
cloud/apps/relay-edge/test/store-parity.test.ts
-1
View File
@@ -10,4 +10,3 @@ vi.mock('../../relay/src/database/database.js', async (importOriginal) => ({
}))
await import('../../relay/src/credentials/credential-store.test.ts')
await import('../../relay/src/credentials/credential-store-cleanup.test.ts')
cloud/apps/relay/src/admin-route-transient-app.test.ts
-17
View File
@@ -53,23 +53,6 @@ describe('admin routes under a database that is briefly out of reach', () => {
expect(response.headers.get('Retry-After')).toBe('7')
expect(await response.json()).toEqual({ error: 'database_temporarily_unavailable' })
})
it('answers admission-selector/status with a retryable 503 instead of a conflict', async () => {
const inspectCellAdmissionSelector = vi.fn(async () => {
throw poolTimeout()
})
const app = appWith({ inspectCellAdmissionSelector })
const response = await app.request(
'/v1/admin/admission-selector/status',
adminRequest({ v: 1, attemptId: '22222222-2222-4222-8222-222222222222' })
)
expect(response.status).toBe(503)
expect(response.headers.get('Retry-After')).toBe('5')
expect(await response.json()).toEqual({ error: 'database_temporarily_unavailable' })
})
})
function config(overrides: Partial<RelayConfig> = {}): RelayConfig {
cloud/apps/relay/src/assignment/assignment-control-supersession-postgres.test.ts
-140
View File
@@ -1,140 +0,0 @@
import { afterAll, beforeAll, describe, expect, it } from 'vitest'
import { RelayAssignmentStore } from './assignment-store.js'
import { openRelayDatabase, type RelayDatabase } from '../database/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])
// A snapshot left by an aborted run rejects the replayed watermark with stale_connection_snapshot.
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])
}
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
]
)
// The store reserves a unit per control lease, so a hand-written pair has to
// carry its own reservation or the fixture starts out of balance.
await databases[0]!.query(
`UPDATE relay_cells SET reserved_requests = reserved_requests + 2 WHERE cell_id = ?`,
[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)
})
cloud/apps/relay/src/assignment/assignment-dead-reserve-cell-postgres.test.ts
-58
View File
@@ -169,7 +169,6 @@ describePostgres('a host leaving a dead reserve cell, on PostgreSQL', () => {
const ledger = new ReserveLedgerWriter(database, { now: () => now, log: () => undefined })
for (const epoch of [6, 7]) ledger.enqueue({ userId: USER, relayHostId: host, cellId: dead.id, epoch })
await ledger.flush()
expect(ledger.pending()).toBe(0)
expect(await rowOf(host)).toEqual({ cell_id: live.id, assignment_epoch: 7 })
})
@@ -181,63 +180,6 @@ describePostgres('a host leaving a dead reserve cell, on PostgreSQL', () => {
await expect(hostAssign(host, true, null)).rejects.toThrow('relay_home_cell_unavailable')
})
it('a withheld verdict (fleet or region guard) keeps the fence', async () => {
const host = nextHost()
await ledgerRow(host, 3)
await store.setCellAdmitMode(dead.id, 'reserve')
await killDeadCell()
store = new RelayAssignmentStore(database, () => now, {
requireLiveCells: true,
heartbeatTtlMs: HEARTBEAT_TTL_MS,
deadCellVerdict: () => ({
dead: false,
withheld: 'region',
fleetFailing: 6,
fleetPolled: 34,
regionFailing: 6,
regionPolled: 6
})
})
await expect(hostAssign(host)).rejects.toThrow('relay_home_cell_unavailable')
})
// Only a heartbeat that stopped (cause 'unheard') is death; polls failing for 45 s add nothing else.
it('a cell heartbeating but not ready keeps its hosts, however long its polls fail', async () => {
const host = nextHost()
await ledgerRow(host, 3)
await store.setCellAdmitMode(dead.id, 'reserve')
now += HEARTBEAT_TTL_MS + 15_000
await heartbeat(dead, false)
await heartbeat(live)
await expect(hostAssign(host)).rejects.toThrow('relay_home_cell_unavailable')
expect(await rowOf(host)).toEqual({ cell_id: dead.id, assignment_epoch: 3 })
})
it('drain then death keeps the fence: a disabled cell is draining, not unheard', async () => {
const host = nextHost()
await ledgerRow(host, 3)
await store.setCellAdmitMode(dead.id, 'reserve')
await database.query(`UPDATE relay_cells SET enabled = 0 WHERE cell_id = ?`, [dead.id])
try {
await killDeadCell()
await expect(hostAssign(host)).rejects.toThrow('relay_home_cell_unavailable')
expect(await rowOf(host)).toEqual({ cell_id: dead.id, assignment_epoch: 3 })
} finally {
await database.query(`UPDATE relay_cells SET enabled = 1 WHERE cell_id = ?`, [dead.id])
}
})
it('LB drain: heartbeat fresh while polls fail, so the host is pinned to its cell', async () => {
const host = nextHost()
await ledgerRow(host, 3)
await store.setCellAdmitMode(dead.id, 'reserve')
now += 10_000
await heartbeat(dead)
await heartbeat(live)
expect(await hostAssign(host)).toMatchObject({ cellId: dead.id, assignmentEpoch: 3 })
expect(await rowOf(host)).toEqual({ cell_id: dead.id, assignment_epoch: 3 })
})
it('DB outage or director freeze: heartbeat stale but polls answered (no verdict), so 503', async () => {
const host = nextHost()
await ledgerRow(host, 3)
cloud/apps/relay/src/assignment/assignment-deployment-status-postgres.test.ts
-23
View File
@@ -84,27 +84,4 @@ describePostgres('PostgreSQL deployment status', () => {
restartBlockingReservedRequests: -1
})
})
it('reads the same counts through the reserve-aware status, and none for a reserve-mode cell', async () => {
const store = new RelayAssignmentStore(database!, () => 100)
const reserveCell = { ...cell, id: `${cell.id}-reserve`, url: 'https://deployment-status-reserve.example.com' }
await store.reconcileCells([cell, reserveCell])
await store.assign({ userId: 'aware-user', relayHostId: 'd1e2f3a4b5c6' })
expect(await store.reserveAwareCellDeploymentStatus(cell.id)).toEqual({
admitMode: 'db',
status: await store.cellDeploymentStatus(cell.id)
})
await store.setCellAdmitMode(reserveCell.id, 'reserve')
await expect(store.cellDeploymentStatus(reserveCell.id)).rejects.toThrow('cell_in_reserve_mode')
const reserve = await store.reserveAwareCellDeploymentStatus(reserveCell.id)
expect(reserve).toMatchObject({
admitMode: 'reserve',
status: { cellId: reserveCell.id, cellUrl: reserveCell.url, enabled: true, capacityRequests: 20 }
})
for (const field of ['assignments', 'activityLeases', 'reservedRequests', 'restartBlockingReservedRequests']) {
expect(reserve.status).not.toHaveProperty(field)
}
await expect(store.reserveAwareCellDeploymentStatus('missing-cell')).rejects.toThrow('cell_not_found')
})
})
cloud/apps/relay/src/assignment/assignment-headroom-scope.test.ts
+13 -49
View File
@@ -3,42 +3,27 @@ import { RelayAssignmentStore } from './assignment-store.js'
import { openInMemoryRelayDatabase, type RelayDatabase } from '../database/database.js'
const NOW = 100_000
const CELL_COUNT = 32
const identity = { userId: 'headroom-user', relayHostId: 'headroomhost0001' }
function observeHeadroom(database: RelayDatabase, rowCounts: number[]): RelayDatabase {
return {
dialect: database.dialect,
async query(sql, params) {
const rows = await database.query(sql, params)
if (sql.includes('FROM relay_cell_connection_limits limits')) rowCounts.push(rows.length)
return rows
},
queryLocked: (sql, params, options) => database.queryLocked(sql, params, options),
transaction: (operation, options) =>
database.transaction((transaction) => operation(observeHeadroom(transaction, rowCounts)), options),
close: () => database.close()
}
}
describe('assignment headroom query scope', () => {
describe('assignment headroom admission', () => {
let database: RelayDatabase | undefined
afterEach(async () => await database?.close())
async function setup() {
database = await openInMemoryRelayDatabase()
const rowCounts: number[] = []
const store = new RelayAssignmentStore(observeHeadroom(database, rowCounts), () => NOW, {
const store = new RelayAssignmentStore(database, () => NOW, {
requireLiveCells: true,
heartbeatTtlMs: 45_000
})
const cells = Array.from({ length: CELL_COUNT }, (_, index) => ({
id: `cell-${String(index).padStart(2, '0')}`,
url: `https://cell-${index}.example.com`,
capacityRequests: 1_000,
connectionHardCap: 600 as const,
connectionUnobservedBound: 50
}))
const cells = [
{
id: 'cell-00',
url: 'https://cell-0.example.com',
capacityRequests: 1_000,
connectionHardCap: 600 as const,
connectionUnobservedBound: 50
}
]
await store.reconcileCells(cells)
for (const cell of cells) {
await store.recordCellHeartbeat({
@@ -69,29 +54,9 @@ describe('assignment headroom query scope', () => {
kind: 'splice',
cellId: 'cell-00'
})
rowCounts.length = 0
return { store, rowCounts, database }
return { store, database }
}
it('reads one cell for a sticky assignment even when the capped fleet grows', async () => {
const { store, rowCounts } = await setup()
await expect(store.assign(identity)).resolves.toMatchObject({
...identity,
cellId: 'cell-00',
assignmentEpoch: 1
})
console.info('sticky headroom rows', { fleetCells: CELL_COUNT, rowCounts })
expect(rowCounts).toEqual([1])
})
it('still considers the full fleet for a new placement', async () => {
const { store, rowCounts } = await setup()
await expect(
store.assign({ userId: 'new-user', relayHostId: 'newheadroomhost1' })
).resolves.toMatchObject({ cellId: 'cell-01', assignmentEpoch: 1 })
expect(rowCounts).toEqual([CELL_COUNT])
})
it.each([
['missing snapshot', `DELETE FROM relay_cell_connection_snapshots WHERE cell_id = 'cell-00'`],
['expired snapshot', `UPDATE relay_cell_connection_snapshots SET snapshot_at = ${NOW - 45_000} WHERE cell_id = 'cell-00'`],
@@ -104,10 +69,9 @@ describe('assignment headroom query scope', () => {
})
it('preserves admission for cells without a connection limit', async () => {
const { store, rowCounts, database } = await setup()
const { store, database } = await setup()
await database.query(`DELETE FROM relay_cell_connection_limits WHERE cell_id = 'cell-00'`)
await expect(store.assign(identity)).resolves.toMatchObject({ cellId: 'cell-00' })
expect(rowCounts).toEqual([0])
})
it('counts outstanding reservations at the admission boundary', async () => {
cloud/apps/relay/src/assignment/assignment-home-cell-unavailable.test.ts
+3 -75
View File
@@ -1,5 +1,5 @@
import { afterEach, describe, expect, it } from 'vitest'
import { RelayAssignmentStore, RelayHomeCellUnavailableError } from './assignment-store.js'
import { RelayAssignmentStore } from './assignment-store.js'
import type { RelayCellConfig } from '../config.js'
import { openInMemoryRelayDatabase, type RelayDatabase } from '../database/database.js'
@@ -7,16 +7,6 @@ const HEARTBEAT_TTL_MS = 45_000
const START_MS = 100
const IDENTITY = { userId: 'user-1', relayHostId: 'host000000000001' }
// A connection-limited cell is what makes the committed fence mandatory, and
// that is the branch which used to answer "capacity exhausted".
const FENCED_CELL: RelayCellConfig = {
id: 'home',
url: 'https://home.example.com',
capacityRequests: 1_000,
connectionHardCap: 600,
connectionUnobservedBound: 50
}
const databases: RelayDatabase[] = []
afterEach(async () => {
@@ -29,7 +19,7 @@ interface Harness {
setNow: (value: number) => void
}
async function setup(cells: RelayCellConfig[] = [FENCED_CELL]): Promise<Harness> {
async function setup(cells: RelayCellConfig[]): Promise<Harness> {
const database = await openInMemoryRelayDatabase()
databases.push(database)
let now = START_MS
@@ -45,76 +35,14 @@ async function setup(cells: RelayCellConfig[] = [FENCED_CELL]): Promise<Harness>
cellIncarnation: `1111111${cells.indexOf(cell)}-1111-4111-8111-111111111111`,
startedAt: 50,
ready,
observedRequests: 0,
...(cell.connectionHardCap === undefined
? {}
: {
totalConnections: 0,
inFlightConnections: 0,
reservedConnectionUnits: 0,
enforcedConnectionUnits: 0,
connectionHardCap: cell.connectionHardCap,
connectionUnobservedBound: cell.connectionUnobservedBound
})
observedRequests: 0
})
}
for (const cell of cells) await heartbeat(cell, true)
return { store, heartbeat, setNow: (value: number) => (now = value) }
}
async function assignFailure(store: RelayAssignmentStore): Promise<unknown> {
return await store.assign(IDENTITY).then(
() => new Error('assign unexpectedly succeeded'),
(error: unknown) => error
)
}
function homeCellError(error: unknown): RelayHomeCellUnavailableError {
expect(error).toBeInstanceOf(RelayHomeCellUnavailableError)
// oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: the assertion above.
return error as RelayHomeCellUnavailableError
}
describe('home cell unavailable', () => {
it('names a readiness failure rather than reporting fleet capacity', async () => {
const { store, heartbeat, setNow } = await setup()
await store.assign(IDENTITY)
setNow(START_MS + 1_000)
await heartbeat(FENCED_CELL, false)
const error = homeCellError(await assignFailure(store))
expect(error.message).toBe('relay_home_cell_unavailable')
expect(error.unavailableCause).toBe('not_ready')
expect(error.cellId).toBe(FENCED_CELL.id)
})
it('names a heartbeat gap as unheard even though the cell last reported ready', async () => {
const { store, setNow } = await setup()
await store.assign(IDENTITY)
setNow(START_MS + HEARTBEAT_TTL_MS + 1)
expect(homeCellError(await assignFailure(store)).unavailableCause).toBe('unheard')
})
it('names a drained cell as draining ahead of its heartbeat gap', async () => {
const { store, setNow } = await setup()
await store.assign(IDENTITY)
await store.configureCell(FENCED_CELL, false)
setNow(START_MS + HEARTBEAT_TTL_MS + 1)
expect(homeCellError(await assignFailure(store)).unavailableCause).toBe('draining')
})
it('still reports capacity exhaustion when the fleet has no headroom', async () => {
const { store } = await setup([{ ...FENCED_CELL, capacityRequests: 1 }])
await store.assign(IDENTITY)
await expect(
store.assign({ userId: 'user-2', relayHostId: 'host000000000002' })
).rejects.toThrow('relay_capacity_exhausted')
})
it('rehomes instead of rejecting when the unavailable cell needs no fence', async () => {
const home: RelayCellConfig = {
id: 'home',
cloud/apps/relay/src/assignment/assignment-inventory-snapshot.test.ts
-108
View File
@@ -1,108 +0,0 @@
import { describe, expect, it } from 'vitest'
import {
formatAssignmentInventorySnapshot,
readAssignmentInventorySnapshot
} from './assignment-inventory-snapshot.js'
import { openInMemoryRelayDatabase } from '../database/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,
hostNotArrivedLast24Hours: 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)
})
})
cloud/apps/relay/src/assignment/assignment-isolated-cell-replacement-postgres.test.ts
+1 -4
View File
@@ -302,7 +302,7 @@ describePostgres('PostgreSQL re-placement off a cell isolated for a roll', () =>
).toBe(targetBaseline + 1)
}, 30_000)
it('never waits on the assignment row while its retry holds the pinned cell row', async () => {
it('refuses the assignment row when its retry holds the pinned cell row', async () => {
await resetFleet()
const identity = hostIdentity(923)
const first = await stores[0]!.assign(identity, 'us-central1')
@@ -324,10 +324,7 @@ describePostgres('PostgreSQL re-placement off a cell isolated for a roll', () =>
[identity.userId, identity.relayHostId]
)
await releaseCell()
const cellFreedAt = performance.now()
expect(await regrant).toBeInstanceOf(RelayAssignmentRowBusyError)
// NOWAIT once the cell row is held: nowhere near the 1s bounded wait.
expect(performance.now() - cellFreedAt).toBeLessThan(300)
} finally {
await releaseAssignment?.()
await releaseCell()
cloud/apps/relay/src/assignment/assignment-lease.test.ts
+2 -140
View File
@@ -1,22 +1,10 @@
import { createHash } from 'node:crypto'
import { EventEmitter } from 'node:events'
import { RELAY_CLOSE_CODE } from '@orca-cloud/relay-contract/close-codes'
import { SignJWT } from 'jose'
import nacl from 'tweetnacl'
import { describe, expect, it, vi } from 'vitest'
import type WebSocket from 'ws'
import { describe, expect, it } from 'vitest'
import {
ASSIGNMENT_LEASE_AUDIENCE,
assignmentLeaseKeyId,
classifyAssignmentLease
} from './assignment-lease.js'
import { AssignmentLeaseShadow } from './assignment-lease-shadow.js'
import type { RelayAssignmentStore } from './assignment-store.js'
import type { RelayConfig } from '../config.js'
import type { RelayCredentialStore } from '../credentials/credential-store.js'
import { HostSessionRegistry } from '../connections/host-session-registry.js'
import type { RelayTokenClaims } from '../authentication/relay-token-verifier.js'
import { ProcessQueuedByteBudget } from '../connections/splice-forwarder.js'
const KEY = new TextEncoder().encode('0123456789abcdef0123456789abcdef')
const OTHER_KEY = new TextEncoder().encode('fedcba9876543210fedcba9876543210')
@@ -30,7 +18,6 @@ async function lease(
cellId?: string
assignmentEpoch?: number
sub?: string
relayHostId?: string
expiresAt?: string | number
} = {}
): Promise<string> {
@@ -40,7 +27,7 @@ async function lease(
cellId: overrides.cellId ?? CELL_ID,
cellUrl: 'https://c7.relay.example.test',
assignmentEpoch: overrides.assignmentEpoch ?? 4,
relayHostId: overrides.relayHostId ?? HOST.relayHostId
relayHostId: HOST.relayHostId
})
.setProtectedHeader(kid === null ? { alg: 'HS256' } : { alg: 'HS256', kid })
.setIssuer('https://relay.example.test')
@@ -73,7 +60,6 @@ describe('assignment lease classification', () => {
expect(await classify(await lease({ kid: assignmentLeaseKeyId(OTHER_KEY) }))).toBe(
'bad-signature'
)
expect(await classify('not.a.jwt')).toBe('bad-signature')
expect(await classify('x'.repeat(5_000))).toBe('bad-signature')
expect(await classify(await lease({ sub: 'user-2' }))).toBe('wrong-host')
expect(await classify(await lease({ cellId: 'production-gce-c8' }))).toBe('wrong-cell')
@@ -81,127 +67,3 @@ describe('assignment lease classification', () => {
expect(await classify(await lease({ assignmentEpoch: 5 }))).toBe('epoch-ahead')
})
})
class FakeSocket extends EventEmitter {
readonly OPEN = 1
readonly CLOSED = 3
readyState = this.OPEN
readonly send = vi.fn()
readonly close = vi.fn((code?: number) => {
this.readyState = this.CLOSED
this.emit('close', code, Buffer.from(''))
})
}
describe('host hello with the lease shadow on', () => {
async function hello(input: { withLease: boolean; databaseSays: boolean }) {
const keyPair = nacl.box.keyPair()
const relayHostId = createHash('sha256')
.update(keyPair.publicKey)
.digest('base64url')
.slice(0, 16)
const identity: RelayTokenClaims = {
sub: HOST.userId,
prof: 'profile-1',
relayHostId,
purpose: 'host-control',
exp: 4_102_444_800
}
const recorded: string[] = []
const shadow = new AssignmentLeaseShadow({ enabled: () => true, key: KEY, cellId: CELL_ID })
const record = shadow.record.bind(shadow)
vi.spyOn(shadow, 'record').mockImplementation((check, dbValid, dbReadMs) => {
void check.then((leaseClass) => recorded.push(`${leaseClass}:${dbValid}`))
record(check, dbValid, dbReadMs)
})
const verifyCellAssignment = vi.fn(async () => input.databaseSays)
const registry = new HostSessionRegistry(
config(),
vi.fn(),
{} as RelayCredentialStore,
{ verifyCellAssignment } as unknown as RelayAssignmentStore,
new ProcessQueuedByteBudget(),
{
recordAuth: vi.fn(),
recordForwardedBytes: vi.fn(),
recordHttp: vi.fn(),
recordReconnect: vi.fn(),
recordSql: vi.fn()
},
Date.now,
Math.random,
'incarnation-1',
shadow
)
const socket = new FakeSocket()
const hostLease = input.withLease ? await lease({ relayHostId }) : undefined
registry.acceptControl(
socket as unknown as WebSocket,
identity,
undefined,
undefined,
hostLease
)
socket.emit(
'message',
Buffer.from(
JSON.stringify({
type: 'host-hello',
v: 1,
relayHostId,
assignmentEpoch: 4,
hostPublicKeyB64: Buffer.from(keyPair.publicKey).toString('base64'),
appVersion: '1.4.200'
})
),
false
)
await vi.waitFor(() => expect(recorded).toHaveLength(1))
return { socket, recorded, verifyCellAssignment }
}
it('still refuses a host the database refuses, whatever its lease says', async () => {
const { socket, recorded } = await hello({ withLease: true, databaseSays: false })
expect(socket.close).toHaveBeenCalledWith(RELAY_CLOSE_CODE.WRONG_CELL, 'wrong assignment epoch')
expect(recorded).toEqual(['valid:false'])
})
it('still admits a host the database admits, with no lease at all', async () => {
const { socket, recorded, verifyCellAssignment } = await hello({
withLease: false,
databaseSays: true
})
expect(verifyCellAssignment).toHaveBeenCalledTimes(1)
expect(socket.close).not.toHaveBeenCalled()
expect(String(socket.send.mock.calls[0]?.[0])).toContain('"type":"host-challenge"')
expect(recorded).toEqual(['absent:true'])
})
})
function config(): RelayConfig {
return {
port: 8080,
publicUrl: 'https://c7.relay.example.test',
cellUrl: 'https://c7.relay.example.test',
authIssuer: 'https://auth.example.test',
authAudience: 'orca-relay',
jwksUrl: 'https://auth.example.test/jwks',
assignmentSigningKey: KEY,
role: 'cell',
cellId: CELL_ID,
cells: [],
adminAudience: 'https://relay.example.test/v1/admin/drain',
deployServiceAccount: 'deploy@example.test',
runtimeServiceAccount: 'relay-cell@example.test',
adminJwksUrl: 'https://auth.example.test/jwks',
databasePoolMax: 10,
publicAssignmentsEnabled: true,
publicAssignmentConcurrency: 2,
publicAssignmentQueueMax: 128,
publicAssignmentWaitMs: 4_000,
publicResolveConcurrency: 1,
publicResolveWaitMs: 5_000,
publicAssignmentRetryAfterSeconds: 5,
dataDir: './data'
}
}
cloud/apps/relay/src/assignment/assignment-placed-marker.test.ts
-48
View File
@@ -1,48 +0,0 @@
import { ASSIGNMENT_LIMITS } from '@orca-cloud/relay-contract/assignment-invariants'
import { afterEach, describe, expect, it } from 'vitest'
import { RelayAssignmentStore } from './assignment-store.js'
import { openInMemoryRelayDatabase, type RelayDatabase } from '../database/database.js'
const CELL = { id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 }
const IDENTITY = { userId: 'user-a', relayHostId: 'host000000000001' }
describe('assign marks only the answers it minted as placed', () => {
let database: RelayDatabase | undefined
afterEach(async () => await database?.close())
async function setup(now: () => number): Promise<RelayAssignmentStore> {
database = await openInMemoryRelayDatabase()
const store = new RelayAssignmentStore(database, now)
await store.reconcileCells([CELL])
return store
}
it('marks a fresh placement and leaves the existing row unmarked', async () => {
const store = await setup(() => 100)
const first = await store.assign(IDENTITY)
const again = await store.assign(IDENTITY)
expect(first.placed).toBe(true)
expect(again).toMatchObject({ cellId: first.cellId, assignmentEpoch: first.assignmentEpoch })
expect(again.placed).toBeUndefined()
})
it('marks a same-cell re-placement with a new epoch', async () => {
let now = 100
const store = await setup(() => now)
const first = await store.assign(IDENTITY)
await store.changeActivity(IDENTITY, 'control', -1)
now += ASSIGNMENT_LIMITS.activityLeaseMs + 1
await store.releaseExpiredActivityLeases()
await store.releaseExpiredActivity()
now += ASSIGNMENT_LIMITS.dormantTtlMs
const replaced = await store.assign(IDENTITY)
expect(replaced).toMatchObject({
cellId: first.cellId,
assignmentEpoch: first.assignmentEpoch + 1
})
expect(replaced.placed).toBe(true)
})
})
cloud/apps/relay/src/assignment/assignment-rejection-log-window.test.ts
-82
View File
@@ -1,82 +0,0 @@
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<string>({
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('keeps distinct keys on independent windows', () => {
const timers = manualTimers()
const window = new AssignmentRejectionLogWindow<string>({
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<string>({
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<number, () => 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()
}
}
}
}
cloud/apps/relay/src/assignment/assignment-rejection-logging.test.ts
-144
View File
@@ -1,144 +0,0 @@
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('../authentication/relay-token-verifier.js', () => ({
createRelayTokenVerifier: () => fakes.verifyRelayToken,
readBearer: (value: string | undefined) => value?.replace(/^Bearer /, '') ?? null
}))
import { createRelayApp } from '../app.js'
import { relayHostLogDigest } from '../observability/relay-host-log-digest.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('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)}`)
})
})
function assignmentRequest(
relayHostId: string,
extra: Record<string, unknown> = {}
): RequestInit & { headers: Record<string, string> } {
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> = {}): 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
}
}
cloud/apps/relay/src/assignment/assignment-reserve-mode-evacuation.test.ts
-11
View File
@@ -62,15 +62,4 @@ describe('evacuation from reserve-mode cells', () => {
await store.setCellAdmitMode('cell-a', 'db')
expect(await store.completeReadyEvacuations()).toBe(1)
})
it('refuses evacuation status for a reserve cell, where activeCellMigrations would count no units', async () => {
const store = await setup(() => 1_000)
await store.setCellAdmitMode('cell-a', 'reserve')
await expect(store.cellEvacuationStatus('cell-a', 'cell-b', true)).rejects.toThrow(
'cell_in_reserve_mode'
)
await expect(store.cellEvacuationStatus('cell-b', 'cell-a', false)).rejects.toThrow(
'cell_in_reserve_mode'
)
})
})
cloud/apps/relay/src/assignment/assignment-reserve-mode-postgres.test.ts
-44
View File
@@ -202,49 +202,5 @@ describePostgres('reserve-mode census on PostgreSQL', () => {
await holder.close()
}
}, 20_000)
// A stall must not double connection demand: concurrent assigns share one reserve-set read,
// and after it fails the last answer stands for the backoff without another query.
it('reads the reserve set once for concurrent assigns during a stall, and not again in the backoff', async () => {
const timed = await openRelayDatabase({ databaseUrl, dataDir: '/tmp', statementTimeoutMs: 300 })
const holder = await openRelayDatabase({ databaseUrl, dataDir: '/tmp' })
let release!: () => void
const released = new Promise<void>((resolve) => (release = resolve))
let held!: () => void
const holding = new Promise<void>((resolve) => (held = resolve))
const lock = holder.transaction(async (transaction) => {
await transaction.query(`LOCK TABLE relay_cell_admit_modes IN ACCESS EXCLUSIVE MODE`)
held()
await released
})
try {
await holding
const store = new RelayAssignmentStore(timed, () => 500_000)
await store.reconcileCells([cellA, cellB], false)
await heartbeat(store, cellA)
await heartbeat(store, cellB)
const query = timed.query.bind(timed)
let reads = 0
timed.query = async (sql, params) => {
if (sql.includes('FROM relay_cell_admit_modes m')) reads += 1
return await query(sql, params)
}
const placed = await Promise.all(
Array.from({ length: 6 }, async (_, index) =>
await store.assign({ userId: USER, relayHostId: `hoststall0000${String(index).padStart(3, '0')}` })
)
)
expect(placed).toHaveLength(6)
expect(reads).toBe(1)
await store.assign({ userId: USER, relayHostId: 'hoststall0000100' })
expect(await store.reserveModeCells()).toBeNull()
expect(reads).toBe(1)
} finally {
release()
await lock
await timed.close()
await holder.close()
}
}, 20_000)
})
cloud/apps/relay/src/assignment/assignment-reserve-mode.test.ts
+1 -54
View File
@@ -59,30 +59,11 @@ describe('assignment reserve-mode behavior', () => {
const { unknown, restore } = await unknownSet(() => now)
expect(await unknown.reserveModeCells()).toBeNull()
expect(await unknown.releaseExpiredActivityLeases()).toBe(0)
expect(await unknown.evacuateDeadCells()).toBe(0)
await restore()
await setReserve(store, [])
expect(await store.releaseExpiredActivityLeases()).toBeGreaterThan(0)
})
it('keeps an idle host sticky on a reserve-mode cell where today it would be re-placed', async () => {
let now = 1_000
const store = await setup(() => now)
const identity = { userId: 'user-a', relayHostId: 'host000000000001' }
const first = await store.assign(identity)
await store.changeActivity(identity, 'control', -1)
now += ASSIGNMENT_LIMITS.activityLeaseMs + ASSIGNMENT_LIMITS.dormantTtlMs + 1
await setReserve(store, [first.cellId])
expect(await store.assign(identity)).toMatchObject({
cellId: first.cellId,
assignmentEpoch: first.assignmentEpoch
})
// Not knowing which cells are switched on is treated the same way.
const { unknown, restore } = await unknownSet(() => now)
expect((await unknown.assign(identity)).assignmentEpoch).toBe(first.assignmentEpoch)
await restore()
})
it('keeps a connected host pinned when its reserve cell goes existing-only after a ledger write', async () => {
let now = 1_000
const store = await setup(() => now)
@@ -125,7 +106,7 @@ describe('assignment reserve-mode behavior', () => {
}
return await query(sql, params)
}
return { counter, release, restore: () => (database!.query = query) }
return { release, restore: () => (database!.query = query) }
}
it('reads afresh after an admit-mode write instead of joining a read begun before it', async () => {
@@ -155,12 +136,10 @@ describe('assignment reserve-mode behavior', () => {
try {
const before = store.reserveModeCells()
await store.setCellAdmitMode('cell-a', 'reserve')
const reads = held.counter.reads
held.release()
await before
// A write cleared the cache: the next caller reads now, not after a 5 s backoff.
expect([...((await store.reserveModeCells()) ?? [])]).toEqual(['cell-a'])
expect(held.counter.reads).toBeGreaterThan(reads)
} finally {
held.release()
held.restore()
@@ -195,17 +174,6 @@ describe('assignment reserve-mode behavior', () => {
}
})
it('never places a fresh host on a reserve-mode cell through the database', async () => {
const store = await setup(() => 1_000, ['cell-a'])
for (let index = 0; index < 4; index += 1) {
const placed = await store.assign({
userId: `user-${index}`,
relayHostId: `host${String(index).padStart(12, '0')}`
})
expect(placed.cellId).toBe('cell-b')
}
})
it('refuses drains, evacuations and status reads that would treat a reserve-mode cell as empty', async () => {
const store = await setup(() => 1_000, ['cell-a'])
await expect(store.cellDeploymentStatus('cell-a')).rejects.toBeInstanceOf(
@@ -232,25 +200,4 @@ describe('assignment reserve-mode behavior', () => {
await expect(unknown.cellDeploymentStatus('cell-b')).rejects.toThrow('reserve_mode_unknown')
await restore()
})
it('refuses the flip to reserve while a migration names the cell', async () => {
const store = await setup(() => 1_000)
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, created_at, updated_at)
VALUES ('user-a', 'host000000000001', 'cell-b', 'cell-a', 1, 2, 0, 1, 9999, 1000, 1000)`
)
await expect(store.setCellAdmitMode('cell-a', 'reserve')).rejects.toThrow(
'cell_has_open_migration'
)
expect(await store.cellAdmitMode('cell-a')).toMatchObject({ admitMode: 'db' })
await store.setCellAdmitMode('cell-a', 'db')
await database!.query(`UPDATE relay_assignment_migrations SET aborted_at = 1001`)
await store.setCellAdmitMode('cell-a', 'reserve')
expect(await store.cellAdmitMode('cell-a')).toMatchObject({
admitMode: 'reserve',
updatedAt: 1_000
})
})
})
cloud/apps/relay/src/assignment/assignment-reserve-spill.test.ts
-152
View File
@@ -1,152 +0,0 @@
import { ASSIGNMENT_LIMITS } from '@orca-cloud/relay-contract/assignment-invariants'
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import {
RelayAssignmentStore,
type AssignOptions,
type RelayAssignment
} from './assignment-store.js'
import type { RelayCellConfig } from '../config.js'
import {
openInMemoryRelayDatabase,
type RelayDatabase,
type RelayLockOptions,
type RelayTransactionOptions,
type SqlRow
} from '../database/database.js'
import { ReservePaceError } from './reserve-placement/reserve-assignment.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 }
]
const IDENTITY = { userId: 'spill-user', relayHostId: 'spillhost0000001' }
const ON_US_RESERVE: RelayAssignment = {
...IDENTITY,
cellId: 'us-reserve',
cellUrl: 'https://us-reserve.relay.example.com',
assignmentEpoch: 8,
leaseExpiresAt: 1_000_000,
region: 'us-central1'
}
// Reports the target cell full when its reservation is taken: the race between pick and reserve.
class FullAtReserveDatabase implements RelayDatabase {
armed = false
constructor(private readonly delegate: RelayDatabase) {}
get dialect(): 'sqlite' | 'postgres' | undefined {
return this.delegate.dialect
}
async query(sql: string, params?: unknown[]): Promise<SqlRow[]> {
return await this.delegate.query(sql, params)
}
async queryLocked(sql: string, params?: unknown[], options?: RelayLockOptions): Promise<SqlRow[]> {
const rows = await this.delegate.queryLocked(sql, params, options)
if (!this.armed || sql.trim() !== 'SELECT * FROM relay_cells WHERE cell_id = ?') {
return rows
}
return rows.map((row) => ({ ...row, reserved_requests: row.capacity_requests }))
}
async transaction<T>(operation: (transaction: RelayDatabase) => Promise<T>, options?: RelayTransactionOptions): Promise<T> {
return await this.delegate.transaction(async (inner) => {
const wrapped = new FullAtReserveDatabase(inner)
wrapped.armed = this.armed
return await operation(wrapped)
}, options)
}
async close(): Promise<void> {
await this.delegate.close()
}
}
describe('reserve spill before relay_capacity_exhausted', () => {
let database: FullAtReserveDatabase
let store: RelayAssignmentStore
let now: number
beforeEach(async () => {
now = 100
database = new FullAtReserveDatabase(await openInMemoryRelayDatabase())
store = new RelayAssignmentStore(database, () => now)
await store.reconcileCells(CELLS)
})
afterEach(async () => {
vi.restoreAllMocks()
await database.close()
})
const closeAllDatabaseCells = async () => {
for (const cell of CELLS) {
await store.configureCell(cell, 'migration-only')
}
}
const assignAsia = async (step5: AssignOptions) =>
await store.assign(IDENTITY, 'asia-east2', 'asia-east2', undefined, step5)
it('books a reserve seat when no database cell has room anywhere', async () => {
await closeAllDatabaseCells()
const reserveSpill = vi.fn(async () => ON_US_RESERVE)
await expect(assignAsia({ reserveSpill })).resolves.toEqual({ ...ON_US_RESERVE, placed: true })
expect(reserveSpill).toHaveBeenCalledWith({ skipOwnRegion: false, rowEpoch: 0 })
expect(await database.query(`SELECT cell_id FROM relay_assignments`)).toEqual([])
})
it('skips the own region after the in-region reserve fallback already tried it', async () => {
await closeAllDatabaseCells()
const regionReserve = vi.fn(async () => null)
const reserveSpill = vi.fn(async () => ON_US_RESERVE)
await expect(assignAsia({ regionReserve, reserveSpill })).resolves.toMatchObject({ cellId: 'us-reserve' })
expect(regionReserve).toHaveBeenCalledOnce()
expect(reserveSpill).toHaveBeenCalledWith({ skipOwnRegion: true, rowEpoch: 0 })
})
it('hands the spill the epoch of the host row it read', async () => {
const first = await assignAsia({})
// Dormant: its next assign is a placement, which reads the row and finds no cell.
await store.changeActivity(IDENTITY, 'control', -1)
now += ASSIGNMENT_LIMITS.activityLeaseMs + 1
await store.releaseExpiredActivityLeases()
await store.releaseExpiredActivity()
now += ASSIGNMENT_LIMITS.dormantTtlMs
await closeAllDatabaseCells()
const reserveSpill = vi.fn(async () => ON_US_RESERVE)
await assignAsia({ reserveSpill })
expect(reserveSpill).toHaveBeenCalledWith({ skipOwnRegion: false, rowEpoch: first.assignmentEpoch })
})
it('refuses with relay_capacity_exhausted when no reserve seat exists either', async () => {
await closeAllDatabaseCells()
await expect(assignAsia({ reserveSpill: async () => null })).rejects.toThrow('relay_capacity_exhausted')
})
it('answers a paced spill with its pace, not a refusal', async () => {
await closeAllDatabaseCells()
const reserveSpill = vi.fn(async () => {
throw new ReservePaceError(4)
})
await expect(assignAsia({ reserveSpill })).rejects.toBeInstanceOf(ReservePaceError)
})
it('answers an in-region reserve pace before any spill', async () => {
await closeAllDatabaseCells()
const reserveSpill = vi.fn(async () => ON_US_RESERVE)
const regionReserve = async () => {
throw new ReservePaceError(2)
}
await expect(assignAsia({ regionReserve, reserveSpill })).rejects.toBeInstanceOf(ReservePaceError)
expect(reserveSpill).not.toHaveBeenCalled()
})
it('does not spill a reservation race: the host retries and re-picks', async () => {
database.armed = true
const reserveSpill = vi.fn(async () => ON_US_RESERVE)
await expect(assignAsia({ reserveSpill })).rejects.toThrow('relay_capacity_exhausted')
expect(reserveSpill).not.toHaveBeenCalled()
})
})
cloud/apps/relay/src/assignment/assignment-roll-replaced-cell.test.ts
+7 -37
View File
@@ -140,21 +140,12 @@ async function setup(cells: RelayCellConfig[] = [SOURCE, TARGET]) {
}
}
// The old process stopped heartbeating at START_MS, so it is unheard from here.
const OLD_UNHEARD_AT = START_MS + HEARTBEAT_TTL_MS + 1_000
// Both routes that refused in production: a hinted assign (placement lane after its
// verification read finds the cell not live) and resolve, which then assigns.
async function hintedAssign(store: RelayAssignmentStore) {
return await store.assign(IDENTITY, 'us-central1', 'us-central1', undefined, {
hostRequested: true
})
}
async function resolveRoute(store: RelayAssignmentStore) {
return (await store.resolve(IDENTITY)) ?? (await store.assign(IDENTITY))
}
async function expectRefused(
attempt: Promise<unknown>,
cause: 'not_ready' | 'unheard'
@@ -164,10 +155,7 @@ async function expectRefused(
expect(error).toMatchObject({ unavailableCause: cause })
}
describe.each([
['hinted assign', hintedAssign],
['resolve', resolveRoute]
] as const)('a host offline through a roll redials via %s', (_route, redial) => {
describe('a host offline through a roll redials via hinted assign', () => {
describe('while the replacement process boots (not ready)', () => {
it('is re-placed when that process started after the roll isolated the cell', async () => {
const { store, first, heartbeat, at, isolateForRoll } = await setup()
@@ -175,7 +163,7 @@ describe.each([
await at(20_000)
await heartbeat(SOURCE, { incarnation: NEW_INCARNATION, startedAt: 15_000, ready: false })
const moved = await redial(store)
const moved = await hintedAssign(store)
expect(moved).toMatchObject({ cellId: TARGET.id })
expect(moved.assignmentEpoch).toBeGreaterThan(first.assignmentEpoch)
})
@@ -186,7 +174,7 @@ describe.each([
await at(20_000)
await heartbeat(SOURCE, { ready: false })
await expectRefused(redial(store), 'not_ready')
await expectRefused(hintedAssign(store), 'not_ready')
})
it('is refused while the new process reports holding connections', async () => {
@@ -200,7 +188,7 @@ describe.each([
totalConnections: 1
})
await expectRefused(redial(store), 'not_ready')
await expectRefused(hintedAssign(store), 'not_ready')
})
it('is refused once the new process itself goes unheard', async () => {
@@ -210,15 +198,7 @@ describe.each([
await heartbeat(SOURCE, { incarnation: NEW_INCARNATION, startedAt: 15_000, ready: false })
await at(20_000 + HEARTBEAT_TTL_MS + 1_000)
await expectRefused(redial(store), 'unheard')
})
it('is refused on an ordinary readiness dip with no roll stamp', async () => {
const { store, heartbeat, at } = await setup()
await at(20_000)
await heartbeat(SOURCE, { incarnation: NEW_INCARNATION, startedAt: 15_000, ready: false })
await expectRefused(redial(store), 'not_ready')
await expectRefused(hintedAssign(store), 'unheard')
})
it('is refused on a migration-only cell parked without a roll stamp', async () => {
@@ -227,20 +207,10 @@ describe.each([
await at(20_000)
await heartbeat(SOURCE, { incarnation: NEW_INCARNATION, startedAt: 15_000, ready: false })
await expectRefused(redial(store), 'not_ready')
await expectRefused(hintedAssign(store), 'not_ready')
})
})
// Residual: the old process is past its TTL and the new one has not heartbeated yet, so
// nothing proves the drained sockets gone. The client's retry lands once the new one is heard.
it('is still refused between the old process going unheard and the new one heartbeating', async () => {
const { store, at, isolateForRoll } = await setup()
await isolateForRoll()
await at(OLD_UNHEARD_AT)
await expectRefused(redial(store), 'unheard')
})
it('is refused when the roll stamp is stale, even with a new process', async () => {
const { store, heartbeat, at, isolateForRoll } = await setup()
await isolateForRoll()
@@ -248,7 +218,7 @@ describe.each([
await at(stale)
await heartbeat(SOURCE, { incarnation: NEW_INCARNATION, startedAt: stale - 1, ready: false })
await expectRefused(redial(store), 'not_ready')
await expectRefused(hintedAssign(store), 'not_ready')
})
})
cloud/apps/relay/src/assignment/assignment-row-lock-order-postgres.test.ts
-134
View File
@@ -1,134 +0,0 @@
import { afterAll, beforeAll, describe, expect, it } from 'vitest'
import { RelayAssignmentStore } from './assignment-store.js'
import { openRelayDatabase, type RelayDatabase } from '../database/database.js'
const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL
const describePostgres = databaseUrl ? describe : describe.skip
const cells = [
{
id: 'row-lock-order-a',
url: 'https://row-lock-order-a.example.com',
capacityRequests: 200,
connectionHardCap: 600 as const,
connectionUnobservedBound: 50
},
{
id: 'row-lock-order-b',
url: 'https://row-lock-order-b.example.com',
capacityRequests: 200,
connectionHardCap: 600 as const,
connectionUnobservedBound: 50
}
]
const userId = 'row-lock-order-user'
const hosts = ['rowlockhost00001', 'rowlockhost00002', 'rowlockhost00003', 'rowlockhost00004'].map(
(relayHostId) => ({ userId, relayHostId })
)
// pg_stat_database exposes cycles hidden by three 40P01 retries; eventual success alone misses deadlocks that exhaust retries under load.
describePostgres('PostgreSQL row lock order', () => {
const databases: RelayDatabase[] = []
beforeAll(async () => {
for (let index = 0; index < 4; index++) {
databases.push(await openRelayDatabase({ databaseUrl, dataDir: '' }))
}
})
async function removeTestRows(database: RelayDatabase): Promise<void> {
await database.query(
`DELETE FROM relay_control_connection_reservations WHERE user_id = ?`,
[userId]
)
for (const table of [
'relay_control_capabilities',
'relay_assignment_activity_leases',
'relay_post_drain_migration_pins',
'relay_assignment_migration_incarnations',
'relay_assignment_migrations',
'relay_assignments'
]) {
await database.query(`DELETE FROM ${table} WHERE user_id = ?`, [userId])
}
for (const cell of cells) {
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])
}
}
}
afterAll(async () => {
if (databases[0]) await removeTestRows(databases[0])
for (const connection of databases) await connection.close()
})
async function deadlockCount(): Promise<number> {
const rows = await databases[0]!.query(
`SELECT deadlocks FROM pg_stat_database WHERE datname = current_database()`
)
return Number(rows[0]!.deadlocks)
}
it('runs the cell accept and the placement retry concurrently without a cycle', async () => {
await removeTestRows(databases[0]!)
const stores = databases.map((database) => new RelayAssignmentStore(database, () => 100))
await stores[0]!.reconcileCells(cells)
for (const cell of cells) {
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 placements = new Map<string, { cellId: string; assignmentEpoch: number }>()
for (const identity of hosts) {
const assignment = await stores[0]!.assign(identity)
placements.set(identity.relayHostId, assignment)
}
const before = await deadlockCount()
// The accept path (host rows, then the cell row last) against the paths
// that must read the inventory first: placement and evacuation.
for (let round = 0; round < 12; round++) {
await Promise.allSettled(
hosts.flatMap((identity, index) => {
const placement = placements.get(identity.relayHostId)!
const store = stores[index % stores.length]!
const other = stores[(index + 1) % stores.length]!
return [
store.activateControl(identity, {
cellId: placement.cellId,
assignmentEpoch: placement.assignmentEpoch,
generation: round + 2
}),
other.assign(identity),
other.startEvacuation(
identity,
placement.cellId === cells[0]!.id ? cells[1]!.id : cells[0]!.id
)
]
})
)
}
const after = await deadlockCount()
expect(after - before).toBe(0)
}, 120_000)
})
cloud/apps/relay/src/assignment/migrations/dead-cell-evacuation-precheck.test.ts
+7 -105
View File
@@ -1,83 +1,14 @@
import pg from 'pg'
import { afterAll, afterEach, describe, expect, it } from 'vitest'
import { afterEach, describe, expect, it } from 'vitest'
import { RelayAssignmentStore } from '../assignment-store.js'
import type { RelayCellConfig } from '../../config.js'
import {
openInMemoryRelayDatabase,
openRelayDatabase,
type RelayDatabase,
type SqlRow
} from '../../database/database.js'
import { openInMemoryRelayDatabase, type RelayDatabase } from '../../database/database.js'
// Both dialects must skip host scans with no actionable cells; stale existing-only rows previously scanned every assignment for nothing.
const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL
const schema = 'relay_dead_cell_precheck_test'
const HOST_QUERY = 'FROM relay_assignments assignment'
const CELLS: RelayCellConfig[] = [
{ id: 'cell-a', url: 'https://relay-a.example.com', capacityRequests: 10 },
{ id: 'cell-b', url: 'https://relay-b.example.com', capacityRequests: 10 }
]
class RecordingDatabase implements RelayDatabase {
readonly statements: string[] = []
constructor(private readonly delegate: RelayDatabase) {}
get dialect() {
return this.delegate.dialect
}
async query(sql: string, params: unknown[] = []): Promise<SqlRow[]> {
this.statements.push(sql)
return await this.delegate.query(sql, params)
}
async queryLocked(...args: Parameters<RelayDatabase['queryLocked']>): Promise<SqlRow[]> {
return await this.delegate.queryLocked(...args)
}
async transaction<T>(operation: (transaction: RelayDatabase) => Promise<T>): Promise<T> {
return await this.delegate.transaction(operation)
}
async close(): Promise<void> {
await this.delegate.close()
}
}
function scopedUrl(): string {
const url = new URL(databaseUrl!)
url.searchParams.set('options', `-c search_path=${schema}`)
return url.toString()
}
async function onAdmin(sql: string): Promise<void> {
const client = new pg.Client({ connectionString: databaseUrl })
await client.connect()
try {
await client.query(sql)
} finally {
await client.end()
}
}
const dialects: [string, () => Promise<RelayDatabase>][] = [
['sqlite', openInMemoryRelayDatabase],
...(databaseUrl
? [
[
'postgres',
async () => {
await onAdmin(`DROP SCHEMA IF EXISTS ${schema} CASCADE`)
await onAdmin(`CREATE SCHEMA ${schema}`)
return await openRelayDatabase({ databaseUrl: scopedUrl(), dataDir: '' })
}
] as [string, () => Promise<RelayDatabase>]
]
: [])
]
describe.each(dialects)('dead-cell evacuation pre-check (%s)', (_dialect, open) => {
describe('dead-cell evacuation pre-check', () => {
let database: RelayDatabase | undefined
afterEach(async () => {
@@ -85,19 +16,14 @@ describe.each(dialects)('dead-cell evacuation pre-check (%s)', (_dialect, open)
database = undefined
})
afterAll(async () => {
if (databaseUrl) await onAdmin(`DROP SCHEMA IF EXISTS ${schema} CASCADE`)
})
async function setup(now: () => number) {
database = await open()
const recording = new RecordingDatabase(database)
const store = new RelayAssignmentStore(recording, now, {
database = await openInMemoryRelayDatabase()
const store = new RelayAssignmentStore(database, now, {
requireLiveCells: true,
heartbeatTtlMs: 45_000
})
await store.reconcileCells(CELLS)
return { store, recording }
return { store }
}
async function heartbeat(store: RelayAssignmentStore, cell: RelayCellConfig): Promise<void> {
@@ -111,31 +37,9 @@ describe.each(dialects)('dead-cell evacuation pre-check (%s)', (_dialect, open)
})
}
it('skips the host query when the only dead cell is unfenced and existing-only', async () => {
let now = 100
const { store, recording } = await setup(() => now)
for (const cell of CELLS) await heartbeat(store, cell)
const identity = { userId: 'user-a', relayHostId: 'host000000000001' }
await store.setCellEnabled('cell-b', false)
expect(await store.assign(identity)).toMatchObject({ cellId: 'cell-a' })
await store.setCellEnabled('cell-b', true)
await store.setCellEnabled('cell-a', false)
now += 45_001
await heartbeat(store, CELLS[1]!)
recording.statements.length = 0
expect(await store.evacuateDeadCells()).toBe(0)
expect(recording.statements.some((sql) => sql.includes(HOST_QUERY))).toBe(false)
expect(
await database!.query(`SELECT cell_id FROM relay_assignments WHERE user_id = ?`, [
identity.userId
])
).toEqual([{ cell_id: 'cell-a' }])
})
it('still evacuates hosts from a dead uncapped cell that admits', async () => {
let now = 100
const { store, recording } = await setup(() => now)
const { store } = await setup(() => now)
for (const cell of CELLS) await heartbeat(store, cell)
const identity = { userId: 'user-a', relayHostId: 'host000000000001' }
await store.setCellEnabled('cell-b', false)
@@ -143,10 +47,8 @@ describe.each(dialects)('dead-cell evacuation pre-check (%s)', (_dialect, open)
await store.setCellEnabled('cell-b', true)
now += 45_001
await heartbeat(store, CELLS[1]!)
recording.statements.length = 0
expect(await store.evacuateDeadCells()).toBe(1)
expect(recording.statements.some((sql) => sql.includes(HOST_QUERY))).toBe(true)
expect((await store.resolve(identity))?.cellId).toBe('cell-b')
})
})
cloud/apps/relay/src/assignment/migrations/drain-migration-row-lookup.test.ts
-54
View File
@@ -1,54 +0,0 @@
import { expect, it } from 'vitest'
import { createDrainMigrationRowLookup } from './drain-migration-row-lookup.js'
import type { SqlRow } from '../../database/database.js'
function text(row: SqlRow, field: string): string {
const value = row[field]
if (typeof value !== 'string') {
throw new Error(`invalid_${field}`)
}
return value
}
it('keeps first assignment matches, lease order, row identity, and separate identity components', () => {
const rows = [
{ user_id: 'a:b', relay_host_id: 'c', value: 1 },
{ user_id: 'a', relay_host_id: 'b:c', value: 2 },
{ user_id: 'a:b', relay_host_id: 'c', value: 3 },
{ user_id: '', relay_host_id: '', value: 4 }
]
const lookup = createDrainMigrationRowLookup(rows, text)
for (const identity of [
{ userId: 'a:b', relayHostId: 'c' },
{ userId: 'a', relayHostId: 'b:c' },
{ userId: '', relayHostId: '' },
{ userId: 'missing', relayHostId: 'c' }
]) {
const expected = rows.filter(
(row) =>
row.user_id === identity.userId && row.relay_host_id === identity.relayHostId
)
expect(lookup.first(identity)).toBe(expected[0])
expect(lookup.all(identity)).toEqual(expected)
lookup.all(identity).forEach((row, index) => expect(row).toBe(expected[index]))
}
})
it('retains lazy validation and short circuiting when an inventory is malformed', () => {
const first = { user_id: 'user', relay_host_id: 'host' }
const identity = { userId: 'user', relayHostId: 'host' }
const lookup = createDrainMigrationRowLookup(
[first, { user_id: null, relay_host_id: 'bad' }],
text
)
expect(lookup.first(identity)).toBe(first)
expect(() => lookup.all(identity)).toThrow('invalid_user_id')
const unrelated = createDrainMigrationRowLookup(
[first, { user_id: 'other', relay_host_id: null }],
text
)
expect(unrelated.all(identity)).toEqual([first])
expect(() => unrelated.first({ userId: 'other', relayHostId: 'host' })).toThrow(
'invalid_relay_host_id'
)
})
cloud/apps/relay/src/assignment/migrations/drain-migration-row-scaling.test.ts
-122
View File
@@ -1,122 +0,0 @@
import { ASSIGNMENT_LIMITS } from '@orca-cloud/relay-contract/assignment-invariants'
import { afterEach, expect, it } from 'vitest'
import { RelayAssignmentStore } from '../assignment-store.js'
import { openInMemoryRelayDatabase, type RelayDatabase, type SqlRow } from '../../database/database.js'
let database: RelayDatabase | undefined
afterEach(async () => await database?.close())
it.each([false, true])(
'looks up a whole drain inventory with linear identity reads (expired: %s)',
async (expired) => {
database = await openInMemoryRelayDatabase()
let measuring = false
let identityReads = 0
let indexedRows = 0
const instrument = (delegate: RelayDatabase): RelayDatabase => ({
dialect: delegate.dialect,
query: (sql, params) => delegate.query(sql, params),
queryLocked: async (sql, params, options) => {
const rows = await delegate.queryLocked(sql, params, options)
if (
!measuring ||
!sql.includes('WHERE EXISTS') ||
!(sql.includes('SELECT assignment.*') || sql.includes('SELECT lease.*'))
) {
return rows
}
indexedRows += rows.length
return rows.map(
(row): SqlRow =>
new Proxy(row, {
get(target, key) {
if (key === 'user_id' || key === 'relay_host_id') {
identityReads++
}
return Reflect.get(target, key)
}
})
)
},
transaction: (operation, options) =>
delegate.transaction((tx) => operation(instrument(tx)), options),
close: () => delegate.close()
})
let now = 100
const store = new RelayAssignmentStore(instrument(database), () => now, {
requireLiveCells: true
})
const cells = ['a', 'b'].map((id) => ({
id: `cell-${id}`,
url: `https://relay-${id}.example.com`,
capacityRequests: 500
}))
await store.reconcileCells(cells)
const incarnation = '11111111-1111-4111-8111-111111111111'
for (const cell of cells) {
await store.recordCellHeartbeat({
cellId: cell.id,
cellUrl: cell.url,
cellIncarnation: incarnation,
startedAt: 50,
ready: true,
observedRequests: 0
})
}
await store.setCellEnabled('cell-b', false)
const identities = Array.from({ length: 50 }, (_, index) => ({
userId: `user-${index % 5}`,
relayHostId: `host${String(index).padStart(12, '0')}`
}))
for (const identity of identities) {
await store.assign(identity)
}
await store.setCellEnabled('cell-b', true)
await store.setCellEnabled('cell-a', false)
for (const identity of identities) {
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: incarnation,
traceValue: '66666666-6666-4666-8666-666666666666',
plannedGraceMs: 120_000
}
await store.prepareCellDrainAttempt(attempt)
if (expired) {
now += ASSIGNMENT_LIMITS.migrationLeaseMs + 1
await store.releaseExpiredActivityLeases()
for (const cell of cells) {
await store.recordCellHeartbeat({
cellId: cell.id,
cellUrl: cell.url,
cellIncarnation: incarnation,
startedAt: 50,
ready: true,
observedRequests: 0
})
}
}
measuring = true
await expect(store.beginCellDrainSend(attempt)).resolves.toMatchObject({
state: 'send-may-have-started',
shouldSend: true
})
expect(indexedRows).toBeGreaterThanOrEqual(100)
expect(identityReads).toBeLessThanOrEqual(indexedRows * 2)
const migrations = await database.query(
'SELECT expires_at FROM relay_assignment_migrations'
)
expect(migrations).toHaveLength(50)
expect(
migrations.every(
(row) => row.expires_at === now + ASSIGNMENT_LIMITS.migrationLeaseMs
)
).toBe(true)
}
)
cloud/apps/relay/src/assignment/migrations/drain-release-cell-row-contention-postgres.test.ts
-282
View File
@@ -1,282 +0,0 @@
import { createHash } from 'node:crypto'
import { afterAll, beforeAll, describe, expect, it } from 'vitest'
import { RelayAssignmentStore } from '../assignment-store.js'
import {
encodeMembership,
type CellAdmissionMembership,
type CellAdmissionState
} from '../../cell/admission/cell-admission-selector.js'
import type { RelayCellConfig } from '../../config.js'
import { openRelayDatabase, type RelayDatabase } from '../../database/database.js'
const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL
const describePostgres = databaseUrl ? describe : describe.skip
const DIRECTOR_INSTANCES = 5
const DIRECTOR_POOL_MAX = 3
const DIRECTOR_APPLICATION = 'drain-release-contention-director'
const USER_PREFIX = 'drain-release-contention'
const NOW = 1_000_000
const CAPPED = {
capacityRequests: 6_000,
connectionHardCap: 3_000,
connectionUnobservedBound: 60
} as const
const SOURCE: RelayCellConfig = {
id: 'drain-release-contention-a-source',
url: 'https://drain-release-contention-source.example.test',
region: 'asia-east2',
...CAPPED
}
const TARGETS: RelayCellConfig[] = ['b', 'c', 'd'].map((suffix) => ({
id: `drain-release-contention-${suffix}-target`,
url: `https://drain-release-contention-${suffix}-target.example.test`,
region: 'asia-east2',
...CAPPED
}))
// Next to the database. Taken only once every Asia neighbour is at its cap.
const US_CELL: RelayCellConfig = {
id: 'drain-release-contention-e-us',
url: 'https://drain-release-contention-us.example.test',
region: 'us-central1',
...CAPPED
}
const CELLS = [SOURCE, ...TARGETS, US_CELL]
type Identity = { userId: string; relayHostId: string }
type Instance = { database: RelayDatabase; store: RelayAssignmentStore }
function hostIdentity(index: number): Identity {
return {
userId: `${USER_PREFIX}-${index}`,
// Relay host ids are fixed-width opaque ids.
relayHostId: `drainrel${String(index).padStart(8, '0')}`
}
}
describePostgres('PostgreSQL drain releases against director placement', () => {
const storeOptions = { requireLiveCells: true, heartbeatTtlMs: 45_000 }
const directors: Instance[] = []
let observer: RelayDatabase
function admin(): Instance {
return directors[0]!
}
async function applySelector(
states: Record<string, CellAdmissionState>,
rollIsolatedCells?: string[]
): Promise<void> {
const current = await admin().store.inspectCellAdmissionSelector()
// Built from relay_cells: the apply requires exact coverage of a shared fleet.
const fleet = await admin().database.query(
`SELECT cell_id FROM relay_cells ORDER BY cell_id ASC`
)
const membership: CellAdmissionMembership = {
existingOnly: [],
migrationOnly: [],
general: []
}
for (const row of fleet) {
const cellId = String(row['cell_id'])
const state =
states[cellId] ??
(current.selector.membership.existingOnly.includes(cellId)
? 'existing-only'
: current.selector.membership.migrationOnly.includes(cellId)
? 'migration-only'
: 'general')
if (state === 'existing-only') membership.existingOnly.push(cellId)
else if (state === 'migration-only') membership.migrationOnly.push(cellId)
else membership.general.push(cellId)
}
await admin().store.applyCellAdmissionSelector({
attemptId: `drain-release-${current.selector.generation}`,
expectedGeneration: current.selector.generation,
...(current.selector.generation === 0
? {
expectedMembershipSha256: createHash('sha256')
.update(encodeMembership(current.selector.membership))
.digest('hex')
}
: {}),
membership,
...(rollIsolatedCells ? { rollIsolatedCells } : {})
})
}
// Other Postgres files write admission through the generation-0 helpers, so
// the selector goes back to generation 0 exactly as found.
async function resetSelectorBoundary(): Promise<void> {
await admin().database.query(
`UPDATE relay_admission_selectors SET generation = 0, attempt_id = NULL
WHERE selector_id = 'general'`
)
await admin().database.query(
`DELETE FROM relay_admission_selector_intents WHERE attempt_id LIKE 'drain-release-%'`
)
await admin().store.reconcileCells([], false)
}
async function deleteHostRows(): Promise<void> {
for (const table of [
'relay_control_capabilities',
'relay_control_connection_reservations',
'relay_assignment_activity_leases',
'relay_assignment_migrations',
'relay_assignment_region_preferences',
'relay_assignments'
]) {
await admin().database.query(`DELETE FROM ${table} WHERE user_id LIKE '${USER_PREFIX}-%'`)
}
}
// `atCap` reports those cells at their connection cap: no headroom left.
async function heartbeatAll(atCap: RelayCellConfig[] = []): Promise<void> {
for (const [index, config] of CELLS.entries()) {
await admin().store.recordCellHeartbeat({
cellId: config.id,
cellUrl: config.url,
cellIncarnation: `2222222${index}-2222-4222-8222-222222222222`,
startedAt: NOW - 1_000,
ready: true,
observedRequests: 0,
region: config.region,
// The cell checks enforced = total + in-flight + reserved.
totalConnections: atCap.includes(config) ? CAPPED.connectionHardCap : 0,
inFlightConnections: 0,
reservedConnectionUnits: 0,
enforcedConnectionUnits: atCap.includes(config) ? CAPPED.connectionHardCap : 0,
connectionHardCap: CAPPED.connectionHardCap,
connectionUnobservedBound: CAPPED.connectionUnobservedBound
})
}
}
// Holds `cellIds`' relay_cells rows in an open transaction until the returned release runs.
async function holdCellRows(cellIds: string[]): Promise<() => Promise<void>> {
let release!: () => void
const released = new Promise<void>((resolve) => (release = resolve))
let held!: () => void
const acquired = new Promise<void>((resolve) => (held = resolve))
const holder = observer.transaction(async (transaction) => {
await transaction.queryLocked(
`SELECT cell_id FROM relay_cells WHERE cell_id IN (${cellIds.map(() => '?').join(', ')})
ORDER BY cell_id ASC`,
cellIds
)
held()
await released
})
await acquired
return async () => {
release()
await holder
}
}
// Every host holds a live control lease on the source, then the source is
// isolated for a roll: the state an Asia drain starts from.
async function seed(hosts: Identity[]): Promise<void> {
await deleteHostRows()
await admin().database.query(
`UPDATE relay_cells SET reserved_requests = 0 WHERE cell_id LIKE '${USER_PREFIX}-%'`
)
// Restoring the source to general also clears the previous run's roll stamp.
await applySelector({
[SOURCE.id]: 'general',
...Object.fromEntries(TARGETS.map((target) => [target.id, 'migration-only' as const]))
})
for (const identity of hosts) {
const grant = await admin().store.assign(identity, 'asia-east2', 'asia-east2')
expect(grant.cellId).toBe(SOURCE.id)
await admin().store.activateControl(identity, {
cellId: SOURCE.id,
assignmentEpoch: grant.assignmentEpoch,
generation: 1
})
}
await applySelector(
Object.fromEntries(TARGETS.map((target) => [target.id, 'general' as const]))
)
await applySelector({ [SOURCE.id]: 'migration-only' }, [SOURCE.id])
}
beforeAll(async () => {
for (let index = 0; index < DIRECTOR_INSTANCES; index += 1) {
const database = await openRelayDatabase({
databaseUrl,
dataDir: '',
poolMax: DIRECTOR_POOL_MAX,
applicationName: DIRECTOR_APPLICATION
})
directors.push({
database,
store: new RelayAssignmentStore(database, () => NOW, storeOptions)
})
}
observer = await openRelayDatabase({ databaseUrl, dataDir: '', poolMax: 1 })
await deleteHostRows()
await admin().store.reconcileCells(CELLS, false)
await resetSelectorBoundary()
await heartbeatAll()
})
afterAll(async () => {
if (directors.length > 0) {
await deleteHostRows()
// Cells before the selector rebuild, or its membership names missing rows.
for (const config of CELLS) {
for (const table of [
'relay_cell_connection_snapshots',
'relay_cell_connection_runtime',
'relay_cell_connection_limits',
'relay_cell_runtime',
'relay_cell_admission',
'relay_cell_regions',
'relay_cells'
]) {
await admin().database.query(`DELETE FROM ${table} WHERE cell_id = ?`, [config.id])
}
}
await resetSelectorBoundary()
}
for (const instance of directors) {
await instance.database.close()
}
await observer?.close()
})
it('places a reconnecting host off the isolated source while that source row is held', async () => {
await heartbeatAll()
const host = hostIdentity(9_000)
await seed([host])
// A drain release holds the source row for its whole transaction; target-only placement never asks for it.
const sourceHeld = await holdCellRows([SOURCE.id])
try {
const grant = await Promise.race([
directors[1]!.store.assign(host, 'asia-east2', 'asia-east2'),
new Promise<'blocked'>((resolve) => setTimeout(() => resolve('blocked'), 3_000))
])
expect(grant).not.toBe('blocked')
expect(TARGETS.map(({ id }) => id)).toContain(grant === 'blocked' ? null : grant.cellId)
} finally {
await sourceHeld()
}
// Control: the same hold on the target rows does block a placement, so the hold above was real.
const other = hostIdentity(9_001)
await seed([other])
const targetsHeld = await holdCellRows(TARGETS.map(({ id }) => id))
const placement = directors[2]!.store.assign(other, 'asia-east2', 'asia-east2')
try {
const raced = await Promise.race([
placement.then(() => 'placed' as const),
new Promise<'blocked'>((resolve) => setTimeout(() => resolve('blocked'), 1_000))
])
expect(raced).toBe('blocked')
} finally {
await targetsHeld()
}
expect((await placement).cellId).not.toBe(SOURCE.id)
}, 30_000)
})
cloud/apps/relay/src/assignment/migrations/migration-recovery-postgres.test.ts
+4 -63
View File
@@ -256,7 +256,7 @@ describePostgres('PostgreSQL migration recovery', () => {
expect(
await database.query(
`SELECT assignment.cell_id, assignment.assignment_epoch,
migration.completed_at, migration.aborted_at
(migration.completed_at IS NOT NULL) AS completed, migration.aborted_at
FROM relay_assignments assignment
JOIN relay_assignment_migrations migration
ON migration.user_id = assignment.user_id
@@ -269,7 +269,7 @@ describePostgres('PostgreSQL migration recovery', () => {
{
cell_id: input.cells.failed.id,
assignment_epoch: '2',
completed_at: String(100 + 45_001 + ASSIGNMENT_LIMITS.migrationLeaseMs + 1),
completed: true,
aborted_at: null
}
])
@@ -309,7 +309,7 @@ describePostgres('PostgreSQL migration recovery', () => {
expect(
await database.query(
`SELECT assignment.cell_id, assignment.assignment_epoch,
migration.completed_at, migration.aborted_at
migration.completed_at, (migration.aborted_at IS NOT NULL) AS aborted
FROM relay_assignments assignment
JOIN relay_assignment_migrations migration
ON migration.user_id = assignment.user_id
@@ -323,9 +323,7 @@ describePostgres('PostgreSQL migration recovery', () => {
cell_id: input.cells.source.id,
assignment_epoch: '3',
completed_at: null,
aborted_at: String(
100 + ASSIGNMENT_LIMITS.migrationLeaseMs + STRANDED_MIGRATION_ABANDON_MS + 1
)
aborted: true
}
])
})
@@ -394,63 +392,6 @@ describePostgres('PostgreSQL migration recovery', () => {
])
})
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<number> => 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)
cloud/apps/relay/src/assignment/migrations/postgres-drain-send-locking.test.ts
+2 -2
View File
@@ -184,7 +184,7 @@ describePostgres('PostgreSQL drain-send locking', () => {
})
now = migration.expiresAt + 1
expect(await store.releaseExpiredActivityLeases()).toBeGreaterThan(0)
await store.releaseExpiredActivityLeases()
await store.recordCellHeartbeat({
cellId: source.id,
cellUrl: source.url,
@@ -207,7 +207,7 @@ describePostgres('PostgreSQL drain-send locking', () => {
cellId: source.id,
cellIncarnation: sourceIncarnation
})
).resolves.toMatchObject({ state: 'send-may-have-started', shouldSend: true })
).resolves.toMatchObject({ shouldSend: true })
await expect(
database.query(
`SELECT activity_kind, cell_id FROM relay_assignment_activity_leases
cloud/apps/relay/src/assignment/migrations/registered-migration-inventory-postgres.test.ts
-139
View File
@@ -1,139 +0,0 @@
import pg from 'pg'
import { afterAll, beforeAll, describe, expect, it } from 'vitest'
import { openRelayDatabase, type RelayDatabase } from '../../database/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
}
]
})
})
})
cloud/apps/relay/src/assignment/migrations/registered-migration-inventory.test.ts
-120
View File
@@ -1,120 +0,0 @@
import { afterEach, describe, expect, it } from 'vitest'
import { openInMemoryRelayDatabase, type RelayDatabase } from '../../database/database.js'
import { REGISTERED_MIGRATION_ABANDON_MS } from './registered-migration-abandonment.js'
import {
formatRegisteredMigrationInventory,
readRegisteredMigrationInventory
} from './registered-migration-inventory.js'
describe('registered migration inventory', () => {
const databases: RelayDatabase[] = []
afterEach(async () => {
for (const database of databases.splice(0)) {
await database.close()
}
})
it('does not restart abandonment when an old source migration lease is refreshed', async () => {
const database = await openInMemoryRelayDatabase()
databases.push(database)
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'
])
await database.query(`UPDATE relay_assignment_migrations SET expires_at = ?, updated_at = ?`, [
now + 1,
now
])
expect((await readRegisteredMigrationInventory(database, now + 1)).abandoned).toBe(1)
})
it('includes unregistered open migrations without logging a host identity', async () => {
const database = await openInMemoryRelayDatabase()
databases.push(database)
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<void> {
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<void> {
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]
)
}
cloud/apps/relay/src/assignment/public-admission/drain-return-admission.test.ts
+6 -81
View File
@@ -1,23 +1,16 @@
import { describe, expect, it } from 'vitest'
import {
DRAIN_RETURN_MIN_RETRY_AFTER_SECONDS,
RelayDrainReturnAdmission,
type DrainReturnDeferral,
type DrainReturnGrant
} from './drain-return-admission.js'
import { RelayPublicAssignmentAdmission } from './public-assignment-admission.js'
type Timer = { at: number; callback: () => void; cancelled: boolean }
function harness(
overrides: {
maxQueued?: number
maxRetryAfterSeconds?: number
onServiceMs?: (durationMs: number) => void
} = {}
) {
let now = 0
const timers: Timer[] = []
const placement = new RelayPublicAssignmentAdmission({
maxConcurrent: 2,
maxQueued: 8,
@@ -26,32 +19,18 @@ function harness(
maxDrainReturnConcurrent: 1,
maxDrainReturnQueued: overrides.maxQueued ?? 0,
drainReturnWaitMs: 3_000,
drainReturnMinIntervalMs: DRAIN_RETURN_MIN_RETRY_AFTER_SECONDS * 1_000,
now: () => now,
schedule: (callback, delayMs) => {
const timer = { at: now + delayMs, callback, cancelled: false }
timers.push(timer)
return () => {
timer.cancelled = true
}
}
drainReturnMinIntervalMs: 2_000,
now: () => 0,
schedule: () => () => undefined
})
const lane = new RelayDrainReturnAdmission(placement, {
maxConcurrent: 1,
maxRetryAfterSeconds: overrides.maxRetryAfterSeconds ?? 300,
now: () => now,
onServiceMs: overrides.onServiceMs
maxRetryAfterSeconds: 300,
now: () => 0
})
return {
lane,
placement,
advance: (ms: number) => {
now += ms
for (const timer of timers.filter((entry) => !entry.cancelled && entry.at <= now)) {
timer.cancelled = true
timer.callback()
}
}
placement
}
}
@@ -60,64 +39,10 @@ function admitted(result: DrainReturnGrant | DrainReturnDeferral): DrainReturnGr
return result
}
function deferred(result: DrainReturnGrant | DrainReturnDeferral): DrainReturnDeferral {
if (result.kind !== 'deferred') throw new Error('expected a deferral')
return result
}
const host = (index: number): string => `host${String(index).padStart(12, '0')}`
describe('drain-return admission', () => {
it('shortens the pacing once the measured service time is short', async () => {
const { lane, advance } = harness()
for (let index = 0; index < 30; index++) {
const grant = admitted(await lane.acquire(host(index)))
advance(50)
grant.lease.release()
advance(2_000)
}
admitted(await lane.acquire(host(100)))
const retries = []
for (let index = 101; index <= 120; index++) {
retries.push(deferred(await lane.acquire(host(index))).retryAfterSeconds)
}
// ~50 ms per re-placement: twenty deferrals span about one second, where the
// 860 ms starting estimate would have spread them over seventeen.
expect(Math.max(...retries)).toBe(3)
})
// The unclamped sample: the EWMA's floor and ceiling would hide the real tail.
it('answers a host’s own early retry with its interval, not a place behind the cohort', async () => {
const { lane } = harness()
admitted(await lane.acquire(host(0))).lease.release()
// A row-busy refusal says "retry in 1 s"; that redial is the host's own.
const repeat = deferred(await lane.acquire(host(0)))
expect(repeat).toMatchObject({
reason: 'host-rate-limited',
retryAfterSeconds: DRAIN_RETURN_MIN_RETRY_AFTER_SECONDS
})
// And it booked nothing: the next overflow still gets the first slot.
admitted(await lane.acquire(host(1)))
expect(deferred(await lane.acquire(host(2))).retryAfterSeconds).toBe(2)
})
it('keeps a host’s reserved slot when it comes back early', async () => {
const { lane, advance } = harness()
admitted(await lane.acquire(host(0)))
for (let index = 1; index <= 4; index++) await lane.acquire(host(index))
expect(deferred(await lane.acquire(host(5))).retryAfterSeconds).toBe(6)
advance(1_000)
// Early and still full: the same instant, not a new slot after everyone else.
expect(deferred(await lane.acquire(host(5))).retryAfterSeconds).toBe(5)
})
it('lets a queued placement go before a queued drain return', async () => {
const { lane, placement } = harness({ maxQueued: 1 })
const drain = admitted(await lane.acquire(host(0)))
cloud/apps/relay/src/assignment/public-admission/public-assignment-admission.test.ts
+3 -42
View File
@@ -1,4 +1,4 @@
import { describe, expect, it, vi } from 'vitest'
import { describe, expect, it } from 'vitest'
import { RelayPublicAssignmentAdmission } from './public-assignment-admission.js'
describe('public assignment admission', () => {
@@ -7,7 +7,7 @@ describe('public assignment admission', () => {
const admission = new RelayPublicAssignmentAdmission({
maxConcurrent: 1,
minIntervalMs: 0,
now: vi.fn(() => 0)
now: () => 0
})
const lease = await admission.acquire('host-a')
expect(lease).not.toBeNull()
@@ -46,19 +46,12 @@ describe('public assignment admission', () => {
})
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
}
}
schedule: () => () => undefined
})
const first = await admission.acquire('host-a')
const second = await admission.acquire('host-b')
@@ -70,12 +63,10 @@ describe('public assignment admission', () => {
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('serializes same-host assignment and reserved recovery without losing priority', async () => {
@@ -100,34 +91,4 @@ describe('public assignment admission', () => {
await expect(admission.acquire('host-b')).resolves.not.toBeNull()
})
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'])
})
})
cloud/apps/relay/src/assignment/public-admission/public-assignment-circuit-breaker.test.ts
-55
View File
@@ -1,55 +0,0 @@
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'
}
}
cloud/apps/relay/src/assignment/public-admission/public-assignment-drain-return-lane.test.ts
+3 -123
View File
@@ -1,6 +1,5 @@
import { afterEach, describe, expect, it, vi } from 'vitest'
import {
RelayAssignmentRowBusyError,
type RelayAssignment,
type ResolvedRelayAssignment
} from '../assignment-store.js'
@@ -21,7 +20,6 @@ vi.mock('../../authentication/relay-token-verifier.js', () => ({
import { createRelayApp } from '../../app.js'
afterEach(() => {
vi.useRealTimers()
vi.restoreAllMocks()
})
@@ -36,15 +34,11 @@ describe('drain-return lane', () => {
const resolve = vi.fn(async ({ relayHostId }: { relayHostId: string }) =>
relayHostId === drained ? isolatedHome(relayHostId) : assignment('cell-o', relayHostId)
)
const outcomes: string[] = []
const servedLanes: 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),
recordAdmissionServiceMs: (lane) => servedLanes.push(lane)
ready: vi.fn(async () => true)
})
// The re-placement holds the drain lane; the sticky lane's only slot is free.
@@ -56,125 +50,12 @@ describe('drain-return lane', () => {
)
expect(reconnect.status).toBe(200)
expect(resolve).toHaveBeenCalledWith(
{ userId: 'user-1', relayHostId: drained },
{ classifyHomeRollIsolation: true }
)
replacement.resolve(assignment('cell-new', drained))
expect((await moving).status).toBe(200)
expect(outcomes).toEqual(['drain-return', 'sticky'])
// The drained host held a sticky slot for its verification read too.
expect(servedLanes).toEqual(['sticky', 'sticky', 'drain-return'])
})
it('defers an overflowing drain return with a paced Retry-After and says why', async () => {
const holder = 'hhhhhhhhhhhhhhhh'
const replacement = deferred<RelayAssignment>()
const assign = vi.fn(async () => await replacement.promise)
const resolve = vi.fn(async ({ relayHostId }: { relayHostId: string }) =>
isolatedHome(relayHostId)
)
const outcomes: string[] = []
const reasons: string[] = []
const retryAfters: number[] = []
const causes: string[] = []
const app = createRelayApp(config({ drainReturnQueueMax: 1, drainReturnWaitMs: 50 }), {
store: {} as never,
assignments: { assign, resolve } as never,
drain: vi.fn(),
ready: vi.fn(async () => true),
recordAssignmentAdmission: (outcome) => outcomes.push(outcome),
recordAssignmentRejectionReason: (lane, reason) => reasons.push(`${lane}:${reason}`),
recordDrainReturnRetryAfter: (seconds) => retryAfters.push(seconds),
recordAssignmentUnavailable: (cause) => causes.push(cause)
})
const warn = vi.spyOn(console, 'warn').mockImplementation(() => {})
const holding = app.request('/v1/assign', assignmentRequest(holder, { reconnect: true }))
await vi.waitFor(() => expect(assign).toHaveBeenCalledTimes(1))
const queued = app.request(
'/v1/assign',
assignmentRequest('qqqqqqqqqqqqqqqq', { reconnect: true })
)
const overflow = await app.request(
'/v1/assign',
assignmentRequest('ffffffffffffffff', { reconnect: true })
)
expect(overflow.status).toBe(503)
expect(overflow.headers.get('retry-after')).toBe('2')
const timedOut = await queued
expect(timedOut.status).toBe(503)
expect(Number(timedOut.headers.get('retry-after'))).toBeGreaterThanOrEqual(2)
expect(reasons).toEqual(['drain-return:queue-full', 'drain-return:wait-timeout'])
expect(retryAfters).toEqual([2, Number(timedOut.headers.get('retry-after'))])
expect(outcomes.filter((outcome) => outcome === 'drain-return-deferred')).toHaveLength(2)
expect(causes).toEqual(['drain-return-deferred', 'drain-return-deferred'])
expect(warn.mock.calls.flat().join('\n')).toMatch(
/lane=drain-return hinted=true reason=queue-full .*retryAfter=2/
)
replacement.resolve(assignment('cell-new', holder))
expect((await holding).status).toBe(200)
})
it('answers a drained host’s row-busy redial with the short interval', async () => {
vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout', 'Date'] })
vi.spyOn(console, 'warn').mockImplementation(() => {})
const holder = 'hhhhhhhhhhhhhhhh'
const host = 'bbbbbbbbbbbbbbbb'
const holding = deferred<RelayAssignment>()
let busy = true
const assign = vi.fn(async ({ relayHostId }: { relayHostId: string }) => {
if (relayHostId === holder) return await holding.promise
if (relayHostId === host && busy) {
busy = false
throw new RelayAssignmentRowBusyError()
}
return assignment('cell-new', relayHostId)
})
const resolve = vi.fn(async ({ relayHostId }: { relayHostId: string }) =>
isolatedHome(relayHostId)
)
const app = createRelayApp(config({ drainReturnQueueMax: 1 }), {
store: {} as never,
assignments: { assign, resolve } as never,
drain: vi.fn(),
ready: vi.fn(async () => true)
})
const request = (relayHostId: string) =>
Promise.resolve(
app.request('/v1/assign', assignmentRequest(relayHostId, { reconnect: true }))
)
const held = request(holder)
await vi.waitFor(() => expect(assign).toHaveBeenCalledTimes(1))
const queued = request(host)
await vi.advanceTimersByTimeAsync(10)
// A cohort behind it books the next ten service slots.
for (let index = 0; index < 10; index++) {
const overflow = await request(hostId('f', index))
expect(overflow.status).toBe(503)
}
await vi.advanceTimersByTimeAsync(2_490)
holding.resolve(assignment('cell-new', holder))
expect((await held).status).toBe(200)
const refused = await queued
expect(refused.headers.get('retry-after')).toBe('1')
await vi.advanceTimersByTimeAsync(1_000)
const redial = await request(host)
// Its own lane interval, not a slot after the ten hosts that came later.
expect(redial.status).toBe(503)
expect(redial.headers.get('retry-after')).toBe('2')
})
})
function hostId(prefix: string, index: number): string {
return `${prefix}${String(index).padStart(15, '0')}`
}
function assignmentRequest(relayHostId: string, extra: { reconnect?: boolean } = {}): RequestInit {
return {
method: 'POST',
@@ -209,7 +90,7 @@ function deferred<T>() {
return { promise, resolve }
}
function config(overrides: Partial<RelayConfig> = {}): RelayConfig {
function config(): RelayConfig {
return {
port: 8080,
publicUrl: 'https://relay.example.test',
@@ -233,7 +114,6 @@ function config(overrides: Partial<RelayConfig> = {}): RelayConfig {
publicResolveConcurrency: 1,
publicResolveWaitMs: 5_000,
publicAssignmentRetryAfterSeconds: 5,
dataDir: './data',
...overrides
dataDir: './data'
}
}
cloud/apps/relay/src/assignment/public-admission/public-assignment-overload.test.ts
+2 -69
View File
@@ -18,68 +18,10 @@ import { createRelayApp } from '../../app.js'
describe('public assignment overload', () => {
it('fairly drains a bounded incident-scale burst without raising database concurrency', async () => {
const pending = new Map<string, ReturnType<typeof deferred<RelayAssignment>>>()
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<RelayAssignment>()
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('refuses an oversized chunked body before parsing it, sparing only the reserve routes', async () => {
const resolve = vi.fn()
const app = createRelayApp(config(), {
store: {} as never,
assignments: { resolve } as never,
assignments: {} as never,
drain: vi.fn(),
ready: vi.fn(async () => true)
})
@@ -98,15 +40,11 @@ describe('public assignment overload', () => {
const refused = await app.request('/v1/resolve', chunked(64 * 1024))
expect(refused.status).toBe(413)
expect(await refused.json()).toEqual({ error: 'request_too_large' })
expect(resolve).not.toHaveBeenCalled()
// Reserve bodies may reach 16 KiB, so the 4 KiB app limit must not answer for them. An
// unavailable route refuses before its own limit, including a body past 16 KiB.
// unavailable route refuses before its own limit.
const reserve = await app.request('/internal/reserve', chunked(6 * 1024))
expect(reserve.status).toBe(404)
expect(await reserve.json()).toEqual({ error: 'reserve_unavailable' })
const oversized = await app.request('/internal/reserve', chunked(20 * 1024))
expect(oversized.status).toBe(404)
expect(await oversized.json()).toEqual({ error: 'reserve_unavailable' })
})
it('reserves one bounded resolve lane while assignment work is saturated', async () => {
@@ -133,8 +71,6 @@ describe('public assignment overload', () => {
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)
@@ -143,8 +79,6 @@ describe('public assignment overload', () => {
expect(excessResolve.status).toBe(503)
expect(excessAssign.status).toBe(503)
expect(assign).toHaveBeenCalledTimes(2)
expect(resolveResume).not.toHaveBeenCalled()
pendingAssignments
.get('aaaaaaaaaaaaaaaa')
@@ -153,7 +87,6 @@ describe('public assignment overload', () => {
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))
cloud/apps/relay/src/assignment/public-admission/public-assignment-reconnect-lane.test.ts
+1 -90
View File
@@ -37,13 +37,11 @@ describe('public assignment reconnect lane', () => {
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)
ready: vi.fn(async () => true)
})
// Saturate the placement lane with two in-flight newcomers, zero queue.
@@ -59,18 +57,6 @@ describe('public assignment reconnect lane', () => {
)
expect(reconnected.status).toBe(200)
expect(resolve).toHaveBeenCalledWith(
{ userId: 'user-1', relayHostId: reconnecting },
{ classifyHomeRollIsolation: true }
)
expect(assign).toHaveBeenCalledWith(
{ userId: 'user-1', relayHostId: reconnecting },
undefined,
undefined,
undefined,
{ hostRequested: true }
)
expect(outcomes).toContain('sticky')
pending.get(newcomerA)?.resolve(assignment('cell-a', newcomerA))
pending.get(newcomerB)?.resolve(assignment('cell-b', newcomerB))
@@ -78,48 +64,6 @@ describe('public assignment reconnect lane', () => {
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('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('refuses a dial whose assignment row is busy at once, and frees the fast lane', async () => {
const busy = 'bbbbbbbbbbbbbbbb'
const next = 'nnnnnnnnnnnnnnnn'
@@ -143,7 +87,6 @@ describe('public assignment reconnect lane', () => {
expect(refused.status).toBe(503)
expect(refused.headers.get('retry-after')).toBe('1')
expect(await refused.json()).toEqual({ error: 'assignment_row_busy' })
expect(warn.mock.calls.flat().join('\n')).toContain('reason=relay_assignment_row_busy')
// The single fast-lane slot went back with the refusal.
expect(
(await app.request('/v1/assign', assignmentRequest(next, { reconnect: true }))).status
@@ -151,38 +94,6 @@ describe('public assignment reconnect lane', () => {
warn.mockRestore()
})
it('asks a busy resume to retry in a second', async () => {
const host = 'rrrrrrrrrrrrrrrr'
const app = createRelayApp(config(), {
store: {
resolveResume: vi.fn(async () => ({ userId: 'user-1', relayDeviceId: 'device-1' }))
} as never,
assignments: {
assign: vi.fn(async () => {
throw new RelayAssignmentRowBusyError()
}),
resolve: vi.fn(async () => null)
} as never,
drain: vi.fn(),
ready: vi.fn(async () => true)
})
const warn = vi.spyOn(console, 'warn').mockImplementation(() => {})
const response = await app.request('/v1/resolve', {
method: 'POST',
headers: { 'content-type': 'application/json' },
body: JSON.stringify({
v: 1,
relayHostId: host,
resumeToken: Buffer.alloc(32, 1).toString('base64url')
})
})
expect(response.status).toBe(503)
expect(response.headers.get('retry-after')).toBe('1')
expect(await response.json()).toEqual({ error: 'relay_assignment_row_busy' })
warn.mockRestore()
})
})
function assignmentRequest(relayHostId: string, extra: { reconnect?: boolean } = {}): RequestInit {
cloud/apps/relay/src/assignment/regional-rehome/idle-regional-rehome-reconciliation.test.ts
-17
View File
@@ -39,23 +39,6 @@ async function setup() {
}
describe('idle cutover durable reconciliation', () => {
it('only permits reopening when the exact source still owns the assignment', async () => {
const { store } = await setup()
expect(await store.reconcileIdleRegionalRehome(request)).toBe('not-committed')
await store.activateControl(request, {
cellId: request.sourceCellId,
assignmentEpoch: request.sourceAssignmentEpoch,
generation: request.sourceGeneration + 1,
cellIncarnation: request.sourceCellIncarnation
})
expect(await store.reconcileIdleRegionalRehome(request)).toBe('stale')
})
it('does not reopen an obsolete source after an assignment change', async () => {
const { store } = await setup()
await store.startEvacuation(request, request.targetCellId)
expect(await store.reconcileIdleRegionalRehome(request)).toBe('stale')
})
it('waits for an outstanding assignment transaction before deciding authority', async () => {
const { database, store } = await setup()
cloud/apps/relay/src/assignment/regional-rehome/idle-regional-rehome-store.test.ts
-132
View File
@@ -111,51 +111,6 @@ async function setup() {
}
describe('constrained idle regional assignment transaction', () => {
it.each(['missing', 'disabled', 'future'] as const)(
'does only one read per tick with %s durable control and sees later enablement',
async (state) => {
const { store, database, safety } = await setup()
const control = (await database.query(
"SELECT * FROM relay_region_rehome_control WHERE control_id = 'global'"
))[0]!
if (state === 'missing') {
await database.query('DELETE FROM relay_region_rehome_control')
} else {
await database.query(
"UPDATE relay_region_rehome_control SET enabled = ?, not_before = ? WHERE control_id = 'global'",
[state === 'disabled' ? 0 : 1, safety.observedAt + (state === 'future' ? 1 : 0)]
)
}
const query = vi.spyOn(database, 'query')
const transaction = vi.spyOn(database, 'transaction')
for (let tick = 0; tick < 3; tick++) {
query.mockClear()
expect(await store.selectIdleRegionalRehomeCandidates(safety)).toEqual([])
expect(query).toHaveBeenCalledTimes(1)
expect(query.mock.calls[0]![0]).toMatch(/^SELECT .*FROM relay_region_rehome_control/s)
expect(transaction).not.toHaveBeenCalled()
}
if (state === 'missing') {
const columns = Object.keys(control)
await database.query(
`INSERT INTO relay_region_rehome_control (${columns.join(',')}) VALUES (${columns.map(() => '?').join(',')})`,
Object.values(control)
)
} else {
await database.query(
"UPDATE relay_region_rehome_control SET enabled = 1, not_before = ? WHERE control_id = 'global'",
[safety.observedAt]
)
}
query.mockClear()
expect(await store.selectIdleRegionalRehomeCandidates(safety)).toHaveLength(1)
expect(query.mock.calls.length).toBeGreaterThan(1)
await database.query("UPDATE relay_region_rehome_control SET enabled = 0 WHERE control_id = 'global'")
query.mockClear()
expect(await store.selectIdleRegionalRehomeCandidates(safety)).toEqual([])
expect(query).toHaveBeenCalledTimes(1)
}
)
it.each([10, 11])('reserves source activity plus assignment at target capacity %i', async (capacity) => {
const { store, database, safety, request } = await setup()
@@ -184,33 +139,6 @@ describe('constrained idle regional assignment transaction', () => {
})
})
it.each(['next_dispatch_at', 'paused_until'] as const)(
'skips the candidate join while %s holds the durable dispatch budget closed',
async (column) => {
const { store, database, safety } = await setup()
const query = vi.spyOn(database, 'query')
// One assignment only: setup leaves both fields at 0, and naming the other
// one too would assign this column twice, which Postgres rejects.
await database.query(
`UPDATE relay_region_rehome_worker_state SET ${column} = ? WHERE worker_id = 'global'`,
[safety.observedAt + 1]
)
for (let tick = 0; tick < 3; tick++) {
query.mockClear()
expect(await store.selectIdleRegionalRehomeCandidates(safety)).toEqual([])
expect(query).toHaveBeenCalledTimes(2)
expect(query.mock.calls[1]![0]).toMatch(/FROM relay_region_rehome_worker_state/s)
}
await database.query(
`UPDATE relay_region_rehome_worker_state SET ${column} = ? WHERE worker_id = 'global'`,
[safety.observedAt]
)
query.mockClear()
expect(await store.selectIdleRegionalRehomeCandidates(safety)).toHaveLength(1)
expect(query.mock.calls.length).toBeGreaterThan(2)
}
)
// sourceRequestUnits reads a reserve cell's missing phone leases as zero units; selection must drop it.
it.each(['source', 'target'])('drops every candidate while the %s cell is recorded as reserve', async (cellId) => {
const { store, safety } = await setup()
@@ -267,51 +195,6 @@ describe('constrained idle regional assignment transaction', () => {
expect(first.some((c) => c.relayHostId === next[0]!.relayHostId)).toBe(false)
})
it('progresses past a full page of busy candidates without writing eligibility state', async () => {
const { store, database, safety } = await setup()
for (const table of [
'relay_assignments',
'relay_assignment_activity_leases',
'relay_control_capabilities',
'relay_region_decisions'
]) {
const template = (
await database.query(`SELECT * FROM ${table} WHERE user_id = ? AND relay_host_id = ?`, [
identity.userId,
identity.relayHostId
])
)[0]!
const columns = Object.keys(template)
for (let index = 0; index < 100; index++) {
const values = columns.map((column) =>
column === 'user_id' || column === 'relay_host_id' ? '?' : column
)
await database.query(
`INSERT INTO ${table} (${columns.join(', ')}) SELECT ${values.join(', ')} FROM ${table}
WHERE user_id = ? AND relay_host_id = ?`,
[
`idle-store-test-${String(index).padStart(3, '0')}`,
`pagehost${String(index).padStart(8, '0')}`,
identity.userId,
identity.relayHostId
]
)
}
}
const first = await store.selectIdleRegionalRehomeCandidates(safety)
const next = await store.selectIdleRegionalRehomeCandidates(safety)
expect(first).toHaveLength(100)
expect(next).toHaveLength(1)
expect(next[0]!.relayHostId).toBe('pagehost00000099')
const restarted = new RelayAssignmentStore(database, () => safety.observedAt, {
regionalRehomeCohortPercent: 100
})
expect(await restarted.selectIdleRegionalRehomeCandidates(safety)).toEqual(first)
expect(await database.query('SELECT * FROM relay_region_rehome_attempts')).toEqual([])
const decisions = await database.query('SELECT last_considered_at FROM relay_region_decisions')
expect(decisions.every((decision) => Number(decision.last_considered_at) === 0)).toBe(true)
})
it.runIf(Boolean(process.env.ORCA_IDLE_REHOME_POSTGRES_URL))(
'rechecks generation when replacement wins after the initial authority lookup',
async () => {
@@ -429,21 +312,6 @@ describe('constrained idle regional assignment transaction', () => {
})
expect(await database.query('SELECT * FROM relay_region_rehome_attempts')).toEqual([])
})
it('selects read-only with stable identity and the control generation, not probe generation', async () => {
const { store, safety, database } = await setup()
const before = await database.query('SELECT * FROM relay_assignments')
const candidates = await store.selectIdleRegionalRehomeCandidates(safety)
expect(candidates).toHaveLength(1)
expect(candidates[0]).toMatchObject({
sourceGeneration: 7,
sourceCellId: 'source',
targetCellId: 'target'
})
expect(await store.selectIdleRegionalRehomeCandidates(safety)).toEqual(candidates)
expect(await database.query('SELECT * FROM relay_assignments')).toEqual(before)
expect(await database.query('SELECT * FROM relay_region_rehome_attempts')).toEqual([])
})
})
function holdStatement(database: RelayDatabase, fragment: string) {
cloud/apps/relay/src/assignment/regional-rehome/idle-regional-rehome-window.test.ts
-85
View File
@@ -151,52 +151,6 @@ describe('idle regional rehome candidate window', () => {
expect(new Set(candidates.map((candidate) => candidate.attemptId)).size).toBe(3)
})
it('drops only the targets without room for the host plus its source activity', async () => {
const { store, database, safety } = await setup()
await seedHost(store, identity)
await database.query(
'UPDATE relay_assignment_activity_leases SET request_units = 4 WHERE user_id = ?',
[identity.userId]
)
// Five units needed: four source units plus the assignment the move reserves.
await database.query("UPDATE relay_cells SET capacity_requests = 4 WHERE cell_id = 'asia-idle'")
const short = await store.selectIdleRegionalRehomeCandidates(safety)
expect(short.map((candidate) => candidate.targetCellId)).toEqual(['asia-mid', 'asia-busy'])
// Exactly enough room is enough; it ranks last because the ratio is per capacity.
await database.query("UPDATE relay_cells SET capacity_requests = 5 WHERE cell_id = 'asia-idle'")
const exact = await store.selectIdleRegionalRehomeCandidates(safety)
expect(exact.map((candidate) => candidate.targetCellId)).toEqual([
'asia-mid',
'asia-busy',
'asia-idle'
])
})
it('reads a bounded window of decisions however many hosts are eligible', async () => {
const { store, database, safety } = await setup()
await seedHost(store, identity)
await cloneHosts(database, identity, 700)
const query = database.query.bind(database)
const decisionRows: number[] = []
vi.spyOn(database, 'query').mockImplementation(async (sql, params) => {
const rows = await query(sql, params)
if (/^SELECT user_id, relay_host_id FROM relay_region_decisions/.test(sql)) {
decisionRows.push(rows.length)
}
return rows
})
const first = await store.selectIdleRegionalRehomeCandidates(safety)
const next = await store.selectIdleRegionalRehomeCandidates(safety)
expect(decisionRows.length).toBeGreaterThan(0)
for (const count of decisionRows) {
expect(count).toBeLessThanOrEqual(500)
}
expect(first).toHaveLength(102)
expect(next).toHaveLength(102)
const firstKeys = new Set(first.map((c) => `${c.userId}/${c.targetCellId}`))
expect(next.every((c) => !firstKeys.has(`${c.userId}/${c.targetCellId}`))).toBe(true)
})
it("keeps the window's last host when a decision turns eligible between the two reads", async () => {
const { store, database, safety, now } = await setup()
await seedHost(store, identity)
@@ -252,7 +206,6 @@ describe('idle regional rehome candidate window', () => {
for (let poll = 0; poll < 20 && !wrapped; poll++) {
const page = await store.selectIdleRegionalRehomeCandidates(safety)
expect(page.length).toBeGreaterThan(0)
expect(page.length).toBeLessThanOrEqual(102)
for (const candidate of page) {
const key = `${candidate.userId}/${candidate.targetCellId}`
if (seen.has(key)) {
@@ -267,42 +220,4 @@ describe('idle regional rehome candidate window', () => {
expect(seen.size).toBe(363)
expect(wrapped).toBe(true)
})
it('excludes a host inside its rehome cooldown and takes it back after', async () => {
const { store, database, safety, now } = await setup()
await seedHost(store, identity)
await database.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, created_at, updated_at)
VALUES ('cooled', ?, ?, 'asia-east2', 'us', ?, 'asia-idle', ?, 0, 9, 0, 0, ?, ?)`,
[identity.userId, identity.relayHostId, incarnations[0], incarnations[2], now - 1_000, now]
)
expect(await store.selectIdleRegionalRehomeCandidates(safety)).toEqual([])
await database.query('UPDATE relay_region_rehome_attempts SET created_at = ?', [
now - 604_800_000 - 1
])
expect(await store.selectIdleRegionalRehomeCandidates(safety)).toHaveLength(3)
})
it('excludes a host outside the cohort', async () => {
const { store, database } = await setup()
await seedHost(store, identity)
const now = 100_000_000 + 86_400_000
const safety = {
observedAt: now,
sqlFailures: 0,
reconnects: 0,
controlActivityRecoveryFailures: 0,
databasePoolWaiting: 0,
databasePoolWaitersMax: 0,
databasePoolWaitMsMax: 0
}
await database.query('UPDATE relay_region_decisions SET cohort_bucket = 40')
const narrow = new RelayAssignmentStore(database, () => now, { regionalRehomeCohortPercent: 40 })
expect(await narrow.selectIdleRegionalRehomeCandidates(safety)).toEqual([])
const wide = new RelayAssignmentStore(database, () => now, { regionalRehomeCohortPercent: 41 })
expect(await wide.selectIdleRegionalRehomeCandidates(safety)).toHaveLength(3)
})
})
cloud/apps/relay/src/assignment/regional-rehome/region-correction-restart.test.ts
+1 -21
View File
@@ -80,20 +80,11 @@ describe('persisted region decisions across director restart', () => {
expect(row.outcome).toBe('inconclusive')
expect(Number(row.expires_at)).toBe(context.window.expiresAt)
const successor = (await context.issue(restarted.store)).window!
expect(successor.generation).toBe(context.window.generation + 1)
expect(successor.generation).toBe(2)
expect(await restarted.store.exchangeRegionCorrection(identity, context.report, epoch))
.toMatchObject({ reportStatus: 'stale' })
})
it('uses server expiry after a restart regardless of an old client report', async () => {
const context = await setup()
context.setNow(context.window.expiresAt)
const restarted = await context.restart()
expect(await restarted.store.exchangeRegionCorrection(identity, context.report, context.window.assignmentEpoch))
.toMatchObject({ reportStatus: 'expired' })
expect(await restarted.store.previewRegionCorrection()).toEqual({ expired: 1 })
})
it('does not interpret a persisted future-policy window using the old policy after rollback', async () => {
const context = await setup()
await context.database.query('UPDATE relay_region_decisions SET policy_version = 2')
@@ -105,15 +96,4 @@ describe('persisted region decisions across director restart', () => {
expect(row.preferred_region).toBeNull()
expect(row.report_json).toBeNull()
})
it('keeps generation ordering when the server clock moves backwards across restart', async () => {
const context = await setup()
context.setNow(1_000_000_000 - 60_000)
const restarted = await context.restart()
const successor = (await context.issue(restarted.store)).window!
expect(successor.generation).toBe(context.window.generation + 1)
expect(successor.expiresAt).toBe(context.window.expiresAt - 60_000)
expect(await restarted.store.exchangeRegionCorrection(identity, context.report, context.window.assignmentEpoch))
.toMatchObject({ reportStatus: 'stale' })
})
})
cloud/apps/relay/src/assignment/regional-rehome/region-correction-store.test.ts
-144
View File
@@ -124,104 +124,7 @@ async function window(context: Awaited<ReturnType<typeof setup>>) {
return result.window!
}
async function regionalMigration(context: Awaited<ReturnType<typeof setup>>) {
const migration = await context.store.startEvacuation(identity, cells[1]!.id)
const attemptId = '33333333-3333-4333-8333-333333333333'
await context.database.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,created_at,updated_at)
VALUES (?,?,?,'asia-east2',?,?,?,?,?,?,60000,1,?,?)`,
[
attemptId,
identity.userId,
identity.relayHostId,
cells[0]!.id,
incarnations[0],
cells[1]!.id,
incarnations[1],
migration.previousEpoch,
migration.assignmentEpoch,
context.now(),
context.now()
]
)
return { migration }
}
describe('ordered region decisions and migration outcomes', () => {
it('reports aggregate migration lifecycle and reservations without identity disclosure or writes', async () => {
const context = await setup()
const { migration } = await regionalMigration(context)
context.advance(1_000)
const before = await context.database.query('SELECT * FROM relay_region_rehome_attempts')
const outcomes = await context.store.regionCorrectionOutcomes()
expect(outcomes).toEqual([
expect.objectContaining({
sourceCellId: cells[0]!.id,
targetCellId: cells[1]!.id,
state: 'registering',
count: 1,
oldestOpenMs: 1_000
})
])
expect(outcomes[0]!.targetReservedUnits).toBeGreaterThan(0)
expect(JSON.stringify(outcomes)).not.toContain(identity.relayHostId)
expect(JSON.stringify(outcomes)).not.toContain(identity.userId)
expect(await context.database.query('SELECT * FROM relay_region_rehome_attempts')).toEqual(
before
)
await context.store.activateControl(identity, {
cellId: cells[1]!.id,
assignmentEpoch: migration.assignmentEpoch,
generation: 1
})
await context.store.markMigrationTargetRegistered(identity, {
cellId: cells[1]!.id,
assignmentEpoch: migration.assignmentEpoch
})
expect(await context.store.regionCorrectionOutcomes()).toEqual([
expect.objectContaining({ state: 'registered' })
])
await context.store.releaseActivity(identity, context.activityId)
expect(await context.store.completeReadyRegionalRehomes()).toBe(1)
expect(await context.store.regionCorrectionOutcomes()).toEqual([
expect.objectContaining({
state: 'completed',
targetReservedUnits: 0,
oldestOpenMs: 0
})
])
})
it('previews the uncapped fleet without writes, claims, or locked reads', async () => {
const context = await setup()
const query = context.database.query.bind(context.database)
const transaction = context.database.transaction.bind(context.database)
const queryLocked = context.database.queryLocked.bind(context.database)
context.database.query = async (sql, params) => {
expect(sql.trim()).toMatch(/^(SELECT|WITH)/i)
return query(sql, params)
}
context.database.transaction = async () => {
throw new Error('preview_must_not_open_mutating_transaction')
}
context.database.queryLocked = async () => {
throw new Error('preview_must_not_lock')
}
try {
const preview = await context.store.previewRegionalRehomeEligibility()
expect(preview.counts['no-verified-decision']).toBeGreaterThanOrEqual(1)
expect(preview.globalSafetyFailure).toBe('process-safety-unavailable')
expect(JSON.stringify(preview)).not.toContain(identity.relayHostId)
expect(JSON.stringify(preview)).not.toContain(identity.userId)
} finally {
context.database.query = query
context.database.transaction = transaction
context.database.queryLocked = queryLocked
}
})
it('allocates distinct ordered generations for concurrent window issuers', async () => {
const context = await setup()
const replies = await Promise.all([window(context), window(context), window(context)])
@@ -244,53 +147,6 @@ describe('ordered region decisions and migration outcomes', () => {
).toMatchObject({ reportStatus: 'stale' })
})
it('compares with assigned region, preserves hints, and never extends a window on report', async () => {
const context = await setup()
await context.store.assign(identity, 'asia-east2')
const issued = await window(context)
expect(issued.incumbentRegion).toBe('us-central1')
context.advance(50)
await context.store.exchangeRegionCorrection(
identity,
{
v: 1,
action: 'report',
generation: issued.generation,
assignmentEpoch: issued.assignmentEpoch,
policyVersion: 1,
outcome: 'conclusive',
measurements: { 'us-central1': 110, 'asia-east2': 90 }
},
issued.assignmentEpoch
)
expect(await context.store.previewRegionCorrection()).toEqual({ ineligible: 1 })
const row = (await context.database.query(`SELECT * FROM relay_region_decisions`))[0]!
expect(Number(row.expires_at)).toBe(issued.expiresAt)
const hint = (
await context.database.query(
`SELECT preferred_region FROM relay_assignment_region_preferences WHERE user_id = ?`,
[identity.userId]
)
)[0]
expect(hint?.preferred_region).toBe('asia-east2')
context.advance(24 * 60 * 60_000)
expect(
await context.store.exchangeRegionCorrection(
identity,
{
v: 1,
action: 'report',
generation: issued.generation,
assignmentEpoch: issued.assignmentEpoch,
policyVersion: 1,
outcome: 'inconclusive',
reason: 'late'
},
issued.assignmentEpoch
)
).toMatchObject({ reportStatus: 'expired' })
})
it('rejects stale assignment basis and requires both thresholds', async () => {
const context = await setup()
for (const measurements of [
cloud/apps/relay/src/assignment/regional-rehome/regional-rehome-constraint-migration-postgres.test.ts
+4 -26
View File
@@ -1,10 +1,6 @@
import pg from 'pg'
import { afterAll, beforeEach, describe, expect, it } from 'vitest'
import {
openRelayDatabase,
REGIONAL_REHOME_DEFAULT_HOST_COOLDOWN_MS,
type RelayDatabase
} from '../../database/database.js'
import { openRelayDatabase, type RelayDatabase } from '../../database/database.js'
const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL
const describePostgres = databaseUrl ? describe : describe.skip
@@ -127,25 +123,15 @@ describePostgres('PostgreSQL regional rehome constraint migration', () => {
await expect(
client.query(INSERT_ATTEMPT, attemptValues('attempt-3', 'europe-west1'))
).rejects.toMatchObject({ code: '23514' })
const constraints = await client.query(
`SELECT conname FROM pg_constraint
WHERE conrelid = 'relay_region_rehome_attempts'::regclass
AND conname LIKE '%preferred_region%'
ORDER BY conname`
)
expect(constraints.rows).toEqual([
{ conname: 'relay_region_rehome_attempts_preferred_region_valid' }
])
// The existing control row keeps its tuning and gains the cooldown.
// Preserve existing control tuning across the upgrade.
const control = await client.query(
`SELECT generation, preference_max_age_ms, host_cooldown_ms
`SELECT generation, preference_max_age_ms
FROM relay_region_rehome_control WHERE control_id = 'global'`
)
expect(control.rows).toEqual([
{
generation: '3',
preference_max_age_ms: '86400000',
host_cooldown_ms: String(REGIONAL_REHOME_DEFAULT_HOST_COOLDOWN_MS)
preference_max_age_ms: '86400000'
}
])
})
@@ -182,14 +168,6 @@ describePostgres('PostgreSQL regional rehome constraint migration', () => {
await withClient(async (client) => {
await client.query(`SET search_path = ${schema}`)
await client.query(INSERT_ATTEMPT, attemptValues('attempt-4', 'us-central1'))
const constraints = await client.query(
`SELECT conname FROM pg_constraint
WHERE conrelid = 'relay_region_rehome_attempts'::regclass
AND conname LIKE '%preferred_region%'`
)
expect(constraints.rows).toEqual([
{ conname: 'relay_region_rehome_attempts_preferred_region_valid' }
])
})
}, 60_000)
})
cloud/apps/relay/src/assignment/regional-rehome/regional-rehome-source-region-filter.test.ts
+2 -14
View File
@@ -1,4 +1,4 @@
import { afterEach, describe, expect, it, vi } from 'vitest'
import { describe, expect, it } from 'vitest'
import { RelayAssignmentStore } from '../assignment-store.js'
import { openIdleRehomeTestDatabase } from './idle-regional-rehome-test-database.js'
@@ -134,10 +134,6 @@ const moves = (candidates: { userId: string; sourceCellId: string; targetCellId:
to: cells.find(({ id }) => id === targetCellId)!.region
}))
afterEach(() => {
vi.restoreAllMocks()
})
describe('rehome source-region filter', () => {
it('keeps US to Asia moves and drops Asia to US moves', async () => {
const context = await setup('us-central1')
@@ -151,13 +147,5 @@ describe('rehome source-region filter', () => {
await context.database.close()
})
it('reports the filtered hosts in the preview, agreeing with selection', async () => {
const context = await setup('us-central1')
const preview = await context.store.previewRegionalRehomeEligibility(context.safety())
expect(preview.counts).toEqual({
'eligible:us-central1-to-asia-east2': 1,
'source-outside-director-region': 1
})
await context.database.close()
})
})
cloud/apps/relay/src/assignment/regional-rehome/regional-rehome-store.test.ts
-169
View File
@@ -30,88 +30,6 @@ const targetIncarnation = '22222222-2222-4222-8222-222222222222'
describe('regional rehome assignment state', () => {
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.tryIdleRehome()
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.database.query(
'SELECT drain_outcome FROM relay_region_rehome_attempts WHERE attempt_id = ?',
[attempt!.attemptId]
)
).toEqual([{ drain_outcome: '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
})
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 drain_receipt_at, completed_at, aborted_at FROM relay_region_rehome_attempts`
)
).toEqual([{ drain_receipt_at: context.now(), completed_at: context.now(), aborted_at: null }])
expect(targetControl).toMatch(/^control:/)
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.tryIdleRehome()
context.advance(6 * 60_000)
await heartbeat(context.store, target, targetIncarnation, 3, 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('completes healthy candidates past a poisoned attempt and logs it', async () => {
const context = await setup()
const poisoned = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' }
@@ -166,52 +84,6 @@ describe('regional rehome assignment state', () => {
).toEqual([{ completed_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.tryIdleRehome()
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()
})
})
type Context = Awaited<ReturnType<typeof setup>>
@@ -281,35 +153,6 @@ function collectCandidateFailureWarnings(
}
}
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<Record<string, number>> {
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 activatePreferredSource(
context: Context,
identity: { userId: string; relayHostId: string }
@@ -344,18 +187,6 @@ async function activatePreferredSource(
return control
}
async function activateSource(
context: Context,
identity: { userId: string; relayHostId: string }
): Promise<string> {
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: IdleRehomeTestStore,
cell: typeof source | typeof target,
cloud/apps/relay/src/assignment/regional-rehome/regional-rehome-target-row-lock-postgres.test.ts
-99
View File
@@ -1,7 +1,6 @@
import { afterAll, beforeAll, beforeEach, describe, expect, it } from 'vitest'
import { RelayAssignmentStore } from '../assignment-store.js'
import {
consumeRelayCellInventoryHold,
openRelayDatabase,
type PostgresDatabase,
type RelayDatabase
@@ -18,21 +17,6 @@ const describePostgres = databaseUrl ? describe : describe.skip
// Roughly the round trip from an Asia cell to the us-central1 database.
const STATEMENT_DELAY_MS = 150
// Probe the formerly fleet-locked tables plus admission/region tables read by commit.
const CELL_TABLES = [
'relay_cells',
'relay_cell_runtime',
'relay_cell_capabilities',
'relay_cell_rehome_safety',
'relay_cell_admission',
'relay_cell_regions'
] as const
// The target-row statement locks exactly these, held from it to COMMIT.
const TARGET_LOCKED = ['target:relay_cells', 'target:relay_cell_admission']
type Trip = { sql: string; lockable: Record<string, boolean>; probeMs: number }
type DelayControl = StatementDelay & { beforeTrip: (sql: string) => Promise<void> }
describePostgres('PostgreSQL regional rehome target-row lock', () => {
@@ -97,27 +81,6 @@ describePostgres('PostgreSQL regional rehome target-row lock', () => {
}
}
async function lockable(table: string, cellId: string): Promise<boolean> {
try {
await observer.transaction(
async (transaction) => {
await transaction.queryLocked(
`SELECT cell_id FROM ${table} WHERE cell_id = ?`,
[cellId],
{
failIfUnavailable: true
}
)
},
{ reportRetries: false }
)
return true
} catch (error) {
if (error instanceof Error && error.message === 'database_lock_unavailable') return false
throw error
}
}
async function reservedRequests(cellId: string): Promise<number> {
const row = (
await primary.query(`SELECT reserved_requests FROM relay_cells WHERE cell_id = ?`, [cellId])
@@ -152,59 +115,6 @@ describePostgres('PostgreSQL regional rehome target-row lock', () => {
)
}
it('holds only the target cell and admission rows until the next COMMIT boundary', async () => {
const context = await fixture()
const request = await context.select()
const trips: Trip[] = []
const probed = {
source: context.source.id,
target: context.target.id,
bystander: context.bystander.id
}
control.beforeTrip = async (sql) => {
const probeStartedAt = performance.now()
const state: Record<string, boolean> = {}
for (const [role, cellId] of Object.entries(probed)) {
for (const table of CELL_TABLES) state[`${role}:${table}`] = await lockable(table, cellId)
}
trips.push({ sql, lockable: state, probeMs: performance.now() - probeStartedAt })
}
consumeRelayCellInventoryHold(delayed)
control.enabled = true
const result = await context.delayedStore.commitIdleRegionalRehome(request, context.safety())
control.enabled = false
expect(result).toEqual({ outcome: 'committed' })
expect(trips.length).toBeGreaterThanOrEqual(15)
expect(trips[trips.length - 1]!.sql).toBe('COMMIT')
for (const trip of trips) {
for (const [key, free] of Object.entries(trip.lockable)) {
if (TARGET_LOCKED.includes(key)) continue
expect({ sql: trip.sql, key, free }).toEqual({ sql: trip.sql, key, free: true })
}
}
for (const key of TARGET_LOCKED) {
expect(trips.filter((trip) => !trip.lockable[key]).map((trip) => trip.sql)).toEqual([
'COMMIT'
])
}
const counts = consumeRelayCellInventoryHold(delayed)
const commitProbeMs = trips.at(-1)!.probeMs
console.info(
JSON.stringify({
event: 'rehome_target_row_hold',
trips: trips.length,
commitProbeMs,
...counts
})
)
expect(counts.rehomeTargetRowHolds).toBe(1)
expect(counts.cellInventoryHoldMaxSite).toBe('rehome-target-row')
expect(counts.rehomeTargetRowHoldMsMax).toBeGreaterThanOrEqual(STATEMENT_DELAY_MS)
expect(await reservedRequests(context.target.id)).toBe(context.targetReservedBefore + 2)
}, 45_000)
it('defers when the target leaves general admission after selection', async () => {
const context = await fixture()
const request = await context.select()
@@ -240,23 +150,14 @@ describePostgres('PostgreSQL regional rehome target-row lock', () => {
})
await lockedPromise
}
consumeRelayCellInventoryHold(delayed)
control.enabled = true
const startedAt = performance.now()
const result = await context.delayedStore.commitIdleRegionalRehome(request, context.safety())
const elapsedMs = performance.now() - startedAt
control.enabled = false
release()
await writer
expect(result).toEqual({ outcome: 'deferred', reason: 'candidate-ineligible' })
expect(consumeRelayCellInventoryHold(delayed)).toMatchObject({
cellInventoryLockUnavailable: 1,
rehomeTargetRowHolds: 0
})
// NOWAIT: no 1s lock_timeout wait and no transaction retry behind the writer.
expect(elapsedMs).toBeLessThan(40 * STATEMENT_DELAY_MS)
await expectNothingCommitted(context)
})
cloud/apps/relay/src/assignment/regional-rehome/regional-rehome-target-selection.test.ts
-202
View File
@@ -1,202 +0,0 @@
import { describe, expect, it } from 'vitest'
import { RelayAssignmentStore } from '../assignment-store.js'
import { openInMemoryRelayDatabase } from '../../database/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, {
regionalRehomeCohortPercent: 100,
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,
hostCooldownMs: 7 * 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,
idleRegionalRehome: true,
cellIncarnation: incarnation(1)
})
const { window } = await store.exchangeRegionCorrection(
identity,
{ v: 1, action: 'issue-window' },
assignment.assignmentEpoch
)
expect(window).toBeDefined()
await store.exchangeRegionCorrection(
identity,
{
v: 1,
action: 'report',
generation: window!.generation,
assignmentEpoch: assignment.assignmentEpoch,
policyVersion: 1,
outcome: 'conclusive',
measurements: { 'us-central1': 180, 'asia-east2': 40 }
},
assignment.assignmentEpoch
)
}
const safety = () => ({
observedAt: clock,
sqlFailures: 0,
reconnects: 0,
controlActivityRecoveryFailures: 0,
databasePoolWaiting: 0,
databasePoolWaitersMax: 0,
databasePoolWaitMsMax: 0
})
return { database, store, beat, activatePreferredSource, safety }
}
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, 3, {
observedRequests: 0,
enforcedConnections: 0,
sqlFailures: 0
})
// Lowest load but the connection hard cap is exhausted.
await context.beat(noHeadroom, 2, 3, {
observedRequests: 0,
enforcedConnections: 999,
sqlFailures: 0
})
await context.beat(unclean, 3, 3, {
observedRequests: 0,
enforcedConnections: 0,
sqlFailures: UNCLEAN
})
await context.beat(highLoad, 4, 3, {
observedRequests: 50,
enforcedConnections: 0,
sqlFailures: 0
})
await context.beat(lowLoad, 5, 3, {
observedRequests: 10,
enforcedConnections: 0,
sqlFailures: 0
})
await context.activatePreferredSource()
const candidates = await context.store.selectIdleRegionalRehomeCandidates(context.safety())
const attempt = candidates[0]
expect(attempt?.targetCellId).toBe(lowLoad.id)
expect(await context.store.commitIdleRegionalRehome(attempt!, context.safety(), 100)).toEqual({
outcome: 'committed'
})
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, 3, {
observedRequests: 0,
enforcedConnections: 0,
sqlFailures: 0
})
await context.beat(noHeadroom, 2, 3, {
observedRequests: 0,
enforcedConnections: 999,
sqlFailures: 0
})
await context.beat(unclean, 3, 3, {
observedRequests: 0,
enforcedConnections: 0,
sqlFailures: UNCLEAN
})
await context.beat(highLoad, 4, 3, {
observedRequests: 50,
enforcedConnections: 0,
sqlFailures: 0
})
await context.beat(lowLoad, 5, 3, {
observedRequests: 10,
enforcedConnections: 0,
sqlFailures: UNCLEAN
})
await context.activatePreferredSource()
const candidates = await context.store.selectIdleRegionalRehomeCandidates(context.safety())
const attempt = candidates[0]
expect(attempt?.targetCellId).toBe(highLoad.id)
expect(await context.store.commitIdleRegionalRehome(attempt!, context.safety(), 100)).toEqual({
outcome: 'committed'
})
await context.database.close()
})
})
cloud/apps/relay/src/assignment/regional-rehome/regional-rehome-worker.test.ts
-189
View File
@@ -1,189 +0,0 @@
import { afterEach, describe, expect, it, vi } from 'vitest'
import type { RelayAssignmentStore } from '../assignment-store.js'
import type { RelayConfig } from '../../config.js'
import { startRegionalRehomeWorker } from './regional-rehome-worker.js'
describe('regional rehome worker', () => {
afterEach(() => vi.restoreAllMocks())
it('bounds empty polling to the six-second cadence and stops its timer', async () => {
vi.useFakeTimers()
const selectIdleRegionalRehomeCandidates = vi.fn().mockResolvedValue([])
const worker = startRegionalRehomeWorker(
config(),
{
selectIdleRegionalRehomeCandidates
} as unknown as RelayAssignmentStore,
{
safetySnapshot: () => safety(Date.now()),
random: () => 0
}
)!
try {
await vi.advanceTimersByTimeAsync(0)
expect(selectIdleRegionalRehomeCandidates).toHaveBeenCalledTimes(1)
await vi.advanceTimersByTimeAsync(5_999)
expect(selectIdleRegionalRehomeCandidates).toHaveBeenCalledTimes(1)
await vi.advanceTimersByTimeAsync(1)
expect(selectIdleRegionalRehomeCandidates).toHaveBeenCalledTimes(2)
worker.stop()
await vi.advanceTimersByTimeAsync(60_000)
expect(selectIdleRegionalRehomeCandidates).toHaveBeenCalledTimes(2)
} finally {
worker.stop()
vi.useRealTimers()
}
})
it('stops walking the page when the source names a deferral no candidate can pass', async () => {
const fetchImpl = respondWith([{ outcome: 'deferred', reason: 'concurrency-limit' }])
const summaries = collectSummaries()
try {
await runOnePoll(fetchImpl, 3, summaries)
} finally {
summaries.restore()
}
expect(fetchImpl).toHaveBeenCalledTimes(1)
expect(summaries.entries).toEqual([
{
event: 'orca_relay_idle_rehome_dispatch_summary',
candidates: 3,
dispatched: 1,
stoppedBy: 'concurrency-limit',
outcomes: { 'deferred:concurrency-limit': 1 }
}
])
})
it('keeps its whole-page walk when the source sends no reason at all', async () => {
const fetchImpl = respondWith([{ outcome: 'deferred' }])
const summaries = collectSummaries()
try {
await runOnePoll(fetchImpl, 3, summaries)
} finally {
summaries.restore()
}
expect(fetchImpl).toHaveBeenCalledTimes(3)
expect(summaries.entries[0]).toMatchObject({
stoppedBy: null,
outcomes: { deferred: 3 }
})
})
it('walks past a deferral that only concerns the one candidate', async () => {
const fetchImpl = respondWith([
{ outcome: 'deferred', reason: 'host-unsupported' },
{ outcome: 'busy' },
{ outcome: 'committed' }
])
const summaries = collectSummaries()
try {
await runOnePoll(fetchImpl, 4, summaries)
} finally {
summaries.restore()
}
expect(fetchImpl).toHaveBeenCalledTimes(3)
expect(summaries.entries[0]).toMatchObject({
dispatched: 3,
stoppedBy: 'committed',
outcomes: { 'deferred:host-unsupported': 1, busy: 1, committed: 1 }
})
})
})
function respondWith(bodies: { outcome: string; reason?: string }[]) {
let index = 0
return vi.fn(async () => {
const body = bodies[Math.min(index++, bodies.length - 1)]!
return new Response(JSON.stringify({ v: 1, ...body }), {
headers: { 'content-type': 'application/json' }
})
})
}
function collectSummaries() {
const entries: Record<string, unknown>[] = []
let arrived: (() => void) | undefined
// Construction immediately polls; await first instead of waiting for a timer tick.
const first = new Promise<void>((resolve) => {
arrived = resolve
})
const original = console.warn
console.warn = (line: unknown, ...rest: unknown[]) => {
try {
const parsed = JSON.parse(line as string) as Record<string, unknown>
if (parsed.event === 'orca_relay_idle_rehome_dispatch_summary') {
entries.push(parsed)
arrived?.()
return
}
} catch {
// Not a JSON log line; fall through to the original writer.
}
original(line as string, ...rest)
}
return { entries, first, restore: () => (console.warn = original) }
}
async function runOnePoll(
fetchImpl: typeof fetch,
candidates: number,
summaries: { first: Promise<void> }
): Promise<void> {
const assignments = {
selectIdleRegionalRehomeCandidates: vi.fn(async () =>
Array.from({ length: candidates }, (_, index) => ({
v: 1 as const,
attemptId: `00000000-0000-4000-8000-00000000000${index}`,
userId: `user-${index}`,
relayHostId: 'abcdefghijklmnop',
sourceCellId: 'us-c1',
sourceCellUrl: 'https://us-c1.relay.example.test',
sourceCellIncarnation: '11111111-1111-4111-8111-111111111111',
sourceAssignmentEpoch: 1,
sourceGeneration: 1,
targetCellId: 'asia-c1'
}))
)
} as unknown as RelayAssignmentStore
const worker = startRegionalRehomeWorker(config(), assignments, {
fetch: fetchImpl,
identityToken: async () => 'token',
safetySnapshot: () => safety(Date.now()),
intervalMs: 60_000
})!
try {
await summaries.first
} finally {
worker.stop()
}
}
function config(overrides: Partial<RelayConfig> = {}): 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
}
}
cloud/apps/relay/src/assignment/relay-region-placement.test.ts
-26
View File
@@ -45,16 +45,6 @@ describe('Relay regional placement', () => {
afterEach(async () => await database.close())
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' })
@@ -219,20 +209,4 @@ describe('Relay regional placement', () => {
})
})
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
await expect(store.releaseExpiredRegionPreferences()).resolves.toBe(0)
now += 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([])
})
})
cloud/apps/relay/src/assignment/reserve-placement/placement-load-band.test.ts
+1 -81
View File
@@ -1,5 +1,5 @@
import { describe, expect, it } from 'vitest'
import { PLACEMENT_INTAKE_SHARE, PlacementLoadBand } from './placement-load-band.js'
import { PlacementLoadBand } from './placement-load-band.js'
import { RelayAssignmentStore } from '../assignment-store.js'
import { openInMemoryRelayDatabase } from '../../database/database.js'
@@ -13,74 +13,8 @@ function seededRandom(seed: number): () => number {
}
}
type SimulatedCell = { cellId: string; reserved: number; capacity: number }
const quiet = () => undefined
// Each pick goes to one of the directors at random, as the load balancer spreads requests.
function simulate(
directors: PlacementLoadBand[],
cells: SimulatedCell[],
picks: number,
at: (index: number) => number,
route: () => number = seededRandom(99)
) {
const counts = new Map(cells.map((cell) => [cell.cellId, 0]))
for (let index = 0; index < picks; index += 1) {
const ranked = cells
.filter((cell) => cell.reserved < cell.capacity)
.map((cell) => ({ cellId: cell.cellId, loadRatio: cell.reserved / cell.capacity, cell }))
.sort(
(left, right) => left.loadRatio - right.loadRatio || left.cellId.localeCompare(right.cellId)
)
const director = directors[Math.floor(route() * directors.length)]
const picked = director?.pick(ranked, at(index))
if (!picked) throw new Error('no candidate')
picked.cell.reserved += 1
counts.set(picked.cellId, (counts.get(picked.cellId) ?? 0) + 1)
}
return counts
}
// c27's drain on 10-08: 2,770 hosts onto one empty Asia cell (2,805 free)
// and four fuller ones (2,311 free between them).
function c27DrainCells(): SimulatedCell[] {
return [
{ cellId: 'cell-empty', reserved: 0, capacity: 4000 },
...['cell-b', 'cell-c', 'cell-d', 'cell-e'].map((cellId, index) => ({
cellId,
reserved: 4000 - 577 - (index === 0 ? 3 : 0),
capacity: 4000
}))
]
}
function drainShare(paceMs: number, seed: number) {
const cells = c27DrainCells()
const directors = Array.from(
{ length: 5 },
(_, index) => new PlacementLoadBand(seededRandom(seed + index), quiet)
)
const hosts = 2770
const counts = simulate(directors, cells, hosts, (index) => (index * paceMs) / hosts)
return { counts, cells, share: (counts.get('cell-empty') ?? 0) / hosts }
}
describe('PlacementLoadBand', () => {
it.each([
['1,200 s', 1_200_000],
['300 s', 300_000]
])('keeps the empty cell near the share across five directors at the %s pace', (_, paceMs) => {
for (const seed of [1, 2, 3]) {
const { counts, cells, share } = drainShare(paceMs, seed * 10)
expect([...counts.values()].reduce((sum, count) => sum + count, 0)).toBe(2770)
expect(share).toBeGreaterThan(0.3)
expect(share).toBeLessThan(0.5)
for (const cell of cells) expect(cell.reserved).toBeLessThanOrEqual(cell.capacity)
}
})
})
describe('leastLoadedCell with the load band', () => {
async function setupRegions() {
const database = await openInMemoryRelayDatabase()
@@ -121,20 +55,6 @@ describe('leastLoadedCell with the load band', () => {
return counts
}
it('spreads within the preferred region and keeps an empty US cell out', async () => {
const { database, store } = await setupRegions()
try {
const counts = await placements(store, 200)
expect(counts.get('us-a')).toBeUndefined()
expect(counts.get('asia-a')).toBeLessThanOrEqual(PLACEMENT_INTAKE_SHARE * 200 + 1)
const used = [...counts.values()].filter((count) => count >= 20)
expect(used.length).toBeGreaterThanOrEqual(3)
} finally {
await database.close()
}
})
it('fills the preferred region to its hard cap before spilling, never past it', async () => {
const { database, store } = await setupRegions()
try {
cloud/apps/relay/src/assignment/reserve-placement/reserve-ledger-writer.test.ts
+2 -31
View File
@@ -26,20 +26,6 @@ describe('reserve ledger writer', () => {
)
)[0]
it('writes a booked join, and only a higher epoch replaces it', async () => {
const { database: db, writer } = await setup()
writer.enqueue({ ...HOST, cellId: 'c1', epoch: 5 })
await writer.flush()
expect(await row(db)).toMatchObject({ cell_id: 'c1', assignment_epoch: 5 })
writer.enqueue({ ...HOST, cellId: 'c2', epoch: 4 })
writer.enqueue({ ...HOST, cellId: 'c3', epoch: 5 })
await writer.flush()
expect(await row(db)).toMatchObject({ cell_id: 'c1', assignment_epoch: 5 })
writer.enqueue({ ...HOST, cellId: 'c2', epoch: 6 })
await writer.flush()
expect(await row(db)).toMatchObject({ cell_id: 'c2', assignment_epoch: 6 })
})
it('drops a moved row’s counters but keeps a same-cell row’s', async () => {
const { database: db, writer } = await setup()
await db.query(
@@ -57,30 +43,15 @@ describe('reserve ledger writer', () => {
expect(await row(db, 'bbbbbbbbbbbbbbbb')).toMatchObject({ cell_id: 'c1', reserved_controls: 2 })
})
it('keeps its queue through a database stall and drops the oldest past the bound', async () => {
it('persists a booked join after a database stall', async () => {
let now = 1_000
const { database: db, writer } = await setup(() => now)
const query = vi.spyOn(db, 'query').mockRejectedValueOnce(new Error('timeout'))
vi.spyOn(db, 'query').mockRejectedValueOnce(new Error('timeout'))
writer.enqueue({ ...HOST, cellId: 'c1', epoch: 5 })
await writer.flush()
expect(writer.pending()).toBe(1)
await writer.flush()
expect(query).toHaveBeenCalledTimes(1)
now += 1_000
await writer.flush()
expect(writer.pending()).toBe(0)
expect(await row(db)).toMatchObject({ cell_id: 'c1', assignment_epoch: 5 })
const bounded = new ReserveLedgerWriter(db, { queueMax: 2, log: () => undefined })
const hosts = ['oldest0000000000', 'middle0000000000', 'newest0000000000']
for (const relayHostId of hosts) {
bounded.enqueue({ ...HOST, relayHostId, cellId: 'c1', epoch: 1 })
}
expect(bounded.pending()).toBe(2)
await bounded.flush()
expect(await row(db, hosts[0])).toBeUndefined()
expect(await row(db, hosts[1])).toMatchObject({ cell_id: 'c1', assignment_epoch: 1 })
expect(await row(db, hosts[2])).toMatchObject({ cell_id: 'c1', assignment_epoch: 1 })
})
it('reconciles the map’s seats and never lowers an epoch', async () => {
cloud/apps/relay/src/assignment/reserve-placement/reserve-placement-simulation.test.ts
-179
View File
@@ -1,179 +0,0 @@
import { describe, expect, it } from 'vitest'
import {
runReservePlacementSimulation,
type SimulationConfig,
type SimulationReport
} from './reserve-placement-simulation.js'
const SEEDS = [11, 23, 47]
const rtt = (region: string) => (region === 'asia-east2' ? 170 : 10)
function fleet(): SimulationConfig['cells'] {
const cells: SimulationConfig['cells'] = []
for (let index = 0; index < 24; index += 1) {
cells.push({
cellId: `us${String(index).padStart(2, '0')}`,
region: 'us-central1',
hardCap: 600,
ceiling: 480,
intakePerSec: 5,
// A third of the fleet stays in database mode, one cell on an image without step 5.
mode: index % 3 === 0 ? 'db' : 'reserve',
old: index === 0,
preseated: 150
})
}
for (let index = 0; index < 6; index += 1) {
cells.push({
cellId: `asia${index}`,
region: 'asia-east2',
hardCap: 600,
ceiling: 480,
intakePerSec: 2,
mode: index === 0 ? 'db' : 'reserve',
preseated: 200
})
}
return cells
}
function expectInvariants(report: SimulationReport): void {
expect(report.violations).toEqual([])
}
describe('step 5 deterministic simulation', () => {
// Invariant 9 (dead-man) never trips here; its own tests below cover it.
it.each(SEEDS)(
'holds invariants 1-8 across mixed versions, lagging polls, restarts and database stalls (seed %i)',
async (seed) => {
const report = await runReservePlacementSimulation({
seed,
durationMs: 25 * 60_000,
directors: 5,
// One director still on the database-only revision.
oldDirectors: 1,
cells: fleet(),
hosts: 6_600,
hostRegion: (index) => (index % 5 === 0 ? 'asia-east2' : 'us-central1'),
arrivalPerSec: 20,
meanSessionMs: 8 * 60_000,
reconnectDelayMs: [500, 5_000],
duplicateAssignShare: 0.05,
halfOpenShare: 0.05,
pollLagMs: (director) => [0, 200, 1_000, 4_000, 10_000][director]!,
rttMs: rtt,
cellRestarts: [{ at: 6 * 60_000, cellId: 'us04' }, { at: 14 * 60_000, cellId: 'asia2' }],
directorRestarts: [{ at: 8 * 60_000, director: 2 }, { at: 16 * 60_000, director: 4 }],
// Drained hosts land elsewhere while their old seat lingers: placement supersedes it.
drains: [{ at: 12 * 60_000, cellId: 'us05', paceMs: 60_000 }],
databaseStalls: Array.from({ length: 3 }, (_, index) => ({
at: (index * 610 + 120) * 1_000,
durationMs: 7_000
}))
})
expectInvariants(report)
expect(report.reservePlacements).toBeGreaterThan(500)
expect(report.stickyAnswers).toBeGreaterThan(500)
expect(report.seatedAtEnd).toBeGreaterThan(4_000)
// The run reached the paths the invariants guard, so a pass is not vacuous.
expect(report.supersedes).toBeGreaterThan(0)
expect(report.rowAheadDemotions).toBeGreaterThan(0)
expect(report.directorRestarts).toBe(2)
expect(report.stickyReserves).toBeGreaterThan(0)
expect(report.databasePlacements).toBeGreaterThan(0)
},
120_000
)
it.each(SEEDS)(
'near full: extra reserve calls stay within 10% and exhausted requests within 1% (seed %i)',
async (seed) => {
const cells = Array.from({ length: 17 }, (_, index) => ({
cellId: `us${index}`,
region: 'us-central1' as const,
hardCap: 700,
ceiling: 600,
intakePerSec: 50,
mode: 'reserve' as const,
preseated: 582
}))
const report = await runReservePlacementSimulation({
seed,
durationMs: 3 * 60_000,
directors: 5,
cells,
hosts: 17 * 582 + 300,
hostRegion: () => 'us-central1',
arrivalPerSec: 10,
meanSessionMs: 0,
reconnectDelayMs: [1_000, 3_000],
duplicateAssignShare: 0,
halfOpenShare: 0,
pollLagMs: () => 0,
rttMs: rtt
})
expectInvariants(report)
expect(report.reservePlacements).toBeGreaterThan(250)
expect(report.reserveCalls / report.reservePlacements - 1).toBeLessThanOrEqual(0.1)
expect(report.exhausted / report.requests).toBeLessThanOrEqual(0.01)
for (const share of Object.values(report.maxOccupancyShare)) expect(share).toBeLessThanOrEqual(1)
},
120_000
)
it.each(SEEDS.flatMap((seed) => [[1_200_000, seed], [300_000, seed]] as const))(
'a drain onto an empty cell at the %i ms pace spreads within each cell intake budget (seed %i)',
async (paceMs, seed) => {
// Production c27 drain: 1,200s pacing toward cells with 2,805/841/777/667/26 free seats.
const free = [2_805, 841, 777, 667, 26]
const cells: SimulationConfig['cells'] = free.map((seats, index) => ({
cellId: `asia${index + 1}`,
region: 'asia-east2',
hardCap: 3_000,
ceiling: 2_830,
intakePerSec: 2,
mode: 'reserve',
preseated: 2_830 - seats
}))
cells.unshift({
cellId: 'asia0',
region: 'asia-east2',
hardCap: 3_000,
ceiling: 2_830,
intakePerSec: 2,
mode: 'reserve',
preseated: 2_770
})
const preseated = cells.reduce((sum, cell) => sum + (cell.preseated ?? 0), 0)
const report = await runReservePlacementSimulation({
seed,
durationMs: 30 * 60_000,
directors: 5,
cells,
hosts: preseated,
hostRegion: () => 'asia-east2',
arrivalPerSec: 1,
meanSessionMs: 0,
reconnectDelayMs: [1_000, 3_000],
duplicateAssignShare: 0.02,
halfOpenShare: 0,
pollLagMs: () => 0,
rttMs: rtt,
drains: [{ at: 10_000, cellId: 'asia0', paceMs }]
})
expectInvariants(report)
// Below every cell's budget nothing waits; above it the cells' buckets pace the drain.
if (paceMs === 300_000) expect(report.intakeRefusals + report.paced).toBeGreaterThan(0)
const moved = Object.entries(report.arrivalsByCell)
const total = moved.reduce((sum, [, count]) => sum + count, 0)
expect(total).toBeGreaterThan(2_000)
// Spread: the emptiest cell takes a share, not the drain.
expect((report.arrivalsByCell.asia1 ?? 0) / total).toBeLessThan(0.6)
expect(moved.filter(([, count]) => count / total > 0.1).length).toBeGreaterThanOrEqual(3)
},
120_000
)
})
cloud/apps/relay/src/assignment/reserve-placement/reserve-placement-simulation.ts
-1357
View File
File diff suppressed because it is too large Load Diff
cloud/apps/relay/src/authentication/stale-if-error-jwks.test.ts
+1 -31
View File
@@ -2,7 +2,7 @@ import { exportJWK, generateKeyPair, SignJWT, type CryptoKey, type JWK } from 'j
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import type { RelayConfig } from '../config.js'
import { createRelayTokenVerifier } from './relay-token-verifier.js'
import { createStaleIfErrorJWKSet, RELAY_JWKS_MAX_STALE_MS } from './stale-if-error-jwks.js'
import { createStaleIfErrorJWKSet } from './stale-if-error-jwks.js'
const config = {
jwksUrl: 'https://auth.example.test/.well-known/jwks.json',
@@ -43,37 +43,7 @@ beforeEach(() => vi.useFakeTimers({ now: Date.parse('2026-10-11T13:00:00Z') }))
afterEach(() => vi.useRealTimers())
describe('relay token keys while auth is down', () => {
it('accepts a valid token from the last served set once the cache is past its 10 min age', async () => {
const key = await signingKey('auth-1')
const authServer = auth(key)
const jwks = createStaleIfErrorJWKSet(new URL(config.jwksUrl), { fetch: authServer.fetch })
const verify = createRelayTokenVerifier(config as RelayConfig, jwks)
expect(await verify(await relayToken(key))).not.toBeNull()
vi.setSystemTime(Date.now() + 20 * MINUTE)
authServer.state.up = false
const fetches = authServer.fetch.mock.calls.length
expect(await verify(await relayToken(key))).toMatchObject({ sub: 'user-1' })
expect(authServer.fetch.mock.calls.length).toBe(fetches + 1)
// Inside the retry window the stale set answers without waiting on another fetch.
expect(await verify(await relayToken(key))).not.toBeNull()
expect(authServer.fetch.mock.calls.length).toBe(fetches + 1)
})
it('rejects a token signed by a key the last served set never had', async () => {
const key = await signingKey('auth-1')
const authServer = auth(key)
const jwks = createStaleIfErrorJWKSet(new URL(config.jwksUrl), { fetch: authServer.fetch })
const verify = createRelayTokenVerifier(config as RelayConfig, jwks)
expect(await verify(await relayToken(key))).not.toBeNull()
vi.setSystemTime(Date.now() + MINUTE)
authServer.state.up = false
expect(await verify(await relayToken(await signingKey('unknown')))).toBeNull()
})
it('bounds the stale set at 30 minutes since the last served set', async () => {
expect(RELAY_JWKS_MAX_STALE_MS).toBe(30 * MINUTE)
const key = await signingKey('auth-1')
const authServer = auth(key)
const verify = createRelayTokenVerifier(
cloud/apps/relay/src/cell-status-reserve-view.test.ts
-314
View File
@@ -1,314 +0,0 @@
import { afterEach, describe, expect, it, vi } from 'vitest'
import type { RelayConfig } from './config.js'
vi.mock('./authentication/admin-token-verifier.js', () => ({
createAdminTokenVerifier: () => async (token: string) => token === 'deploy-token',
createReadOnlyAdminTokenVerifier: () => async () => false,
createRegionalRehomeControlApplyTokenVerifier: () => async () => false,
createRegionalRehomeRuntimeTokenVerifier: () => async () => false,
createRegionalRehomeTokenVerifier: () => async () => false,
createRuntimeTokenVerifier: () => async () => false
}))
vi.mock('./authentication/relay-token-verifier.js', () => ({
createRelayTokenVerifier: () => async () => null,
readBearer: (value: string | undefined) => value?.replace(/^Bearer /, '') ?? null
}))
import { createRelayApp } from './app.js'
import { RelayAssignmentStore } from './assignment/assignment-store.js'
import type { CellReserveClient } from './cell/reservations/cell-reserve-client.js'
import { openInMemoryRelayDatabase, type RelayDatabase } from './database/database.js'
import { ReserveAssignment, type ReservePlacementMode } from './assignment/reserve-placement/reserve-assignment.js'
import { ReservePlacer } from './assignment/reserve-placement/reserve-placement.js'
import {
ShadowSeatDirectory,
type SeatFeedCell,
type SeatFeedResponse
} from './cell/seats/shadow-seat-directory.js'
const INCARNATION = '11111111-1111-4111-8111-111111111111'
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 PER_HOST_FIELDS = [
'reservedRequests',
'assignments',
'activityLeases',
'activityRequestUnits',
'restartBlockingActivityLeases',
'restartBlockingActivityRequestUnits',
'restartBlockingReservedRequests',
'outgoingMigrations',
'incomingMigrations'
]
function feed(cellId: string, overrides: Partial<SeatFeedResponse> = {}): SeatFeedResponse {
return {
v: 1,
cellId,
incarnation: INCARNATION,
seq: 1,
at: 0,
draining: false,
counts: { seats: 40, controls: 42, bookings: 3, units: 45, ceiling: 500 },
intake: { perSec: 5, burst: 20, tokens: 20 },
flagsApplied: { generation: 3, flags: { admitMode: 'reserve' } },
admitModeEffective: 'reserve',
admitModeRaw: 'reserve',
full: [],
...overrides
}
}
describe('reserve-aware cell status', () => {
let database: RelayDatabase | undefined
afterEach(async () => {
await database?.close()
database = undefined
})
async function setup(
options: { mode?: ReservePlacementMode; startedAt?: number; placement?: boolean } = {}
) {
const now = { value: 100_000 }
database = await openInMemoryRelayDatabase()
const store = new RelayAssignmentStore(database, () => now.value, { heartbeatTtlMs: 45_000 })
await store.reconcileCells(CELLS)
for (const cell of CELLS) {
await store.recordCellHeartbeat({
cellId: cell.id,
cellUrl: cell.url,
cellIncarnation: INCARNATION,
startedAt: 50,
ready: true,
observedRequests: 0
})
}
await store.assign({ userId: 'user-a', relayHostId: 'host000000000001' })
await store.setCellAdmitMode('cell-a', 'reserve')
const cells: SeatFeedCell[] = CELLS.map((cell) => ({
cellId: cell.id,
cellUrl: cell.url,
region: 'us-central1',
heartbeatExpiresAt: Number.MAX_SAFE_INTEGER,
requiredForComplete: true,
general: true
}))
const directory = new ShadowSeatDirectory()
directory.setCells(cells.map((cell) => cell.cellId))
const reservePlacement = new ReserveAssignment({
mode: options.mode ?? 'on',
directory,
cells: () => cells,
startedAt: options.startedAt ?? 0,
placer: new ReservePlacer(
() => now.value,
() => 0.25
),
// oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: the status route never books or demotes.
client: {} as CellReserveClient,
readRow: async () => null,
now: () => now.value
})
const app = createRelayApp(config(), {
store: {} as never,
assignments: store,
...(options.placement === false ? {} : { reservePlacement }),
drain: vi.fn(),
ready: vi.fn(async () => true)
})
const status = async (body: Record<string, unknown>) =>
await app.request('/v1/admin/cell-status', {
method: 'POST',
headers: { authorization: 'Bearer deploy-token', 'content-type': 'application/json' },
body: JSON.stringify({ v: 1, ...body })
})
return { store, directory, now, status }
}
it('reads a database cell exactly as today, and leaves a reserve cell without per-host counts', async () => {
const { store } = await setup()
expect(await store.reserveAwareCellDeploymentStatus('cell-b')).toEqual({
admitMode: 'db',
status: await store.cellDeploymentStatus('cell-b')
})
await expect(store.cellDeploymentStatus('cell-a')).rejects.toThrow('cell_in_reserve_mode')
const reserve = await store.reserveAwareCellDeploymentStatus('cell-a')
expect(reserve).toMatchObject({
admitMode: 'reserve',
status: {
cellId: 'cell-a',
runtime: { cellIncarnation: INCARNATION, ready: true, heartbeatFresh: true }
}
})
expect(Object.keys(reserve.status).filter((key) => PER_HOST_FIELDS.includes(key))).toEqual([])
await database!.query(
'ALTER TABLE relay_cell_admit_modes RENAME TO relay_cell_admit_modes_gone'
)
const unknown = new RelayAssignmentStore(database!, () => 100_000)
await expect(unknown.reserveAwareCellDeploymentStatus('cell-b')).rejects.toThrow(
'reserve_mode_unknown'
)
})
it('keeps the 404 for a reserve cell without the opt-in, and 400s an unknown view', async () => {
const { status } = await setup()
const refused = await status({ cellId: 'cell-a' })
expect(refused.status).toBe(404)
expect(await refused.json()).toEqual({ error: 'cell_in_reserve_mode' })
expect((await status({ cellId: 'cell-a', view: 'all' })).status).toBe(400)
})
it('answers a reserve cell from the map, with perHost null and never zeros', async () => {
const { directory, now, status } = await setup()
directory.apply('cell-a', feed('cell-a', { draining: true }), now.value, now.value - 1_500)
directory.apply(
'cell-b',
feed('cell-b', { admitModeEffective: 'db', admitModeRaw: 'db' }),
now.value,
now.value
)
now.value += 20_000
const response = await status({ cellId: 'cell-a', view: 'reserve-aware' })
expect(response.status).toBe(200)
const body = await statusBody(response)
expect(body.status).toMatchObject({
cellId: 'cell-a',
admissionState: 'general',
runtime: { cellIncarnation: INCARNATION },
perHost: null,
reserve: {
admitMode: 'reserve',
feedStatus: 'live',
// Reported, not enforced: the caller rejects a stale map.
polledAgeMs: 21_500,
incarnation: INCARNATION,
draining: true,
controls: 42,
seats: 40,
bookings: 3,
units: 45,
ceiling: 500,
intakePerSec: 5,
free: 500 - (45 - 3) - 3
}
})
expect(Object.keys(body.status).filter((key) => PER_HOST_FIELDS.includes(key))).toEqual([])
})
it('reports nulls, not zeros, for a restarted cell that has not applied its reserve switch', async () => {
const { directory, now, status } = await setup()
directory.apply(
'cell-a',
feed('cell-a', {
counts: { seats: 0, controls: 0 },
intake: undefined,
flagsApplied: { generation: 0, flags: {} },
admitModeEffective: 'db',
admitModeRaw: 'db'
}),
now.value,
now.value
)
const body = await statusBody(await status({ cellId: 'cell-a', view: 'reserve-aware' }))
expect(body.status.reserve).toMatchObject({
admitMode: 'db',
bookings: null,
units: null,
ceiling: null,
intakePerSec: null,
free: null
})
})
it('answers a retryable 503 while this director cannot vouch for the cell', async () => {
const unavailable = async (response: Response) => ({
status: response.status,
retryAfter: response.headers.get('Retry-After'),
body: await response.json()
})
const expected = (reason: string) => ({
status: 503,
retryAfter: '5',
body: { error: 'reserve_status_unavailable', reason }
})
for (const options of [{ mode: 'dry-run' as const }, { placement: false }]) {
const { directory, now, status } = await setup(options)
directory.apply('cell-a', feed('cell-a'), now.value, now.value)
expect(await unavailable(await status({ cellId: 'cell-a', view: 'reserve-aware' }))).toEqual(
expected('placement-off')
)
await database?.close()
}
// A starting director with one cell unheard.
const starting = await setup({ startedAt: 99_000 })
starting.directory.apply('cell-a', feed('cell-a'), starting.now.value, starting.now.value)
expect(
await unavailable(await starting.status({ cellId: 'cell-a', view: 'reserve-aware' }))
).toEqual(expected('map-incomplete'))
await database?.close()
const { directory, now, status } = await setup()
directory.apply('cell-b', feed('cell-b'), now.value, now.value)
directory.markNoFeed('cell-a', now.value)
expect(await unavailable(await status({ cellId: 'cell-a', view: 'reserve-aware' }))).toEqual(
expected('no-feed')
)
// A feed without the drain flag or counts is an older image: no feed, never an empty cell.
directory.apply('cell-a', feed('cell-a', { draining: undefined }), now.value, now.value)
expect(await unavailable(await status({ cellId: 'cell-a', view: 'reserve-aware' }))).toEqual(
expected('no-feed')
)
directory.apply('cell-a', feed('cell-a', { counts: undefined }), now.value, now.value)
expect(await unavailable(await status({ cellId: 'cell-a', view: 'reserve-aware' }))).toEqual(
expected('no-feed')
)
// The process that answered the feed is not the one heartbeating the runtime row.
directory.apply(
'cell-a',
feed('cell-a', { incarnation: '22222222-2222-4222-8222-222222222222' }),
now.value,
now.value
)
expect(await unavailable(await status({ cellId: 'cell-a', view: 'reserve-aware' }))).toEqual(
expected('incarnation-mismatch')
)
})
})
async function statusBody(response: Response) {
// oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: the route answers { v, status } on 200.
return (await response.json()) as { v: number; status: Record<string, unknown> }
}
function config(overrides: Partial<RelayConfig> = {}): RelayConfig {
return {
port: 8080,
publicUrl: 'https://relay.example.test',
cellUrl: 'https://relay.example.test',
region: 'us-central1',
authIssuer: 'https://auth.example.test',
authAudience: 'orca-relay',
jwksUrl: 'https://auth.example.test/jwks',
assignmentSigningKey: new Uint8Array(32),
role: 'director',
cellId: 'director',
cells: [],
adminAudience: 'https://relay.example.test/v1/admin/drain',
deployServiceAccount: 'deploy@example.test',
runtimeServiceAccount: 'relay-cell@example.test',
adminJwksUrl: 'https://auth.example.test/jwks',
databasePoolMax: 10,
publicAssignmentsEnabled: true,
publicAssignmentConcurrency: 2,
publicAssignmentQueueMax: 128,
publicAssignmentWaitMs: 4_000,
publicResolveConcurrency: 1,
publicResolveWaitMs: 5_000,
publicAssignmentRetryAfterSeconds: 5,
dataDir: './data',
...overrides
}
}
cloud/apps/relay/src/cell/admission/cell-admission-selector.test.ts
-170
View File
@@ -1,4 +1,3 @@
import { ASSIGNMENT_LIMITS } from '@orca-cloud/relay-contract/assignment-invariants'
import { createHash } from 'node:crypto'
import { afterEach, describe, expect, it } from 'vitest'
import { RelayAssignmentStore } from '../../assignment/assignment-store.js'
@@ -27,41 +26,6 @@ const BASE_MEMBERSHIP = {
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<SqlRow[]> {
return await this.delegate.query(sql, params)
}
async queryLocked(
sql: string,
params?: unknown[],
options?: RelayLockOptions
): Promise<SqlRow[]> {
return await this.delegate.queryLocked(sql, params, options)
}
async transaction<T>(operation: (transaction: RelayDatabase) => Promise<T>): Promise<T> {
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<void> {
await this.delegate.close()
}
}
class AfterTransactionDatabase implements RelayDatabase {
private afterTransaction: (() => Promise<void>) | undefined
@@ -154,34 +118,6 @@ describe('relay cell admission selector', () => {
})
})
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)
@@ -283,94 +219,6 @@ describe('relay cell admission selector', () => {
])
})
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({
@@ -416,24 +264,6 @@ describe('relay cell admission selector', () => {
).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({
cloud/apps/relay/src/cell/admission/cell-admission-startup-postgres.test.ts
-83
View File
@@ -1,83 +0,0 @@
import pg from 'pg'
import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from 'vitest'
import { RelayAssignmentStore } from '../../assignment/assignment-store.js'
import { reconcileCellAdmissionAtStartup } from './cell-admission-startup.js'
import type { RelayCellConfig } from '../../config.js'
import { openRelayDatabase, type RelayDatabase } from '../../database/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<void>((resolve) => (releaseLock = resolve))
const lockAcquired = new Promise<void>((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)
})
cloud/apps/relay/src/cell/admission/cell-admission-startup.test.ts
-44
View File
@@ -1,44 +0,0 @@
import { afterEach, describe, expect, it, vi } from 'vitest'
import {
reconcileCellAdmissionAtStartup
} 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('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)
})
})
cloud/apps/relay/src/cell/admission/cell-connection-cap-postgres.test.ts
-119
View File
@@ -1,119 +0,0 @@
import { afterAll, beforeAll, describe, expect, it } from 'vitest'
import { RelayAssignmentStore } from '../../assignment/assignment-store.js'
import type { RelayCellConfig } from '../../config.js'
import { openRelayDatabase, type RelayDatabase } from '../../database/database.js'
const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL
const describePostgres = databaseUrl ? describe : describe.skip
const CELL_PREFIX = 'connection-cap-min-postgres'
const FLEET_SIZE = 16
function cellConfig(id: string, connectionHardCap: 1_000 | 3_000): RelayCellConfig {
return {
id,
url: `https://${id}.example.com`,
capacityRequests: 1_000,
connectionHardCap,
connectionUnobservedBound: 60
}
}
async function removeCells(database: RelayDatabase): Promise<void> {
for (const table of [
'relay_cell_connection_snapshots',
'relay_cell_connection_runtime',
'relay_cell_connection_limits',
'relay_cell_capabilities',
'relay_cell_rehome_safety',
'relay_cell_runtime',
'relay_cell_admission',
'relay_cell_regions',
'relay_cells'
]) {
await database.query(`DELETE FROM ${table} WHERE cell_id LIKE ?`, [`${CELL_PREFIX}-%`])
}
}
function storeOn(database: RelayDatabase): RelayAssignmentStore {
return new RelayAssignmentStore(database, () => 100, {
requireLiveCells: true,
heartbeatTtlMs: 45_000
})
}
async function heartbeat(
store: RelayAssignmentStore,
cell: RelayCellConfig,
connectionHardCap: number,
units: number
): Promise<void> {
await store.recordCellHeartbeat({
cellId: cell.id,
cellUrl: cell.url,
cellIncarnation: '11111111-1111-4111-8111-111111111111',
startedAt: 50,
ready: true,
observedRequests: 0,
totalConnections: units,
inFlightConnections: 0,
reservedConnectionUnits: 0,
enforcedConnectionUnits: units,
connectionHardCap,
connectionUnobservedBound: 60
})
}
describePostgres('PostgreSQL effective connection cap', () => {
let database: RelayDatabase
let racer: RelayDatabase
beforeAll(async () => {
database = await openRelayDatabase({ databaseUrl, dataDir: '' })
racer = await openRelayDatabase({ databaseUrl, dataDir: '' })
await removeCells(database)
})
afterAll(async () => {
await removeCells(database)
await racer.close()
await database.close()
})
it('keeps 16 raised cells placeable while a heartbeat races the raise', async () => {
const ids = Array.from({ length: FLEET_SIZE }, (_, index) => `${CELL_PREFIX}-fleet-${index}`)
const lower = ids.map((id) => cellConfig(id, 1_000))
const director = storeOn(database)
const otherInstance = storeOn(racer)
await director.reconcileCells(lower, true)
for (const cell of lower) await heartbeat(director, cell, 1_000, 10)
// The heartbeat lands before or after the raise; either order must leave the same state.
await Promise.all([
director.reconcileCells(
ids.map((id) => cellConfig(id, 3_000)),
true
),
heartbeat(otherInstance, lower[0]!, 1_000, 11)
])
for (const id of ids) {
await expect(director.cellDeploymentStatus(id)).resolves.toMatchObject({
connectionCapacity: {
hardCap: 3_000,
reportedHardCap: 1_000,
effectiveHardCap: 1_000,
heartbeatFresh: true
}
})
const rows = await database.query(
`SELECT cell_id FROM relay_cell_connection_snapshots WHERE cell_id = ?`,
[id]
)
expect(rows).toHaveLength(1)
}
// The old processes keep heartbeating their 1,000 under the raised config.
for (const cell of lower) {
await expect(heartbeat(otherInstance, cell, 1_000, 12)).resolves.toBeUndefined()
}
})
})
cloud/apps/relay/src/cell/admission/cell-connection-cap.test.ts
+2 -83
View File
@@ -1,14 +1,11 @@
import { createHash } from 'node:crypto'
import { afterEach, describe, expect, it, vi } from 'vitest'
import { RelayAssignmentStore } from '../../assignment/assignment-store.js'
import { encodeMembership, type CellAdmissionMembership } from './cell-admission-selector.js'
import type { RelayCellConfig } from '../../config.js'
import { openInMemoryRelayDatabase, type RelayDatabase } from '../../database/database.js'
const HEARTBEAT_TTL_MS = 45_000
const BOUND = 60
const OLD_INCARNATION = '11111111-1111-4111-8111-111111111111'
const NEW_INCARNATION = '22222222-2222-4222-8222-222222222222'
function cell(index: number, connectionHardCap: 1_000 | 3_000): RelayCellConfig {
const id = `us-c${String(index).padStart(2, '0')}`
@@ -40,7 +37,7 @@ type Report = {
async function setup(cells: RelayCellConfig[]) {
const database = await openInMemoryRelayDatabase()
databases.push(database)
let now = 100
const now = 100
const store = new RelayAssignmentStore(database, () => now, {
requireLiveCells: true,
heartbeatTtlMs: HEARTBEAT_TTL_MS
@@ -81,32 +78,13 @@ async function setup(cells: RelayCellConfig[]) {
}
return rows
}
const isolateForRoll = async (cellId: string): Promise<void> => {
const current = await store.inspectCellAdmissionSelector()
const membership: CellAdmissionMembership = {
existingOnly: [],
migrationOnly: [cellId],
general: cells.map(({ id }) => id).filter((id) => id !== cellId)
}
await store.applyCellAdmissionSelector({
attemptId: `isolate-${cellId}`,
expectedGeneration: current.selector.generation,
expectedMembershipSha256: createHash('sha256')
.update(encodeMembership(current.selector.membership))
.digest('hex'),
membership,
rollIsolatedCells: [cellId]
})
}
return {
database,
store,
heartbeat,
capacity,
reportedHardCap,
telemetryRows,
isolateForRoll,
setNow: (value: number) => (now = value)
telemetryRows
}
}
@@ -215,63 +193,4 @@ describe('effective connection cap = min(configured, reported)', () => {
harness.store.assign({ userId: 'user-a', relayHostId: 'host000000000001' })
).rejects.toThrow('relay_capacity_exhausted')
})
it('places a mixed fleet against each cell’s own reported cap', async () => {
const small = cell(1, 3_000)
const large = cell(2, 3_000)
const harness = await setup([small, large])
await harness.heartbeat(small, { connectionHardCap: 1_000, units: 840 })
await harness.heartbeat(large, { connectionHardCap: 3_000, units: 840 })
expect(await harness.capacity(small.id)).toMatchObject({ normalAdmissionPause: 840 })
expect(await harness.capacity(large.id)).toMatchObject({ normalAdmissionPause: 2_840 })
for (let index = 0; index < 4; index++) {
await expect(
harness.store.assign({ userId: `user-${index}`, relayHostId: `host00000000000${index}` })
).resolves.toMatchObject({ cellId: large.id })
}
})
it('has no unavailable window while a cell rolls 1,000 -> 3,000', async () => {
const cells = [cell(1, 1_000), cell(2, 1_000)]
const [rolled, neighbour] = cells
const harness = await setup(cells)
const identity = { userId: 'user-a', relayHostId: 'host000000000001' }
await harness.heartbeat(rolled!, { connectionHardCap: 1_000 })
await harness.heartbeat(neighbour!, { connectionHardCap: 1_000, units: 500 })
const home = await harness.store.assign(identity)
expect(home.cellId).toBe(rolled!.id)
await harness.store.reconcileCells(
cells.map((_, index) => cell(index + 1, 3_000)),
true
)
// The old process keeps heartbeating its 1,000 well past the TTL; the host stays sticky.
for (const at of [30_000, 60_000, 90_000]) {
harness.setNow(at)
await harness.heartbeat(rolled!, { connectionHardCap: 1_000 })
await harness.heartbeat(neighbour!, { connectionHardCap: 1_000, units: 500 })
await expect(harness.store.assign(identity)).resolves.toMatchObject({
cellId: rolled!.id,
assignmentEpoch: home.assignmentEpoch
})
}
// The roll isolates the live old process: the isolated-incumbent lane re-places the host.
await harness.isolateForRoll(rolled!.id)
const moved = await harness.store.assign(identity)
expect(moved).toMatchObject({ cellId: neighbour!.id, placed: true })
// The replacement boots and reports the new cap.
harness.setNow(120_000)
await harness.heartbeat(rolled!, {
connectionHardCap: 3_000,
incarnation: NEW_INCARNATION,
startedAt: 110_000
})
expect(await harness.capacity(rolled!.id)).toMatchObject({
effectiveHardCap: 3_000,
heartbeatFresh: true
})
})
})
cloud/apps/relay/src/cell/cell-heartbeat-client.test.ts
-55
View File
@@ -1,55 +0,0 @@
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
})
})
})
cloud/apps/relay/src/cell/control-flags/relay-control-flag-channel.test.ts
+3 -37
View File
@@ -5,7 +5,7 @@ import {
cellFlagObjectName,
cellFlagParser
} from './cell-flags.js'
import { CONTROL_FLAG_POLL_MS, startControlFlagChannel } from './relay-control-flag-channel.js'
import { startControlFlagChannel } from './relay-control-flag-channel.js'
const CELL_ID = 'production-gce-c7'
const ACCESS_TOKEN = 'ya29.secret-access-token'
@@ -46,7 +46,7 @@ function cellObject(generation: number, flags: Record<string, unknown>, cellId =
}
}
function cellChannel(fetchImpl: typeof fetch, autoStart = false) {
function cellChannel(fetchImpl: typeof fetch) {
return startControlFlagChannel({
objectName: cellFlagObjectName(CELL_ID),
defaults: CELL_FLAG_DEFAULTS,
@@ -54,7 +54,7 @@ function cellChannel(fetchImpl: typeof fetch, autoStart = false) {
appliedEvent: CELL_FLAGS_APPLIED_EVENT,
fetch: fetchImpl,
random: () => 0.5,
autoStart
autoStart: false
})
}
@@ -77,25 +77,6 @@ describe('control flag channel', () => {
.filter((line) => line.includes(`"event":"${CELL_FLAGS_APPLIED_EVENT}"`))
.map((line): unknown => JSON.parse(line))
it('flips a flag on and back off, logging each applied generation once', async () => {
const google = fakeGoogle(cellObject(5, { readinessLocal: true, ticketCheck: 'shadow' }))
const channel = cellChannel(google.fetchImpl)
await channel.poll()
await channel.poll()
expect(google.state.storageRequests[1]?.searchParams.get('ifGenerationNotMatch')).toBe('5')
google.state.object = cellObject(7, { readinessLocal: false, ticketCheck: 'off' })
await channel.poll()
expect(channel.applied()).toEqual({ generation: 7, flags: CELL_FLAG_DEFAULTS })
expect(appliedLines()).toEqual([
{
event: CELL_FLAGS_APPLIED_EVENT,
generation: 5,
flags: { ...CELL_FLAG_DEFAULTS, readinessLocal: true, ticketCheck: 'shadow' }
},
{ event: CELL_FLAGS_APPLIED_EVENT, generation: 7, flags: CELL_FLAG_DEFAULTS }
])
})
it('keeps the applied values through every kind of lost contact, logging once per cause', async () => {
const google = fakeGoogle(cellObject(5, { readinessLocal: true }))
const channel = cellChannel(google.fetchImpl)
@@ -230,19 +211,4 @@ describe('control flag channel', () => {
expect(channel.applied()).toMatchObject({ generation: 5, flags: { readinessLocal: true } })
})
it('reads one at a time and on a jittered 5 s cadence', async () => {
vi.useFakeTimers()
const google = fakeGoogle(cellObject(5, { readinessLocal: true }))
const channel = cellChannel(google.fetchImpl)
await Promise.all([channel.poll(), channel.poll()])
expect(google.state.storageRequests).toHaveLength(1)
const running = cellChannel(google.fetchImpl, true)
await vi.advanceTimersByTimeAsync(0)
expect(google.state.storageRequests).toHaveLength(2)
await vi.advanceTimersByTimeAsync(CONTROL_FLAG_POLL_MS)
expect(google.state.storageRequests).toHaveLength(3)
running.stop()
await vi.advanceTimersByTimeAsync(CONTROL_FLAG_POLL_MS * 3)
expect(google.state.storageRequests).toHaveLength(3)
})
})
cloud/apps/relay/src/cell/reservations/cell-admit-effective-writer.test.ts
+4 -55
View File
@@ -4,69 +4,34 @@ import {
CellAdmitEffectiveWriter,
DEAD_MAN_SLOT_CLAIM_RETRY_MS
} from './cell-admit-effective-writer.js'
import { DEAD_MAN_SLOT_HOLD_MS } from './cell-dead-man-trip-slot.js'
function setup(
options: { fail?: boolean; busy?: boolean; busyWhen?: () => boolean; slotFree?: () => boolean } = {}
options: { slotFree?: () => boolean } = {}
) {
const clock = { value: 1_000 }
const writes: ('db' | 'reserve' | 'claim')[] = []
const pending: Array<() => void> = []
const logs: string[] = []
const writer = new CellAdmitEffectiveWriter({
write: async (mode) => {
writes.push(mode)
await new Promise<void>((resolve) => pending.push(resolve))
if (options.fail) throw new Error('connection refused')
},
claimTrip: async (slots) => {
writes.push('claim')
await new Promise<void>((resolve) => pending.push(resolve))
return slots > 0 && (options.slotFree?.() ?? true)
},
databaseBusy: () => options.busy === true || options.busyWhen?.() === true,
now: () => clock.value,
log: (line) => logs.push(line)
databaseBusy: () => false,
now: () => clock.value
})
const settle = async () => {
while (pending.length > 0) pending.shift()!()
await new Promise((resolve) => setTimeout(resolve, 0))
}
return { clock, writes, logs, writer, settle }
return { clock, writes, writer, settle }
}
describe('CellAdmitEffectiveWriter', () => {
it('writes on the first tick, on every change, and every 15 s, one write at a time', async () => {
const { clock, writes, writer, settle } = setup()
writer.tick('reserve', true)
writer.tick('db', true)
expect(writes).toEqual(['reserve'])
await settle()
// Tripped: the switch still says reserve, the raw mode db.
writer.tick('db', true)
await settle()
writer.tick('db', true)
expect(writes).toEqual(['reserve', 'db'])
clock.value += CELL_ADMIT_EFFECTIVE_REFRESH_MS
writer.tick('db', true)
expect(writes).toEqual(['reserve', 'db', 'db'])
})
// Step 5 off: a db switch never writes, so the dark roll adds no database write.
// Best-effort: a failed write never throws into the cell, logs at most once a minute, and retries.
it('retries a failed write on the next tick and rate-limits its log line', async () => {
const { clock, writes, logs, writer, settle } = setup({ fail: true })
writer.tick('db', true)
await settle()
clock.value += 1_000
writer.tick('db', true)
await settle()
expect(writes).toEqual(['db', 'db'])
expect(logs).toHaveLength(1)
expect(JSON.parse(logs[0]!)).toMatchObject({ event: 'orca_relay_cell_admit_effective_write_failed', mode: 'db' })
})
it('claims a held trip slot in place of its refresh, and permits the trip only once it lands', async () => {
let free = false
const { clock, writes, writer, settle } = setup({ slotFree: () => free })
@@ -97,20 +62,4 @@ describe('CellAdmitEffectiveWriter', () => {
expect(writes).toEqual(['reserve', 'claim', 'reserve', 'claim'])
})
it('refreshes a tripped row past a saturated pool before its slot hold runs out', async () => {
let busy = false
const { clock, writes, writer, settle } = setup({ busyWhen: () => busy })
writer.tick('db', true)
await settle()
busy = true
clock.value += 4 * 60_000
writer.tick('db', true)
expect(writes).toEqual(['db'])
clock.value += 60_000
writer.tick('db', true)
await settle()
expect(writes).toEqual(['db', 'db'])
expect(DEAD_MAN_SLOT_HOLD_MS).toBeGreaterThan(5 * 60_000)
})
})
cloud/apps/relay/src/cell/reservations/cell-reserve-admission.blackbox.test.ts
-163
View File
@@ -8,8 +8,6 @@ import { afterEach, describe, expect, it, vi } from 'vitest'
import WebSocket from 'ws'
import type { RawData } from 'ws'
import { CELL_FLAG_DEFAULTS, type CellFlags } from '../control-flags/cell-flags.js'
import { CELL_SEAT_DRIFT_SWEEP_MS } from '../seats/cell-seat-drift.js'
import type { CellSeatLog } from '../seats/cell-seat-log.js'
import type { RelayConfig } from '../../config.js'
import { PostgresDatabase } from '../../database/database.js'
import type { AppliedControlFlags } from '../control-flags/relay-control-flag-channel.js'
@@ -73,20 +71,6 @@ function cellConfig(port: number, issuer: string): RelayConfig {
}
}
// The registry's private feed log, to plant the drift the cell must repair.
function seatLogOf(sessions: object): CellSeatLog {
// oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: HostSessionRegistry holds its feed log as `seatLog`.
return (sessions as { seatLog: CellSeatLog }).seatLog
}
// The JSON lines a test's console.warn spy saw for one event.
function warned(event: string): Array<Record<string, unknown>> {
return vi
.mocked(console.warn)
.mock.calls.map(([line]) => (typeof line === 'string' && line.startsWith('{') ? JSON.parse(line) : null))
.filter((entry) => entry?.event === event)
}
function nextMessage(socket: WebSocket): Promise<Record<string, unknown>> {
return new Promise((resolve, reject) => {
socket.once('message', (data: RawData) => {
@@ -265,9 +249,6 @@ describe('a reserve-mode cell admits from memory with its database wedged', () =
book,
setFlags: (next: Partial<CellFlags>) => {
applied = { generation: applied.generation + 1, flags: { ...CELL_FLAG_DEFAULTS, ...next } }
},
advance: (ms: number) => {
elapsedMs += ms
}
}
}
@@ -338,97 +319,6 @@ describe('a reserve-mode cell admits from memory with its database wedged', () =
return (page as { full: Array<{ relayHostId: string; epoch: number }> }).full.map((entry) => entry.epoch)
}
it('repairs at once the seat a director names when no session backs it', async () => {
const cell = await startCell({ admitMode: 'reserve' })
const identity = await cell.host()
cell.book(identity, 9)
await cell.connect(identity, 9)
// A stale seat at a generation above the live session's: its changes could never replace it.
const log = seatLogOf(cell.relay.sessions)
log.append({ kind: 'join', userId: 'user-1', relayHostId: identity.hostId, epoch: 8, generation: 5, at: 4_242 })
const before = cell.relay.sessions.seatFeed(null)
const seq = 'full' in before ? before.seq : 0
const seat = { v: 1 as const, userId: 'user-1', relayHostId: identity.hostId, epoch: 8, joinedAt: 4_242 }
expect(cell.relay.sessions.demote(seat)).toBe('not-seated')
expect(cell.relay.sessions.seatFeed(seq)).toMatchObject({
changes: [
{ kind: 'leave', epoch: 8, generation: 5, closeCode: 4409 },
{ kind: 'join', epoch: 9, generation: 1 }
]
})
expect(seatEpochs(cell)).toEqual([9])
expect(warned('orca_relay_cell_seat_log_drift')).toMatchObject([
{ kind: 'unbacked-seat', via: 'demote', reason: 'epoch', seatEpoch: 8, sessionEpoch: 9, relisted: true }
])
expect(JSON.stringify(warned('orca_relay_cell_seat_log_drift'))).not.toContain(identity.hostId)
// Asked again, there is nothing left to repair.
expect(cell.relay.sessions.demote(seat)).toBe('not-seated')
expect(cell.relay.sessions.seatFeed(seq + 2)).toMatchObject({ changes: [] })
})
it('leaves a seat no session backs once two sweeps in a row find it, with no director asking', async () => {
const cell = await startCell({ admitMode: 'db' })
const ghost = await cell.host()
seatLogOf(cell.relay.sessions).append({
kind: 'join',
userId: 'user-1',
relayHostId: ghost.hostId,
epoch: 10,
generation: 3,
at: 1_000
})
expect(seatEpochs(cell)).toEqual([10])
cell.advance(CELL_SEAT_DRIFT_SWEEP_MS)
expect(seatEpochs(cell)).toEqual([])
expect(warned('orca_relay_cell_seat_log_drift')).toMatchObject([
{ kind: 'unbacked-seat', via: 'sweep', reason: 'no-session', seatEpoch: 10, seatGeneration: 3, relisted: false }
])
// Left as moved: the cell will not seat it again at 10 from memory.
expect(cell.relay.sessions.seatFeed(0)).toMatchObject({ changes: [{}, { kind: 'leave', closeCode: 4409 }] })
})
it('lists again a live session the feed lost, once two sweeps in a row find it', async () => {
const cell = await startCell({ admitMode: 'reserve' })
const identity = await cell.host()
cell.book(identity, 9)
await cell.connect(identity, 9)
seatLogOf(cell.relay.sessions).append({
kind: 'leave',
userId: 'user-1',
relayHostId: identity.hostId,
epoch: 9,
generation: 1,
at: 5_000
})
expect(seatEpochs(cell)).toEqual([])
cell.advance(CELL_SEAT_DRIFT_SWEEP_MS)
expect(seatEpochs(cell)).toEqual([9])
expect(warned('orca_relay_cell_seat_log_drift')).toMatchObject([{ kind: 'unlisted-session', sessionEpoch: 9 }])
})
// A drained session's seat stays listed, with no session, until its socket's close lands.
it('repairs nothing for a seat that one sweep finds unbacked and is gone by the next', async () => {
const cell = await startCell({ admitMode: 'reserve' })
const identity = await cell.host()
cell.book(identity, 9)
const live = await cell.connect(identity, 9)
// oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: ws keeps its net.Socket as `_socket`.
;(live.socket as unknown as { _socket: { pause(): void } })._socket.pause()
cell.relay.sessions.drainHost({
attemptId: 'attempt-1',
userId: 'user-1',
relayHostId: identity.hostId,
sourceAssignmentEpoch: 9,
graceMs: 0
})
await new Promise((resolve) => setTimeout(resolve, 20))
expect(seatEpochs(cell)).toEqual([9])
await new Promise((resolve) => setTimeout(resolve, RELAY_WEBSOCKET_FORCE_CLOSE_MS + 300))
cell.advance(CELL_SEAT_DRIFT_SWEEP_MS)
expect(seatEpochs(cell)).toEqual([])
expect(warned('orca_relay_cell_seat_log_drift')).toEqual([])
}, 10_000)
// Drained sockets can close seconds after map removal; reusing their generation let late leaves erase reconnected live seats.
it('keeps a session that reconnected while its drained predecessor was still closing', async () => {
const cell = await startCell({ admitMode: 'reserve' })
@@ -456,19 +346,6 @@ describe('a reserve-mode cell admits from memory with its database wedged', () =
expect(seatEpochs(cell)).toEqual([10])
}, 10_000)
// Repeated fresh joins for one identity replace session and seat together, leaving no drift to repair.
it('stays on the live session when the same identity connects afresh each turn', async () => {
const cell = await startCell({ admitMode: 'reserve' })
const identity = await cell.host()
for (const epoch of [9, 10, 11]) {
cell.book(identity, epoch)
expect((await cell.connect(identity, epoch)).ack).toMatchObject({ type: 'host-hello-ack' })
expect(seatEpochs(cell)).toEqual([epoch])
cell.advance(CELL_SEAT_DRIFT_SWEEP_MS)
}
expect(seatEpochs(cell)).toEqual([11])
expect(warned('orca_relay_cell_seat_log_drift')).toEqual([])
})
})
it('gives every booking unit back when the cell starts to drain', async () => {
@@ -542,44 +419,4 @@ describe('a reserve-mode cell admits from memory with its database wedged', () =
}).toEqual(before)
}, 30_000)
it('never crosses the hard cap under seeded bookings, hellos, rebooks and closes', async () => {
const hardCap = 6
const cell = await startCell({ admitMode: 'reserve' }, { hardCap, controlReserve: 1 })
const hosts = await Promise.all(Array.from({ length: 12 }, () => cell.host()))
const epochs = new Map<string, number>()
const open: WebSocket[] = []
let seed = 7
const random = () => {
seed = (seed * 1_103_515_245 + 12_345) % 2_147_483_648
return seed / 2_147_483_648
}
let peak = 0
for (let step = 0; step < 240; step += 1) {
const identity = hosts[Math.floor(random() * hosts.length)]!
const roll = random()
if (roll < 0.4) {
// A director's fresh booking, or a re-book at a newer epoch.
const epoch = (epochs.get(identity.hostId) ?? 1) + (random() < 0.5 ? 0 : 1)
if (cell.book(identity, epoch)[0]?.outcome === 'ok') epochs.set(identity.hostId, epoch)
} else if (roll < 0.8) {
const reply = await cell.connect(identity, epochs.get(identity.hostId) ?? 1)
if (reply.ack?.type === 'host-hello-ack') open.push(reply.socket!)
} else if (open.length > 0) {
const socket = open.splice(Math.floor(random() * open.length), 1)[0]!
// A newer control for the same host may already have replaced it.
if (socket.readyState !== WebSocket.CLOSED) {
const closed = new Promise((resolve) => socket.once('close', resolve))
socket.close()
await closed
}
}
const counts = cell.relay.connectionSnapshot()!
peak = Math.max(peak, counts.physicalConnections)
expect(counts.physicalConnections).toBeLessThanOrEqual(hardCap)
expect(counts.enforcedConnectionUnits).toBeLessThanOrEqual(hardCap)
}
// The run did reach the cap, so the bound was exercised.
expect(peak).toBeGreaterThanOrEqual(hardCap - 1)
}, 60_000)
})
cloud/apps/relay/src/cell/reservations/cell-reserve-client.test.ts
-41
View File
@@ -1,41 +0,0 @@
import { describe, expect, it, vi } from 'vitest'
import { CellReserveClient } from './cell-reserve-client.js'
const target = { cellId: 'c1', cellUrl: 'https://c1.example' }
const item = (userId: string, epoch = 3) => ({
userId,
relayHostId: 'abcdefghijklmnop',
epoch,
ttlMs: 30_000
})
describe('cell reserve client', () => {
it('batches a cell’s bookings into one request and maps each answer back', async () => {
const fetchImpl = vi.fn(async (_url: string | URL | Request, init?: RequestInit) => {
const body = JSON.parse(String(init?.body)) as { items: Array<{ userId: string }> }
return Response.json({
v: 1,
results: body.items.map((entry) =>
entry.userId === 'b' ? { outcome: 'seated-newer', epoch: 9 } : { outcome: 'ok' }
)
})
})
const client = new CellReserveClient({
directorId: 'director-a',
identityToken: async () => 'token',
fetch: fetchImpl
})
const answers = await Promise.all([
client.reserve(target, item('a')),
client.reserve(target, item('b')),
client.reserve(target, item('c'), true)
])
expect(answers).toEqual([{ outcome: 'ok' }, { outcome: 'seated-newer', epoch: 9 }, { outcome: 'ok' }])
// One booking batch and one dry-run batch.
expect(fetchImpl).toHaveBeenCalledTimes(2)
const first = JSON.parse(String(fetchImpl.mock.calls[0]![1]?.body)) as Record<string, unknown>
expect(first).toMatchObject({ v: 1, directorId: 'director-a', items: [{ userId: 'a' }, { userId: 'b' }] })
expect(String(fetchImpl.mock.calls[0]![0])).toBe('https://c1.example/internal/reserve')
})
})
cloud/apps/relay/src/cell/reservations/cell-reserve-dead-man-slots.blackbox.test.ts
-9
View File
@@ -128,13 +128,4 @@ describe('dead-man trip slots', () => {
expect(seated.socket.readyState).toBe(WebSocket.OPEN)
}, 30_000)
it('without the flag, trips with every slot taken, as before', async () => {
let now = Date.now()
const { database, reserve } = await cell(() => now)
await trippedPeer(database, 'production-gce-c90', now)
await reserve.bookedHost(3)
expect(reserve.relay.sessions.inReserveMode()).toBe(true)
now += reserveDeadManWindowMs(reserve.config.cellId) + 1_000
expect(reserve.relay.sessions.inReserveMode()).toBe(false)
}, 30_000)
})
cloud/apps/relay/src/cell/reservations/cell-reserve-dead-man.test.ts
-9
View File
@@ -55,15 +55,6 @@ describe('ReserveDeadMan', () => {
])
})
it('keeps the plain window for a later reserve write once a director has been in touch', () => {
const { clock, deadMan } = setup(3 * WINDOW_MS)
deadMan.mode(reserve(1))
deadMan.contact()
expect(deadMan.mode(reserve(2))).toBe('reserve')
clock.value += WINDOW_MS + 1
expect(deadMan.mode(reserve(2))).toBe('db')
})
// A claim answered after the director came back must not trip the cell later without a slot.
it('drops a permit for a silence that contact has ended', () => {
const { clock, deadMan, events } = setup()
cloud/apps/relay/src/cell/reservations/cell-reserve-flip-back-postgres.test.ts
-114
View File
@@ -1,114 +0,0 @@
import { mkdtempSync, rmSync } from 'node:fs'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
import { afterEach, describe, expect, it, vi } from 'vitest'
import WebSocket from 'ws'
import { openRelayDatabase, type RelayDatabase } from '../../database/database.js'
import { openDelayedPostgresDatabase, type StatementDelay } from '../../test-fixtures/delayed-postgres-database.js'
import { RESERVE_MODE_CELL_ID, startReserveModeCell } from '../../test-fixtures/reserve-mode-cell.js'
// PostgreSQL flip-back keeps controls connected, reserves most pool capacity for DB work and folds cell deltas per tick even across cross-region round trips (E-mix).
const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL
const describePostgres = databaseUrl ? describe : describe.skip
async function until(check: () => Promise<boolean> | boolean, timeoutMs: number): Promise<void> {
const deadline = Date.now() + timeoutMs
while (!(await check())) {
if (Date.now() > deadline) throw new Error('condition not reached')
await new Promise((resolve) => setTimeout(resolve, 50))
}
}
describePostgres('flipping a reserve-mode cell back on PostgreSQL', () => {
const cleanup: Array<() => Promise<void> | void> = []
afterEach(async () => {
for (const close of cleanup.splice(0).reverse()) await close()
vi.restoreAllMocks()
// The database is shared with the other PostgreSQL suites, some of which read every cell.
const database = await openRelayDatabase({ databaseUrl, dataDir: tmpdir() })
try {
const tables = await database.query(
`SELECT table_name FROM information_schema.columns
WHERE table_schema = current_schema() AND column_name = 'cell_id' AND table_name LIKE 'relay_%'`
)
for (const row of tables) {
await database.query(`DELETE FROM ${String(row.table_name)} WHERE cell_id = ?`, [RESERVE_MODE_CELL_ID])
}
} finally {
await database.close()
}
})
async function cellOn(database: RelayDatabase, databasePoolMax = 10) {
vi.spyOn(console, 'log').mockImplementation(() => undefined)
vi.spyOn(console, 'warn').mockImplementation(() => undefined)
const dataDir = mkdtempSync(join(tmpdir(), 'orca-relay-flip-back-pg-'))
cleanup.push(() => rmSync(dataDir, { recursive: true, force: true }))
return await startReserveModeCell({ database, dataDir, cleanup, databasePoolMax })
}
it('re-registers every memory-admitted control, ten flips in a row', async () => {
const database = await openRelayDatabase({ databaseUrl, dataDir: tmpdir() })
cleanup.push(() => database.close())
for (let flip = 0; flip < 10; flip += 1) {
const cell = await cellOn(database)
const hosts = []
for (let index = 0; index < 3; index += 1) hosts.push(await cell.bookedHost(5 + index))
const ids = new Set(hosts.map((entry) => entry.identity.hostId))
const leased = async () =>
(await cell.controlLeases()).filter((row) => ids.has(String(row.relay_host_id)))
expect(await leased()).toEqual([])
cell.setAdmitMode('db')
await until(async () => (await leased()).length === 3, 10_000)
await until(() => cell.relay.sessions.reregistrationPending() === 0, 10_000)
expect(hosts.every((entry) => entry.socket.readyState === WebSocket.OPEN)).toBe(true)
for (const entry of hosts) entry.socket.close()
}
}, 120_000)
it('keeps re-registration to a third of the pool and one cell-row write per tick, at 170 ms a trip', async () => {
const delay: StatementDelay = { enabled: false, delayMs: 170 }
let cellRowWrites = 0
let leaseWritesInFlight = 0
let peakInFlight = 0
delay.beforeTrip = async (sql) => {
if (/UPDATE relay_cells\b/.test(sql)) cellRowWrites += 1
}
const database = openDelayedPostgresDatabase(databaseUrl!, delay, 10)
cleanup.push(() => database.close())
// The schema and the cell row, without the delay.
const setup = await openRelayDatabase({ databaseUrl, dataDir: tmpdir() })
await setup.close()
const cell = await cellOn(database, 10)
const deferred = cell.relay.assignments.activateControlDeferringCell.bind(cell.relay.assignments)
vi.spyOn(cell.relay.assignments, 'activateControlDeferringCell').mockImplementation(async (...args) => {
leaseWritesInFlight += 1
peakInFlight = Math.max(peakInFlight, leaseWritesInFlight)
try {
return await deferred(...args)
} finally {
leaseWritesInFlight -= 1
}
})
const hosts = []
for (let index = 0; index < 9; index += 1) hosts.push(await cell.bookedHost(5 + index))
const ids = new Set(hosts.map((entry) => entry.identity.hostId))
delay.enabled = true
const startedAt = Date.now()
cell.setAdmitMode('db')
await until(() => cell.relay.sessions.reregistrationPending() > 0, 5_000)
await until(() => cell.relay.sessions.reregistrationPending() === 0, 45_000)
const elapsedMs = Date.now() - startedAt
delay.enabled = false
const leased = (await cell.controlLeases()).filter((row) => ids.has(String(row.relay_host_id)))
expect(leased).toHaveLength(9)
expect(hosts.every((entry) => entry.socket.readyState === WebSocket.OPEN)).toBe(true)
// Pool of 10: at most 3 activations at once, leaving 7 for hellos and renewals.
expect(peakInFlight).toBeLessThanOrEqual(3)
// Far fewer cell-row writes than activations: deltas fold into one write per tick.
expect(cellRowWrites).toBeLessThan(9)
process.stderr.write(`flip_back_latency ${JSON.stringify({ elapsedMs, peakInFlight, cellRowWrites })}\n`)
}, 60_000)
})
cloud/apps/relay/src/cell/reservations/cell-reserve-rollback.blackbox.test.ts
+2 -44
View File
@@ -3,7 +3,6 @@ import { tmpdir } from 'node:os'
import { join } from 'node:path'
import { afterEach, describe, expect, it, vi } from 'vitest'
import WebSocket from 'ws'
import { CELL_ADMIT_EFFECTIVE_REFRESH_MS } from './cell-admit-effective-writer.js'
import { RESERVE_DEAD_MAN_MS, reserveDeadManWindowMs } from './cell-reserve-dead-man.js'
import { openRelayDatabase } from '../../database/database.js'
import {
@@ -34,7 +33,7 @@ describe('flipping a reserve-mode cell back to the database', () => {
vi.restoreAllMocks()
})
async function cell(now?: () => number, options: { flagsUnread?: boolean } = {}) {
async function cell(now?: () => number) {
vi.spyOn(console, 'log').mockImplementation(() => undefined)
vi.spyOn(console, 'warn').mockImplementation(() => undefined)
vi.spyOn(console, 'error').mockImplementation(() => undefined)
@@ -42,7 +41,7 @@ describe('flipping a reserve-mode cell back to the database', () => {
cleanup.push(() => rmSync(dataDir, { recursive: true, force: true }))
const database = await openRelayDatabase({ dataDir })
cleanup.push(() => database.close())
return await startReserveModeCell({ database, dataDir, cleanup, ...options, ...(now ? { now } : {}) })
return await startReserveModeCell({ database, dataDir, cleanup, ...(now ? { now } : {}) })
}
it('registers every memory-admitted control without disconnecting any, matching the seat report', async () => {
@@ -159,47 +158,6 @@ describe('flipping a reserve-mode cell back to the database', () => {
expect(host.socket.readyState).toBe(WebSocket.OPEN)
}, 60_000)
it('takes its flip-back pace from the switch file', async () => {
const reserve = await cell()
for (let index = 0; index < 4; index += 1) await reserve.bookedHost(5 + index)
const real = reserve.relay.assignments.activateControlDeferringCell.bind(reserve.relay.assignments)
let inFlight = 0
let peak = 0
vi.spyOn(reserve.relay.assignments, 'activateControlDeferringCell').mockImplementation(async (...args) => {
inFlight += 1
peak = Math.max(peak, inFlight)
await new Promise((resolve) => setTimeout(resolve, 50))
try {
return await real(...args)
} finally {
inFlight -= 1
}
})
reserve.setAdmitMode('db', { reregisterInFlight: 1 })
await until(async () => (await reserve.controlLeases()).length === 4)
expect(peak).toBe(1)
}, 30_000)
// Step 5 off adds no writes before/after flag reads; leaving reserve owes exactly one db write.
it('writes an admit-effective row only while its switch says reserve, plus one db write on the way out', async () => {
let now = Date.now()
const booting = await cell(() => now, { flagsUnread: true })
await new Promise((resolve) => setTimeout(resolve, 2_500))
expect(await booting.admitEffectiveRow()).toBeUndefined()
booting.setAdmitMode('db')
now += 3 * CELL_ADMIT_EFFECTIVE_REFRESH_MS
await new Promise((resolve) => setTimeout(resolve, 2_500))
expect(await booting.admitEffectiveRow()).toBeUndefined()
booting.setAdmitMode('reserve')
await until(async () => (await booting.admitEffectiveRow())?.mode === 'reserve')
booting.setAdmitMode('db')
await until(async () => (await booting.admitEffectiveRow())?.mode === 'db')
const last = (await booting.admitEffectiveRow())!.updated_at
now += 3 * CELL_ADMIT_EFFECTIVE_REFRESH_MS
await new Promise((resolve) => setTimeout(resolve, 2_500))
expect((await booting.admitEffectiveRow())!.updated_at).toEqual(last)
}, 30_000)
it('flips itself back to db when no placing director has been in touch for the dead-man window', async () => {
let now = Date.now()
const reserve = await cell(() => now)
cloud/apps/relay/src/cell/seats/cell-poll-reachability.test.ts
+1 -9
View File
@@ -50,7 +50,7 @@ const answer = (status: number) => new Response('{}', { status })
// Every cell serves a feed page on the first round; `respond` answers every round after it.
function poller(
respond: (cellId: string) => Response | Promise<Response>,
options: { identityFails?: boolean; cells?: SeatFeedCell[]; neverServed?: ReadonlySet<string> } = {}
options: { cells?: SeatFeedCell[]; neverServed?: ReadonlySet<string> } = {}
) {
let now = 1_000
let round = 0
@@ -66,7 +66,6 @@ function poller(
return round === 0 ? page(cellId) : await respond(cellId)
}),
identityToken: vi.fn(async () => {
if (options.identityFails && round > 0) throw new Error('metadata_unavailable')
return 'token-1'
}),
now: () => now,
@@ -107,13 +106,6 @@ describe('seat-feed poller reachability', () => {
}
})
it('does not blame a cell for a token this director cannot mint', async () => {
const cells = poller((cellId) => page(cellId), { identityFails: true })
const at = await cells.run(CELL_POLL_DEAD_STREAK_MS)
cells.stop()
expect(cells.deadCellVerdict('c2', at)).toBeNull()
})
it('counts only cells that served a feed: legacy 404 rows neither fill the sample nor dilute it', async () => {
// Staging's shape: a few real cells and 16 legacy rows. Every real cell fails.
const cells = poller(() => answer(502), { cells: [...CELLS.slice(0, 3), ...LEGACY] })
cloud/apps/relay/src/cell/seats/cell-seat-feed.test.ts
-336
View File
@@ -1,336 +0,0 @@
import { EventEmitter } from 'node:events'
import { RELAY_CLOSE_CODE } from '@orca-cloud/relay-contract/close-codes'
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import type WebSocket from 'ws'
import { z } from 'zod'
import type { RelayAssignmentStore } from '../../assignment/assignment-store.js'
import type { CellFlags } from '../control-flags/cell-flags.js'
import type { RelayConfig } from '../../config.js'
import type { RelayCredentialStore } from '../../credentials/credential-store.js'
import { HostSessionRegistry, type HostSession } from '../../connections/host-session-registry.js'
import type { AppliedControlFlags } from '../control-flags/relay-control-flag-channel.js'
import type { RelayRuntimeObserver } from '../../observability/relay-observability.js'
import type { RelayTokenClaims } from '../../authentication/relay-token-verifier.js'
import { ProcessQueuedByteBudget } from '../../connections/splice-forwarder.js'
const adminTokenChecks = vi.hoisted(() => ({ count: 0 }))
vi.mock('../../authentication/admin-token-verifier.js', () => ({
createAdminTokenVerifier: () => async (token: string) => {
adminTokenChecks.count += 1
return token === 'deploy-token'
},
createReadOnlyAdminTokenVerifier: () => async () => false,
createRegionalRehomeControlApplyTokenVerifier: () => async () => false,
createRegionalRehomeRuntimeTokenVerifier: () => async () => false,
createRegionalRehomeTokenVerifier: () => async (token: string) => token === 'rehome-token',
createRuntimeTokenVerifier: () => async () => false
}))
import { createRelayApp } from '../../app.js'
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 incarnation = '11111111-1111-4111-8111-111111111111'
// The published v1 shape, so a field rename fails here before it reaches a director.
const SeatChangeSchema = z.object({
seq: z.number(),
kind: z.enum(['join', 'leave', 'drain-only', 'active']),
userId: z.string(),
relayHostId: z.string(),
epoch: z.number(),
generation: z.number(),
state: z.enum(['active', 'drain-only']).optional(),
closeCode: z.number().optional(),
at: z.number()
})
const SeatFeedReplySchema = z.object({
v: z.literal(1),
cellId: z.string(),
incarnation: z.string(),
at: z.number(),
draining: z.boolean(),
counts: z.object({
controls: z.number(),
seats: z.number()
}),
flagsApplied: z
.object({
generation: z.number(),
flags: z.object({
readinessLocal: z.boolean(),
ticketCheck: z.enum(['off', 'shadow', 'enforce']),
rejectionFence: z.boolean(),
admitMode: z.enum(['db', 'reserve']),
reserveDryRun: z.boolean(),
phonePath: z.enum(['db', 'memory'])
}),
ignoredKeys: z.array(z.string()).optional()
})
.optional(),
seq: z.number(),
changes: z.array(SeatChangeSchema).optional(),
more: z.boolean().optional(),
full: z
.array(
z.object({
userId: z.string(),
relayHostId: z.string(),
epoch: z.number(),
generation: z.number(),
state: z.enum(['active', 'drain-only']),
joinedAt: z.number()
})
)
.optional()
})
function config(overrides: Partial<RelayConfig> = {}): RelayConfig {
return {
port: 8080,
publicUrl: 'https://c7.relay.example.test',
cellUrl: 'https://c7.relay.example.test',
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
}
}
function hostIdentity(index: number, exp = 4_102_444_800): RelayTokenClaims {
return {
sub: `user-${index}`,
prof: 'profile-1',
relayHostId: `host${String(index).padStart(12, '0')}`,
purpose: 'host-control',
exp
}
}
// What a host's auth-refresh token verifies as.
const verifyRelayToken = vi.fn<(token: string) => Promise<RelayTokenClaims | null>>()
function createCell(relayConfig = config(), cellFlags?: () => AppliedControlFlags<CellFlags>) {
const assignments = {
activateControl: vi.fn().mockResolvedValue('control:production-gce-c7:1'),
markMigrationTargetRegistered: vi.fn().mockResolvedValue(undefined),
releaseActivity: vi.fn().mockResolvedValue(true),
renewControlActivities: vi.fn().mockResolvedValue([])
} as unknown as RelayAssignmentStore
const observer = {
recordAuth: vi.fn(),
recordForwardedBytes: vi.fn(),
recordHttp: vi.fn(),
recordReconnect: vi.fn(),
recordSql: vi.fn()
} satisfies RelayRuntimeObserver
const registry = new HostSessionRegistry(
relayConfig,
verifyRelayToken,
{} as RelayCredentialStore,
assignments,
new ProcessQueuedByteBudget(),
observer,
Date.now,
Math.random,
incarnation
)
const activate = (
registry as unknown as {
activate: (
socket: WebSocket,
identity: RelayTokenClaims,
existing: HostSession | null,
generation: number,
rebind: boolean,
assignmentEpoch: number,
appVersion: string
) => Promise<void>
}
).activate.bind(registry)
const connect = async (
index: number,
options: { generation?: number; epoch?: number; rebind?: boolean; exp?: number } = {}
): Promise<FakeSocket> => {
const socket = new FakeSocket()
const identity = hostIdentity(index, options.exp)
await activate(
socket as unknown as WebSocket,
identity,
registry.get({ userId: identity.sub, relayHostId: identity.relayHostId }),
options.generation ?? 1,
options.rebind ?? false,
options.epoch ?? 1,
'1.4.200'
)
return socket
}
const app = createRelayApp(relayConfig, {
store: {} as never,
assignments: {} as never,
drain: (graceMs, options) => registry.drain(graceMs, options ?? {}),
cellIncarnation: incarnation,
cellSeatFeed: (sinceSeq) => registry.seatFeed(sinceSeq),
cellFlags,
cellPhonePathEffective: () => registry.phonePathEffective(),
isDraining: () => registry.isDraining(),
runtimeCounts: () => ({
totalConnections: 0,
preAuthConnections: 0,
queuedBytes: 0,
...registry.runtimeCounts()
}),
ready: vi.fn(async () => true)
})
const poll = async (since?: string, token = 'rehome-token') =>
await app.request(
`/v1/admin/cell-seats${since === undefined ? '' : `?since=${encodeURIComponent(since)}`}`,
{ headers: { authorization: `Bearer ${token}` } }
)
const read = async (since?: string) => SeatFeedReplySchema.parse(await (await poll(since)).json())
return { registry, connect, app, poll, read }
}
describe('cell seat feed', () => {
beforeEach(() => vi.useFakeTimers())
afterEach(() => {
vi.clearAllTimers()
vi.useRealTimers()
})
it('keeps a closing socket seated until its close lands, as the change log does', async () => {
const { connect, read } = createCell()
const socket = await connect(1)
// The ws close handshake: no longer OPEN, close event not yet emitted.
socket.readyState = 2
const closing = await read()
expect(closing.counts).toMatchObject({ controls: 0, seats: 1 })
expect(closing.full).toHaveLength(1)
socket.readyState = socket.CLOSED
socket.emit('close', 1000, Buffer.from(''))
expect((await read()).counts).toMatchObject({ controls: 0, seats: 0 })
})
it('reports join, drain-only and leave as changes after the cursor', async () => {
const { registry, connect, read } = createCell()
const socket = await connect(1)
const { seq } = await read()
registry.drain(60_000)
socket.close(RELAY_CLOSE_CODE.DRAINING, 'relay draining')
await connect(2)
const body = await read(`${incarnation}:${seq}`)
expect(body.full).toBeUndefined()
expect(body.draining).toBe(true)
expect(body.changes?.map((change) => change.kind)).toEqual(['drain-only', 'leave'])
expect(body.changes?.[1]).toMatchObject({
userId: 'user-1',
closeCode: RELAY_CLOSE_CODE.DRAINING,
generation: 1
})
// The second connect landed after the drain, so the cell refused it: no join.
expect(body.seq).toBe(seq + 2)
})
it('reports an auth expiry as drain-only and the refresh that clears it as active', async () => {
vi.spyOn(console, 'warn').mockImplementation(() => {})
const { connect, read } = createCell()
const exp = Math.floor(Date.now() / 1000) + 5
const socket = await connect(1, { exp })
const { seq } = await read()
// The 15 s heartbeat finds the token expired, inside its grace.
await vi.advanceTimersByTimeAsync(15_000)
verifyRelayToken.mockResolvedValueOnce(hostIdentity(1, exp + 3_600))
socket.emit(
'message',
Buffer.from(JSON.stringify({ type: 'auth-refresh', relayJwt: 'refreshed' })),
false
)
await vi.waitFor(() => expect(verifyRelayToken).toHaveBeenCalled())
const body = await read(`${incarnation}:${seq}`)
expect(body.changes?.map((change) => change.kind)).toEqual(['drain-only', 'active'])
expect(body.counts.seats).toBe(1)
expect((await read()).full).toEqual([expect.objectContaining({ state: 'active' })])
})
it('does not report a rebind as a leave, and joins the rebound socket', async () => {
const { connect, read } = createCell()
await connect(1)
const { seq } = await read()
await connect(1, { rebind: true, epoch: 2 })
const body = await read(`${incarnation}:${seq}`)
expect(body.changes).toEqual([
expect.objectContaining({ kind: 'join', userId: 'user-1', epoch: 2, state: 'active' })
])
})
it('reports a superseded generation leaving after its successor joined', async () => {
const { connect, read } = createCell()
await connect(1)
const { seq } = await read()
await connect(1, { generation: 2 })
const body = await read(`${incarnation}:${seq}`)
// Readers drop the leave: its generation is older than the seat's.
expect(body.changes?.map((change) => [change.kind, change.generation])).toEqual([
['leave', 1],
['join', 2]
])
})
it('resyncs a cursor from another incarnation with a full snapshot', async () => {
const { connect, poll, read } = createCell()
await connect(1)
const body = await read('22222222-2222-4222-8222-222222222222:1')
expect(body.full).toHaveLength(1)
expect((await poll('not a cursor')).status).toBe(400)
})
it('accepts only the directors rehome identity, verified once per poll', async () => {
const { poll } = createCell()
adminTokenChecks.count = 0
expect((await poll(undefined, 'deploy-token')).status).toBe(401)
expect((await poll(undefined, 'wrong')).status).toBe(401)
expect((await poll()).status).toBe(200)
// The admin middleware leaves this route to its own rehome check.
expect(adminTokenChecks.count).toBe(0)
})
it('is unavailable on a cell without the rehome pair, and on directors', async () => {
const withoutPair = createCell(
config({ rehomeAudience: undefined, rehomeDirectorServiceAccount: undefined })
)
const response = await withoutPair.poll()
expect(response.status).toBe(404)
expect(await response.json()).toEqual({ error: 'seat_feed_unavailable' })
const director = createCell(config({ role: 'director', cellId: 'director' }))
expect((await director.poll()).status).toBe(404)
})
})
cloud/apps/relay/src/cell/seats/cell-seat-log.test.ts
+1 -26
View File
@@ -1,5 +1,5 @@
import { describe, expect, it } from 'vitest'
import { CELL_SEAT_FEED_PAGE_MAX, CellSeatLog } from './cell-seat-log.js'
import { CellSeatLog } from './cell-seat-log.js'
function append(log: CellSeatLog, count: number): void {
for (let index = 0; index < count; index += 1) {
@@ -24,16 +24,6 @@ describe('cell seat log', () => {
expect(log.read(5)).toMatchObject({ seq: 10, full: expect.any(Array) })
})
it('pages a long backlog and resumes from the returned cursor', () => {
const log = new CellSeatLog(CELL_SEAT_FEED_PAGE_MAX * 2)
append(log, CELL_SEAT_FEED_PAGE_MAX + 5)
const first = log.read(0)
expect(first).toMatchObject({ seq: CELL_SEAT_FEED_PAGE_MAX, more: true })
expect('changes' in first && first.changes).toHaveLength(CELL_SEAT_FEED_PAGE_MAX)
const second = log.read(first.seq)
expect(second).toMatchObject({ seq: CELL_SEAT_FEED_PAGE_MAX + 5, more: false })
expect('changes' in second && second.changes).toHaveLength(5)
})
})
describe('cell seat log snapshot', () => {
@@ -85,20 +75,5 @@ describe('cell seat cursor', () => {
expect(log.recentlyLeftOf(host.userId, host.relayHostId)).toBeUndefined()
})
it('sends only the newest page of leavers in a full snapshot', () => {
const log = new CellSeatLog(CELL_SEAT_FEED_PAGE_MAX * 4, () => 0)
const leavers = CELL_SEAT_FEED_PAGE_MAX + 5
for (let index = 0; index < leavers; index += 1) {
const host = { userId: `user-${index}`, relayHostId: 'abcdefghijklmnop', epoch: 1, generation: 1 }
log.append({ kind: 'join', ...host, at: 0 })
log.append({ kind: 'leave', ...host, closeCode: 1006, at: 0 })
}
const page = log.read(null)
const left = 'full' in page ? page.recentlyLeft : []
expect(left).toHaveLength(CELL_SEAT_FEED_PAGE_MAX)
expect(left.at(-1)?.userId).toBe(`user-${leavers - 1}`)
// The cell itself still remembers the oldest.
expect(log.recentlyLeftOf('user-0', 'abcdefghijklmnop')).toBeDefined()
})
})
cloud/apps/relay/src/cell/seats/shadow-directory-compare.test.ts
-166
View File
@@ -1,166 +0,0 @@
import { describe, expect, it, vi } from 'vitest'
import type { RelayAssignment } from '../../assignment/assignment-store.js'
import type { RelayConfig } from '../../config.js'
import {
ShadowDirectoryCompare
} from './shadow-directory-compare.js'
import { ShadowSeatDirectory, type SeatFeedResponse } from './shadow-seat-directory.js'
const fakes = vi.hoisted(() => ({
verifyRelayToken: vi.fn(async (token: string) => ({ sub: 'user-a', relayHostId: token }))
}))
vi.mock('../../authentication/relay-token-verifier.js', () => ({
createRelayTokenVerifier: () => fakes.verifyRelayToken,
readBearer: (value: string | undefined) => value?.replace(/^Bearer /, '') ?? null
}))
import { createRelayApp } from '../../app.js'
const HOST = 'hhhhhhhhhhhhhhhh'
const IDENTITY = { userId: 'user-a', relayHostId: HOST }
function feed(cellId: string, overrides: Partial<SeatFeedResponse> = {}): SeatFeedResponse {
return { v: 1, cellId, incarnation: `inc-${cellId}`, seq: 1, at: 1, full: [], ...overrides }
}
function seat(epoch: number, state = 'active', relayHostId = HOST, joinedAt = 1) {
return { userId: 'user-a', relayHostId, epoch, generation: 1, state, joinedAt }
}
function directoryWith(
seats: Record<string, ReturnType<typeof seat>[]>,
noFeed: string[] = [],
unlive: string[] = []
) {
const directory = new ShadowSeatDirectory()
const cellIds = [...Object.keys(seats), ...noFeed]
directory.setCells(cellIds, cellIds, {
readAt: 0,
expiresAt: new Map(
cellIds.filter((cellId) => !unlive.includes(cellId)).map((cellId) => [cellId, 1e12])
)
})
for (const [cellId, full] of Object.entries(seats)) {
directory.apply(cellId, feed(cellId, { full }), 10)
}
for (const cellId of noFeed) directory.markNoFeed(cellId, 10)
return directory
}
describe('unexplained samples', () => {
it('log only the cell and epoch of the real database record, never raw ids', () => {
const lines: Record<string, unknown>[] = []
const compare = new ShadowDirectoryCompare(
directoryWith({ 'cell-a': [seat(3)] }),
() => 20,
(line) => lines.push(JSON.parse(line))
)
compare.compare('sticky-verify', IDENTITY, { ...assignment('cell-a'), assignmentEpoch: 2 })
compare.flush(20)
const sample = lines.find((line) => line.event === 'orca_relay_shadow_compare_unexplained')
expect(sample).toMatchObject({
class: 'epoch-mismatch',
db: { cellId: 'cell-a', assignmentEpoch: 2 }
})
expect(JSON.stringify(lines)).not.toContain(HOST)
expect(JSON.stringify(lines)).not.toContain('user-a')
})
})
describe('compare hooks on the director routes', () => {
async function respond(compareShadowSeats?: () => void) {
const resolved = assignment('cell-a')
const app = createRelayApp(config(), {
// SAFETY: /v1/resolve reads only resolveResume.
store: { resolveResume: vi.fn(async () => ({ userId: 'user-a' })) } as never,
// SAFETY: both routes read only resolve and assign.
assignments: {
resolve: vi.fn(async () => resolved),
assign: vi.fn(async () => resolved)
} as never,
drain: vi.fn(),
ready: vi.fn(async () => true),
compareShadowSeats
})
const sticky = await app.request('/v1/assign', {
method: 'POST',
headers: { authorization: `Bearer ${HOST}`, 'content-type': 'application/json' },
body: JSON.stringify({ v: 1, relayHostId: HOST, reconnect: true })
})
const resolve = await app.request('/v1/resolve', {
method: 'POST',
headers: { 'content-type': 'application/json' },
body: JSON.stringify({
v: 1,
relayHostId: HOST,
resumeToken: Buffer.alloc(32, 1).toString('base64url')
})
})
return [
{ status: sticky.status, body: await sticky.text() },
{ status: resolve.status, body: await resolve.text() }
]
}
it('compares after each database answer without changing a response byte', async () => {
vi.useFakeTimers({ now: 1_000_000 })
try {
const baseline = await respond()
const routes: unknown[] = []
const compared = await respond((...args: unknown[]) => {
routes.push(args[0])
})
const throwing = await respond(() => {
throw new Error('compare broke')
})
expect(baseline[0]?.status).toBe(200)
expect(baseline[1]?.status).toBe(200)
expect(compared).toEqual(baseline)
expect(throwing).toEqual(baseline)
expect(routes).toEqual(['sticky-verify', 'resolve'])
} finally {
vi.useRealTimers()
}
})
})
function assignment(cellId: string): RelayAssignment {
return {
userId: 'user-a',
relayHostId: HOST,
cellId,
cellUrl: `https://${cellId}.relay.example.test`,
assignmentEpoch: 3,
leaseExpiresAt: 1_300_000
}
}
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: true,
publicAssignmentConcurrency: 2,
publicAssignmentQueueMax: 128,
publicAssignmentWaitMs: 4_000,
publicResolveConcurrency: 1,
publicResolveWaitMs: 5_000,
publicAssignmentRetryAfterSeconds: 5,
dataDir: './data'
}
}
cloud/apps/relay/src/cell/seats/shadow-seat-directory.test.ts
-94
View File
@@ -3,7 +3,6 @@ import { RelayAssignmentStore } from '../../assignment/assignment-store.js'
import { openInMemoryRelayDatabase } from '../../database/database.js'
import {
SHADOW_SEAT_OUTRANK_SKEW_MS,
SHADOW_SEAT_RECENTLY_LEFT_MAX_HOSTS,
ShadowSeatDirectory,
startShadowSeatPoller,
type SeatFeedCell,
@@ -141,7 +140,6 @@ describe('startShadowSeatPoller', () => {
{ cellId: 'cell-a', epoch: 3, generation: 1, state: 'active' }
])
expect(cell.directory.cellState('cell-a')).toMatchObject({ status: 'live', seq: 5 })
expect(cell.identityToken).toHaveBeenCalledTimes(1)
})
it('resyncs from a full snapshot when the cell restarts with a new incarnation', async () => {
@@ -284,7 +282,6 @@ describe('startShadowSeatPoller', () => {
await cell.tick()
await cell.tick()
cell.stop()
expect(call).toBe(3)
expect(cell.requests[2]?.searchParams.get('since')).toBe(`${INCARNATION}:1`)
expect(cell.directory.seatsOf('user-a', 'host-2')).toHaveLength(1)
})
@@ -336,31 +333,6 @@ describe('startShadowSeatPoller', () => {
expect(cell.directory.isComplete()).toBe(true)
})
it('forgets the switches a cell reported once its image no longer has the feed', async () => {
let rolledBack = false
const cell = poller({
'cell-a': () => json(feed({ seq: 1, full: [] })),
'cell-b': () =>
rolledBack
? new Response('Not Found', { status: 404 })
: json(
feed({
cellId: 'cell-b',
seq: 1,
full: [],
flagsApplied: { generation: 3, flags: { readinessLocal: true } }
})
)
})
await cell.tick()
expect(cell.directory.cellState('cell-b')?.flagsApplied?.generation).toBe(3)
rolledBack = true
await cell.tick()
cell.stop()
expect(cell.directory.cellState('cell-b')).toMatchObject({ status: 'no-feed' })
expect(cell.directory.cellState('cell-b')?.flagsApplied).toBeUndefined()
})
it('keeps a reserve switch through a 404 until a feed says db or the process is replaced', () => {
const directory = new ShadowSeatDirectory()
directory.setCells(['cell-b'])
@@ -379,58 +351,6 @@ describe('startShadowSeatPoller', () => {
expect(directory.cellState('cell-b')?.flagsApplied).toBeUndefined()
})
it('completes without a cell that is dead, empty or isolated', async () => {
const cell = poller({ 'cell-a': () => json(feed({ seq: 1, full: [] })) }, 'all', async () => [
CELLS[0]!,
{ ...CELLS[1]!, requiredForComplete: false }
])
await cell.tick()
cell.stop()
expect(cell.directory.cellState('cell-b')).toMatchObject({ status: 'pending' })
expect(cell.directory.isComplete()).toBe(true)
})
it('backs off cell-list reads while the database fails', async () => {
const listCells = vi.fn(async (): Promise<SeatFeedCell[]> => {
throw new Error('pool timeout')
})
const cell = poller({}, 'all', listCells)
for (let second = 0; second < 40; second += 1) {
await cell.tick()
cell.advance(1_000)
}
cell.stop()
// Retries at +1, +2, +4, +8, +16 s, then at most every 30 s.
expect(listCells.mock.calls.length).toBeLessThanOrEqual(7)
expect(listCells.mock.calls.length).toBeGreaterThanOrEqual(5)
})
it('keeps polling the last cell list while a cell-list read hangs', async () => {
const readEntered = Promise.withResolvers<void>()
const readDone = Promise.withResolvers<SeatFeedCell[]>()
const listCells = vi
.fn<() => Promise<SeatFeedCell[]>>()
.mockResolvedValueOnce(CELLS)
.mockImplementation(() => {
readEntered.resolve()
return readDone.promise
})
const cell = poller({ 'cell-a': () => json(feed({ seq: 1, full: [] })) }, ['cell-a'], listCells)
await cell.tick()
cell.advance(31_000)
const round = cell.tick()
await readEntered.promise
await Promise.resolve()
expect(listCells).toHaveBeenCalledTimes(2)
expect(cell.requests.map((url) => url.hostname)).toEqual([
'cell-a.example.test',
'cell-a.example.test'
])
readDone.resolve(CELLS)
await round
cell.stop()
})
it('treats a malformed or mislabelled body as unverifiable', async () => {
const cell = poller({
'cell-a': () => json({ v: 2 }),
@@ -483,20 +403,6 @@ describe('startShadowSeatPoller', () => {
expect(cell.directory.seatsOf('user-a', 'host-1')).toEqual([])
})
it('bounds the recently-left memory', () => {
const directory = new ShadowSeatDirectory()
directory.setCells(['cell-a'])
const hosts = SHADOW_SEAT_RECENTLY_LEFT_MAX_HOSTS + 10
directory.apply(
'cell-a',
feed({ seq: 1, full: Array.from({ length: hosts }, (_, index) => seat(`host-${index}`)) }),
10
)
directory.apply('cell-a', feed({ incarnation: 'restarted', seq: 1, full: [] }), 20)
expect(directory.summary(20).recentlyLeftHosts).toBe(SHADOW_SEAT_RECENTLY_LEFT_MAX_HOSTS)
expect(directory.recentlyLeftOf('user-a', 'host-0', 20)).toEqual([])
expect(directory.recentlyLeftOf('user-a', `host-${hosts - 1}`, 20)).toHaveLength(1)
})
})
describe('RelayAssignmentStore.seatFeedCells', () => {
cloud/apps/relay/src/connections/cell-phone-path.blackbox.test.ts
+8 -253
View File
@@ -2,11 +2,9 @@ import { mkdtempSync, rmSync } from 'node:fs'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
import { afterEach, describe, expect, it, vi } from 'vitest'
import { reserveDeadManWindowMs } from '../cell/reservations/cell-reserve-dead-man.js'
import { openRelayDatabase } from '../database/database.js'
import {
nextFrame,
RESERVE_MODE_CELL_ID,
startReserveModeCell
} from '../test-fixtures/reserve-mode-cell.js'
import { StatementRecordingDatabase } from '../test-fixtures/statement-recording-database.js'
@@ -22,19 +20,6 @@ const OTHER_CELL = {
capacityRequests: 900
}
const RESUME_DEVICE_READ = expect.stringMatching(/^SELECT .* FROM relay_devices\b/)
const HOST_ASSIGNMENT_READ = expect.stringMatching(
/^SELECT assignment\.\*.* FROM relay_assignments\b/
)
async function until(check: () => Promise<boolean>, timeoutMs = 10_000): Promise<void> {
const deadline = Date.now() + timeoutMs
while (!(await check())) {
if (Date.now() > deadline) throw new Error('condition not reached')
await new Promise((resolve) => setTimeout(resolve, 50))
}
}
describe('the memory phone path', () => {
const cleanup: Array<() => Promise<void> | void> = []
@@ -58,11 +43,7 @@ describe('the memory phone path', () => {
cleanup,
...(now ? { now } : {})
})
// The admit-effective writer ticks on its own; every other statement is the phone's.
const phoneStatements = () =>
database.statements.filter((sql) => !sql.includes('relay_cell_admit_effective'))
// The test's own reads and writes go around the recorder.
return { ...reserve, database, raw: inner, phoneStatements }
return { ...reserve, database, raw: inner }
}
async function seatedPhoneHost(phonePath: 'db' | 'memory') {
@@ -70,27 +51,23 @@ describe('the memory phone path', () => {
reserve.setAdmitMode('reserve', { phonePath })
const seated = await reserve.bookedHost(3)
await reserve.installDevice(seated.identity.hostId, TOKEN)
reserve.database.statements.length = 0
return { reserve, seated }
}
// A phone resumed and spliced to its desktop, waiting to confirm.
async function splicedPhone(
phonePath: 'db' | 'memory',
device: { expiresAt?: number } = {},
clock: { now?: () => number; beforeDial?: () => void } = {}
device: { expiresAt?: number } = {}
) {
const reserve = await cell(clock.now)
const reserve = await cell()
reserve.setAdmitMode('reserve', { phonePath })
const seated = await reserve.bookedHost(3)
await reserve.installDevice(seated.identity.hostId, TOKEN, device)
clock.beforeDial?.()
const connOpen = nextFrame(seated.socket, 'conn-open')
const phone = await reserve.dialPhone(seated.identity.hostId, TOKEN)
const basisConnId = String((await connOpen).connId)
await reserve.attachData(await connOpen, seated.ack.generation)
expect(await phone.hello).toMatchObject({ ok: true })
reserve.database.statements.length = 0
// The desktop's confirm, sent in reply to the phone over the end-to-end channel.
const confirm = async (reqId: string, basis = basisConnId) => {
const answer = Promise.race([
@@ -108,54 +85,7 @@ describe('the memory phone path', () => {
const deviceRow = async (reserve: Awaited<ReturnType<typeof cell>>) =>
(await reserve.raw.query(`SELECT current_expires_at FROM relay_devices`))[0]
const warnings = (event: string) =>
vi
.mocked(console.warn)
.mock.calls.map(([line]) => String(line))
.filter((line) => line.includes(event))
.map((line) => JSON.parse(line))
it('admits a seated phone resume from one device read and writes no activity lease', async () => {
const { reserve, seated } = await seatedPhoneHost('memory')
const connOpen = nextFrame(seated.socket, 'conn-open')
const phone = await reserve.dialPhone(seated.identity.hostId, TOKEN)
expect(await connOpen).toMatchObject({ kind: 'resume', relayDeviceId: 'device-1' })
expect(reserve.phoneStatements()).toEqual([RESUME_DEVICE_READ])
await reserve.attachData(await connOpen, seated.ack.generation)
expect(await phone.hello).toMatchObject({
type: 'relay-hello',
ok: true,
credentialKind: 'resume'
})
expect(reserve.phoneStatements()).toEqual([RESUME_DEVICE_READ])
expect(await reserve.activityLeases()).toEqual([])
})
it('persists the connection basis and activity leases while phonePath is db', async () => {
const { reserve, seated } = await seatedPhoneHost('db')
const connOpen = nextFrame(seated.socket, 'conn-open')
const phone = await reserve.dialPhone(seated.identity.hostId, TOKEN)
await reserve.attachData(await connOpen, seated.ack.generation)
expect(await phone.hello).toMatchObject({ type: 'relay-hello', ok: true })
expect((await reserve.activityLeases()).map((lease) => lease.activity_kind).sort()).toEqual([
'confirmation',
'splice'
])
expect(
await reserve.raw.query(
`SELECT basis_conn_id, relay_device_id, credential_kind, active FROM relay_connection_bases`
)
).toEqual([
{
basis_conn_id: String((await connOpen).connId),
relay_device_id: 'device-1',
credential_kind: 'resume',
active: 1
}
])
})
it('closes a revoked or expired device 4401 and never tries it as an invite', async () => {
it('closes a revoked or expired device 4401', async () => {
const reserve = await cell()
reserve.setAdmitMode('reserve', { phonePath: 'memory' })
const seated = await reserve.bookedHost(3)
@@ -165,32 +95,12 @@ describe('the memory phone path', () => {
expiresAt: Date.now() - 1_000
})
for (const refused of [token(2), token(3)]) {
reserve.database.statements.length = 0
const phone = await reserve.dialPhone(seated.identity.hostId, refused)
expect(await phone.hello).toEqual({ type: 'relay-hello', ok: false, code: 4401 })
expect((await phone.closed).code).toBe(4401)
expect(reserve.phoneStatements()).toEqual([RESUME_DEVICE_READ])
}
})
it('closes 4503 when the device read fails, without trying the invite path', async () => {
const { reserve, seated } = await seatedPhoneHost('memory')
vi.spyOn(reserve.relay.store, 'lookupResume').mockRejectedValue(
new Error('connection terminated')
)
const invite = vi.spyOn(reserve.relay.store, 'resolveInviteForMove')
for (let dial = 0; dial < 2; dial += 1) {
const phone = await reserve.dialPhone(seated.identity.hostId, TOKEN)
expect((await phone.closed).code).toBe(4503)
}
expect(invite).not.toHaveBeenCalled()
expect(reserve.phoneStatements()).toEqual([])
// One line opens the window; the repeat waits for its closing line.
expect(warnings('orca_relay_client_accept_db_error')).toEqual([
{ event: 'orca_relay_client_accept_db_error', cause: 'error', cellId: reserve.config.cellId }
])
})
it('falls through to the invite path only when no device row matches', async () => {
const { reserve, seated } = await seatedPhoneHost('memory')
const created = nextFrame(seated.socket, 'invite-created')
@@ -208,7 +118,7 @@ describe('the memory phone path', () => {
])
})
it('answers a phone whose host has no seat here from one assignment read', async () => {
it('distinguishes moved and offline hosts when a phone finds no seat here', async () => {
const reserve = await cell()
reserve.setAdmitMode('reserve', { phonePath: 'memory' })
await reserve.relay.assignments.reconcileCells([...reserve.config.cells, OTHER_CELL])
@@ -222,31 +132,19 @@ describe('the memory phone path', () => {
)
}
await reserve.installDevice(host.hostId, TOKEN)
reserve.database.statements.length = 0
const phone = await reserve.dialPhone(host.hostId, TOKEN)
const closed = await phone.closed
expect(reserve.phoneStatements()).toEqual([RESUME_DEVICE_READ, HOST_ASSIGNMENT_READ])
return closed.code
}
expect(await dial('elsewhere')).toBe(4409)
expect(await dial('none')).toBe(4404)
expect(warnings('orca_relay_client_miss_row_names_here')).toEqual([])
expect(await dial('here')).toBe(4404)
expect(warnings('orca_relay_client_miss_row_names_here')).toEqual([
expect.objectContaining({
cellId: reserve.config.cellId,
relayHostIdDigest: expect.any(String)
})
])
vi.spyOn(reserve.relay.assignments, 'resolve').mockRejectedValueOnce(
new Error('connection terminated')
)
const host = await reserve.host()
await reserve.installDevice(host.hostId, TOKEN)
expect((await (await reserve.dialPhone(host.hostId, TOKEN)).closed).code).toBe(4409)
expect(warnings('orca_relay_client_miss_read_error')).toEqual([
{ event: 'orca_relay_client_miss_read_error', cause: 'error', cellId: reserve.config.cellId }
])
})
it('takes the database path for a phone whose switch turned off during its device read', async () => {
@@ -265,7 +163,7 @@ describe('the memory phone path', () => {
])
})
it('confirms from the memory basis with one fresh read, renewing a device due for it', async () => {
it('confirms from the memory basis and renews a device due for it', async () => {
const { reserve, confirm } = await splicedPhone('memory')
const before = Number((await deviceRow(reserve))?.current_expires_at)
const answer = await confirm('confirm-1')
@@ -278,12 +176,6 @@ describe('the memory phone path', () => {
const after = Number((await deviceRow(reserve))?.current_expires_at)
expect(after).toBeGreaterThan(before)
expect(answer.resumeExpiresAt).toBe(after)
expect(reserve.phoneStatements()).toEqual([
RESUME_DEVICE_READ,
expect.stringMatching(
/^UPDATE relay_devices SET current_expires_at = \?.*RETURNING current_expires_at$/
)
])
expect(await reserve.activityLeases()).toEqual([])
// Its audit row leaves with the next batch.
await reserve.relay.store.resumeAudit.flush()
@@ -293,10 +185,8 @@ describe('the memory phone path', () => {
expect(audits.map((row) => JSON.parse(String(row.detail_json)))).toEqual([
{ reqId: 'confirm-1', basisConnId: expect.any(String), renewed: true }
])
// A retry of the same request answers from memory; another basis cannot reuse its reqId.
reserve.database.statements.length = 0
// A retry keeps its answer; another basis cannot reuse its reqId.
expect(await confirm('confirm-1')).toEqual(answer)
expect(reserve.phoneStatements()).toEqual([])
expect(await confirm('confirm-1', 'another-basis')).toMatchObject({
type: 'control-error',
code: 'confirmation_tuple_mismatch'
@@ -310,7 +200,6 @@ describe('the memory phone path', () => {
renewed: true,
resumeExpiresAt: fresh
})
expect(notDue.reserve.phoneStatements()).toEqual([RESUME_DEVICE_READ])
const due = Date.now() + 20 * 24 * 60 * 60 * 1000
const failing = await splicedPhone('memory', { expiresAt: due })
@@ -320,74 +209,20 @@ describe('the memory phone path', () => {
resumeExpiresAt: due
})
expect(Number((await deviceRow(failing.reserve))?.current_expires_at)).toBe(due)
// Never silent: a renewal left failing ages phones toward re-pairing.
expect(warnings('orca_relay_resume_renewal_failed')).toEqual([
{
event: 'orca_relay_resume_renewal_failed',
cause: 'error',
cellId: failing.reserve.config.cellId
}
])
})
it('replays a stored answer after its splice closed, and drops the basis itself', async () => {
const { reserve, phone, confirm } = await splicedPhone('memory')
const { phone, confirm } = await splicedPhone('memory')
const answer = await confirm('confirm-1')
phone.socket.close()
await phone.closed
await new Promise((resolve) => setTimeout(resolve, 50))
reserve.database.statements.length = 0
expect(await confirm('confirm-1')).toEqual(answer)
expect(await confirm('confirm-1', 'another-basis')).toMatchObject({
code: 'confirmation_tuple_mismatch'
})
expect(reserve.phoneStatements()).toEqual([])
// A new request on the closed basis finds none in memory, so it goes to the database.
expect(await confirm('confirm-2')).toMatchObject({ code: 'confirmation_not_active' })
expect(reserve.phoneStatements()[0]).toBe('BEGIN')
})
it('answers a grace-version confirm renewed:false with its expiry unchanged, as today', async () => {
const { reserve, confirm } = await splicedPhone('memory')
const stored = Number((await deviceRow(reserve))?.current_expires_at)
// An install rotated the credential after the accept read, so the accepted version is grace.
await reserve.raw.query(
`UPDATE relay_devices SET grace_hash = current_hash, grace_version = 1, grace_expires_at = ?,
current_hash = 'rotated', current_version = 2`,
[Date.now() + 60_000]
)
expect(await confirm('confirm-1')).toMatchObject({
type: 'device-resume-confirmed',
currentVersion: 2,
acceptedAs: 'current',
renewed: false,
resumeExpiresAt: stored
})
expect(reserve.phoneStatements()).toEqual([RESUME_DEVICE_READ])
})
it('answers every device change since the accept read as the database confirm does', async () => {
const changes: Record<string, [string, unknown[]]> = {
revoked: ['UPDATE relay_devices SET revoked_at = ?', [Date.now() - 1]],
expired: ['UPDATE relay_devices SET current_expires_at = ?', [Date.now() - 1]],
retired: [
"UPDATE relay_devices SET current_hash = 'rotated-twice', current_version = 3, grace_version = 2",
[]
],
unchanged: ['UPDATE relay_devices SET updated_at = updated_at', []]
}
const outcome = async (phonePath: 'db' | 'memory', [sql, params]: [string, unknown[]]) => {
const { reserve, confirm } = await splicedPhone(phonePath)
await reserve.raw.query(sql, params)
const answer = await confirm('confirm-1')
return answer.type === 'control-error'
? answer.code
: `${answer.type}:${answer.acceptedAs}:${answer.renewed}`
}
for (const [name, change] of Object.entries(changes)) {
expect(await outcome('memory', change), name).toBe(await outcome('db', change))
}
expect(await outcome('memory', changes.revoked!)).toBe('reject-revoked')
})
it('durably confirms and renews a database-path resume, releasing only its confirmation lease', async () => {
@@ -411,82 +246,6 @@ describe('the memory phone path', () => {
expect(await confirm('confirm-1')).toEqual(answer)
})
describe('rolling a memory cell back', () => {
// Flip back by the switch, or by the dead-man; either way the confirm finishes from memory.
type Cell = Awaited<ReturnType<typeof cell>>
type Advance = (ms: number) => void
type Rollback = {
beforeDial?: (advance: Advance) => void
roll: (reserve: Cell, advance: Advance) => void
}
// The dead-man trips just after the phone's accept: its window ran from the host's booking.
const window = reserveDeadManWindowMs(RESERVE_MODE_CELL_ID)
const rollbacks: Record<string, Rollback> = {
'admitMode=db': { roll: (reserve) => reserve.setAdmitMode('db', { phonePath: 'memory' }) },
'phonePath=db': { roll: (reserve) => reserve.setAdmitMode('reserve', { phonePath: 'db' }) },
'a dead-man trip': {
beforeDial: (advance) => advance(window - 5_000),
roll: (_reserve, advance) => advance(6_000)
}
}
for (const [name, rollback] of Object.entries(rollbacks)) {
it(`finishes an in-flight confirm on its memory basis after ${name}, writing and leaking no lease`, async () => {
let now = Date.now()
const advance: Advance = (ms) => {
now += ms
}
const spliced = await splicedPhone(
'memory',
{},
{
now: () => now,
beforeDial: () => rollback.beforeDial?.(advance)
}
)
const { reserve, seated } = spliced
// The flag tool records reserve in Postgres before the cell admits from memory.
await reserve.relay.assignments.reconcileCells([...reserve.config.cells, OTHER_CELL])
await reserve.relay.assignments.setCellAdmitMode(reserve.config.cellId, 'reserve')
rollback.roll(reserve, advance)
expect(reserve.relay.sessions.phonePathEffective()).toBe('db')
reserve.database.statements.length = 0
expect(await spliced.confirm('confirm-1')).toMatchObject({
type: 'device-resume-confirmed',
renewed: true
})
expect(
reserve.phoneStatements().filter((sql) => !sql.includes('relay_assignment'))
).toEqual([
RESUME_DEVICE_READ,
expect.stringMatching(/^UPDATE relay_devices SET current_expires_at/)
])
// Only the control lease a flip back re-registers; nothing for the phone.
const leases = async () =>
(await reserve.activityLeases()).map((lease) => lease.activity_kind)
expect((await leases()).filter((kind) => kind !== 'control')).toEqual([])
// Postgres still records reserve, so evacuation leaves the cell alone.
await expect(
reserve.relay.assignments.startEvacuation(
{ userId: 'user-1', relayHostId: seated.identity.hostId },
OTHER_CELL.id
)
).rejects.toThrow('cell_in_reserve_mode')
// The splice still runs, and closing it releases nothing it never took.
expect(spliced.phone.socket.readyState).toBe(spliced.phone.socket.OPEN)
// Let a flip back finish leasing its control (and fold its relay_cells write) first.
await until(async () => reserve.relay.sessions.reregistrationPending() === 0)
await new Promise((resolve) => setTimeout(resolve, 250))
reserve.database.statements.length = 0
spliced.phone.socket.close()
await spliced.phone.closed
await new Promise((resolve) => setTimeout(resolve, 50))
expect(reserve.phoneStatements()).toEqual([])
expect((await leases()).filter((kind) => kind !== 'control')).toEqual([])
})
}
})
// E-dep-6 in miniature: the staging drill blocks one memory cell's database for 10 minutes.
it('rides out a database outage: splices keep forwarding, new work fails cleanly and recovers', async () => {
let now = Date.now()
@@ -537,13 +296,9 @@ describe('the memory phone path', () => {
// Flipped back during the block: the in-flight confirm still finishes on its memory basis.
reserve.setAdmitMode('db', { phonePath: 'memory' })
reserve.database.failing = null
reserve.database.statements.length = 0
expect(await confirm('confirm-c', inFlight.basisConnId)).toMatchObject({
type: 'device-resume-confirmed'
})
expect(reserve.phoneStatements().find((sql) => !sql.includes('relay_assignment'))).toEqual(
RESUME_DEVICE_READ
)
// And new resumes take the database path again.
const connOpen = nextFrame(seated.socket, 'conn-open')
await reserve.dialPhone(seated.identity.hostId, token(1))
cloud/apps/relay/src/connections/host-close-reason-memory.test.ts
-37
View File
@@ -1,37 +0,0 @@
import { ASSIGNMENT_LIMITS } from '@orca-cloud/relay-contract/assignment-invariants'
import { RELAY_HOST_CLOSE_REASON } from '@orca-cloud/relay-contract/host-close-reason'
import { describe, expect, it } from 'vitest'
import { HostCloseReasonMemory } from './host-close-reason-memory.js'
function memoryAt(clock: { now: number }): HostCloseReasonMemory {
return new HostCloseReasonMemory(() => clock.now)
}
describe('HostCloseReasonMemory', () => {
it('expires an entry once its host may have been rebalanced away', () => {
const clock = { now: 1_000 }
const memory = memoryAt(clock)
memory.record('a', RELAY_HOST_CLOSE_REASON.SIGNED_OUT)
clock.now += ASSIGNMENT_LIMITS.dormantTtlMs - 1
expect(memory.read('a')).toBe(RELAY_HOST_CLOSE_REASON.SIGNED_OUT)
clock.now += 1
expect(memory.read('a')).toBeNull()
expect(memory.size()).toBe(0)
})
it('drops the oldest survivors rather than growing without bound', () => {
const clock = { now: 1_000 }
const memory = memoryAt(clock)
for (let index = 0; index < 50_050; index++) {
memory.record(`host-${index}`, RELAY_HOST_CLOSE_REASON.SIGNED_OUT)
}
expect(memory.size()).toBe(50_000)
expect(memory.read('host-0')).toBeNull()
expect(memory.read('host-50049')).toBe(RELAY_HOST_CLOSE_REASON.SIGNED_OUT)
})
})
cloud/apps/relay/src/connections/host-control-proof-cleanup.test.ts
+1 -33
View File
@@ -1,4 +1,3 @@
import nacl from 'tweetnacl'
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import {
answerProof,
@@ -16,24 +15,7 @@ afterEach(() => {
describe('host control proof cleanup', () => {
it('releases the challenge timer and listeners when its peer closes before proof', async () => {
const h = fixture()
const socket = await openProof(h)
expect(vi.getTimerCount()).toBe(1)
expect(socket.listenerCount('message')).toBe(1)
socket.peerClose()
expect({
timers: vi.getTimerCount(),
message: socket.listenerCount('message'),
close: socket.listenerCount('close')
}).toEqual({ timers: 0, message: 0, close: 0 })
await vi.advanceTimersByTimeAsync(10_000)
expect(socket.close).not.toHaveBeenCalled()
expect(h.activate).not.toHaveBeenCalled()
expect(h.database.query).not.toHaveBeenCalled()
})
it('does no challenge crypto, send, timer or registration after a closed peer finishes verification', async () => {
it('does not challenge or register a closed peer after verification finishes', async () => {
const h = fixture()
let finish!: (valid: boolean) => void
h.verify.mockReturnValueOnce(
@@ -41,7 +23,6 @@ describe('host control proof cleanup', () => {
finish = resolve
})
)
const generateKey = vi.spyOn(nacl.box, 'keyPair')
const socket = new ProofSocket()
h.registry.acceptControl(socket.registrySocket(), h.identity)
socket.emit('message', Buffer.from(h.hello), false)
@@ -50,11 +31,6 @@ describe('host control proof cleanup', () => {
finish(true)
await vi.advanceTimersByTimeAsync(0)
expect(socket.send).not.toHaveBeenCalled()
expect(generateKey).not.toHaveBeenCalled()
expect(vi.getTimerCount()).toBe(0)
expect(socket.listenerCount('message')).toBe(0)
expect(h.activate).not.toHaveBeenCalled()
expect(h.database.query).not.toHaveBeenCalled()
expect(
h.registry.get({ userId: h.identity.sub, relayHostId: h.identity.relayHostId })
).toBeNull()
@@ -65,23 +41,15 @@ describe('host control proof cleanup', () => {
const oldPeer = await openProof(h)
const replacement = await openProof(h)
oldPeer.peerClose()
expect(vi.getTimerCount()).toBe(1)
answerProof(replacement, h.keyPair)
await vi.advanceTimersByTimeAsync(0)
expect(h.activate).toHaveBeenCalledOnce()
expect(h.recordAuth).toHaveBeenCalledExactlyOnceWith(true)
expect(
h.registry.get({ userId: h.identity.sub, relayHostId: h.identity.relayHostId })?.socket
).toBe(replacement)
expect(replacement.send).toHaveBeenCalledTimes(2)
expect(replacement.listenerCount('message')).toBe(1)
expect(replacement.listenerCount('close')).toBe(2)
expect(vi.getTimerCount()).toBe(1)
await vi.advanceTimersByTimeAsync(10_000)
expect(oldPeer.close).not.toHaveBeenCalled()
expect(replacement.close).not.toHaveBeenCalled()
h.registry.drain(0)
await vi.advanceTimersByTimeAsync(0)
expect(vi.getTimerCount()).toBe(0)
})
})
cloud/apps/relay/src/connections/host-control-resume-identity.test.ts
-2
View File
@@ -1,4 +1,3 @@
import { CONTROL_CONTINUITY_LIMITS } from '@orca-cloud/relay-contract/control-continuity'
import { RELAY_CLOSE_CODE } from '@orca-cloud/relay-contract/close-codes'
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import type { RelayTokenClaims } from '../authentication/relay-token-verifier.js'
@@ -46,7 +45,6 @@ async function lapsedHost() {
)
const session = h.registry.get({ userId: h.identity.sub, relayHostId: h.identity.relayHostId })
expect(session?.state).toBe('orphaned')
expect(CONTROL_CONTINUITY_LIMITS.orphanGraceMs).toBeGreaterThan(5_000)
await vi.advanceTimersByTimeAsync(5_000)
return { h, ack, session: session!, resumeHello: resumeHelloFor(h.hello, ack) }
}
cloud/apps/relay/src/connections/host-session-client-accept.test.ts
-57
View File
@@ -229,7 +229,6 @@ describe('client accept abandoned mid-DB-phase', () => {
capacity
)
await vi.advanceTimersByTimeAsync(0)
expect(h.acquireActivity).toHaveBeenCalledOnce()
// The phone's 12s bound fires while the cell still waits on Postgres.
client.close(1000, 'client bound')
capacity.release()
@@ -247,54 +246,12 @@ describe('client accept abandoned mid-DB-phase', () => {
{ userId: identity.sub, relayHostId: identity.relayHostId },
expect.stringMatching(/^confirmation:/)
)
expect(h.observer.recordClientAcceptAbandoned).toHaveBeenCalledWith(
'activity',
expect.any(Number)
)
const line = warn.mock.calls
.map((call) => String(call[0]))
.find((entry) => entry.includes('orca_relay_client_accept_abandoned'))
expect(line).toBeDefined()
expect(JSON.parse(line!)).toMatchObject({ stage: 'activity' })
expect(line).not.toContain(identity.relayHostId)
} finally {
warn.mockRestore()
h.registry.drain(0)
vi.advanceTimersByTime(0)
}
})
it('stops after a slow credential reservation without acquiring an activity lease', async () => {
const h = harness()
await activeHost(h)
const slowReserve = deferred<CredentialReservation>()
h.store.reserveCredential.mockReturnValueOnce(slowReserve.promise)
const client = new FakeSocket()
const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined)
try {
const accepting = h.registry.acceptClient(
client as unknown as WebSocket,
identity.relayHostId,
'credential'
)
await vi.advanceTimersByTimeAsync(0)
client.close(1000, 'client bound')
slowReserve.resolve(reservation)
await accepting
expect(h.acquireActivity).not.toHaveBeenCalled()
expect(h.store.failReservation).toHaveBeenCalledWith(reservation)
expect(h.observer.recordClientAcceptAbandoned).toHaveBeenCalledWith(
'credential',
expect.any(Number)
)
} finally {
warn.mockRestore()
h.registry.drain(0)
vi.advanceTimersByTime(0)
}
})
})
// The attach resolves its owning session once and hands it to the unfenced leg;
@@ -360,7 +317,6 @@ describe('host data attach ownership', () => {
1
)
await vi.advanceTimersByTimeAsync(0)
expect(h.store.recordConnectionBasis).toHaveBeenCalledOnce()
expect(client.send).not.toHaveBeenCalledWith(expect.stringContaining('relay-hello'))
basis.resolve()
expect(await attaching).toBe(true)
@@ -420,17 +376,4 @@ describe('paced drain and the phones of a host not yet told', () => {
expect(served.close).not.toHaveBeenCalled()
expect(untold.send).toHaveBeenCalledWith(expect.stringContaining('conn-open'))
})
it("refuses that host's phones as soon as its own drain is sent", async () => {
const { h, untold } = await twoHostCell()
h.registry.drain(0, { paceWindowMs: 40_000 })
await vi.advanceTimersByTimeAsync(40_000)
// Told by its zero-grace close, a timer nested 1 ms after the send.
await vi.advanceTimersByTimeAsync(1)
expect(untold.send).toHaveBeenCalledWith(expect.stringContaining('"type":"drain"'))
const refused = await dial(h, laterHostId)
expect(refused.close).toHaveBeenCalledWith(RELAY_CLOSE_CODE.DRAINING, expect.any(String))
})
})
cloud/apps/relay/src/connections/host-session-drain-release-order.test.ts
-45
View File
@@ -1,6 +1,5 @@
import { RELAY_CLOSE_CODE } from '@orca-cloud/relay-contract/close-codes'
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import { DRAIN_RELEASE_BEFORE_CLOSE_MS } from './host-session-registry.js'
import {
answerProof,
hostControlProofFixture,
@@ -70,48 +69,4 @@ describe('a zero-grace drain releases the host before telling it', () => {
socket.close.mock.invocationCallOrder[0]!
)
})
it('closes on the bound when the release has not landed, as before this ordering', async () => {
const { h, socket } = await connectedHost()
vi.spyOn(h.assignments, 'releaseActivity').mockReturnValue(pendingRelease().promise)
h.registry.drain(0)
await vi.advanceTimersByTimeAsync(DRAIN_RELEASE_BEFORE_CLOSE_MS - 1)
expect(socket.close).not.toHaveBeenCalled()
await vi.advanceTimersByTimeAsync(2)
expect(frames(socket)).toEqual([{ type: 'drain', graceMs: 0, recovery: 'resolve-director' }])
expect(socket.close).toHaveBeenCalledWith(
RELAY_CLOSE_CODE.DRAINING,
'resolve configured director'
)
const slow = vi
.mocked(console.warn)
.mock.calls.map(([line]) => String(line))
.filter((line) => line.includes('"orca_relay_host_drain_release_slow"'))
expect(slow).toHaveLength(1)
expect(slow[0]).not.toContain(h.identity.relayHostId)
})
it('still tells a host with grace at once, and releases it before the closing deadline', async () => {
const { h, socket } = await connectedHost()
const release = pendingRelease()
vi.spyOn(h.assignments, 'releaseActivity').mockReturnValue(release.promise)
h.registry.drain(10_000)
expect(frames(socket)).toEqual([
{ type: 'drain', graceMs: 10_000, recovery: 'resolve-director' }
])
await vi.advanceTimersByTimeAsync(10_000)
expect(socket.close).not.toHaveBeenCalled()
release.settle('released')
await vi.advanceTimersByTimeAsync(0)
expect(socket.close).toHaveBeenCalledWith(
RELAY_CLOSE_CODE.DRAINING,
'resolve configured director'
)
// Told once, by the drain itself; the close adds no second frame.
expect(frames(socket)).toHaveLength(1)
})
})
cloud/apps/relay/src/connections/host-session-registry.test.ts
-534
View File
@@ -1,534 +0,0 @@
import { EventEmitter } from 'node:events'
import { CONTROL_CONTINUITY_LIMITS } from '@orca-cloud/relay-contract/control-continuity'
import { RELAY_CLOSE_CODE } from '@orca-cloud/relay-contract/close-codes'
import { RELAY_HOST_CAPABILITY_IDLE_REGIONAL_REHOME } from '@orca-cloud/relay-contract/control-messages'
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import type WebSocket from 'ws'
import type { RelayAssignmentStore } from '../assignment/assignment-store.js'
import type { RelayConfig } from '../config.js'
import type { RelayCredentialStore } from '../credentials/credential-store.js'
import { CONTROL_RENEWAL_BATCH_INTERVAL_MS } from '../control-leases/control-renewal-batch.js'
import type {
ControlRenewalOutcome
} from '../control-leases/control-renewal-statement.js'
import {
HostSessionRegistry,
IDLE_REHOME_MIN_CONTROL_AGE_MS,
type HostSession
} from './host-session-registry.js'
import type { RelayRuntimeObserver } from '../observability/relay-observability.js'
import type { RelayTokenClaims } from '../authentication/relay-token-verifier.js'
import { ProcessQueuedByteBudget } from './splice-forwarder.js'
// A due renewal leaves the heartbeat as a batch enqueue, so the store only sees
// the tick once the batch window closes.
async function closeRenewalWindow(): Promise<void> {
await vi.advanceTimersByTimeAsync(CONTROL_RENEWAL_BATCH_INTERVAL_MS)
}
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<void>
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<T>(): {
promise: Promise<T>
resolve(value: T): void
} {
let resolve!: (value: T) => void
const promise = new Promise<T>((complete) => {
resolve = complete
})
return { promise, resolve }
}
function createRegistry(
activateControl: RelayAssignmentStore['activateControl'],
store: Partial<RelayCredentialStore> = {},
verifyRelayToken: (token: string) => Promise<RelayTokenClaims | null> = vi.fn()
): {
registry: HostSessionRegistry
activate: ActivateSession
acquireActivity: ReturnType<typeof vi.fn>
renewControlActivities: ReturnType<typeof vi.fn<RelayAssignmentStore['renewControlActivities']>>
releaseActivity: ReturnType<typeof vi.fn>
observer: {
recordAuth: ReturnType<typeof vi.fn>
recordControlClose: ReturnType<typeof vi.fn>
recordPostProofClose: ReturnType<typeof vi.fn>
recordSpliceClose: ReturnType<typeof vi.fn>
}
} {
const acquireActivity = vi.fn().mockResolvedValue(undefined)
const releaseActivity = vi.fn().mockResolvedValue(true)
const renewControlActivities = vi
.fn<RelayAssignmentStore['renewControlActivities']>()
.mockImplementation(async (rows) => rows.map(() => 'renewed'))
const assignments = {
activateControl,
markMigrationTargetRegistered: vi.fn().mockResolvedValue(undefined),
resolve: vi.fn().mockResolvedValue({ cellId: config.cellId }),
acquireActivity,
renewControlActivities,
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(),
recordPostProofClose: vi.fn(),
recordSpliceClose: vi.fn()
} satisfies RelayRuntimeObserver
const registry = new HostSessionRegistry(
config,
verifyRelayToken,
store as RelayCredentialStore,
assignments,
new ProcessQueuedByteBudget(),
observer,
Date.now,
Math.random,
'incarnation-1'
)
// 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<void>
}
).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,
renewControlActivities,
releaseActivity,
observer
}
}
describe('host session cleanup races', () => {
beforeEach(() => vi.useFakeTimers())
afterEach(() => {
vi.clearAllTimers()
vi.useRealTimers()
})
// c34 ghost seat came from activation outliving orphan grace and reviving a removed session outside the map.
it('refuses a rebind whose session left the map while its activation was in flight', async () => {
const blocked = deferred<string>()
const activateControl = vi
.fn<RelayAssignmentStore['activateControl']>()
.mockResolvedValueOnce('control:production-gce-c3:1')
.mockReturnValueOnce(blocked.promise)
const { registry, activate, releaseActivity } = createRegistry(activateControl)
const key = { userId: identity.sub, relayHostId: identity.relayHostId }
const originalSocket = new FakeSocket()
// oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: FakeSocket implements the registry's WebSocket event and lifecycle surface.
await activate(originalSocket as unknown as WebSocket, identity, null, 1, false, 1)
const original = registry.get(key)
originalSocket.close(1006, 'network reset')
expect(original?.state).toBe('orphaned')
const rebindSocket = new FakeSocket()
// oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: FakeSocket implements the registry's WebSocket event and lifecycle surface.
const rebinding = activate(rebindSocket as unknown as WebSocket, identity, original, 1, true, 1)
await vi.advanceTimersByTimeAsync(0)
expect(activateControl).toHaveBeenCalledTimes(2)
// The grace runs out while the database answers slowly.
await vi.advanceTimersByTimeAsync(CONTROL_CONTINUITY_LIMITS.orphanGraceMs)
expect(registry.get(key)).toBeNull()
blocked.resolve('control:production-gce-c3:1')
await rebinding
expect(rebindSocket.close).toHaveBeenCalledWith(
RELAY_CLOSE_CODE.PEER_DROPPED,
'control activation superseded'
)
expect(registry.get(key)).toBeNull()
expect(original?.state).toBe('closed')
const page = registry.seatFeed(null)
expect('full' in page && page.full).toEqual([])
expect(releaseActivity).toHaveBeenCalledWith(key, 'control:production-gce-c3:1')
})
it('skips a closed queued control so its live retry avoids abandoned database work', async () => {
const stalled = deferred<string>()
const activateControl = vi
.fn<RelayAssignmentStore['activateControl']>()
.mockReturnValueOnce(stalled.promise)
.mockImplementation(async () => {
await new Promise<void>((resolve) => setTimeout(resolve, 4_000))
return 'control:production-gce-c3:1'
})
const { registry, activate, releaseActivity } = createRegistry(activateControl)
const firstSocket = new FakeSocket()
// oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: FakeSocket implements the registry's WebSocket event and lifecycle surface.
const first = activate(firstSocket as unknown as WebSocket, identity, null, 1, false, 1)
await vi.advanceTimersByTimeAsync(0)
expect(activateControl).toHaveBeenCalledOnce()
const abandonedSocket = new FakeSocket()
// oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: FakeSocket implements the registry's WebSocket event and lifecycle surface.
const abandoned = activate(abandonedSocket as unknown as WebSocket, identity, null, 1, false, 1)
const liveSocket = new FakeSocket()
const startedAt = Date.now()
let liveCompletedAt: number | undefined
// oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: FakeSocket implements the registry's WebSocket event and lifecycle surface.
const live = activate(liveSocket as unknown as WebSocket, identity, null, 1, false, 1).then(
() => {
liveCompletedAt = Date.now()
}
)
firstSocket.close()
abandonedSocket.close()
stalled.resolve('control:production-gce-c3:1')
await vi.advanceTimersByTimeAsync(8_000)
await Promise.all([first, abandoned, live])
console.log(
JSON.stringify({
scenario: 'closed queued control before a live retry',
activationCalls: activateControl.mock.calls.length,
activityReleases: releaseActivity.mock.calls.length,
liveReadyMs: liveCompletedAt === undefined ? null : liveCompletedAt - startedAt
})
)
expect(activateControl).toHaveBeenCalledTimes(2)
expect(releaseActivity).toHaveBeenCalledOnce()
expect(liveCompletedAt! - startedAt).toBe(4_000)
expect(registry.get({ userId: identity.sub, relayHostId: identity.relayHostId })?.socket).toBe(
liveSocket
)
})
it('ignores a denial belonging to the socket before a same-generation rebind', async () => {
const h = createRegistry(vi.fn().mockResolvedValue('control:production-gce-c3:1'))
const oldSocket = new FakeSocket()
await h.activate(oldSocket as unknown as WebSocket, identity, null, 1, false, 1)
const session = h.registry.get({ userId: identity.sub, relayHostId: identity.relayHostId })!
let finish!: (outcomes: ControlRenewalOutcome[]) => void
h.renewControlActivities.mockReturnValueOnce(
new Promise<ControlRenewalOutcome[]>((resolve) => {
finish = resolve
})
)
await vi.advanceTimersByTimeAsync(15_000)
await closeRenewalWindow()
const replacement = new FakeSocket()
await h.activate(replacement as unknown as WebSocket, identity, session, 1, true, 1)
finish(['activity_cell_not_authoritative'])
await vi.advanceTimersByTimeAsync(0)
expect(replacement.close).not.toHaveBeenCalled()
expect(session.socket).toBe(replacement)
expect(session.generation).toBe(1)
})
})
describe('source-owned idle cutover', () => {
beforeEach(() => vi.useFakeTimers())
afterEach(() => {
vi.clearAllTimers()
vi.useRealTimers()
})
const request = {
attemptId: 'idle-1',
userId: identity.sub,
relayHostId: identity.relayHostId,
sourceAssignmentEpoch: 1,
sourceGeneration: 1,
sourceCellIncarnation: 'incarnation-1',
targetCellId: 'target'
}
async function source(
store: Partial<RelayCredentialStore> = {},
controlAgeMs = IDLE_REHOME_MIN_CONTROL_AGE_MS
) {
const h = createRegistry(vi.fn().mockResolvedValue('control:1'), store)
const socket = new FakeSocket()
h.registry.acceptControl(
socket as unknown as WebSocket,
identity,
undefined,
new Set([RELAY_HOST_CAPABILITY_IDLE_REGIONAL_REHOME])
)
socket.removeAllListeners('message')
await h.activate(socket as unknown as WebSocket, identity, null, 1, false, 1)
// Moves the clock without firing the heartbeat timers that would close the session.
vi.setSystemTime(Date.now() + controlAgeMs)
return { ...h, socket, session: h.registry.get(request)! }
}
it('rejects new accepts and replacements synchronously while a commit awaits', async () => {
const h = await source()
const result = deferred<{ outcome: 'deferred' }>()
const commit = vi.fn().mockReturnValue(result.promise)
const moving = h.registry.idleRehome(
request,
commit,
vi.fn().mockResolvedValue('not-committed')
)
const duplicate = h.registry.idleRehome(
request,
commit,
vi.fn().mockResolvedValue('not-committed')
)
const client = new FakeSocket()
const release = vi.fn()
await h.registry.acceptClient(
client as unknown as WebSocket,
identity.relayHostId,
'credential',
{ release } as never
)
expect(client.close).toHaveBeenCalledWith(RELAY_CLOSE_CODE.WRONG_CELL, expect.any(String))
expect(release).toHaveBeenCalledOnce()
const replacement = new FakeSocket()
await h.activate(replacement as unknown as WebSocket, identity, h.session, 2, false, 1)
expect(replacement.close).toHaveBeenCalledWith(RELAY_CLOSE_CODE.WRONG_CELL, expect.any(String))
result.resolve({ outcome: 'deferred' })
await moving
await duplicate
expect(commit).toHaveBeenCalledOnce()
expect(h.socket.close).not.toHaveBeenCalled()
expect(
await h.registry.idleRehome(
{ ...request, attemptId: 'next' },
vi.fn().mockResolvedValue({ outcome: 'committed' }),
vi.fn().mockResolvedValue('not-committed')
)
).toEqual({ outcome: 'committed' })
})
it.each(['ambiguous', 'deferred'])(
'keeps %s outcomes fenced until locked reconciliation succeeds',
async (claim) => {
const h = await source()
const reconcile = vi
.fn()
.mockRejectedValueOnce(new Error('database unavailable'))
.mockRejectedValueOnce(new Error('database unavailable'))
.mockResolvedValue('not-committed')
const moving = h.registry.idleRehome(
request,
claim === 'ambiguous'
? vi.fn().mockRejectedValue(new Error('lost commit reply'))
: vi.fn().mockResolvedValue({ outcome: 'deferred' }),
reconcile
)
await vi.advanceTimersByTimeAsync(50)
expect(
await h.registry.idleRehome(
{ ...request, attemptId: 'other' },
vi.fn(),
vi.fn().mockResolvedValue('not-committed')
)
).toEqual({ outcome: 'busy' })
expect(h.socket.close).not.toHaveBeenCalled()
await vi.advanceTimersByTimeAsync(300)
expect(await moving).toEqual({ outcome: 'deferred' })
expect(reconcile).toHaveBeenCalledTimes(3)
expect(h.socket.close).not.toHaveBeenCalled()
}
)
it('holds attach ownership through basis failure reservation cleanup', async () => {
const basis = deferred<void>()
const cleanup = deferred<void>()
const h = await source({
recordConnectionBasis: vi.fn().mockImplementation(async () => {
await basis.promise
throw new Error('basis failed')
}),
failReservation: vi.fn().mockReturnValue(cleanup.promise)
})
const client = new FakeSocket()
h.session.pendingConns.set('conn', {
connId: 'conn',
connTicket: 'ticket',
client: client as unknown as WebSocket,
reservation: {
userId: identity.sub,
relayHostId: identity.relayHostId,
credentialKind: 'invite',
relayDeviceId: 'device-1',
tokenHash: 'invite-hash',
reservationId: 'reservation-1',
leaseExpiresAt: Date.now() + 1000
},
attachTimer: setTimeout(() => {}, 1000),
credentialActivityId: null,
phonePath: 'db',
timing: {
startedAt: Date.now(),
connOpenAt: Date.now(),
stageMs: { assignment: 0, credential: 0, activity: 0 }
}
})
const attached = h.registry.acceptHostData(
new FakeSocket() as unknown as WebSocket,
'conn',
'ticket',
1
)
const commit = vi.fn().mockResolvedValue({ outcome: 'deferred' })
expect(await h.registry.idleRehome(request, commit, vi.fn())).toEqual({ outcome: 'busy' })
basis.resolve()
await vi.advanceTimersByTimeAsync(0)
expect(h.session.activeConnIds.size).toBe(0)
expect(await h.registry.idleRehome(request, commit, vi.fn())).toEqual({ outcome: 'busy' })
expect(commit).not.toHaveBeenCalled()
cleanup.resolve()
await attached
})
})
describe('paced drain', () => {
beforeEach(() => vi.useFakeTimers())
afterEach(() => {
vi.clearAllTimers()
vi.useRealTimers()
})
async function connectHosts(count: number): Promise<{
registry: HostSessionRegistry
sockets: FakeSocket[]
}> {
const activateControl = vi
.fn<RelayAssignmentStore['activateControl']>()
.mockResolvedValue('control:production-gce-c3:1')
const { registry, activate } = createRegistry(activateControl)
const sockets: FakeSocket[] = []
for (let index = 0; index < count; index += 1) {
const socket = new FakeSocket()
sockets.push(socket)
await activate(
socket as unknown as WebSocket,
{ ...identity, sub: `user-${index}` },
null,
1,
false,
1
)
socket.send.mockClear()
}
return { registry, sockets }
}
function drainsSent(sockets: FakeSocket[]): number {
return sockets.filter((socket) =>
socket.send.mock.calls.some(([payload]) => String(payload).includes('"type":"drain"'))
).length
}
// Windows here stay under the 75s control-silence watchdog, which would otherwise close
// a test socket that never heartbeats before its paced send is due.
it('gives each host its own grace after its own send, not after the call', async () => {
const { registry, sockets } = await connectHosts(2)
registry.drain(10_000, { paceWindowMs: 40_000 })
await vi.advanceTimersByTimeAsync(10_000)
expect(sockets[0]!.readyState).toBe(sockets[0]!.CLOSED)
expect(sockets[1]!.readyState).toBe(sockets[1]!.OPEN)
// Its own send at 40s plus its own 10s grace, not 10s from the drain call.
await vi.advanceTimersByTimeAsync(39_999)
expect(sockets[1]!.readyState).toBe(sockets[1]!.OPEN)
await vi.advanceTimersByTimeAsync(1)
expect(sockets[1]!.readyState).toBe(sockets[1]!.CLOSED)
})
it('lets an emergency drain supersede the sends still queued by a paced one', async () => {
const { registry, sockets } = await connectHosts(4)
registry.drain(0, { paceWindowMs: 40_000 })
await vi.advanceTimersByTimeAsync(0)
expect(drainsSent(sockets)).toBe(1)
registry.drain(0)
await vi.advanceTimersByTimeAsync(0)
expect(drainsSent(sockets)).toBe(4)
const sendsAfterEmergency = sockets.map((socket) => socket.send.mock.calls.length)
await vi.advanceTimersByTimeAsync(40_000)
expect(sockets.map((socket) => socket.send.mock.calls.length)).toEqual(sendsAfterEmergency)
})
})
cloud/apps/relay/src/connections/host-signed-out-rejection.test.ts
-178
View File
@@ -1,178 +0,0 @@
import { EventEmitter } from 'node:events'
import { CONTROL_CONTINUITY_LIMITS } from '@orca-cloud/relay-contract/control-continuity'
import { RELAY_CLOSE_CODE } from '@orca-cloud/relay-contract/close-codes'
import { RELAY_HOST_CLOSE_REASON } from '@orca-cloud/relay-contract/host-close-reason'
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import type WebSocket from 'ws'
import type { RelayAssignmentStore } from '../assignment/assignment-store.js'
import type { RelayConfig } from '../config.js'
import type { RelayCredentialStore } from '../credentials/credential-store.js'
import { HostSessionRegistry } from './host-session-registry.js'
import type { RelayRuntimeObserver } from '../observability/relay-observability.js'
import type { RelayTokenClaims } from '../authentication/relay-token-verifier.js'
import { ProcessQueuedByteBudget } from './splice-forwarder.js'
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 ?? ''))
})
readonly terminate = vi.fn(() => {
this.readyState = this.CLOSED
this.emit('close', 1006, Buffer.alloc(0))
})
}
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: []
} as unknown as RelayConfig
const identity = {
sub: 'user-1',
prof: 'profile-1',
org: 'org-1',
relayHostId: 'AbCdEf0123_-xyZ9'
} as unknown as RelayTokenClaims
const reservation = {
userId: identity.sub,
relayHostId: identity.relayHostId,
credentialKind: 'resume',
relayDeviceId: 'device-1',
leaseExpiresAt: Date.now() + 60_000
}
function createRegistry() {
const store = {
resolveResume: vi.fn().mockResolvedValue({ userId: identity.sub }),
reserveCredential: vi.fn().mockResolvedValue(reservation),
failReservation: vi.fn().mockResolvedValue(undefined)
}
const assignments = {
activateControl: vi.fn().mockResolvedValue('control:production-gce-c3:1'),
markMigrationTargetRegistered: vi.fn().mockResolvedValue(undefined),
resolve: vi.fn().mockResolvedValue({ cellId: config.cellId }),
acquireActivity: vi.fn().mockResolvedValue(undefined),
renewControlActivity: vi.fn().mockResolvedValue(undefined),
releaseActivity: vi.fn().mockResolvedValue(true)
} 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,
vi.fn(),
store as unknown as RelayCredentialStore,
assignments,
new ProcessQueuedByteBudget(),
observer
)
const activate = (socket: WebSocket, generation: number): Promise<void> =>
(
registry as unknown as {
activate: (
socket: WebSocket,
identity: RelayTokenClaims,
existing: null,
generation: number,
rebind: boolean,
assignmentEpoch: number,
appVersion: string
) => Promise<void>
}
).activate(socket, identity, null, generation, false, 1, '1.4.173')
return { registry, activate }
}
async function dialPhone(registry: HostSessionRegistry): Promise<FakeSocket> {
const phone = new FakeSocket()
await registry.acceptClient(phone as unknown as WebSocket, identity.relayHostId, 'credential')
return phone
}
// The 4404 hello body is unchanged: every shipped phone parses it with a strict
// schema, so the cause has to ride the close frame instead.
const HOST_OFFLINE_HELLO = JSON.stringify({
type: 'relay-hello',
ok: false,
code: RELAY_CLOSE_CODE.HOST_OFFLINE
})
describe('host sign-out reason on phone rejection', () => {
beforeEach(() => vi.useFakeTimers())
afterEach(() => {
vi.clearAllTimers()
vi.useRealTimers()
})
it('names the sign-out to a phone that arrives after the host is gone', async () => {
const { registry, activate } = createRegistry()
const control = new FakeSocket()
await activate(control as unknown as WebSocket, 1)
control.close(1000, RELAY_HOST_CLOSE_REASON.SIGNED_OUT)
vi.advanceTimersByTime(CONTROL_CONTINUITY_LIMITS.orphanGraceMs + 1)
const phone = await dialPhone(registry)
expect(phone.send).toHaveBeenCalledWith(HOST_OFFLINE_HELLO)
expect(phone.close).toHaveBeenCalledWith(
RELAY_CLOSE_CODE.HOST_OFFLINE,
RELAY_HOST_CLOSE_REASON.SIGNED_OUT
)
})
it('ignores a close reason the host invented', async () => {
const { registry, activate } = createRegistry()
const control = new FakeSocket()
await activate(control as unknown as WebSocket, 1)
control.close(1000, 'signed-out-ish')
vi.advanceTimersByTime(CONTROL_CONTINUITY_LIMITS.orphanGraceMs + 1)
const phone = await dialPhone(registry)
expect(phone.close).toHaveBeenCalledWith(
RELAY_CLOSE_CODE.HOST_OFFLINE,
'relay connection rejected'
)
})
it('forgets the sign-out once the host proves itself again', async () => {
const { registry, activate } = createRegistry()
const control = new FakeSocket()
await activate(control as unknown as WebSocket, 1)
control.close(1000, RELAY_HOST_CLOSE_REASON.SIGNED_OUT)
vi.advanceTimersByTime(CONTROL_CONTINUITY_LIMITS.orphanGraceMs + 1)
const reconnected = new FakeSocket()
await activate(reconnected as unknown as WebSocket, 2)
// Drop it abruptly, as a network death would, so only the stale memory
// could still name a cause.
reconnected.terminate()
vi.advanceTimersByTime(CONTROL_CONTINUITY_LIMITS.orphanGraceMs + 1)
const phone = await dialPhone(registry)
expect(phone.close).toHaveBeenCalledWith(
RELAY_CLOSE_CODE.HOST_OFFLINE,
'relay connection rejected'
)
})
})
cloud/apps/relay/src/connections/relay-connection-hard-cap.blackbox.test.ts
+1 -34
View File
@@ -122,15 +122,12 @@ describe('relay connection hard cap', () => {
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' }
{ ...publicJwk, kid: 'test-key', alg: 'ES256', use: 'sig' }
]
})
)
@@ -188,36 +185,6 @@ describe('relay connection hard cap', () => {
await new Promise<void>((resolve) => relay.server.once('listening', resolve))
cleanup.push(() => new Promise<void>((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}`)
cloud/apps/relay/src/connections/relay-connection-ledger.test.ts
-33
View File
@@ -13,38 +13,6 @@ function socket(): { webSocket: WebSocket; close: () => void } {
describe('RelayConnectionLedger', () => {
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('transfers a pending host-data unit without increasing enforced capacity', () => {
const ledger = new RelayConnectionLedger(6, 2)
const phone = ledger.tryReservePhone()!
@@ -128,4 +96,3 @@ describe('RelayConnectionLedger', () => {
expect(ledger.canReserveBooking(100)).toBe(false)
})
})
cloud/apps/relay/src/connections/relay-first-frame-close.blackbox.test.ts
+8 -104
View File
@@ -1,7 +1,6 @@
import { connect, type Socket } from 'node:net'
import { afterEach, expect, it, vi } from 'vitest'
import { RELAY_CLOSE_CODE } from '@orca-cloud/relay-contract/close-codes'
import { RELAY_PROTOCOL_LIMITS } from '@orca-cloud/relay-contract/protocol-limits'
import { loadRelayConfig } from '../config.js'
import type { RelayDatabase } from '../database/database.js'
import { createRelayServer } from '../relay-server.js'
@@ -19,10 +18,7 @@ const IDLE_LEDGER = {
enforcedConnectionUnits: 0
}
async function fixture(
options: { role?: 'cell' | 'director'; hardCap?: number } = {}
) {
const role = options.role ?? 'cell'
async function fixture() {
const database: RelayDatabase = {
query: vi.fn(async () => []),
queryLocked: vi.fn(async () => []),
@@ -35,27 +31,20 @@ async function fixture(
ORCA_RELAY_AUTH_ISSUER: 'https://auth.example.test',
ORCA_RELAY_JWKS_URL: 'https://auth.example.test/jwks',
ORCA_RELAY_ASSIGNMENT_SIGNING_KEY: 'synthetic-assignment-key-for-test-only',
ORCA_RELAY_ROLE: role,
ORCA_RELAY_ROLE: 'cell',
ORCA_RELAY_ADMIN_AUDIENCE: 'https://auth.example.test/admin',
ORCA_RELAY_DEPLOY_SERVICE_ACCOUNT: 'deploy@example.test',
ORCA_RELAY_CELL_CONNECTION_HARD_CAP: '600',
ORCA_RELAY_CELL_CONNECTION_UNOBSERVED_BOUND: '60',
...(role === 'director'
? {
ORCA_RELAY_CELLS_JSON: JSON.stringify([
{ id: 'cell-1', url: 'https://cell-1.example.test', capacityRequests: 900 }
])
}
: {})
ORCA_RELAY_CELL_CONNECTION_UNOBSERVED_BOUND: '60'
})
const relay = createRelayServer(config, database, {
connectionLedgerLimits: { hardCap: options.hardCap ?? 3, controlReserve: 1 }
connectionLedgerLimits: { hardCap: 3, controlReserve: 1 }
})
await new Promise<void>((resolve) => relay.server.listen(0, '127.0.0.1', resolve))
cleanups.push(() => new Promise<void>((resolve) => relay.server.close(() => resolve())))
const address = relay.server.address()
if (!address || typeof address === 'string') throw new Error('missing test port')
return { relay, port: address.port, database }
return { relay, port: address.port }
}
type RawPeer = {
@@ -108,9 +97,6 @@ async function silentUpgrade(port: number, target: string): Promise<RawPeer> {
}
}
// A 43-character base64url credential is the shortest value RelayAuthSchema accepts.
const WELL_FORMED_CREDENTIAL = 'abcdefghijklmnopqrstuvwxyzABCDEFGH012345678'
function maskedTextFrame(payload: string): Buffer {
const body = Buffer.from(payload)
const mask = Buffer.from([1, 2, 3, 4])
@@ -118,17 +104,6 @@ function maskedTextFrame(payload: string): Buffer {
return Buffer.concat([Buffer.from([0x81, 0x80 | body.length]), mask, masked])
}
function relayAuthFrame(): Buffer {
return maskedTextFrame(
JSON.stringify({
type: 'relay-auth',
v: 1,
mode: 'connect',
credential: WELL_FORMED_CREDENTIAL
})
)
}
// Rejected peer’s last unmasked frame: 0x88, length, big-endian status, then UTF-8 reason.
function readCloseFrame(received: Buffer): { code: number; reason: string } | null {
const start = received.lastIndexOf(0x88)
@@ -149,67 +124,16 @@ async function expectIdleLedger(
const PHONE_TARGET = '/v1/connect/abcdefghijklmnop'
it.each([
{ label: 'first-frame timeout', target: PHONE_TARGET, opcode: undefined },
{ label: 'binary first frame', target: PHONE_TARGET, opcode: 0x82 },
{ label: 'invalid phone auth', target: PHONE_TARGET, opcode: 0x81 },
{ label: 'invalid host-data auth', target: '/v1/host/data/connection-1', opcode: 0x81 }
])(
'releases admission after $label even when the peer ignores close',
async ({ target, opcode }) => {
const { relay, port, database } = await fixture()
const peer = await silentUpgrade(port, target)
const firstFrameDeadline = opcode === undefined ? RELAY_PROTOCOL_LIMITS.firstFrameDeadlineMs : 0
if (opcode !== undefined) peer.socket.write(Buffer.from([opcode, 0x80, 0, 0, 0, 0]))
await expectIdleLedger(relay, firstFrameDeadline + 2_000)
expect(relay.runtimeCounts().preAuthConnections).toBe(0)
expect(database.query).not.toHaveBeenCalled()
expect(database.queryLocked).not.toHaveBeenCalled()
await silentUpgrade(port, PHONE_TARGET)
expect(relay.connectionSnapshot()?.enforcedConnectionUnits).toBe(2)
}
)
it('releases admission after a rejected director invite when the peer ignores close', async () => {
const { relay, port } = await fixture({ role: 'director' })
it('releases admission after a binary first frame even when the peer ignores close', async () => {
const { relay, port } = await fixture()
const peer = await silentUpgrade(port, PHONE_TARGET)
peer.socket.write(relayAuthFrame())
peer.socket.write(Buffer.from([0x82, 0x80, 0, 0, 0, 0]))
await expectIdleLedger(relay, 2_000)
expect(readCloseFrame(peer.received())).toEqual({
code: RELAY_CLOSE_CODE.BAD_OUTER_CREDENTIAL,
reason: 'invalid invite'
})
expect(relay.runtimeCounts().preAuthConnections).toBe(0)
await silentUpgrade(port, PHONE_TARGET)
expect(relay.connectionSnapshot()?.enforcedConnectionUnits).toBe(2)
})
it('releases admission after a director move redirect when the peer ignores close', async () => {
const { relay, port } = await fixture({ role: 'director' })
const identity = { userId: 'user-1', relayHostId: 'abcdefghijklmnop' }
vi.spyOn(relay.store, 'resolveInviteForMove').mockResolvedValue({
userId: identity.userId,
relayDeviceId: 'device-1'
})
vi.spyOn(relay.assignments, 'resolve').mockResolvedValue({
...identity,
cellId: 'cell-1',
cellUrl: 'https://cell-1.example.test',
assignmentEpoch: 1,
leaseExpiresAt: Date.now() + 60_000
})
const peer = await silentUpgrade(port, PHONE_TARGET)
peer.socket.write(relayAuthFrame())
await expectIdleLedger(relay, 2_000)
expect(peer.received().toString('utf8')).toContain('"type":"relay-moved"')
expect(readCloseFrame(peer.received())).toEqual({
code: RELAY_CLOSE_CODE.DRAINING,
reason: 'connect to assigned cell'
})
await silentUpgrade(port, PHONE_TARGET)
expect(relay.connectionSnapshot()?.enforcedConnectionUnits).toBe(2)
})
// This responsive peer proves close bytes precede FIN, not delivery under backpressure: terminate can drop queued frames after 1s.
it('delivers the rejection code and a graceful FIN to a peer that keeps reading', async () => {
const { relay, port } = await fixture()
@@ -224,23 +148,3 @@ it('delivers the rejection code and a graceful FIN to a peer that keeps reading'
expect(peer.transport().error).toBeNull()
await expectIdleLedger(relay, 2_000)
})
const BINARY_FIRST_FRAME = Buffer.from([0x82, 0x80, 0, 0, 0, 0])
const CONCURRENT_PEERS_PER_SOURCE = 4
// Ledger does not clamp at zero, so a disconnect racing rejection exposes double release as a negative count.
it('releases exactly once when a rejected peer disconnects at the same moment', async () => {
const { relay, port } = await fixture({ hardCap: 2 * CONCURRENT_PEERS_PER_SOURCE + 1 })
const peers = await Promise.all(
Array.from({ length: CONCURRENT_PEERS_PER_SOURCE }, () => silentUpgrade(port, PHONE_TARGET))
)
for (const peer of peers) {
peer.socket.write(BINARY_FIRST_FRAME)
peer.socket.destroy()
}
await expectIdleLedger(relay, 3_000)
expect(relay.connectionSnapshot()).toMatchObject(IDLE_LEDGER)
expect(relay.runtimeCounts().preAuthConnections).toBe(0)
await silentUpgrade(port, PHONE_TARGET)
expect(relay.connectionSnapshot()?.enforcedConnectionUnits).toBe(2)
})
cloud/apps/relay/src/connections/relay-host-hello-shed.test.ts
+1 -7
View File
@@ -147,7 +147,7 @@ describe('host hello shedding under database pool pressure', () => {
})
}
it('admits hellos behind a deep queue that is still moving', async () => {
it('admits hellos behind a deep queue below the age threshold', async () => {
const cell = await startCell(120)
cell.advance(HOST_HELLO_SHED_OLDEST_WAIT_MS - 300)
@@ -167,12 +167,6 @@ describe('host hello shedding under database pool pressure', () => {
cell.advance(HOST_HELLO_SHED_OLDEST_WAIT_MS)
expect(await dial(cell)).toEqual({ status: 503, retryAfter: '2' })
// Shipped desktops lack unexpected-response listeners; refusal must reach their ordinary connect-error backoff.
const desktop = new WebSocket(cell.url, {
headers: { authorization: `Bearer ${cell.token}` }
})
const error = await new Promise<Error>((resolve) => desktop.once('error', resolve))
expect(error.message).toBe('Unexpected server response: 503')
})
})
cloud/apps/relay/src/connections/relay-host-proof-failure.blackbox.test.ts
+3 -8
View File
@@ -28,7 +28,7 @@ describe('relay host proof failures', () => {
})
// A verifyCellAssignment connect-timeout rejection once escaped beginProof and crashed whole cells.
it.each(['cell', 'combined'] as const)('contains %s host hello database timeouts', async (role) => {
it('contains cell host hello database timeouts', async () => {
const keys = await generateKeyPair('ES256')
const publicJwk = await exportJWK(keys.publicKey)
const jwksServer: Server = createServer((_request, response) => {
@@ -64,7 +64,7 @@ describe('relay host proof failures', () => {
authAudience: 'orca-relay',
jwksUrl: issuer,
assignmentSigningKey: new Uint8Array(32),
role,
role: 'cell',
cellId: 'production-gce-c3',
cells: [{ id: 'production-gce-c3', url: relayUrl, capacityRequests: 4_000 }],
adminAudience: `${relayUrl}/admin`,
@@ -88,7 +88,7 @@ describe('relay host proof failures', () => {
await new Promise<void>((resolve) => relay.server.once('listening', resolve))
cleanup.push(() => new Promise<void>((resolve) => relay.server.close(() => resolve())))
vi.spyOn(console, 'log').mockImplementation(() => undefined)
const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined)
vi.spyOn(console, 'warn').mockImplementation(() => undefined)
const unhandled: unknown[] = []
const onUnhandled = (reason: unknown): void => {
unhandled.push(reason)
@@ -140,11 +140,6 @@ describe('relay host proof failures', () => {
)
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<void>((resolve) => setImmediate(resolve))
expect(unhandled).toEqual([])
})
cloud/apps/relay/src/connections/splice-forwarder.test.ts
-37
View File
@@ -12,8 +12,6 @@ const { Receiver } = createRequire(new URL('../splice-forwarder.test.ts', import
Receiver: new (options: { isServer: true }) => Writable
}
const FRAME_OVERHEAD = RELAY_ADMISSION_BUDGETS.queuedFrameOverheadBytes
class FakeSocket extends EventEmitter {
readonly OPEN = 1
readyState = 1
@@ -48,7 +46,6 @@ describe('splice forwarder', () => {
onClose: vi.fn()
})
client.emit('message', Buffer.alloc(8 * 1024 * 1024), true)
expect(budget.current()).toBe(8 * 1024 * 1024 + FRAME_OVERHEAD)
client.emit('message', Buffer.alloc(512 * 1024), true)
expect(client.closes[0]?.code).toBe(4429)
expect(host.closes[0]?.code).toBe(4429)
@@ -73,7 +70,6 @@ describe('splice forwarder', () => {
client.emit('message', frame, true)
expect(client.closes).toEqual([])
expect(host.closes).toEqual([])
expect(budget.current()).toBe(8 * 1024 * 1024 + FRAME_OVERHEAD)
host.bufferedAmount = 0
vi.advanceTimersByTime(25)
@@ -109,9 +105,6 @@ describe('splice forwarder', () => {
receiver.write(batch)
expect(client._socket.pause).toHaveBeenCalledTimes(
RELAY_ADMISSION_BUDGETS.spliceHardQueuedFrames
)
expect(closed).toHaveBeenCalledWith({
code: 4429,
reason: 'relay queue limit exceeded',
@@ -122,34 +115,4 @@ describe('splice forwarder', () => {
expect(host.closes[0]?.code).toBe(4429)
expect(budget.current()).toBe(0)
})
it('charges empty queued frames against the shared byte budget until drained or closed', () => {
vi.useFakeTimers()
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()
})
for (let index = 0; index < 3; index += 1) {
client.emit('message', Buffer.alloc(0), true)
}
expect(budget.current()).toBe(3 * FRAME_OVERHEAD)
host.bufferedAmount = 0
vi.advanceTimersByTime(25)
expect(host.sent).toHaveLength(3)
expect(budget.current()).toBe(0)
host.bufferedAmount = 1024 * 1024
client.emit('message', Buffer.alloc(0), true)
client.emit('message', Buffer.from('x'), true)
expect(budget.current()).toBe(2 * FRAME_OVERHEAD + 1)
client.close(1000, 'done')
expect(budget.current()).toBe(0)
})
})
cloud/apps/relay/src/control-leases/control-accept-cell-row-lock-postgres.test.ts
+1 -68
View File
@@ -1,9 +1,6 @@
import pg from 'pg'
import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest'
import { afterAll, beforeAll, describe, expect, it } from 'vitest'
import { RelayAssignmentStore } from '../assignment/assignment-store.js'
import type { RelayConfig } from '../config.js'
import { openRelayDatabase, type RelayDatabase } from '../database/database.js'
import { createRelayServer } from '../relay-server.js'
const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL
const describePostgres = databaseUrl ? describe : describe.skip
@@ -174,41 +171,6 @@ describePostgres('PostgreSQL control accept without a held cell row', () => {
expect(Number(cells[0]!.reserved_requests)).toBe(Number(units[0]!.units))
}, 20_000)
// Drive the production store: wrappers can omit optional fused commit and silently restore a separate COMMIT round trip.
it('commits every counter write in one message through the production store', async () => {
await removeTestRows(databases[0]!)
const store = new RelayAssignmentStore(databases[0]!, () => 100)
await prepareCell(store)
const { assignments } = createRelayServer(cellConfig(), databases[0]!, { now: () => 100 })
const assignment = await store.assign(first)
const firstControl = await assignments.activateControl(first, {
cellId: cell.id,
assignmentEpoch: assignment.assignmentEpoch,
generation: 1
})
const sent = vi.spyOn(pg.Client.prototype, 'query')
let texts: string[] = []
try {
await assignments.releaseActivity(first, firstControl)
await assignments.activateControl(first, {
cellId: cell.id,
assignmentEpoch: assignment.assignmentEpoch,
generation: 2
})
await assignments.acquireActivity(first, {
activityId: 'invite:fused',
kind: 'invite',
cellId: cell.id
})
} finally {
texts = sent.mock.calls.map(([text]) => String(text))
sent.mockRestore()
}
expect(texts.filter((text) => text.endsWith('; COMMIT'))).toHaveLength(3)
expect(texts.filter((text) => text === 'COMMIT')).toEqual([])
})
async function prepareCell(store: RelayAssignmentStore): Promise<void> {
await store.reconcileCells([cell])
await store.recordCellHeartbeat({
@@ -227,33 +189,4 @@ describePostgres('PostgreSQL control accept without a held cell row', () => {
connectionUnobservedBound: 50
})
}
})
function cellConfig(): RelayConfig {
return {
port: 0,
publicUrl: cell.url,
cellUrl: cell.url,
authIssuer: 'https://auth.example.test',
authAudience: 'orca-relay',
jwksUrl: 'https://auth.example.test/jwks',
assignmentSigningKey: new Uint8Array(32),
role: 'cell',
cellId: cell.id,
cells: [],
adminAudience: 'https://relay.example.test/v1/admin/drain',
deployServiceAccount: 'deploy@example.test',
runtimeServiceAccount: 'relay-cell@example.test',
adminJwksUrl: 'https://auth.example.test/jwks',
databasePoolMax: 10,
publicAssignmentsEnabled: true,
publicAssignmentConcurrency: 2,
publicAssignmentQueueMax: 128,
publicAssignmentWaitMs: 4_000,
publicResolveConcurrency: 1,
publicResolveWaitMs: 5_000,
publicAssignmentRetryAfterSeconds: 5,
dataDir: './data'
}
}
cloud/apps/relay/src/control-leases/control-accept-cell-row-lock.test.ts
-131
View File
@@ -1,131 +0,0 @@
import { afterEach, beforeEach, describe, expect, it } from 'vitest'
import type { RelayCellConfig } from '../config.js'
import { RelayAssignmentStore } from '../assignment/assignment-store.js'
import {
openInMemoryRelayDatabase,
type RelayDatabase,
type RelayLockOptions,
type RelayTransactionOptions,
type SqlRow
} from '../database/database.js'
const CELL: RelayCellConfig = {
id: 'accept-cell-a',
url: 'https://accept-a.example.com',
capacityRequests: 2
}
const host = { userId: 'accept-user', relayHostId: 'acceptho00000001' }
const second = { userId: 'accept-user', relayHostId: 'acceptho00000002' }
const third = { userId: 'accept-user', relayHostId: 'acceptho00000003' }
type Statement = { sql: string; locked: boolean }
// Records the statements a transaction issues, in order, so what the accept
// path does with the shared cell row can be asserted rather than described.
class RecordingDatabase implements RelayDatabase {
constructor(
private readonly inner: RelayDatabase,
readonly statements: Statement[] = []
) {}
async query(sql: string, params: unknown[] = []): Promise<SqlRow[]> {
this.statements.push({ sql, locked: false })
return await this.inner.query(sql, params)
}
async queryLocked(
sql: string,
params: unknown[] = [],
options: RelayLockOptions = {}
): Promise<SqlRow[]> {
this.statements.push({ sql, locked: true })
return await this.inner.queryLocked(sql, params, options)
}
async transaction<T>(
operation: (transaction: RelayDatabase) => Promise<T>,
options: RelayTransactionOptions = {}
): Promise<T> {
return await this.inner.transaction(
async (transaction) =>
await operation(new RecordingDatabase(transaction, this.statements)),
options
)
}
async close(): Promise<void> {
await this.inner.close()
}
}
describe('control accept cell-row lock span', () => {
let recorder: RecordingDatabase
let store: RelayAssignmentStore
let assignmentEpoch: number
beforeEach(async () => {
recorder = new RecordingDatabase(await openInMemoryRelayDatabase())
store = new RelayAssignmentStore(recorder, () => 100)
await store.reconcileCells([CELL])
assignmentEpoch = (await store.assign(host)).assignmentEpoch
})
afterEach(async () => {
await recorder?.close()
})
async function accept(generation: number): Promise<string> {
recorder.statements.length = 0
return await store.activateControl(host, {
cellId: CELL.id,
assignmentEpoch,
generation
})
}
function cellStatements(): Statement[] {
return recorder.statements.filter((statement) => /relay_cells/.test(statement.sql))
}
async function reservedRequests(): Promise<number> {
const row = (
await recorder.query(`SELECT reserved_requests FROM relay_cells WHERE cell_id = ?`, [
CELL.id
])
)[0]!
return Number(row.reserved_requests)
}
async function cellLeaseUnits(): Promise<number> {
const row = (
await recorder.query(
`SELECT COALESCE(SUM(request_units), 0) AS units
FROM relay_assignment_activity_leases WHERE cell_id = ?`,
[CELL.id]
)
)[0]!
return Number(row.units)
}
it('leaves the cell row untouched when a rebind retires and installs one control', async () => {
await accept(1)
await accept(2)
expect(cellStatements()).toEqual([])
expect(await reservedRequests()).toBe(1)
expect(await cellLeaseUnits()).toBe(1)
})
it('still refuses an accept that would exceed the cell capacity', async () => {
const control = await accept(1)
await store.assign(second)
await store.releaseActivity(host, control)
await store.assign(third)
expect(await reservedRequests()).toBe(CELL.capacityRequests)
await expect(accept(2)).rejects.toThrow('relay_capacity_exhausted')
expect(await reservedRequests()).toBe(CELL.capacityRequests)
expect(await cellLeaseUnits()).toBe(CELL.capacityRequests)
})
})

Some files were not shown because too many files have changed in this diff Show More