diff --git a/api/Dockerfile b/api/Dockerfile index e7dbcc4e..c065c4a2 100644 --- a/api/Dockerfile +++ b/api/Dockerfile @@ -255,6 +255,7 @@ RUN dnf install -y --setopt=install_weak_deps=False \ libkrun \ socat \ curl \ + openssl \ util-linux \ grep \ gawk \ @@ -268,8 +269,9 @@ COPY --from=nsjail-builder /usr/local/bin/sandbox-rootfs-setup /sandbox-rootfs-s COPY launcher/entrypoint.sh /usr/local/bin/launcher-entrypoint.sh COPY docker/start-direct-sandbox.sh /usr/local/bin/start-direct-sandbox.sh COPY docker/sandbox-runner-healthcheck.sh /usr/local/bin/sandbox-runner-healthcheck.sh +COPY docker/sandbox-runner-drain.sh /usr/local/bin/sandbox-runner-drain.sh COPY docker/sandbox-entrypoint.sh /usr/local/bin/entrypoint.sh -RUN chmod +x /usr/local/bin/launcher-entrypoint.sh /usr/local/bin/start-direct-sandbox.sh /usr/local/bin/sandbox-runner-healthcheck.sh /usr/local/bin/entrypoint.sh +RUN chmod +x /usr/local/bin/launcher-entrypoint.sh /usr/local/bin/start-direct-sandbox.sh /usr/local/bin/sandbox-runner-healthcheck.sh /usr/local/bin/sandbox-runner-drain.sh /usr/local/bin/entrypoint.sh EXPOSE 2000/tcp ENTRYPOINT ["/usr/local/bin/entrypoint.sh"] diff --git a/api/src/api/drain.admission.test.ts b/api/src/api/drain.admission.test.ts new file mode 100644 index 00000000..b99458cb --- /dev/null +++ b/api/src/api/drain.admission.test.ts @@ -0,0 +1,133 @@ +import { afterEach, describe, expect, test } from 'bun:test'; +import express from 'express'; +import { request, type Server } from 'node:http'; +import { drainSandbox, resetSandboxDrainForTests, sandboxDrainSnapshot } from '../drain'; +import { withSandboxDrain } from './drain'; + +afterEach(resetSandboxDrainForTests); + +async function serve(handler: Parameters[1]) { + const app = express(); + app.post('/execute', withSandboxDrain(express.json(), handler)); + app.use((error: Error, _req: express.Request, res: express.Response, _next: express.NextFunction) => { + res.status(400).json({ error: error.message }); + }); + const server: Server = app.listen(0, '127.0.0.1'); + await new Promise((resolve) => server.once('listening', resolve)); + const address = server.address(); + if (!address || typeof address === 'string') throw new Error('Missing test port'); + return { server, url: `http://127.0.0.1:${address.port}/execute` }; +} + +async function close(server: Server) { + server.closeAllConnections(); + await new Promise((resolve) => server.close(() => resolve())); +} + +async function until(check: () => boolean) { + for (let i = 0; i < 100; i++) { + if (check()) return; + await Bun.sleep(2); + } + throw new Error('Condition did not become true'); +} + +describe('drain admission across request lifecycle', () => { + test('waits for response completion after the async handler returns', async () => { + let returned = false; + let finishResponse: (() => void) | undefined; + const { server, url } = await serve(async (_req, res) => { + res.write('partial result'); + finishResponse = () => res.end(' complete'); + returned = true; + }); + const pendingResponse = fetch(url, { + method: 'POST', headers: { 'Content-Type': 'application/json' }, body: '{}', + }); + try { + await until(() => returned); + let drained = false; + const pendingDrain = drainSandbox(1000, 0).then(result => { drained = true; return result; }); + await Bun.sleep(10); + expect(sandboxDrainSnapshot().activeExecutions).toBe(1); + expect(drained).toBe(false); + finishResponse?.(); + expect(await (await pendingResponse).text()).toBe('partial result complete'); + expect((await pendingDrain).drained).toBe(true); + } finally { + finishResponse?.(); + await close(server); + } + }); + + test('counts an unfinished upload before JSON parsing and waits for its handler', async () => { + let entered = false; + let finish!: () => void; + const pendingHandler = new Promise((resolve) => { finish = resolve; }); + const { server, url } = await serve(async (_req, res) => { + entered = true; + await pendingHandler; + res.json({ ok: true }); + }); + const req = request(url, { method: 'POST', headers: { 'Content-Type': 'application/json' } }); + req.on('error', () => {}); + try { + req.write('{'); + await until(() => sandboxDrainSnapshot().activeExecutions === 1); + let drained = false; + const pending = drainSandbox(1000, 0).then(() => { drained = true; }); + await Bun.sleep(5); + expect(entered).toBe(false); + expect(drained).toBe(false); + req.end('}'); + await until(() => entered); + // Losing the response does not stop an already admitted execution or its cleanup. + req.destroy(); + await Bun.sleep(5); + expect(sandboxDrainSnapshot().activeExecutions).toBe(1); + expect(drained).toBe(false); + finish(); + await pending; + expect(sandboxDrainSnapshot().activeExecutions).toBe(0); + } finally { + finish(); + req.destroy(); + await close(server); + } + }); + + test('releases an aborted upload without invoking the handler', async () => { + let calls = 0; + const { server, url } = await serve(async (_req, res) => { calls++; res.end(); }); + const req = request(url, { method: 'POST', headers: { 'Content-Type': 'application/json' } }); + req.on('error', () => {}); + try { + req.write('{'); + await until(() => sandboxDrainSnapshot().activeExecutions === 1); + const pending = drainSandbox(1000, 0); + req.destroy(); + expect((await pending).drained).toBe(true); + expect(calls).toBe(0); + } finally { + req.destroy(); + await close(server); + } + }); + + test('parser failures release admission and drain rejection happens before parsing', async () => { + let calls = 0; + const { server, url } = await serve(async (_req, res) => { calls++; res.end(); }); + try { + const bad = await fetch(url, { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: '{' }); + expect(bad.status).toBe(400); + expect(sandboxDrainSnapshot().activeExecutions).toBe(0); + await drainSandbox(1000, 0); + const rejected = await fetch(url, { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: '{' }); + expect(rejected.status).toBe(503); + expect((await rejected.json()).error).toBe('runner_draining'); + expect(calls).toBe(0); + } finally { + await close(server); + } + }); +}); diff --git a/api/src/api/drain.routes.test.ts b/api/src/api/drain.routes.test.ts new file mode 100644 index 00000000..12ec7afb --- /dev/null +++ b/api/src/api/drain.routes.test.ts @@ -0,0 +1,115 @@ +import { afterAll, afterEach, beforeAll, describe, expect, test } from 'bun:test'; +import express from 'express'; +import type { Server } from 'node:http'; +import { generateKeyPairSync, sign } from 'node:crypto'; +import { mkdtemp, readFile, rm } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { config } from '../config'; +import { resetSandboxDrainForTests, sandboxDrainSnapshot, startSandboxDrain } from '../drain'; +import { validateDrainConfiguration } from './drain'; +import { RUNNER_DRAIN_MESSAGE } from '../../../shared/drain-contract'; +import v2Router from './v2'; + +let server: Server; +let baseUrl: string; +const originalMinDelay = config.drain_min_delay_ms; +const originalEnabled = config.graceful_drain_enabled; +const originalPublicKey = config.drain_public_key; +let stateDirectory: string; +const hostScript = join(__dirname, '../../../docker/sandbox-runner-drain.sh'); + +async function runHostHook(mode?: string): Promise { + const child = Bun.spawn(['bash', hostScript, ...(mode ? [mode] : [])], { + env: { ...process.env, SANDBOX_DRAIN_STATE_DIR: stateDirectory, PORT: new URL(baseUrl || 'http://localhost:2000').port }, + stdout: 'pipe', stderr: 'pipe', + }); + const output = await new Response(child.stdout).text(); + const error = await new Response(child.stderr).text(); + if (await child.exited !== 0) throw new Error(error); + return output; +} + +beforeAll(async () => { + config.drain_min_delay_ms = 0; + config.graceful_drain_enabled = true; + stateDirectory = await mkdtemp(join(tmpdir(), 'runner-drain-')); + await runHostHook('init'); + config.drain_public_key = await readFile(join(stateDirectory, 'public.b64'), 'utf8'); + validateDrainConfiguration(); + const app = express(); + app.use('/api/v2', v2Router); + server = app.listen(0, '127.0.0.1'); + await new Promise((resolve) => server.once('listening', resolve)); + const address = server.address(); + if (!address || typeof address === 'string') throw new Error('Expected TCP test server'); + baseUrl = `http://127.0.0.1:${address.port}/api/v2`; +}); + +afterEach(() => { + resetSandboxDrainForTests(); + config.graceful_drain_enabled = true; + config.drain_min_delay_ms = 0; +}); + +afterAll(async () => { + config.drain_min_delay_ms = originalMinDelay; + config.graceful_drain_enabled = originalEnabled; + config.drain_public_key = originalPublicKey; + resetSandboxDrainForTests(); + await new Promise((resolve, reject) => server.close((error) => error ? reject(error) : resolve())); + await rm(stateDirectory, { recursive: true, force: true }); +}); + +describe('sandbox drain routes', () => { + test('does not activate drain when the feature is disabled', async () => { + config.graceful_drain_enabled = false; + expect((await fetch(`${baseUrl}/drain`, { method: 'POST' })).status).toBe(404); + const response = await fetch(`${baseUrl}/execute`, { + method: 'POST', headers: { 'Content-Type': 'application/json' }, body: '{}', + }); + expect(response.status).not.toBe(503); + }); + test('rejects unsigned and foreign-pod drain requests without changing admission', async () => { + expect((await fetch(`${baseUrl}/drain`)).status).toBe(404); + expect((await fetch(`${baseUrl}/drain`, { method: 'POST' })).status).toBe(403); + const foreign = generateKeyPairSync('ed25519'); + const signature = sign(null, Buffer.from(RUNNER_DRAIN_MESSAGE), foreign.privateKey).toString('base64'); + expect((await fetch(`${baseUrl}/drain`, { + method: 'POST', headers: { 'X-CodeAPI-Drain-Signature': signature }, + })).status).toBe(403); + expect(sandboxDrainSnapshot().draining).toBe(false); + }); + test('advertises enough delay for endpoint propagation and validates the retry bound', async () => { + config.drain_min_delay_ms = 5000; + validateDrainConfiguration(); + startSandboxDrain(); + const response = await fetch(`${baseUrl}/execute`, { + method: 'POST', headers: { 'Content-Type': 'application/json' }, body: '{}', + }); + expect(response.headers.get('retry-after')).toBe('5'); + config.drain_min_delay_ms = 5001; + expect(validateDrainConfiguration).toThrow('5-second'); + }); + test('fails readiness and rejects execution before validation after drain begins', async () => { + // Run the actual host signer and curl hook against the real guest route. + expect(JSON.parse(await runHostHook())).toEqual({ status: 'drained', activeExecutions: 0 }); + + const readyResponse = await fetch(`${baseUrl}/ready`); + expect(readyResponse.status).toBe(503); + expect(await readyResponse.json()).toEqual({ status: 'draining', activeExecutions: 0 }); + + const executeResponse = await fetch(`${baseUrl}/execute`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: '{}', + }); + expect(executeResponse.status).toBe(503); + expect(executeResponse.headers.get('retry-after')).toBe(String(config.drain_retry_after_seconds)); + expect(executeResponse.headers.get('connection')).toBe('close'); + expect(await executeResponse.json()).toEqual({ + error: 'runner_draining', + message: 'Sandbox runner is draining for shutdown', + }); + }); +}); diff --git a/api/src/api/drain.ts b/api/src/api/drain.ts new file mode 100644 index 00000000..beb21e3f --- /dev/null +++ b/api/src/api/drain.ts @@ -0,0 +1,147 @@ +import express, { type Request, type Response, type RequestHandler } from 'express'; +import { createPublicKey, verify } from 'node:crypto'; +import { MAX_RUNNER_DRAIN_RETRY_DELAY_MS, RUNNER_DRAIN_MESSAGE } from '../../../shared/drain-contract'; +import { config } from '../config'; +import { logger } from '../logger'; +import { + sandboxDrainSnapshot, + admitSandboxExecution, + drainSandbox, +} from '../drain'; +import { checkSandboxWorkspaceHealth } from '../workspace-isolation'; + +const router = express.Router(); + +export function validateDrainConfiguration(): void { + if (!config.graceful_drain_enabled) return; + if (config.drain_min_delay_ms > config.drain_timeout_ms || config.drain_timeout_ms > 2_147_483_647) { + throw new Error('Drain timeout must cover propagation and fit the runtime timer bound'); + } + if (config.drain_min_delay_ms > MAX_RUNNER_DRAIN_RETRY_DELAY_MS + || config.drain_retry_after_seconds * 1000 > MAX_RUNNER_DRAIN_RETRY_DELAY_MS) { + throw new Error('Drain propagation and retry delays must not exceed the 5-second transport retry bound'); + } + const key = createPublicKey({ key: Buffer.from(config.drain_public_key, 'base64'), format: 'der', type: 'spki' }); + if (key.asymmetricKeyType !== 'ed25519') throw new Error('Drain requires an Ed25519 public verifier'); +} + +/** Count uploads before the parser runs, and hold admission through handler cleanup. + * A disconnect during execution is not a release: the job may still have side effects. */ +export function withSandboxDrain( + parse: RequestHandler, + handler: (req: Request, res: Response) => Promise, +): RequestHandler { + return (req, res, next) => { + const release = admitSandboxExecution(); + if (!release) { + res.set('Retry-After', String(Math.max(config.drain_retry_after_seconds, Math.ceil(config.drain_min_delay_ms / 1000)))); + // A pooled connection stays pinned to this pod after Service endpoint removal. + res.set('Connection', 'close'); + res.status(503).json({ + error: 'runner_draining', + message: 'Sandbox runner is draining for shutdown', + }); + return; + } + let uploadAborted = false; + let handlerComplete = false; + let responseComplete = res.writableFinished || res.destroyed; + const maybeRelease = (): void => { + if (!handlerComplete || !responseComplete) return; + res.removeListener('finish', settledResponse); + res.removeListener('close', settledResponse); + release(); + }; + const settledResponse = (): void => { + responseComplete = true; + maybeRelease(); + }; + res.once('finish', settledResponse); + res.once('close', settledResponse); + const abortUpload = (): void => { + uploadAborted = true; + handlerComplete = true; + responseComplete = true; + maybeRelease(); + }; + req.once('aborted', abortUpload); + const parsed = (error?: unknown): void => { + req.removeListener('aborted', abortUpload); + if (error || uploadAborted) { + handlerComplete = true; + if (error) next(error); + maybeRelease(); + return; + } + Promise.resolve().then(() => handler(req, res)).catch(next).finally(() => { + handlerComplete = true; + maybeRelease(); + }); + }; + try { + parse(req, res, parsed); + } catch (error) { + parsed(error); + } + }; +} + +router.get('/ready', async (_req: Request, res: Response) => { + const snapshot = sandboxDrainSnapshot(); + if (snapshot.draining) { + return res.status(503).json({ status: 'draining', activeExecutions: snapshot.activeExecutions }); + } + + try { + const health = await checkSandboxWorkspaceHealth(); + const current = sandboxDrainSnapshot(); + if (current.draining) { + return res.status(503).json({ status: 'draining', activeExecutions: current.activeExecutions }); + } + return res.status(200).json(health); + } catch (error) { + logger.error({ err: error }, 'Sandbox workspace readiness check failed'); + return res.status(503).json({ + status: 'unhealthy', + error: 'workspace_unavailable', + message: 'Sandbox workspace root is unavailable', + }); + } +}); + +const drainHandler = async (_req: Request, res: Response): Promise => { + if (!config.graceful_drain_enabled) { + return res.status(404).json({ message: 'Not Found' }); + } + try { + const signature = _req.get('X-CodeAPI-Drain-Signature') ?? ''; + if (!/^[A-Za-z0-9+/]{86}==$/.test(signature) + || !verify(null, Buffer.from(RUNNER_DRAIN_MESSAGE), + createPublicKey({ key: Buffer.from(config.drain_public_key, 'base64'), format: 'der', type: 'spki' }), + Buffer.from(signature, 'base64'))) { + return res.status(403).json({ message: 'Invalid drain signature' }); + } + } catch { + return res.status(403).json({ message: 'Invalid drain signature' }); + } + const result = await drainSandbox(config.drain_timeout_ms, config.drain_min_delay_ms); + + if (!result.drained) { + logger.warn( + { activeExecutions: result.activeExecutions, timeoutMillis: config.drain_timeout_ms }, + 'Sandbox drain timed out with executions still active', + ); + } else { + logger.info('Sandbox drain completed'); + } + + return res.status(200).json({ + status: result.drained ? 'drained' : 'drain_timeout', + activeExecutions: result.activeExecutions, + }); +}; + +// Only the host has the per-pod private key; the guest receives a public verifier. +router.post('/drain', drainHandler); + +export default router; diff --git a/api/src/api/v2-cleanup.test.ts b/api/src/api/v2-cleanup.test.ts index d1efee5d..fd64a713 100644 --- a/api/src/api/v2-cleanup.test.ts +++ b/api/src/api/v2-cleanup.test.ts @@ -5,6 +5,7 @@ import type { Server } from 'http'; import { tmpdir } from 'os'; import { join } from 'path'; import { config } from '../config'; +import { drainSandbox, resetSandboxDrainForTests } from '../drain'; import { Job } from '../job'; import { loadPackage } from '../runtime'; import { ValidationError } from '../validation'; @@ -42,6 +43,7 @@ afterAll(async () => { await rm(directory, { recursive: true, force: true }); }); afterEach(() => { + resetSandboxDrainForTests(); Job.prototype.prime = originalPrime; Job.prototype.execute = originalExecute; Job.prototype.uploadGeneratedFiles = originalUpload; @@ -92,6 +94,8 @@ for (const outcome of [ events.push('cleanup-end'); }; let responseArrived = false; + let drainCompleted = false; + let pendingDrain: Promise | undefined; const pending = fetch(url, { method: 'POST', headers: { 'Content-Type': 'application/json' }, @@ -111,9 +115,14 @@ for (const outcome of [ throw new Error(`Response preceded cleanup: ${response.status}`); }), ]); + pendingDrain = drainSandbox(1000, 0).then(result => { + drainCompleted = true; + expect(result.drained).toBe(true); + }); // Give an incorrectly early res.json() time to reach the client. await new Promise(resolve => setTimeout(resolve, 30)); expect(responseArrived).toBe(false); + expect(drainCompleted).toBe(false); } finally { releaseCleanup(); const response = await pending; @@ -125,6 +134,7 @@ for (const outcome of [ : 500 ); await response.text(); + await pendingDrain; } expect(events.filter(event => event === 'cleanup-start')).toHaveLength(1); expect(events[events.length - 1]).toBe('cleanup-end'); diff --git a/api/src/api/v2.ts b/api/src/api/v2.ts index cba8e80a..4fc3a9c9 100644 --- a/api/src/api/v2.ts +++ b/api/src/api/v2.ts @@ -26,6 +26,7 @@ import { parseSessionBindingFromHeader, } from '../session-workspace'; import { streamSessionCheckpoint, restoreSessionCheckpoint } from '../session-checkpoint'; +import drainRouter, { withSandboxDrain } from './drain'; import { SESSION_INPUT_CACHE_MAX_OBJECTS, hasCachedInput, @@ -78,6 +79,9 @@ export function deduplicateFilesByDestination(files: TFile[]): TFile[] { } return [...byDestination.values()]; } +// Register lifecycle routes before the content-type guard so the signed +// preStop request can use a bodyless POST. +router.use(drainRouter); function existingDestinationConflictMessage(existing: string, destination: string): string { return existing === destination @@ -475,7 +479,7 @@ router.use('/execute', (_req: Request, res: Response, next: NextFunction) => { * `index.ts` for the routing-order rationale). Keep the limit configurable: * analyst workflows may send large scripts or replay payloads, while file * bytes should still move through the gateway/file path. */ -router.post('/execute', express.json({ limit: config.execute_body_limit }), async (req: Request, res: Response) => { +router.post('/execute', withSandboxDrain(express.json({ limit: config.execute_body_limit }), async (req: Request, res: Response) => { const started = performance.now(); let job: Job | undefined; let cleanedUp = false; @@ -683,7 +687,7 @@ router.post('/execute', express.json({ limit: config.execute_body_limit }), asyn durationSeconds: (performance.now() - started) / 1000, }); } -}); +})); router.post('/workspace/execute', express.json({ limit: '256kb' }), async (req: Request, res: Response) => { let binding; diff --git a/api/src/config.ts b/api/src/config.ts index a5d0c993..84d455d8 100644 --- a/api/src/config.ts +++ b/api/src/config.ts @@ -64,6 +64,11 @@ export const config = { compile_memory_limit: Number(process.env.SANDBOX_COMPILE_MEMORY_LIMIT ?? -1), run_memory_limit: Number(process.env.SANDBOX_RUN_MEMORY_LIMIT ?? -1), max_concurrent_jobs: safeInt(process.env.SANDBOX_MAX_CONCURRENT_JOBS, 8), + drain_timeout_ms: safeInt(process.env.SANDBOX_DRAIN_TIMEOUT_MS, 45_000), + graceful_drain_enabled: process.env.SANDBOX_GRACEFUL_DRAIN_ENABLED === 'true', + drain_public_key: process.env.SANDBOX_DRAIN_PUBLIC_KEY ?? '', + drain_min_delay_ms: safeInt(process.env.SANDBOX_DRAIN_MIN_DELAY_MS, 5_000, 0), + drain_retry_after_seconds: safeInt(process.env.SANDBOX_DRAIN_RETRY_AFTER_SECONDS, 1), per_job_uids: (process.env.SANDBOX_PER_JOB_UIDS ?? 'true') === 'true', /* Image-level enable for persistent session workspaces (stateful sessions). * Only the Lambda MicroVM runner target sets this true; the K8s diff --git a/api/src/drain.test.ts b/api/src/drain.test.ts new file mode 100644 index 00000000..dc4c6403 --- /dev/null +++ b/api/src/drain.test.ts @@ -0,0 +1,65 @@ +import { afterEach, describe, expect, test } from 'bun:test'; +import { + admitSandboxExecution, + drainSandbox, + resetSandboxDrainForTests, + sandboxDrainSnapshot, + startSandboxDrain, + waitForSandboxDrain, +} from './drain'; + +afterEach(() => resetSandboxDrainForTests()); + +describe('sandbox drain state', () => { + test('preStop and signal reuse one bounded drain including after timeout', async () => { + const release = admitSandboxExecution(); + const first = drainSandbox(5, 0); + expect(drainSandbox(1000, 1000)).toBe(first); + expect((await first).drained).toBe(false); + expect(drainSandbox(1000, 1000)).toBe(first); + release?.(); + }); + + test('waits for the endpoint-removal floor even with no active work', async () => { + let completed = false; + const pending = drainSandbox(1000, 20).then(() => { completed = true; }); + await Bun.sleep(5); + expect(completed).toBe(false); + await pending; + expect(completed).toBe(true); + }); + test('rejects new executions once draining starts', () => { + expect(startSandboxDrain()).toBe(true); + expect(startSandboxDrain()).toBe(false); + expect(admitSandboxExecution()).toBeUndefined(); + expect(sandboxDrainSnapshot()).toEqual({ draining: true, activeExecutions: 0 }); + }); + + test('waits for admitted executions and releases each execution once', async () => { + const releaseFirst = admitSandboxExecution(); + const releaseSecond = admitSandboxExecution(); + expect(releaseFirst).toBeFunction(); + expect(releaseSecond).toBeFunction(); + + startSandboxDrain(); + const pending = waitForSandboxDrain(1_000); + releaseFirst?.(); + releaseFirst?.(); + expect(sandboxDrainSnapshot().activeExecutions).toBe(1); + releaseSecond?.(); + + expect(await pending).toEqual({ draining: true, activeExecutions: 0, drained: true }); + }); + + test('reports active executions when the drain times out', async () => { + const release = admitSandboxExecution(); + startSandboxDrain(); + + expect(await waitForSandboxDrain(5)).toEqual({ + draining: true, + activeExecutions: 1, + drained: false, + }); + release?.(); + }); +}); diff --git a/api/src/drain.ts b/api/src/drain.ts new file mode 100644 index 00000000..80fbf8e1 --- /dev/null +++ b/api/src/drain.ts @@ -0,0 +1,85 @@ +type DrainWaiter = () => void; + +let draining = false; +let activeExecutions = 0; +const drainWaiters = new Set(); +let pendingDrain: Promise | undefined; + +export interface SandboxDrainSnapshot { + draining: boolean; + activeExecutions: number; +} + +export interface SandboxDrainResult extends SandboxDrainSnapshot { + drained: boolean; +} + +export function startSandboxDrain(): boolean { + if (draining) return false; + draining = true; + return true; +} + +export function sandboxDrainSnapshot(): SandboxDrainSnapshot { + return { draining, activeExecutions }; +} + +/** + * Admits an execution synchronously so a drain cannot begin between checking + * readiness and accounting for work. The returned release function is + * idempotent because error paths may converge in nested finally blocks. + */ +export function admitSandboxExecution(): (() => void) | undefined { + if (draining) return undefined; + activeExecutions += 1; + let released = false; + + return () => { + if (released) return; + released = true; + activeExecutions = Math.max(0, activeExecutions - 1); + if (activeExecutions === 0) { + for (const resolve of drainWaiters) resolve(); + drainWaiters.clear(); + } + }; +} + +export async function waitForSandboxDrain(timeoutMillis: number): Promise { + if (activeExecutions === 0) { + return { ...sandboxDrainSnapshot(), drained: true }; + } + + let timeout: ReturnType | undefined; + let resolveWaiter: DrainWaiter | undefined; + const drained = await new Promise((resolve) => { + resolveWaiter = () => resolve(true); + drainWaiters.add(resolveWaiter); + timeout = setTimeout(() => resolve(false), timeoutMillis); + (timeout as { unref?: () => void }).unref?.(); + }); + + if (resolveWaiter) drainWaiters.delete(resolveWaiter); + if (timeout) clearTimeout(timeout); + return { ...sandboxDrainSnapshot(), drained }; +} + +/** preStop and SIGTERM share a deadline; a timeout must not buy a second drain budget. */ +export function drainSandbox(timeoutMillis: number, minDelayMillis: number): Promise { + if (!pendingDrain) { + startSandboxDrain(); + pendingDrain = Promise.all([ + waitForSandboxDrain(timeoutMillis), + new Promise((resolve) => setTimeout(resolve, minDelayMillis)), + ]).then(([result]) => result); + } + return pendingDrain; +} + +export function resetSandboxDrainForTests(): void { + draining = false; + activeExecutions = 0; + pendingDrain = undefined; + for (const resolve of drainWaiters) resolve(); + drainWaiters.clear(); +} diff --git a/api/src/index.ts b/api/src/index.ts index 00d10e9d..cc6da827 100644 --- a/api/src/index.ts +++ b/api/src/index.ts @@ -10,6 +10,8 @@ import { startWarmupCommand } from './warmup'; import { hostedAppSupervisor, validateHostedAppStartup } from './hosted-app'; import v2Router from './api/v2'; import lifecycleRouter, { LIFECYCLE_HOOK_BASE_PATH } from './api/lifecycle'; +import { drainSandbox } from './drain'; +import { validateDrainConfiguration } from './api/drain'; const app = express(); @@ -66,6 +68,7 @@ app.use((err: HttpError, _req: express.Request, res: express.Response, _next: ex }); async function main(): Promise { + validateDrainConfiguration(); validateHostedAppStartup(); validateHardenedSandboxStartup(); await initializeSandboxWorkspaceIsolation(); @@ -110,6 +113,15 @@ async function main(): Promise { const shutdown = async (): Promise => { if (shuttingDown) return; shuttingDown = true; + if (config.graceful_drain_enabled) { + const drainResult = await drainSandbox(config.drain_timeout_ms, config.drain_min_delay_ms); + if (!drainResult.drained) { + logger.warn( + { activeExecutions: drainResult.activeExecutions, timeoutMillis: config.drain_timeout_ms }, + 'Signal-triggered sandbox drain timed out', + ); + } + } stopWorkspaceReaper(); await closeHttpServerWithTimeout(); await hostedAppSupervisor.shutdown().catch((err) => { diff --git a/api/src/metrics.ts b/api/src/metrics.ts index 328b4871..41966256 100644 --- a/api/src/metrics.ts +++ b/api/src/metrics.ts @@ -54,6 +54,8 @@ function routeLabel(req: Request): string { if (req.path === '/metrics') return '/metrics'; if (req.path === '/api/v2/execute') return '/api/v2/execute'; if (req.path === '/api/v2/health') return '/api/v2/health'; + if (req.path === '/api/v2/ready') return '/api/v2/ready'; + if (req.path === '/api/v2/drain') return '/api/v2/drain'; if (req.path === '/api/v2/runtimes') return '/api/v2/runtimes'; return 'unmatched'; } diff --git a/docker/sandbox-entrypoint.sh b/docker/sandbox-entrypoint.sh index 0ea2bb66..b11ef563 100644 --- a/docker/sandbox-entrypoint.sh +++ b/docker/sandbox-entrypoint.sh @@ -3,6 +3,11 @@ set -e KVM_ENABLED="${KVM_ENABLED:-true}" +if [[ "${SANDBOX_GRACEFUL_DRAIN_ENABLED:-false}" == true ]]; then + /usr/local/bin/sandbox-runner-drain.sh init + export SANDBOX_DRAIN_PUBLIC_KEY="$(cat "${SANDBOX_DRAIN_STATE_DIR:-/run/codeapi-drain}/public.b64")" +fi + if [ "$KVM_ENABLED" = "true" ]; then exec /usr/local/bin/launcher-entrypoint.sh "$@" fi diff --git a/docker/sandbox-runner-drain.sh b/docker/sandbox-runner-drain.sh new file mode 100644 index 00000000..da47b088 --- /dev/null +++ b/docker/sandbox-runner-drain.sh @@ -0,0 +1,20 @@ +#!/usr/bin/env bash +set -euo pipefail + +# This directory stays on the host; it is not mounted into the guest or jobs. +state_dir="${SANDBOX_DRAIN_STATE_DIR:-/run/codeapi-drain}" +if [[ "${1:-}" == init ]]; then + umask 077 + mkdir -p "$state_dir" + openssl genpkey -algorithm ED25519 -out "$state_dir/private.pem" + openssl pkey -in "$state_dir/private.pem" -pubout -outform DER | base64 | tr -d '\n' > "$state_dir/public.b64" + printf %s codeapi-runner-drain-v1 > "$state_dir/message" + exit 0 +fi + +timeout_ms="${SANDBOX_DRAIN_TIMEOUT_MS:-45000}" +[[ "$timeout_ms" =~ ^[1-9][0-9]{0,8}$ ]] || { echo 'Invalid SANDBOX_DRAIN_TIMEOUT_MS' >&2; exit 1; } +signature="$(openssl pkeyutl -sign -rawin -inkey "$state_dir/private.pem" -in "$state_dir/message" | base64 | tr -d '\n')" +exec curl --fail --silent --show-error --max-time "$(( (timeout_ms + 999) / 1000 + 5 ))" \ + -X POST -H "X-CodeAPI-Drain-Signature: $signature" \ + "http://127.0.0.1:${PORT:-2000}/api/v2/drain" diff --git a/docker/sandbox-runner-healthcheck.sh b/docker/sandbox-runner-healthcheck.sh index c639848b..73d6a273 100755 --- a/docker/sandbox-runner-healthcheck.sh +++ b/docker/sandbox-runner-healthcheck.sh @@ -8,6 +8,9 @@ timeout_seconds="${SANDBOX_RUNNER_HEALTHCHECK_TIMEOUT_SECONDS:-5}" manifest_clock_tolerance_seconds=30 port="${PORT:-2000}" url="${SANDBOX_RUNNER_HEALTHCHECK_URL:-http://127.0.0.1:${port}/api/v2/health}" +if [[ "${1:-}" == "--ready" ]]; then + url="http://127.0.0.1:${port}/api/v2/ready" +fi validate_non_negative_integer() { local name="$1" diff --git a/helm/codeapi/Chart.yaml b/helm/codeapi/Chart.yaml index dfcdd45b..6de60a24 100644 --- a/helm/codeapi/Chart.yaml +++ b/helm/codeapi/Chart.yaml @@ -3,7 +3,7 @@ apiVersion: v2 name: codeapi description: A Helm chart for Code Interpreter API - scalable code execution service type: application -version: 0.3.2 # Chart version (bump this when you change the chart) +version: 0.3.3 # Chart version (bump this when you change the chart) appVersion: "2.0.0" # App version (bump this when you change the app) # Keywords for searching diff --git a/helm/codeapi/README.md b/helm/codeapi/README.md index f63950ae..a5baa448 100644 --- a/helm/codeapi/README.md +++ b/helm/codeapi/README.md @@ -49,6 +49,28 @@ use it outside local development. ## Production deployment notes +**Runner shutdown.** Set `workerSandbox.sandboxRunner.gracefulDrain.enabled=true` +to stop accepting executions before termination. Readiness fails during drain; +liveness keeps the existing workspace, file-descriptor, and guest-clock checks. +The preStop hook and SIGTERM share one drain deadline. Defaults are a 5-second +minimum wait, a 45-second execution deadline, and a 60-second pod grace period. +Use the `gracefulDrain` settings rather than environment overrides; the chart +requires the grace period to exceed the drain deadline by more than 5 seconds. +The worker retries only a `503` with `error: runner_draining`, at most twice; +network failures and ambiguous execution results are not retried. +The advertised retry delay covers the propagation floor; both settings must +fit the worker's 5-second maximum delay. With defaults, `Retry-After` is 5. +The preStop exec hook signs a request with a per-pod Ed25519 private key on the +host. Only its public verifier enters the guest. Remote requests without that +signature, including signatures from another pod, cannot activate drain. + +Validate first in a development cluster: start an execution with a known result, +terminate its runner while it is active, and confirm the result arrives once, +readiness fails, new work reaches a ready runner, and the old pod exits within +its grace period. Repeat with an execution exceeding the drain deadline to +verify bounded shutdown. Promote to staging, repeat these checks, then promote +to production. A successful health check alone does not validate active drain. + This chart deploys the full service stack on a single cluster and is the supported path for self-hosting. A few things are intentionally left to your platform rather than templated here: external ingress/service mesh, KEDA-style diff --git a/helm/codeapi/templates/worker-sandbox-deployment.yaml b/helm/codeapi/templates/worker-sandbox-deployment.yaml index 2819d998..58ee9c77 100644 --- a/helm/codeapi/templates/worker-sandbox-deployment.yaml +++ b/helm/codeapi/templates/worker-sandbox-deployment.yaml @@ -7,6 +7,24 @@ Split runtime deployments: {{- fail "workerSandbox.enabled requires egressGateway.enabled=true because sandbox-runner may only egress to the egress gateway" }} {{- end }} {{- if .Values.workerSandbox.enabled }} +{{- $gracefulDrain := .Values.workerSandbox.sandboxRunner.gracefulDrain }} +{{- if $gracefulDrain.enabled }} +{{- range $key := list "minDelaySeconds" "timeoutSeconds" "retryAfterSeconds" "terminationGracePeriodSeconds" }} +{{- $value := toString (get $gracefulDrain $key) }} +{{- if or (not (regexMatch "^(0|[1-9][0-9]{0,5})$" $value)) (and (ne $key "minDelaySeconds") (eq $value "0")) }} +{{- fail (printf "gracefulDrain.%s must be a bounded integer (%s)" $key $value) }} +{{- end }} +{{- end }} +{{- if gt (int $gracefulDrain.minDelaySeconds) (int $gracefulDrain.timeoutSeconds) }} +{{- fail "gracefulDrain.minDelaySeconds must not exceed timeoutSeconds" }} +{{- end }} +{{- if or (gt (int $gracefulDrain.minDelaySeconds) 5) (gt (int $gracefulDrain.retryAfterSeconds) 5) }} +{{- fail "gracefulDrain propagation and retry delays must not exceed the 5-second transport retry bound" }} +{{- end }} +{{- if le (int $gracefulDrain.terminationGracePeriodSeconds) (add (int $gracefulDrain.timeoutSeconds) 5) }} +{{- fail "gracefulDrain.terminationGracePeriodSeconds must exceed timeoutSeconds by more than 5 seconds" }} +{{- end }} +{{- end }} {{- $legacyAutoscaling := default dict .Values.workerSandbox.autoscaling }} {{- if get $legacyAutoscaling "enabled" }} {{- fail "workerSandbox.autoscaling.enabled was replaced by workerSandbox.sandboxRunner.autoscaling.enabled; move autoscaling settings under workerSandbox.sandboxRunner.autoscaling" }} @@ -45,6 +63,9 @@ Split runtime deployments: {{- end }} {{- range .Values.workerSandbox.sandboxExtraEnv }} {{- $name := .name | default "" }} +{{- if or (hasPrefix "SANDBOX_DRAIN_" $name) (eq $name "SANDBOX_GRACEFUL_DRAIN_ENABLED") }} +{{- fail "configure gracefulDrain through workerSandbox.sandboxRunner.gracefulDrain, not sandboxExtraEnv" }} +{{- end }} {{- if and $name (or (regexMatch "(?i)(SECRET|TOKEN|PASSWORD|KEY)" $name) (hasPrefix "REDIS_" $name) (hasPrefix "AWS_" $name) (hasPrefix "S3_" $name) (hasPrefix "MINIO_" $name) (hasPrefix "CODEAPI_" $name) (eq $name "FILE_SERVER_URL") (eq $name "TOOL_CALL_SERVER_URL")) }} {{- fail (printf "workerSandbox.sandboxExtraEnv may not set sensitive/control-plane env var %s; sandbox-runner must only receive runtime knobs and the chart-managed manifest public verifier" $name) }} {{- end }} @@ -246,6 +267,9 @@ spec: serviceAccountName: {{ $sandboxRunnerServiceAccount }} {{- end }} automountServiceAccountToken: {{ .Values.workerSandbox.sandboxServiceAccount.automountServiceAccountToken }} + {{- if $gracefulDrain.enabled }} + terminationGracePeriodSeconds: {{ $gracefulDrain.terminationGracePeriodSeconds }} + {{- end }} {{- with .Values.global.imagePullSecrets }} imagePullSecrets: {{- toYaml . | nindent 8 }} @@ -282,6 +306,16 @@ spec: value: {{ $clockSkewLimit | quote }} - name: SANDBOX_RUNNER_CLOCK_SKEW_LIVENESS_JITTER_SECONDS value: {{ $clockSkewJitter | quote }} + {{- if $gracefulDrain.enabled }} + - name: SANDBOX_GRACEFUL_DRAIN_ENABLED + value: "true" + - name: SANDBOX_DRAIN_MIN_DELAY_MS + value: {{ mul $gracefulDrain.minDelaySeconds 1000 | quote }} + - name: SANDBOX_DRAIN_TIMEOUT_MS + value: {{ mul $gracefulDrain.timeoutSeconds 1000 | quote }} + - name: SANDBOX_DRAIN_RETRY_AFTER_SECONDS + value: {{ $gracefulDrain.retryAfterSeconds | quote }} + {{- end }} - name: EGRESS_GATEWAY_URL value: "http://{{ include "codeapi.fullname" . }}-egress-gateway:{{ .Values.egressGateway.service.port }}" - name: SANDBOX_LOG_LEVEL @@ -388,6 +422,13 @@ spec: - SETFCAP - MKNOD {{- end }} + {{- if $gracefulDrain.enabled }} + lifecycle: + preStop: + exec: + command: + - /usr/local/bin/sandbox-runner-drain.sh + {{- end }} livenessProbe: exec: command: @@ -399,6 +440,9 @@ spec: exec: command: - /usr/local/bin/sandbox-runner-healthcheck.sh + {{- if $gracefulDrain.enabled }} + - --ready + {{- end }} initialDelaySeconds: 30 periodSeconds: 10 timeoutSeconds: 7 diff --git a/helm/codeapi/values.yaml b/helm/codeapi/values.yaml index 33c32af0..6c94f197 100644 --- a/helm/codeapi/values.yaml +++ b/helm/codeapi/values.yaml @@ -216,6 +216,12 @@ workerSandbox: affinity: {} sandboxRunner: + gracefulDrain: + enabled: false + minDelaySeconds: 5 + timeoutSeconds: 45 + retryAfterSeconds: 1 + terminationGracePeriodSeconds: 60 # Defaults to workerSandbox.replicaCount when unset. replicaCount: null # Restart a sandbox-runner before launcher/VMM host-side fd retention can diff --git a/launcher/src/main.rs b/launcher/src/main.rs index 84e48347..c3379249 100644 --- a/launcher/src/main.rs +++ b/launcher/src/main.rs @@ -421,6 +421,11 @@ fn is_allowed_guest_env_key(key: &str, egress_gateway_enabled: bool) -> bool { "SANDBOX_COMPILE_TIMEOUT", "SANDBOX_DATA_DIRECTORY", "SANDBOX_DISABLE_NETWORKING", + "SANDBOX_GRACEFUL_DRAIN_ENABLED", + "SANDBOX_DRAIN_MIN_DELAY_MS", + "SANDBOX_DRAIN_TIMEOUT_MS", + "SANDBOX_DRAIN_RETRY_AFTER_SECONDS", + "SANDBOX_DRAIN_PUBLIC_KEY", "SANDBOX_HTTP_INPUT_CACHE_ENABLED", "SANDBOX_HTTP_INPUT_CACHE_MAX_INFLIGHT", "SANDBOX_HTTP_INPUT_CACHE_MAX_OBJECTS", @@ -877,6 +882,23 @@ mod tests { } } + #[test] + fn guest_env_allowlist_forwards_drain_settings_in_both_egress_modes() { + for egress_gateway_enabled in [false, true] { + for key in [ + "SANDBOX_GRACEFUL_DRAIN_ENABLED", + "SANDBOX_DRAIN_MIN_DELAY_MS", + "SANDBOX_DRAIN_TIMEOUT_MS", + "SANDBOX_DRAIN_RETRY_AFTER_SECONDS", + "SANDBOX_DRAIN_PUBLIC_KEY", + ] { + assert!(is_allowed_guest_env_key(key, egress_gateway_enabled), "{key}"); + } + assert!(!is_allowed_guest_env_key("SANDBOX_DRAIN_SECRET", egress_gateway_enabled)); + assert!(!is_allowed_guest_env_key("SANDBOX_DRAIN_UNDECLARED", egress_gateway_enabled)); + } + } + #[test] fn guest_env_allowlist_forwards_input_file_limit_in_both_egress_modes() { for egress_gateway_enabled in [false, true] { diff --git a/service/src/sandbox-backend/http.test.ts b/service/src/sandbox-backend/http.test.ts index 403772be..4bb0fa81 100644 --- a/service/src/sandbox-backend/http.test.ts +++ b/service/src/sandbox-backend/http.test.ts @@ -14,7 +14,16 @@ type CapturedRequest = { let server: ReturnType; let captured: CapturedRequest[] = []; -let nextResponse: { status: number; body: unknown; delayMs?: number; location?: string } = { status: 200, body: {} }; +type TestResponse = { + status: number; + body: unknown; + delayMs?: number; + location?: string; + headers?: Record; +}; + +let nextResponse: TestResponse = { status: 200, body: {} }; +let responseSequence: TestResponse[] = []; const savedEndpoint = env.SANDBOX_ENDPOINT; @@ -28,12 +37,13 @@ beforeAll(() => { rawBody: await req.text(), headers: Object.fromEntries(req.headers.entries()), }); - if (nextResponse.delayMs) { - await new Promise((resolve) => setTimeout(resolve, nextResponse.delayMs)); + const response = responseSequence.shift() ?? nextResponse; + if (response.delayMs !== undefined) { + await new Promise((resolve) => setTimeout(resolve, response.delayMs)); } - return new Response(JSON.stringify(nextResponse.body), { - status: nextResponse.status, - headers: { 'Content-Type': 'application/json', ...(nextResponse.location ? { Location: nextResponse.location } : {}) }, + return new Response(JSON.stringify(response.body), { + status: response.status, + headers: { 'Content-Type': 'application/json', ...response.headers, ...(response.location !== undefined ? { Location: response.location } : {}) }, }); }, }); @@ -48,6 +58,7 @@ afterAll(() => { afterEach(() => { captured = []; nextResponse = { status: 200, body: {} }; + responseSequence = []; }); function payloadBody(): t.PayloadBody { @@ -151,6 +162,58 @@ describe('HttpSandboxBackend', () => { } }); + test('retries an explicit runner_draining response before execution', async () => { + responseSequence = [ + { + status: 503, + body: { error: 'runner_draining', message: 'Sandbox runner is draining for shutdown' }, + headers: { 'Retry-After': '0' }, + }, + { + status: 200, + body: { session_id: 'sess_exec_1', language: 'python', version: '3.14.4', files: [] }, + }, + ]; + + const result = await new HttpSandboxBackend().execute(request(), context()); + + expect(captured).toHaveLength(2); + expect(result).toEqual({ + session_id: 'sess_exec_1', + language: 'python', + version: '3.14.4', + files: [], + }); + }); + + test('does not retry a generic sandbox 503', async () => { + nextResponse = { status: 503, body: { error: 'sandbox_execution_failed' } }; + + await expect(new HttpSandboxBackend().execute(request(), context())).rejects.toMatchObject({ + response: { status: 503 }, + }); + expect(captured).toHaveLength(1); + }); + + test('bounds retries when every runner rejects admission', async () => { + nextResponse = { status: 503, body: { error: 'runner_draining' }, headers: { 'Retry-After': '0' } }; + await expect(new HttpSandboxBackend().execute(request(), context())).rejects.toMatchObject({ + response: { status: 503 }, + }); + expect(captured).toHaveLength(3); + }); + + test('cancellation interrupts the retry delay without dispatching another execution', async () => { + nextResponse = { status: 503, body: { error: 'runner_draining' }, headers: { 'Retry-After': '999' } }; + const controller = new AbortController(); + const pending = new HttpSandboxBackend().execute(request(), context({ signal: controller.signal })); + for (let i = 0; i < 100 && captured.length === 0; i++) await Bun.sleep(2); + await Bun.sleep(10); + controller.abort(new Error('cancelled during drain retry')); + await expect(pending).rejects.toThrow('cancelled during drain retry'); + expect(captured).toHaveLength(1); + }); + test('propagates the worker abort signal as an axios cancellation', async () => { nextResponse = { status: 200, body: { session_id: 'x' }, delayMs: 5_000 }; const controller = new AbortController(); diff --git a/service/src/sandbox-backend/http.ts b/service/src/sandbox-backend/http.ts index 1ebdfcbe..2c263034 100644 --- a/service/src/sandbox-backend/http.ts +++ b/service/src/sandbox-backend/http.ts @@ -1,12 +1,44 @@ -import axios from 'axios'; +import axios, { type AxiosError, type AxiosResponse } from 'axios'; import type { SandboxBackend, SandboxExecuteContext, SandboxRawResponse, SandboxTransportRequest } from './types'; import { injectTraceHeaders, withSpan } from '../telemetry'; import { Jobs } from '../enum'; import { env } from '../config'; +import { MAX_RUNNER_DRAIN_RETRIES, MAX_RUNNER_DRAIN_RETRY_DELAY_MS } from '../../../shared/drain-contract'; -/** Current behavior: POST the signed request to SANDBOX_ENDPOINT. - * Axios errors are rethrown untouched so the worker's existing - * abort/timeout/sandbox-error mapping stays byte-identical. */ +function isRunnerDrainingError(error: unknown): error is AxiosError { + if (!axios.isAxiosError(error) || error.response?.status !== 503) return false; + const data = error.response.data; + return typeof data === 'object' + && data !== null + && (data as { error?: unknown }).error === 'runner_draining'; +} + +function retryAfterMillis(error: AxiosError): number { + const raw = error.response?.headers['retry-after']; + const seconds = Number(Array.isArray(raw) ? raw[0] : raw); + if (!Number.isFinite(seconds) || seconds < 0) return 1_000; + return Math.min(seconds * 1_000, MAX_RUNNER_DRAIN_RETRY_DELAY_MS); +} + +async function delayWithSignal(timeoutMillis: number, signal: AbortSignal): Promise { + if (signal.aborted) throw signal.reason ?? new Error('Aborted'); + await new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + signal.removeEventListener('abort', onAbort); + resolve(); + }, timeoutMillis); + const onAbort = (): void => { + clearTimeout(timeout); + reject(signal.reason ?? new Error('Aborted')); + }; + signal.addEventListener('abort', onAbort, { once: true }); + }); +} + +/** POST the signed request to SANDBOX_ENDPOINT. An explicit runner_draining + * response is safe to retry because the runner rejects it before execution. + * All ambiguous failures remain untouched so code with side effects is never + * replayed after a connection loss. */ export class HttpSandboxBackend implements SandboxBackend { readonly name = 'http' as const; @@ -16,15 +48,26 @@ export class HttpSandboxBackend implements SandboxBackend { 'url.path': `/${Jobs.execute}`, 'codeapi.language': ctx.language, 'codeapi.sandbox.backend': this.name, - }, () => axios.post( - `${env.SANDBOX_ENDPOINT}/${Jobs.execute}`, - req.body, - { - maxRedirects: 0, - headers: injectTraceHeaders(req.headers), - signal: ctx.signal, + }, async (): Promise> => { + for (let attempt = 0; ; attempt += 1) { + try { + return await axios.post( + `${env.SANDBOX_ENDPOINT}/${Jobs.execute}`, + req.body, + { + maxRedirects: 0, + headers: injectTraceHeaders(req.headers), + signal: ctx.signal, + }, + ); + } catch (error) { + if (!isRunnerDrainingError(error) || attempt >= MAX_RUNNER_DRAIN_RETRIES) { + throw error; + } + await delayWithSignal(retryAfterMillis(error), ctx.signal); + } } - ), 'CLIENT'); + }, 'CLIENT'); if (response.status !== 200) { throw new Error('Error from sandbox'); diff --git a/shared/drain-contract.ts b/shared/drain-contract.ts new file mode 100644 index 00000000..5082e60c --- /dev/null +++ b/shared/drain-contract.ts @@ -0,0 +1,4 @@ +/** Retry bounds keep a signed execution fresh across runner replacement. */ +export const MAX_RUNNER_DRAIN_RETRIES = 2; +export const MAX_RUNNER_DRAIN_RETRY_DELAY_MS = 5_000; +export const RUNNER_DRAIN_MESSAGE = 'codeapi-runner-drain-v1'; diff --git a/tests/graceful_drain.sh b/tests/graceful_drain.sh new file mode 100644 index 00000000..6995d407 --- /dev/null +++ b/tests/graceful_drain.sh @@ -0,0 +1,72 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +SOURCE_CHART="${1:-$ROOT/helm/codeapi}" +PREFIX="${2:-}" +TMP_DIR="$(mktemp -d)" +trap 'rm -rf "$TMP_DIR"' EXIT + +# Render without optional database subcharts or network dependency downloads. +mkdir "$TMP_DIR/chart" +cp "$SOURCE_CHART/values.yaml" "$TMP_DIR/chart/values.yaml" +cp -R "$SOURCE_CHART/templates" "$TMP_DIR/chart/templates" +awk '/^dependencies:/{exit} {print}' "$SOURCE_CHART/Chart.yaml" > "$TMP_DIR/chart/Chart.yaml" + +render() { + local args=( + --set "${PREFIX}workerSandbox.enabled=true" + --set "${PREFIX}egressGateway.enabled=true" + --set "${PREFIX}executionManifest.publicKey=test" + ) + if [[ -n "$PREFIX" ]]; then + args+=(--set "${PREFIX}executionManifest.privateKeySecretName=test") + else + args+=(--set executionManifest.privateKey=test) + fi + helm template drain "$TMP_DIR/chart" "${args[@]}" "$@" +} + +runner() { + awk ' + function emit() { if (kind && name) printf "%s", document } + /^---$/ { emit(); document=""; kind=0; name=0; next } + { document=document $0 "\n" } + /^kind: Deployment$/ { kind=1 } + /^ name: .*sandbox-runner$/ { name=1 } + END { emit() } + ' +} + +render | runner > "$TMP_DIR/off.yaml" +if grep -Eq 'preStop:|terminationGracePeriodSeconds:|SANDBOX_DRAIN_|SANDBOX_GRACEFUL_DRAIN_ENABLED|--ready' "$TMP_DIR/off.yaml"; then + echo 'disabled drain must preserve the previous lifecycle' >&2 + exit 1 +fi +render --set "${PREFIX}workerSandbox.sandboxRunner.gracefulDrain.enabled=true" | runner > "$TMP_DIR/on.yaml" +for expected in 'terminationGracePeriodSeconds: 60' 'preStop:' '/usr/local/bin/sandbox-runner-drain.sh' 'SANDBOX_GRACEFUL_DRAIN_ENABLED' 'SANDBOX_DRAIN_TIMEOUT_MS' 'value: "45000"' 'value: "5000"' 'value: "1"' '- --ready'; do + grep -Fq -- "$expected" "$TMP_DIR/on.yaml" || { echo "missing drain contract: $expected" >&2; exit 1; } +done +# The existing host FD/guest clock check remains both liveness and readiness. +[[ "$(grep -Fc '/usr/local/bin/sandbox-runner-healthcheck.sh' "$TMP_DIR/on.yaml")" == 2 ]] +[[ "$(grep -Fc -- '- --ready' "$TMP_DIR/on.yaml")" == 1 ]] + +for invalid in timeoutSeconds=0 timeoutSeconds=-1 timeoutSeconds=1.5 minDelaySeconds=46 minDelaySeconds=6 retryAfterSeconds=0 retryAfterSeconds=6 terminationGracePeriodSeconds=50; do + if render --set "${PREFIX}workerSandbox.sandboxRunner.gracefulDrain.enabled=true" \ + --set "${PREFIX}workerSandbox.sandboxRunner.gracefulDrain.$invalid" > /dev/null 2> "$TMP_DIR/error"; then + echo "accepted invalid drain budget: $invalid" >&2 + exit 1 + fi + grep -Fq gracefulDrain "$TMP_DIR/error" +done +if [[ -n "$PREFIX" ]]; then + override="${PREFIX}workerSandbox.sandboxEnv.SANDBOX_DRAIN_TIMEOUT_MS=1" +else + override='workerSandbox.sandboxExtraEnv[0].name=SANDBOX_DRAIN_TIMEOUT_MS' +fi +if render --set "$override" > /dev/null 2> "$TMP_DIR/error"; then + echo 'accepted a drain environment override outside the validated settings' >&2 + exit 1 +fi +grep -Fq gracefulDrain "$TMP_DIR/error" +echo 'Graceful drain chart regressions passed' diff --git a/tests/sandbox_runner_healthcheck.sh b/tests/sandbox_runner_healthcheck.sh index 583ff7f9..a84cb553 100755 --- a/tests/sandbox_runner_healthcheck.sh +++ b/tests/sandbox_runner_healthcheck.sh @@ -75,9 +75,14 @@ run_check() { TEST_GUEST_DATE='Tue, 04 Aug 2026 08:00:00 GMT' \ TEST_EXPECTED_URL='http://127.0.0.1:2000/api/v2/health' \ "$@" \ - bash "$ROOT/docker/sandbox-runner-healthcheck.sh" + bash -c 'exec bash "$1" "${TEST_PROBE_ARGUMENT:-}"' _ "$ROOT/docker/sandbox-runner-healthcheck.sh" } +run_check \ + SANDBOX_RUNNER_CLOCK_SKEW_LIVENESS_LIMIT_SECONDS=0 \ + TEST_PROBE_ARGUMENT=--ready \ + TEST_EXPECTED_URL='http://127.0.0.1:2000/api/v2/ready' + run_check \ SANDBOX_RUNNER_CLOCK_SKEW_LIVENESS_LIMIT_SECONDS=10 \ TEST_GUEST_SECONDS=91