diff --git a/cloud/apps/relay/src/assignment-isolated-cell-replacement.test.ts b/cloud/apps/relay/src/assignment-isolated-cell-replacement.test.ts index 306c25043d9..8e5c3717f9e 100644 --- a/cloud/apps/relay/src/assignment-isolated-cell-replacement.test.ts +++ b/cloud/apps/relay/src/assignment-isolated-cell-replacement.test.ts @@ -934,4 +934,37 @@ describe('classifying a reconnect whose home is isolated for a roll', () => { expect(counter.sql.filter((sql) => sql.includes('FROM relay_assignments'))).toHaveLength(1) expect(counter.count('relay_cell_admission')).toBe(1) }) + + it('still treats a host its draining cell already released as a drain return', async () => { + // A drain releases the host's control lease before it tells the host, so the redial + // finds no activity: the stamp, not a live lease, has to carry the classification. + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + const { store, database, isolateForRoll } = await setup() + const first = await store.assign(IDENTITY, 'us-central1') + await store.activateControl(IDENTITY, { + cellId: first.cellId, + assignmentEpoch: first.assignmentEpoch, + generation: 1 + }) + await isolateForRoll(first.cellId) + expect(await store.releaseActivity(IDENTITY, `control:${first.cellId}:1`)).toBe(true) + expect( + await database.query( + `SELECT activity_id FROM relay_assignment_activity_leases WHERE user_id = ?`, + [IDENTITY.userId] + ) + ).toEqual([]) + + expect(await store.resolve(IDENTITY, classify)).toMatchObject({ + cellId: first.cellId, + homeCellRollIsolated: true + }) + warn.mockClear() + const moved = await store.assign(IDENTITY, 'us-central1') + expect(moved.cellId).not.toBe(first.cellId) + expect(moved.assignmentEpoch).toBe(first.assignmentEpoch + 1) + expect(jsonEvents(warn, 'orca_relay_sticky_replaced_off_isolated_cell')).toMatchObject([ + { fromCellId: first.cellId, toCellId: moved.cellId } + ]) + }) }) diff --git a/cloud/apps/relay/src/floating-promise-census.test.ts b/cloud/apps/relay/src/floating-promise-census.test.ts index 10574467484..5853558008b 100644 --- a/cloud/apps/relay/src/floating-promise-census.test.ts +++ b/cloud/apps/relay/src/floating-promise-census.test.ts @@ -18,7 +18,9 @@ const NEVER_REJECTS = new Map([ ['resume-confirmed-audit-buffer.ts#flush', 'catches the insert; a failed batch goes back to the buffer'], ['control-renewal-batch.ts#flush', 'rejects the waiters, never itself'], ['regional-rehome-worker.ts#run', 'one try/catch around the whole poll'], - ['relay-control-flag-channel.ts#poll', 'one catch on the whole read'] + ['relay-control-flag-channel.ts#poll', 'one catch on the whole read'], + ['host-session-registry.ts#releaseActivitySettled', 'one try/catch around the release'], + ['host-session-registry.ts#settledWithin', 'races the settled work against a timer'] ]) const sourceDirectory = fileURLToPath(new URL('.', import.meta.url)) diff --git a/cloud/apps/relay/src/host-session-client-accept.test.ts b/cloud/apps/relay/src/host-session-client-accept.test.ts index 7448631383a..5dca44107d3 100644 --- a/cloud/apps/relay/src/host-session-client-accept.test.ts +++ b/cloud/apps/relay/src/host-session-client-accept.test.ts @@ -823,6 +823,8 @@ describe('paced drain and the phones of a host not yet told', () => { 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) diff --git a/cloud/apps/relay/src/host-session-drain-release-order.test.ts b/cloud/apps/relay/src/host-session-drain-release-order.test.ts new file mode 100644 index 00000000000..4ca22d72fa9 --- /dev/null +++ b/cloud/apps/relay/src/host-session-drain-release-order.test.ts @@ -0,0 +1,171 @@ +import { RELAY_CLOSE_CODE } from '@orca-cloud/relay-contract' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { DRAIN_RELEASE_BEFORE_CLOSE_MS } from './host-session-registry.js' +import { + answerProof, + hostControlProofFixture, + openProof, + ProofSocket +} from './test-fixtures/host-control-proof.js' + +// A drained host redials the director the moment it reads `drain` or the close. While this +// cell's release still holds the host's assignment row, that redial is refused row_busy. + +beforeEach(() => vi.useFakeTimers()) +afterEach(() => { + vi.clearAllTimers() + vi.useRealTimers() + vi.restoreAllMocks() +}) + +async function connectedHost() { + vi.spyOn(console, 'warn').mockImplementation(() => undefined) + const h = hostControlProofFixture() + const socket = await openProof(h, new ProofSocket()) + answerProof(socket, h.keyPair) + await vi.advanceTimersByTimeAsync(0) + const key = { userId: h.identity.sub, relayHostId: h.identity.relayHostId } + expect(h.registry.get(key)?.controlActivityId).toBe('control:1') + socket.send.mockClear() + return { h, socket, key } +} + +function frames(socket: ProofSocket): unknown[] { + return socket.send.mock.calls.map(([frame]) => JSON.parse(frame) as unknown) +} + +function pendingRelease() { + let settle!: (outcome: 'released' | 'failed') => void + const promise = new Promise((resolve, reject) => { + settle = (outcome) => (outcome === 'released' ? resolve(true) : reject(new Error('down'))) + }) + return { promise, settle } +} + +describe('a zero-grace drain releases the host before telling it', () => { + it('tells and closes the host only once its control release has landed', async () => { + const { h, socket, key } = await connectedHost() + const release = pendingRelease() + const releaseActivity = vi + .spyOn(h.assignments, 'releaseActivity') + .mockReturnValue(release.promise) + + h.registry.drain(0) + await vi.advanceTimersByTimeAsync(1) + + expect(releaseActivity).toHaveBeenCalledWith(key, 'control:1') + // Fenced at once: no longer a session here, and not yet told to leave. + expect(h.registry.get(key)).toBeNull() + expect(socket.send).not.toHaveBeenCalled() + expect(socket.close).not.toHaveBeenCalled() + + release.settle('released') + await vi.advanceTimersByTimeAsync(0) + expect(frames(socket)).toEqual([{ type: 'drain', graceMs: 0, recovery: 'resolve-director' }]) + expect(socket.close).toHaveBeenCalledWith( + RELAY_CLOSE_CODE.DRAINING, + 'resolve configured director' + ) + expect(socket.send.mock.invocationCallOrder[0]).toBeLessThan( + 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("starts the host's splice releases before its control release", async () => { + // The store runs one host's releases in the order they start, so waiting on the + // control release covers the splices the drain closed ahead of it. + const { h, key } = await connectedHost() + const order: string[] = [] + vi.spyOn(h.assignments, 'releaseActivity').mockImplementation(async (_identity, id) => { + order.push(id) + return true + }) + const session = h.registry.get(key)! + session.activeSplices.set('conn-a', () => { + session.activeSplices.delete('conn-a') + void h.assignments.releaseActivity(key, 'splice:conn-a') + }) + + h.registry.drain(0) + await vi.advanceTimersByTimeAsync(1) + + expect(order).toEqual(['splice:conn-a', 'control:1']) + }) + + it('closes at once when the release fails, so a database outage never holds the host', async () => { + const { h, socket } = await connectedHost() + const release = pendingRelease() + vi.spyOn(h.assignments, 'releaseActivity').mockReturnValue(release.promise) + + h.registry.drain(0) + await vi.advanceTimersByTimeAsync(1) + release.settle('failed') + await vi.advanceTimersByTimeAsync(0) + + expect(socket.close).toHaveBeenCalledWith( + RELAY_CLOSE_CODE.DRAINING, + 'resolve configured director' + ) + }) + + it('does not wait for a host that holds no lease here', async () => { + // A reserve-mode cell admits hosts from memory, with no activity row to release. + const { h, socket, key } = await connectedHost() + h.registry.get(key)!.controlActivityId = null + const releaseActivity = vi.spyOn(h.assignments, 'releaseActivity') + + h.registry.drain(0) + await vi.advanceTimersByTimeAsync(1) + + expect(releaseActivity).not.toHaveBeenCalled() + expect(frames(socket)).toEqual([{ type: 'drain', graceMs: 0, recovery: 'resolve-director' }]) + expect(socket.close).toHaveBeenCalledWith( + RELAY_CLOSE_CODE.DRAINING, + 'resolve configured director' + ) + }) + + 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) + }) +}) diff --git a/cloud/apps/relay/src/host-session-registry.test.ts b/cloud/apps/relay/src/host-session-registry.test.ts index 0db469ad457..50fcd6aedb5 100644 --- a/cloud/apps/relay/src/host-session-registry.test.ts +++ b/cloud/apps/relay/src/host-session-registry.test.ts @@ -1356,6 +1356,7 @@ describe('source-owned idle cutover', () => { await h.registry.idleRehome(request, commit, vi.fn().mockResolvedValue('not-committed')) ).toEqual({ outcome }) expect(commit).toHaveBeenCalledTimes(outcome === 'busy' ? 0 : 1) + await vi.advanceTimersByTimeAsync(0) expect(h.socket.close).toHaveBeenCalledTimes(outcome === 'busy' ? 0 : 1) }) it('restarts the age when the host rebinds its control socket', async () => { @@ -1414,6 +1415,7 @@ describe('source-owned idle cutover', () => { expect( await h.registry.idleRehome(request, commit, vi.fn().mockResolvedValue('not-committed')) ).toEqual({ outcome: 'committed' }) + await vi.advanceTimersByTimeAsync(0) expect(h.socket.close).toHaveBeenCalledWith(RELAY_CLOSE_CODE.DRAINING, expect.any(String)) expect(h.releaseActivity).toHaveBeenCalled() }) @@ -1961,6 +1963,8 @@ describe('paced drain', () => { it('sends every drain at once when no window is given', async () => { const { registry, sockets } = await connectHosts(4) registry.drain(0) + // A zero-grace drain is told by its close, once the host's release has landed. + await vi.advanceTimersByTimeAsync(0) expect(drainsSent(sockets)).toBe(4) }) @@ -1969,13 +1973,19 @@ describe('paced drain', () => { it('spreads the sends evenly across the window', async () => { const { registry, sockets } = await connectHosts(5) registry.drain(0, { paceWindowMs: 40_000 }) - // The first host is sent synchronously; the last lands on the window's closing edge. + // Each host is told by its zero-grace close, which a timer nested in the send runs 1 ms on. + const advance = async (ms: number): Promise => { + await vi.advanceTimersByTimeAsync(ms) + await vi.advanceTimersByTimeAsync(1) + } + // The first host is told at once; the last lands on the window's closing edge. + await advance(0) expect(drainsSent(sockets)).toBe(1) - await vi.advanceTimersByTimeAsync(10_000) + await advance(10_000) expect(drainsSent(sockets)).toBe(2) - await vi.advanceTimersByTimeAsync(20_000) + await advance(20_000) expect(drainsSent(sockets)).toBe(4) - await vi.advanceTimersByTimeAsync(10_000) + await advance(10_000) expect(drainsSent(sockets)).toBe(5) }) @@ -1987,6 +1997,7 @@ describe('paced drain', () => { const socket = new FakeSocket() registry.acceptControl(socket as unknown as WebSocket, { ...identity, sub: 'user-late' }) expect(socket.close).toHaveBeenCalledWith(RELAY_CLOSE_CODE.DRAINING, 'relay draining') + await vi.advanceTimersByTimeAsync(0) expect(drainsSent(sockets)).toBe(1) }) @@ -2032,8 +2043,10 @@ describe('paced drain', () => { 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) diff --git a/cloud/apps/relay/src/host-session-registry.ts b/cloud/apps/relay/src/host-session-registry.ts index 1ccbabb98e7..2a048f08c08 100644 --- a/cloud/apps/relay/src/host-session-registry.ts +++ b/cloud/apps/relay/src/host-session-registry.ts @@ -189,6 +189,22 @@ function send(socket: WebSocket, type: string, message: object): void { socket.send(JSON.stringify({ type, ...message })) } +// Whether `work` settled within `ms`; never rejects, and unref'd so it never holds exit. +async function settledWithin(work: Promise, ms: number): Promise { + let timer: ReturnType | undefined + const elapsed = new Promise((resolve) => { + timer = setTimeout(() => resolve(false), ms) + timer.unref?.() + }) + const settled = await Promise.race([work.then(settledTrue, settledTrue), elapsed]) + clearTimeout(timer) + return settled +} + +function settledTrue(): true { + return true +} + function readControlFrame( socket: WebSocket, timeoutMs: number, @@ -235,6 +251,12 @@ const preAuthEnded = // stalled predecessor only accumulates doomed sockets. const ACTIVATION_QUEUE_WAIT_MS = 30_000 +// A drained host redials the director the moment it is told, and a redial that meets this +// cell's release still holding its assignment row is refused row_busy. From asia-east2 one +// release holds the row ~1 s (c34/c35 drains, 2026-10), and a host's releases run one after +// another, so 3 s covers a control plus one phone. Past that, it is told anyway. +export const DRAIN_RELEASE_BEFORE_CLOSE_MS = 3_000 + // Why: this lease bounds how long a host lingers on a cell after a missed drain, // and rebinding it is the only passive rebalancing we have, so it has to stay // finite. 6h keeps both properties while cutting control-activation traffic on @@ -1499,8 +1521,11 @@ export class HostSessionRegistry { session.authorityRevision += 1 this.markSeatDrainOnly(session) this.drainSentHosts.add(session.relayHostId) - if (session.socket) send(session.socket, 'drain', { graceMs, recovery: 'resolve-director' }) - this.scheduleDrainTimer(graceMs, () => this.closeDrainedSession(session)) + // With no grace the host is told by the close, after its release, not ahead of it. + if (graceMs > 0 && session.socket) { + send(session.socket, 'drain', { graceMs, recovery: 'resolve-director' }) + } + this.scheduleDrainTimer(graceMs, () => this.closeDrainedSession(session, graceMs === 0)) } // Unref'd so a drain in flight never holds the process open past its own work. @@ -1944,7 +1969,7 @@ export class HostSessionRegistry { this.rejectClient(pending.client, RELAY_CLOSE_CODE.PEER_DROPPED) } existing.socket?.close(RELAY_CLOSE_CODE.PEER_DROPPED, 'replaced by a newer generation') - this.releaseControlActivity(existing) + void this.releaseControlActivity(existing) } const session: HostSession = { identity, @@ -2120,7 +2145,7 @@ export class HostSessionRegistry { for (const close of session.activeSplices.values()) close() const key = this.key(session.identity.sub, session.relayHostId) if (this.sessions.get(key) === session) this.sessions.delete(key) - this.releaseControlActivity(session) + void this.releaseControlActivity(session) }, CONTROL_CONTINUITY_LIMITS.orphanGraceMs) }) if (session.heartbeatTimer) clearInterval(session.heartbeatTimer) @@ -2315,7 +2340,7 @@ export class HostSessionRegistry { }) } - private closeDrainedSession(session: HostSession): void { + private closeDrainedSession(session: HostSession, sendDrainFirst = false): void { if (session.state === 'closed') return const forcedConnections = session.activeConnIds.size + session.pendingConns.size if (forcedConnections > 0) { @@ -2346,29 +2371,62 @@ export class HostSessionRegistry { for (const close of session.activeSplices.values()) { close(RELAY_CLOSE_CODE.DRAINING, 'relay draining') } + const releases: Promise[] = [] for (const pending of session.pendingConns.values()) { clearTimeout(pending.attachTimer) pending.capacityReservation?.release() this.failReservationBestEffort(pending.reservation) if (pending.credentialActivityId) { - this.releaseActivityBestEffort( - { - userId: pending.reservation.userId, - relayHostId: pending.reservation.relayHostId - }, - pending.credentialActivityId + releases.push( + this.releaseActivitySettled( + { + userId: pending.reservation.userId, + relayHostId: pending.reservation.relayHostId + }, + pending.credentialActivityId + ) ) } this.rejectClient(pending.client, RELAY_CLOSE_CODE.DRAINING) } session.pendingConns.clear() session.state = 'closed' - if (session.socket) { - closeRelayWebSocket(session.socket, RELAY_CLOSE_CODE.DRAINING, 'resolve configured director') - } const key = this.key(session.identity.sub, session.relayHostId) if (this.sessions.get(key) === session) this.sessions.delete(key) - this.releaseControlActivity(session) + const controlRelease = this.releaseControlActivity(session) + if (controlRelease) { + releases.push(controlRelease) + } + const socket = session.socket + if (!socket) { + return + } + const close = (): void => { + if (sendDrainFirst && socket.readyState === socket.OPEN) { + send(socket, 'drain', { graceMs: 0, recovery: 'resolve-director' }) + } + closeRelayWebSocket(socket, RELAY_CLOSE_CODE.DRAINING, 'resolve configured director') + } + if (releases.length === 0) { + close() + return + } + // The host redials on this close, so its row is released first; a stuck release still closes. + // The splice releases above are queued ahead of these, so they have landed too. + const startedAt = this.now() + void settledWithin(Promise.all(releases), DRAIN_RELEASE_BEFORE_CLOSE_MS).then((settled) => { + if (!settled) { + console.warn( + JSON.stringify({ + event: 'orca_relay_host_drain_release_slow', + ...this.logIdentity(), + relayHostIdDigest: relayHostLogDigest(session.relayHostId), + waitedMs: this.now() - startedAt + }) + ) + } + close() + }) } private reassertRegionalDrain(session: HostSession): void { @@ -2576,20 +2634,32 @@ export class HostSessionRegistry { closeRelayWebSocket(socket, code, closeReason ?? 'relay connection rejected') } - private releaseControlActivity(session: HostSession): void { - if (!session.controlActivityId) return - this.releaseActivityBestEffort( + private releaseControlActivity(session: HostSession): Promise | undefined { + if (!session.controlActivityId) { + return undefined + } + return this.releaseActivitySettled( { userId: session.identity.sub, relayHostId: session.relayHostId }, session.controlActivityId ) } private releaseActivityBestEffort(identity: RelayIdentityKey, activityId: string): void { - void this.assignments.releaseActivity(identity, activityId).catch(() => { + void this.releaseActivitySettled(identity, activityId) + } + + // Never rejects, so a caller may wait on it without handling the failure again. + private async releaseActivitySettled( + identity: RelayIdentityKey, + activityId: string + ): Promise { + try { + await this.assignments.releaseActivity(identity, activityId) + } catch { // Why: expiry cleanup is the durable fallback; a transient SQL failure while // closing a socket must not become an unhandled rejection that kills the cell. console.warn('[orca-relay] activity release deferred to lease cleanup') - }) + } } private failReservationBestEffort(reservation: CredentialReservation): void {