From 0e7948fa6dca7aec99dbce94d3472c1370b95799 Mon Sep 17 00:00:00 2001 From: Jinwoo Hong <73622457+Jinwoo-H@users.noreply.github.com> Date: Thu, 17 Sep 2026 16:45:25 -0400 Subject: [PATCH] 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. --- ...d-deploy-relay-production-same-cap-job.yml | 9 +- cloud/apps/relay/src/app.ts | 16 ++- .../src/host-session-client-accept.test.ts | 70 +++++++++ .../relay/src/host-session-registry.test.ts | 133 +++++++++++++++++- cloud/apps/relay/src/host-session-registry.ts | 59 ++++++-- .../relay/src/regional-host-drain-app.test.ts | 72 ++++++++++ cloud/apps/relay/src/relay-server.ts | 2 +- ...epare-relay-production-capacity-canary.mjs | 59 ++++++-- ...-relay-production-capacity-canary.test.mjs | 74 +++++++++- .../relay-same-cap-script-census.test.mjs | 11 +- 10 files changed, 462 insertions(+), 43 deletions(-) diff --git a/.github/workflows/cloud-deploy-relay-production-same-cap-job.yml b/.github/workflows/cloud-deploy-relay-production-same-cap-job.yml index fa2d55d3d40..f6a577d2b42 100644 --- a/.github/workflows/cloud-deploy-relay-production-same-cap-job.yml +++ b/.github/workflows/cloud-deploy-relay-production-same-cap-job.yml @@ -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' }} diff --git a/cloud/apps/relay/src/app.ts b/cloud/apps/relay/src/app.ts index d60e2fd264c..978e544f5ff 100644 --- a/cloud/apps/relay/src/app.ts +++ b/cloud/apps/relay/src/app.ts @@ -58,6 +58,8 @@ const RelayCellConnectionHardCapSchema = z.custom( 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) { 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 0b7febaa7b4..7806e8d25b3 100644 --- a/cloud/apps/relay/src/host-session-client-accept.test.ts +++ b/cloud/apps/relay/src/host-session-client-accept.test.ts @@ -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 + 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, hostId: string): Promise { + 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)) + } + }) +}) diff --git a/cloud/apps/relay/src/host-session-registry.test.ts b/cloud/apps/relay/src/host-session-registry.test.ts index 7c50285454c..7fcff6a79f4 100644 --- a/cloud/apps/relay/src/host-session-registry.test.ts +++ b/cloud/apps/relay/src/host-session-registry.test.ts @@ -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() + .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) + }) +}) diff --git a/cloud/apps/relay/src/host-session-registry.ts b/cloud/apps/relay/src/host-session-registry.ts index 0d0cf4a7940..35e3935fb7d 100644 --- a/cloud/apps/relay/src/host-session-registry.ts +++ b/cloud/apps/relay/src/host-session-registry.ts @@ -179,6 +179,10 @@ export class HostSessionRegistry { private readonly hostCloseReasons = new HostCloseReasonMemory(() => this.now()) private readonly hostCapabilities = new WeakMap>() private draining = false + private readonly drainTimers = new Set>() + // 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() private readonly idleWork = new Map() private readonly idleAttempts = new Map< @@ -330,7 +334,10 @@ export class HostSessionRegistry { credential: string, capacityReservation?: PendingHostDataReservation ): Promise { - 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 = setTimeout(() => { + this.drainTimers.delete(timer) + run() + }, delayMs) + timer.unref?.() + this.drainTimers.add(timer) + } + drainHost(input: { attemptId: string userId: string diff --git a/cloud/apps/relay/src/regional-host-drain-app.test.ts b/cloud/apps/relay/src/regional-host-drain-app.test.ts index cf7e3798155..38ebdd3f709 100644 --- a/cloud/apps/relay/src/regional-host-drain-app.test.ts +++ b/cloud/apps/relay/src/regional-host-drain-app.test.ts @@ -697,6 +697,78 @@ async function postPath( }) } +describe('cell drain endpoint pacing', () => { + function appWithDrain(): { + app: ReturnType + drain: ReturnType + } { + const drain = vi.fn() + const app = createRelayApp(config(), { + store: {} as never, + assignments: {} as never, + drain, + cellIncarnation, + ready: vi.fn(async () => true) + } as Parameters[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 { return { port: 8080, diff --git a/cloud/apps/relay/src/relay-server.ts b/cloud/apps/relay/src/relay-server.ts index 50077b41116..3a0668bbcef 100644 --- a/cloud/apps/relay/src/relay-server.ts +++ b/cloud/apps/relay/src/relay-server.ts @@ -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)() diff --git a/cloud/dev/scripts/prepare-relay-production-capacity-canary.mjs b/cloud/dev/scripts/prepare-relay-production-capacity-canary.mjs index 7967d164e4b..9787dbdd62a 100644 --- a/cloud/dev/scripts/prepare-relay-production-capacity-canary.mjs +++ b/cloud/dev/scripts/prepare-relay-production-capacity-canary.mjs @@ -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) diff --git a/cloud/dev/scripts/prepare-relay-production-capacity-canary.test.mjs b/cloud/dev/scripts/prepare-relay-production-capacity-canary.test.mjs index 274a60d2198..c566229b678 100644 --- a/cloud/dev/scripts/prepare-relay-production-capacity-canary.test.mjs +++ b/cloud/dev/scripts/prepare-relay-production-capacity-canary.test.mjs @@ -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 () => { diff --git a/cloud/dev/scripts/relay-same-cap-script-census.test.mjs b/cloud/dev/scripts/relay-same-cap-script-census.test.mjs index 770e20b6353..bd49f7600a7 100644 --- a/cloud/dev/scripts/relay-same-cap-script-census.test.mjs +++ b/cloud/dev/scripts/relay-same-cap-script-census.test.mjs @@ -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)