fix(relay): release a drained host before telling it to leave (#31572)

* fix(relay): release a drained host before telling it to leave

A zero-grace drain sent `drain` and closed with 4503 in the same tick, then
released the host's control lease. The desktop redials the director as soon
as it reads either, and from asia-east2 the release holds the host's
assignment row for ~1 s, so the redial met it and was refused row_busy.

The release now runs first and the host is told once it lands, bounded at
2 s; a failed or stuck release still closes. Wire frames are unchanged.

* fix(relay): give a drained host's release 3 s and log when it runs out

One host's releases run one after another, so 3 s covers a control lease
plus one open phone connection from asia-east2.

* style(relay): brace the drain release branches

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 }
+    ])
+  })
 })

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))

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)

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<boolean>((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)
+  })
+})

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<void> => {
+      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)

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<unknown>, ms: number): Promise<boolean> {
+  let timer: ReturnType<typeof setTimeout> | undefined
+  const elapsed = new Promise<false>((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<void>[] = []
     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<void> | 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<void> {
+    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 {