mirror of
https://github.com/stablyai/orca.git
synced 2026-09-21 16:02:20 +00:00
feat(relay): pace the drain send during a same-cap cell roll (#21284)
* feat(relay): pace the drain send during a same-cap cell roll A same-cap roll drains a cell with graceMs 0, which sends `drain` to all ~800 controls in one pass. Every desktop re-dials on receipt regardless of graceMs, so the whole cell reconnects inside a second. On 2026-09-16 that stampede hit a Cloud SQL stall: attaches timed out, each leaving 10 minutes of late-arrival debt on connection headroom, and placement answered relay_capacity_exhausted fleet-wide for ~13 minutes. Spreading the sends spreads the re-dials. `HostSessionRegistry.drain` takes an optional pacing window and schedules each session's send evenly across it; admission is fenced for every session up front, and each host keeps its own full grace after its own send. /v1/admin/drain accepts `paceWindowMs` (<= 5 min) and echoes what it applied. The same-cap job asks for 120 s, and the drain-completion wait grew by the same amount. A cell still on an older image rejects the field, so the deploy script falls back to an unpaced drain rather than failing the roll. * fix(relay): scope the drain fence to the hosts already told Review of the paced drain found two problems, both from treating "this cell is draining" as one instant when pacing makes it a window. Timers: the sends queued by a paced drain were neither cleared when a later drain superseded them nor unref'd. A SIGTERM mid-window left up to 800 no-op timers holding the event loop open until systemd escalated to SIGKILL. Drain timers are now tracked, cleared on the next drain, and unref'd, so a retry re-arms a session's teardown instead of stacking a second one. Phones: the client fence read the global draining flag, so every phone was refused for the whole window even though its own host had not been told yet and was still serving. The director keeps pointing phones at this cell until their host moves, so they would have looped for up to two minutes. A session is now fenced when its drain is sent, not when the drain starts, and the client paths key off that. New control connections and re-attaches stay fenced globally: nothing new should land on a cell that is going away.
This commit is contained in:
Jinwoo Hong
@@ -57,6 +57,9 @@ jobs:
|
||||
GATE_OVERRIDE_REASON: ${{ inputs.gate-override-reason }}
|
||||
GATE_OVERRIDE_CONFIRMATION: ${{ inputs.gate-override-confirmation }}
|
||||
OUTPUT_DIRECTORY: ${{ github.workspace }}/relay-monitor-evidence
|
||||
# ~800 controls over 2 min is ~7 re-dials/s per cell, well under the director's
|
||||
# 5 x 80 in-flight assign cap. A cell on an older image ignores it and drains at once.
|
||||
DRAIN_PACE_WINDOW_MS: '120000'
|
||||
steps:
|
||||
- name: Require exact reusable-workflow configuration
|
||||
working-directory: .
|
||||
@@ -479,14 +482,16 @@ jobs:
|
||||
echo "SELECTOR_GENERATION_AFTER_ISOLATE=${ISOLATE_GENERATION}" >> "${GITHUB_ENV}"
|
||||
node dev/scripts/prepare-relay-production-capacity-canary.mjs \
|
||||
--director-origin "${DIRECTOR_ORIGIN}" --cell-origin "${CELL_ORIGIN}" \
|
||||
--cell-id "${TARGET_CELL_ID}" --approved-cells same-cap --mode drain
|
||||
--cell-id "${TARGET_CELL_ID}" --approved-cells same-cap --mode drain \
|
||||
--pace-window-ms "${DRAIN_PACE_WINDOW_MS}"
|
||||
# The wait has to outlast the pacing window as well as the leases it waits on.
|
||||
node dev/scripts/verify-relay-capacity-transition.mjs \
|
||||
--director-origin "${DIRECTOR_ORIGIN}" --cell-origin "${CELL_ORIGIN}" \
|
||||
--cell-id "${TARGET_CELL_ID}" --hard-cap "${EXPECTED_HARD_CAP}" \
|
||||
--unobserved-bound "${EXPECTED_UNOBSERVED_BOUND}" \
|
||||
--heartbeat either --admission migration-only --draining required \
|
||||
--activity restart-safe --expected-image-digests "${CURRENT_IMAGE_DIGEST}" \
|
||||
--timeout-ms 900000
|
||||
--timeout-ms 1020000
|
||||
|
||||
- id: capacity-auth
|
||||
if: ${{ inputs.mode != 'verify' }}
|
||||
|
||||
@@ -58,6 +58,8 @@ const RelayCellConnectionHardCapSchema = z.custom<RelayCellConnectionHardCap>(
|
||||
|
||||
const ASSIGNMENT_REJECTION_LOG_WINDOW_MS = 10_000
|
||||
const REGION_CATALOG_CACHE_MS = 30_000
|
||||
// A drain that outlives the roll step it belongs to is an outage, not a pacing win.
|
||||
const DRAIN_PACE_WINDOW_MAX_MS = 5 * 60 * 1_000
|
||||
|
||||
type AdmissionRejectionLogEntry = {
|
||||
route: 'assign' | 'resolve'
|
||||
@@ -72,7 +74,7 @@ export function createRelayApp(
|
||||
operations: {
|
||||
store: RelayCredentialStore
|
||||
assignments: RelayAssignmentStore
|
||||
drain: (graceMs: number) => void
|
||||
drain: (graceMs: number, options?: { paceWindowMs?: number }) => void
|
||||
idleRehome?: (input: IdleRegionalRehomeRequest & {
|
||||
cohortPercent: number
|
||||
directorSafety: RegionalRehomeSafetySnapshot
|
||||
@@ -488,12 +490,18 @@ export function createRelayApp(
|
||||
return context.json({ error: 'invalid_token' }, 401)
|
||||
}
|
||||
const body = z
|
||||
.object({ v: z.literal(1), graceMs: z.number().int().nonnegative().max(60 * 60 * 1000) })
|
||||
.object({
|
||||
v: z.literal(1),
|
||||
graceMs: z.number().int().nonnegative().max(60 * 60 * 1000),
|
||||
// Spreads the drain sends, and so the re-dials, over this window.
|
||||
paceWindowMs: z.number().int().nonnegative().max(DRAIN_PACE_WINDOW_MAX_MS).optional()
|
||||
})
|
||||
.strict()
|
||||
.safeParse(await context.req.json().catch(() => null))
|
||||
if (!body.success) return context.json({ error: 'invalid_request' }, 400)
|
||||
operations.drain(body.data.graceMs)
|
||||
return context.json({ ok: true })
|
||||
const paceWindowMs = body.data.paceWindowMs ?? 0
|
||||
operations.drain(body.data.graceMs, { paceWindowMs })
|
||||
return context.json({ ok: true, paceWindowMs })
|
||||
})
|
||||
app.post('/v1/admin/host-idle-rehome', async (context) => {
|
||||
if (config.role !== 'cell' || !operations.idleRehome) {
|
||||
|
||||
@@ -729,3 +729,73 @@ describe('control lease jitter', () => {
|
||||
vi.advanceTimersByTime(0)
|
||||
})
|
||||
})
|
||||
|
||||
describe('paced drain and the phones of a host not yet told', () => {
|
||||
beforeEach(() => vi.useFakeTimers())
|
||||
afterEach(() => {
|
||||
vi.clearAllTimers()
|
||||
vi.useRealTimers()
|
||||
})
|
||||
|
||||
const laterHostId = 'qrstuvwxyz012345'
|
||||
const laterIdentity = { ...identity, sub: 'user-2', relayHostId: laterHostId }
|
||||
|
||||
async function twoHostCell(): Promise<{
|
||||
h: ReturnType<typeof harness>
|
||||
told: FakeSocket
|
||||
untold: FakeSocket
|
||||
}> {
|
||||
const h = harness()
|
||||
const told = await activeHost(h)
|
||||
const untold = new FakeSocket()
|
||||
await h.activate(untold as unknown as WebSocket, laterIdentity, null, 1, false, 1, '1.4.197')
|
||||
// Both hosts now dial in, so the credential mocks have to answer for either.
|
||||
h.store.resolveResume.mockImplementation(async (hostId: string) => ({
|
||||
userId: hostId === laterHostId ? laterIdentity.sub : identity.sub
|
||||
}))
|
||||
h.store.reserveCredential.mockImplementation(async (hostId: string) => ({
|
||||
...reservation,
|
||||
userId: hostId === laterHostId ? laterIdentity.sub : identity.sub,
|
||||
relayHostId: hostId
|
||||
}))
|
||||
return { h, told, untold }
|
||||
}
|
||||
|
||||
async function dial(h: ReturnType<typeof harness>, hostId: string): Promise<FakeSocket> {
|
||||
const client = new FakeSocket()
|
||||
await h.registry.acceptClient(client as unknown as WebSocket, hostId, 'credential')
|
||||
return client
|
||||
}
|
||||
|
||||
it('serves a host whose drain has not been sent and refuses one whose has', async () => {
|
||||
const { h, told, untold } = await twoHostCell()
|
||||
h.registry.drain(0, { paceWindowMs: 40_000 })
|
||||
|
||||
const refused = await dial(h, identity.relayHostId)
|
||||
expect(refused.close).toHaveBeenCalledWith(RELAY_CLOSE_CODE.DRAINING, expect.any(String))
|
||||
expect(told.send).not.toHaveBeenCalledWith(expect.stringContaining('conn-open'))
|
||||
|
||||
const served = await dial(h, laterHostId)
|
||||
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)
|
||||
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))
|
||||
})
|
||||
|
||||
it('keeps an unpaced drain refusing every phone at once', async () => {
|
||||
const { h } = await twoHostCell()
|
||||
h.registry.drain(0)
|
||||
for (const hostId of [identity.relayHostId, laterHostId]) {
|
||||
const refused = await dial(h, hostId)
|
||||
expect(refused.close).toHaveBeenCalledWith(RELAY_CLOSE_CODE.DRAINING, expect.any(String))
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
@@ -336,15 +336,15 @@ describe('host session cleanup races', () => {
|
||||
session.activeSplices.set('conn-a', () => session.activeSplices.delete('conn-a'))
|
||||
|
||||
// POST /v1/admin/drain has no idempotency guard, and SIGTERM then SIGINT both
|
||||
// reach drain(), so a second teardown can be scheduled for the same session.
|
||||
// reach drain(), so a retry re-sends to every session. It must re-arm the pending
|
||||
// teardown rather than stack a second one: across a paced cell that is 800 orphaned
|
||||
// timers per retry, each one holding the loop open for the rest of the window.
|
||||
registry.drain(0)
|
||||
const scheduled = vi.getTimerCount()
|
||||
registry.drain(0)
|
||||
// Pin the premise: if drain ever gains an idempotency guard, the retry schedules no
|
||||
// second teardown and the assertion below stops defending the write-once snapshot
|
||||
// while still passing. Compare against the count before the retry rather than an
|
||||
// absolute, since the session's heartbeat interval is also pending.
|
||||
expect(vi.getTimerCount()).toBe(scheduled + 1)
|
||||
// Compare against the count before the retry rather than an absolute, since the
|
||||
// session's heartbeat interval is also pending.
|
||||
expect(vi.getTimerCount()).toBe(scheduled)
|
||||
vi.advanceTimersByTime(1)
|
||||
|
||||
// Asserting registry state, not the log line: FakeSocket closes synchronously, so
|
||||
@@ -1696,3 +1696,124 @@ describe('host data attach owner lookup', () => {
|
||||
expect(h.owner.activeConnIds.size).toBe(0)
|
||||
})
|
||||
})
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
it('sends every drain at once when no window is given', async () => {
|
||||
const { registry, sockets } = await connectHosts(4)
|
||||
registry.drain(0)
|
||||
expect(drainsSent(sockets)).toBe(4)
|
||||
})
|
||||
|
||||
// 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('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.
|
||||
expect(drainsSent(sockets)).toBe(1)
|
||||
await vi.advanceTimersByTimeAsync(10_000)
|
||||
expect(drainsSent(sockets)).toBe(2)
|
||||
await vi.advanceTimersByTimeAsync(20_000)
|
||||
expect(drainsSent(sockets)).toBe(4)
|
||||
await vi.advanceTimersByTimeAsync(10_000)
|
||||
expect(drainsSent(sockets)).toBe(5)
|
||||
})
|
||||
|
||||
it('fences admission for every session before the first paced send lands', async () => {
|
||||
const { registry, sockets } = await connectHosts(3)
|
||||
registry.drain(0, { paceWindowMs: 40_000 })
|
||||
expect(registry.isDraining()).toBe(true)
|
||||
// A host whose drain has not been sent yet must already be non-authoritative.
|
||||
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')
|
||||
expect(drainsSent(sockets)).toBe(1)
|
||||
})
|
||||
|
||||
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(10_001)
|
||||
expect(sockets[1]!.readyState).toBe(sockets[1]!.CLOSED)
|
||||
})
|
||||
|
||||
it('leaves no timer behind once an emergency drain cuts a window short', async () => {
|
||||
const { registry } = await connectHosts(4)
|
||||
registry.drain(0, { paceWindowMs: 40_000 })
|
||||
registry.drain(0)
|
||||
await vi.advanceTimersByTimeAsync(0)
|
||||
// Every session is closed, so anything still pending is an orphan of the cut window.
|
||||
expect(vi.getTimerCount()).toBe(0)
|
||||
})
|
||||
|
||||
it('keeps the first teardown snapshot when a regional drain fires before the fleet one', async () => {
|
||||
const { registry, sockets } = await connectHosts(1)
|
||||
const session = registry.get({ userId: 'user-0', relayHostId: identity.relayHostId })!
|
||||
session.activeSplices.set('conn-a', () => session.activeSplices.delete('conn-a'))
|
||||
registry.drainHost({
|
||||
attemptId: 'attempt',
|
||||
userId: 'user-0',
|
||||
relayHostId: identity.relayHostId,
|
||||
sourceAssignmentEpoch: 1,
|
||||
graceMs: 0
|
||||
})
|
||||
registry.drain(10)
|
||||
await vi.advanceTimersByTimeAsync(11)
|
||||
expect(session.closingCounts).toEqual({ splices: 1, pending: 0 })
|
||||
expect(sockets[0]!.readyState).toBe(sockets[0]!.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 })
|
||||
expect(drainsSent(sockets)).toBe(1)
|
||||
registry.drain(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)
|
||||
})
|
||||
})
|
||||
|
||||
@@ -179,6 +179,10 @@ export class HostSessionRegistry {
|
||||
private readonly hostCloseReasons = new HostCloseReasonMemory(() => this.now())
|
||||
private readonly hostCapabilities = new WeakMap<WebSocket, ReadonlySet<string>>()
|
||||
private draining = false
|
||||
private readonly drainTimers = new Set<ReturnType<typeof setTimeout>>()
|
||||
// Hosts whose drain has been sent. Paced sends land minutes apart, so "this cell is
|
||||
// draining" is not the same question as "this host has been told to leave".
|
||||
private readonly drainSentHosts = new Set<string>()
|
||||
|
||||
private readonly idleWork = new Map<string, number>()
|
||||
private readonly idleAttempts = new Map<
|
||||
@@ -330,7 +334,10 @@ export class HostSessionRegistry {
|
||||
credential: string,
|
||||
capacityReservation?: PendingHostDataReservation
|
||||
): Promise<void> {
|
||||
if (this.draining) {
|
||||
// Not `this.draining`: a paced drain tells hosts minutes apart, and the director keeps
|
||||
// pointing phones here until their own host has moved. Refusing them for the whole
|
||||
// window would turn a 2 min drain into a 2 min outage for hosts not yet told.
|
||||
if (this.drainSentHosts.has(hostId)) {
|
||||
capacityReservation?.release()
|
||||
this.rejectClient(socket, RELAY_CLOSE_CODE.DRAINING)
|
||||
return
|
||||
@@ -450,7 +457,7 @@ export class HostSessionRegistry {
|
||||
}
|
||||
// Admission may have crossed a drain or control replacement while persisting activity.
|
||||
if (
|
||||
this.draining ||
|
||||
this.drainSentHosts.has(hostId) ||
|
||||
this.sessions.get(sessionKey) !== session ||
|
||||
session.state !== 'active' ||
|
||||
session.socket !== admittingSocket ||
|
||||
@@ -592,7 +599,7 @@ export class HostSessionRegistry {
|
||||
}
|
||||
// Already admitted attachments may finish a regional drain, but never a retired generation.
|
||||
if (
|
||||
this.draining ||
|
||||
this.drainSentHosts.has(identity.relayHostId) ||
|
||||
this.sessions.get(this.key(identity.userId, identity.relayHostId)) !== session ||
|
||||
this.get(identity)?.state === 'closed' ||
|
||||
!session.activeConnIds.has(connId) ||
|
||||
@@ -846,17 +853,49 @@ export class HostSessionRegistry {
|
||||
return { controls, splices, pendingSplices }
|
||||
}
|
||||
|
||||
drain(graceMs: number): void {
|
||||
drain(graceMs: number, options: { paceWindowMs?: number } = {}): void {
|
||||
this.draining = true
|
||||
for (const session of this.sessions.values()) {
|
||||
if (session.state === 'closed') continue
|
||||
session.authorityRevision += 1
|
||||
session.state = 'drain-only'
|
||||
if (session.socket) send(session.socket, 'drain', { graceMs, recovery: 'resolve-director' })
|
||||
setTimeout(() => this.closeDrainedSession(session), graceMs)
|
||||
// A later drain (an emergency one, or shutdown) owns every session again, so nothing
|
||||
// queued by an earlier paced drain may still fire: it would re-send and, worse, keep
|
||||
// the event loop alive for the rest of a window the operator just cut short.
|
||||
for (const timer of this.drainTimers) clearTimeout(timer)
|
||||
this.drainTimers.clear()
|
||||
const paceWindowMs = Math.max(0, Math.trunc(options.paceWindowMs ?? 0))
|
||||
const targets = [...this.sessions.values()].filter((session) => session.state !== 'closed')
|
||||
// The desktop re-dials the director as soon as it reads `drain`, whatever graceMs says,
|
||||
// so spreading the send is the only thing that spreads the reconnect load.
|
||||
const step = paceWindowMs > 0 && targets.length > 1 ? paceWindowMs / (targets.length - 1) : 0
|
||||
for (const [index, session] of targets.entries()) {
|
||||
const delay = Math.round(step * index)
|
||||
if (delay === 0) {
|
||||
this.sendDrain(session, graceMs)
|
||||
continue
|
||||
}
|
||||
this.scheduleDrainTimer(delay, () => this.sendDrain(session, graceMs))
|
||||
}
|
||||
}
|
||||
|
||||
// A session is only fenced when it is told, not when the drain starts: until its send
|
||||
// lands it is an ordinary live host, and its phones have to keep being able to reach it.
|
||||
private sendDrain(session: HostSession, graceMs: number): void {
|
||||
if (session.state === 'closed') return
|
||||
session.authorityRevision += 1
|
||||
session.state = 'drain-only'
|
||||
this.drainSentHosts.add(session.relayHostId)
|
||||
if (session.socket) send(session.socket, 'drain', { graceMs, recovery: 'resolve-director' })
|
||||
this.scheduleDrainTimer(graceMs, () => this.closeDrainedSession(session))
|
||||
}
|
||||
|
||||
// Unref'd so a drain in flight never holds the process open past its own work.
|
||||
private scheduleDrainTimer(delayMs: number, run: () => void): void {
|
||||
const timer: ReturnType<typeof setTimeout> = setTimeout(() => {
|
||||
this.drainTimers.delete(timer)
|
||||
run()
|
||||
}, delayMs)
|
||||
timer.unref?.()
|
||||
this.drainTimers.add(timer)
|
||||
}
|
||||
|
||||
drainHost(input: {
|
||||
attemptId: string
|
||||
userId: string
|
||||
|
||||
@@ -697,6 +697,78 @@ async function postPath(
|
||||
})
|
||||
}
|
||||
|
||||
describe('cell drain endpoint pacing', () => {
|
||||
function appWithDrain(): {
|
||||
app: ReturnType<typeof createRelayApp>
|
||||
drain: ReturnType<typeof vi.fn>
|
||||
} {
|
||||
const drain = vi.fn()
|
||||
const app = createRelayApp(config(), {
|
||||
store: {} as never,
|
||||
assignments: {} as never,
|
||||
drain,
|
||||
cellIncarnation,
|
||||
ready: vi.fn(async () => true)
|
||||
} as Parameters<typeof createRelayApp>[1])
|
||||
return { app, drain }
|
||||
}
|
||||
|
||||
it('drains everything at once when the caller asks for no pacing', async () => {
|
||||
const { app, drain } = appWithDrain()
|
||||
const response = await postPath(app, '/v1/admin/drain', 'deploy-token', { v: 1, graceMs: 0 })
|
||||
expect(response.status).toBe(200)
|
||||
expect(await response.json()).toEqual({ ok: true, paceWindowMs: 0 })
|
||||
expect(drain).toHaveBeenCalledWith(0, { paceWindowMs: 0 })
|
||||
})
|
||||
|
||||
it('passes the requested window through and echoes what it accepted', async () => {
|
||||
const { app, drain } = appWithDrain()
|
||||
const response = await postPath(app, '/v1/admin/drain', 'deploy-token', {
|
||||
v: 1,
|
||||
graceMs: 0,
|
||||
paceWindowMs: 120_000
|
||||
})
|
||||
expect(response.status).toBe(200)
|
||||
expect(await response.json()).toEqual({ ok: true, paceWindowMs: 120_000 })
|
||||
expect(drain).toHaveBeenCalledWith(0, { paceWindowMs: 120_000 })
|
||||
})
|
||||
|
||||
it('refuses a window that is negative, fractional, or past the cap', async () => {
|
||||
for (const paceWindowMs of [-1, 1.5, 300_001]) {
|
||||
const { app, drain } = appWithDrain()
|
||||
const response = await postPath(app, '/v1/admin/drain', 'deploy-token', {
|
||||
v: 1,
|
||||
graceMs: 0,
|
||||
paceWindowMs
|
||||
})
|
||||
expect(response.status).toBe(400)
|
||||
expect(drain).not.toHaveBeenCalled()
|
||||
}
|
||||
})
|
||||
|
||||
it('accepts the cap itself', async () => {
|
||||
const { app, drain } = appWithDrain()
|
||||
const response = await postPath(app, '/v1/admin/drain', 'deploy-token', {
|
||||
v: 1,
|
||||
graceMs: 0,
|
||||
paceWindowMs: 300_000
|
||||
})
|
||||
expect(response.status).toBe(200)
|
||||
expect(drain).toHaveBeenCalledWith(0, { paceWindowMs: 300_000 })
|
||||
})
|
||||
|
||||
it('still rejects an unauthenticated pacing request', async () => {
|
||||
const { app, drain } = appWithDrain()
|
||||
const response = await postPath(app, '/v1/admin/drain', 'wrong-token', {
|
||||
v: 1,
|
||||
graceMs: 0,
|
||||
paceWindowMs: 120_000
|
||||
})
|
||||
expect(response.status).toBe(401)
|
||||
expect(drain).not.toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
|
||||
function config(overrides: Partial<RelayConfig> = {}): RelayConfig {
|
||||
return {
|
||||
port: 8080,
|
||||
|
||||
@@ -136,7 +136,7 @@ export function createRelayServer(
|
||||
const app = createRelayApp(config, {
|
||||
store,
|
||||
assignments,
|
||||
drain: (graceMs) => sessions.drain(graceMs),
|
||||
drain: (graceMs, options) => sessions.drain(graceMs, options ?? {}),
|
||||
drainHost: (input) => sessions.drainHost(input),
|
||||
idleRehome: (input) => {
|
||||
const now = (options.now ?? Date.now)()
|
||||
|
||||
@@ -35,6 +35,9 @@ function cellOrigin(cellId) {
|
||||
// The same-cap roll covers the Asia cells the US-only capacity rollout never touches.
|
||||
const APPROVED_CELL_LISTS = { 'same-cap': SAME_CAP_CELLS }
|
||||
|
||||
// Matches the cell's own cap on /v1/admin/drain.
|
||||
const MAX_PACE_WINDOW_MS = 5 * 60 * 1_000
|
||||
|
||||
export function parseProductionCapacityCellArguments(argv) {
|
||||
const values = {}
|
||||
for (let index = 0; index < argv.length; index += 2) {
|
||||
@@ -64,11 +67,22 @@ export function parseProductionCapacityCellArguments(argv) {
|
||||
) {
|
||||
throw new Error('production capacity target origin is not exact')
|
||||
}
|
||||
const paceWindowMs = values['pace-window-ms'] === undefined
|
||||
? 0
|
||||
: Number(values['pace-window-ms'])
|
||||
if (
|
||||
!Number.isSafeInteger(paceWindowMs) ||
|
||||
paceWindowMs < 0 ||
|
||||
paceWindowMs > MAX_PACE_WINDOW_MS
|
||||
) {
|
||||
throw new Error('--pace-window-ms must be an integer between 0 and 300000')
|
||||
}
|
||||
return {
|
||||
directorOrigin: DIRECTOR_ORIGIN,
|
||||
cellOrigin: expectedCellOrigin,
|
||||
cellId,
|
||||
mode: values.mode
|
||||
mode: values.mode,
|
||||
paceWindowMs
|
||||
}
|
||||
}
|
||||
|
||||
@@ -82,24 +96,39 @@ export async function prepareProductionCapacityCell(config, overrides = {}) {
|
||||
const fetchImpl = overrides.fetch ?? fetch
|
||||
const token = overrides.token ?? process.env.ORCA_RELAY_ADMIN_ID_TOKEN
|
||||
if (!token || token.length > 8_192) throw new Error('admin identity token is unavailable')
|
||||
const postAt = async (origin, path, body) =>
|
||||
await responseJson(
|
||||
await fetchAdminOnceMore(
|
||||
fetchImpl,
|
||||
`${origin}${path}`,
|
||||
{
|
||||
method: 'POST',
|
||||
headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' },
|
||||
body: JSON.stringify(body)
|
||||
},
|
||||
{ wait: overrides.wait }
|
||||
),
|
||||
path
|
||||
const postRaw = async (origin, path, body) =>
|
||||
await fetchAdminOnceMore(
|
||||
fetchImpl,
|
||||
`${origin}${path}`,
|
||||
{
|
||||
method: 'POST',
|
||||
headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' },
|
||||
body: JSON.stringify(body)
|
||||
},
|
||||
{ wait: overrides.wait }
|
||||
)
|
||||
const postAt = async (origin, path, body) =>
|
||||
await responseJson(await postRaw(origin, path, body), path)
|
||||
const post = async (path, body) => await postAt(config.directorOrigin, path, body)
|
||||
if (config.mode === 'drain') {
|
||||
const paceWindowMs = config.paceWindowMs ?? 0
|
||||
if (paceWindowMs > 0) {
|
||||
const paced = await postRaw(config.cellOrigin, '/v1/admin/drain', {
|
||||
v: 1,
|
||||
graceMs: 0,
|
||||
paceWindowMs
|
||||
})
|
||||
if (paced.ok) {
|
||||
await paced.json().catch(() => ({}))
|
||||
return { changed: false, drained: true, paceWindowMs }
|
||||
}
|
||||
// A cell still on an image without paced drain rejects the unknown field outright.
|
||||
// An unpaced drain is the behaviour that cell already has, so fall back to it.
|
||||
if (paced.status !== 400) throw new Error(`/v1/admin/drain returned ${paced.status}`)
|
||||
await paced.json().catch(() => ({}))
|
||||
}
|
||||
await postAt(config.cellOrigin, '/v1/admin/drain', { v: 1, graceMs: 0 })
|
||||
return { changed: false, drained: true }
|
||||
return { changed: false, drained: true, paceWindowMs: 0 }
|
||||
}
|
||||
const before = await inspectAdmissionSelector(post)
|
||||
const state = selectorCellState(before.selector, config.cellId)
|
||||
|
||||
@@ -90,7 +90,8 @@ describe('production Relay capacity cell admission', () => {
|
||||
directorOrigin: 'https://relay.onorca.dev',
|
||||
cellOrigin: 'https://c7.relay.onorca.dev',
|
||||
cellId: 'production-gce-c7',
|
||||
mode: 'isolate'
|
||||
mode: 'isolate',
|
||||
paceWindowMs: 0
|
||||
})
|
||||
assert.throws(() => parseProductionCapacityCellArguments([
|
||||
'--director-origin', 'https://relay.onorca.dev',
|
||||
@@ -125,7 +126,8 @@ describe('production Relay capacity cell admission', () => {
|
||||
directorOrigin: 'https://relay.onorca.dev',
|
||||
cellOrigin: `https://${hostname}.relay.onorca.dev`,
|
||||
cellId,
|
||||
mode: 'isolate'
|
||||
mode: 'isolate',
|
||||
paceWindowMs: 0
|
||||
})
|
||||
}
|
||||
for (const cellId of ['production-gce-c17', 'production-gce-c18', 'production-gce-c30']) {
|
||||
@@ -168,13 +170,77 @@ describe('production Relay capacity cell admission', () => {
|
||||
{ ...config, mode: 'drain' },
|
||||
{ fetch: fake.fetch, token: 'token' }
|
||||
)
|
||||
assert.deepEqual(result, { changed: false, drained: true })
|
||||
assert.deepEqual(result, { changed: false, drained: true, paceWindowMs: 0 })
|
||||
assert.deepEqual(fake.calls, [{
|
||||
path: '/v1/admin/drain',
|
||||
body: { v: 1, graceMs: 0 }
|
||||
}])
|
||||
})
|
||||
|
||||
it('paces the drain send when the roll asks for a window', async () => {
|
||||
const fake = canaryFetch()
|
||||
const result = await prepareProductionCapacityCell(
|
||||
{ ...config, mode: 'drain', paceWindowMs: 120_000 },
|
||||
{ fetch: fake.fetch, token: 'token' }
|
||||
)
|
||||
assert.deepEqual(result, { changed: false, drained: true, paceWindowMs: 120_000 })
|
||||
assert.deepEqual(fake.calls, [{
|
||||
path: '/v1/admin/drain',
|
||||
body: { v: 1, graceMs: 0, paceWindowMs: 120_000 }
|
||||
}])
|
||||
})
|
||||
|
||||
it('drains unpaced when the cell image rejects the pacing field', async () => {
|
||||
const bodies = []
|
||||
const result = await prepareProductionCapacityCell(
|
||||
{ ...config, mode: 'drain', paceWindowMs: 120_000 },
|
||||
{
|
||||
token: 'token',
|
||||
wait: async () => {},
|
||||
fetch: async (url, init) => {
|
||||
assert.equal(new URL(url).pathname, '/v1/admin/drain')
|
||||
const body = JSON.parse(init.body)
|
||||
bodies.push(body)
|
||||
if (body.paceWindowMs !== undefined) return response({ error: 'invalid_request' }, 400)
|
||||
return response({ v: 1, draining: true })
|
||||
}
|
||||
}
|
||||
)
|
||||
assert.deepEqual(result, { changed: false, drained: true, paceWindowMs: 0 })
|
||||
assert.deepEqual(bodies, [
|
||||
{ v: 1, graceMs: 0, paceWindowMs: 120_000 },
|
||||
{ v: 1, graceMs: 0 }
|
||||
])
|
||||
})
|
||||
|
||||
it('fails a paced drain that the cell rejects for any other reason', async () => {
|
||||
await assert.rejects(
|
||||
prepareProductionCapacityCell(
|
||||
{ ...config, mode: 'drain', paceWindowMs: 120_000 },
|
||||
{
|
||||
token: 'token',
|
||||
wait: async () => {},
|
||||
fetch: async () => response({ error: 'invalid_token' }, 401)
|
||||
}
|
||||
),
|
||||
/returned 401/
|
||||
)
|
||||
})
|
||||
|
||||
it('refuses a pacing window that is not a bounded integer', () => {
|
||||
const argv = (value) => [
|
||||
'--director-origin', 'https://relay.onorca.dev',
|
||||
'--cell-origin', 'https://c26.relay.onorca.dev',
|
||||
'--cell-id', 'production-gce-c26',
|
||||
'--mode', 'drain',
|
||||
'--pace-window-ms', value
|
||||
]
|
||||
for (const value of ['-1', '300001', '1.5', 'soon']) {
|
||||
assert.throws(() => parseProductionCapacityCellArguments(argv(value)), /pace-window-ms/)
|
||||
}
|
||||
assert.equal(parseProductionCapacityCellArguments(argv('300000')).paceWindowMs, 300_000)
|
||||
})
|
||||
|
||||
it('restores only the selected cell to general admission', async () => {
|
||||
const fake = canaryFetch()
|
||||
await prepareProductionCapacityCell(
|
||||
@@ -228,7 +294,7 @@ describe('production Relay capacity cell admission', () => {
|
||||
}
|
||||
)
|
||||
assert.equal(calls, 2)
|
||||
assert.deepEqual(result, { changed: false, drained: true })
|
||||
assert.deepEqual(result, { changed: false, drained: true, paceWindowMs: 0 })
|
||||
})
|
||||
|
||||
it('fails when both drain attempts return a transient 503', async () => {
|
||||
|
||||
@@ -159,7 +159,8 @@ describe('same-cap roll scripts accept every same-cap cell', () => {
|
||||
directorOrigin: 'https://relay.onorca.dev',
|
||||
cellOrigin: `https://${hostname(cellId)}.relay.onorca.dev`,
|
||||
cellId,
|
||||
mode
|
||||
mode,
|
||||
paceWindowMs: 0
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -192,6 +193,14 @@ describe('same-cap roll scripts accept every same-cap cell', () => {
|
||||
}
|
||||
})
|
||||
|
||||
it('paces the drain it sends to the selected cell', () => {
|
||||
const drain = workflow.split('--mode drain')[1] ?? ''
|
||||
assert.match(drain.split('\n').slice(0, 2).join(' '), /--pace-window-ms "\$\{DRAIN_PACE_WINDOW_MS\}"/)
|
||||
assert.match(workflow, /DRAIN_PACE_WINDOW_MS: '120000'/)
|
||||
// The transition wait has to outlast the pacing window on top of the leases it waits on.
|
||||
assert.match(workflow, /--activity restart-safe[\s\S]*?--timeout-ms 1020000/)
|
||||
})
|
||||
|
||||
it('passes this cell\'s rehome protocol and pool on every plan validation the job runs', () => {
|
||||
const invocations = workflow.split('validate-relay-capacity-plan.mjs').slice(1)
|
||||
assert.equal(invocations.length, 2)
|
||||
|
||||
Reference in New Issue
Block a user