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)