fix(worker): make user-box lifecycle concurrency-safe

This commit is contained in:
Gabriel Brown
2026-07-10 16:54:55 -04:00
parent daec4d6dfa
commit 9e7a8722f3
6 changed files with 203 additions and 41 deletions
+2
View File
@@ -1,5 +1,6 @@
import { env } from './env'; import { env } from './env';
import { startWorkerServer } from './server'; import { startWorkerServer } from './server';
import { reconcileExistingBoxes } from './user-container';
import { startWorker } from './worker'; import { startWorker } from './worker';
// Dev-only watchdog: the dev runner chain (turbo → with-env → dotenv → bash) // Dev-only watchdog: the dev runner chain (turbo → with-env → dotenv → bash)
@@ -26,4 +27,5 @@ if (env.devWatchdog) {
} }
startWorkerServer(); startWorkerServer();
await reconcileExistingBoxes();
await startWorker(); await startWorker();
+1
View File
@@ -362,6 +362,7 @@ export const ensureUserContainer = async (args: {
[ [
'run', 'run',
'-d', '-d',
'--init',
'--name', '--name',
name, name,
'--memory', '--memory',
+12 -6
View File
@@ -4,9 +4,10 @@ import type { Server } from 'node:http';
import type { WebSocket } from 'ws'; import type { WebSocket } from 'ws';
import { WebSocketServer } from 'ws'; import { WebSocketServer } from 'ws';
import type { BoxHandle } from './user-container';
import { env } from './env'; import { env } from './env';
import { verifyTerminalToken } from './terminal-token'; import { verifyTerminalToken } from './terminal-token';
import { acquireUserBox, releaseUserBox } from './user-container'; import { acquireUserBox } from './user-container';
import { getTerminalWorkspace } from './worker'; import { getTerminalWorkspace } from './worker';
const clampDimension = (value: unknown) => { const clampDimension = (value: unknown) => {
@@ -63,7 +64,7 @@ const bridge = async (ws: WebSocket, jobId: string) => {
else pendingInput.push(data); else pendingInput.push(data);
}); });
let acquired = false; const handleHolder: { current?: BoxHandle } = {};
let released = false; let released = false;
// Read through a function so TS doesn't narrow `released` to a constant — the // Read through a function so TS doesn't narrow `released` to a constant — the
// cleanup handler flips it asynchronously when the socket closes. // cleanup handler flips it asynchronously when the socket closes.
@@ -72,7 +73,7 @@ const bridge = async (ws: WebSocket, jobId: string) => {
if (released) return; if (released) return;
released = true; released = true;
procHolder.current?.kill(); procHolder.current?.kill();
if (acquired) releaseUserBox(workspace.username); handleHolder.current?.release();
}; };
ws.on('close', cleanup); ws.on('close', cleanup);
ws.on('error', cleanup); ws.on('error', cleanup);
@@ -81,12 +82,13 @@ const bridge = async (ws: WebSocket, jobId: string) => {
// the terminal share the exact same container (Phase 2). // the terminal share the exact same container (Phase 2).
let boxName: string; let boxName: string;
try { try {
boxName = await acquireUserBox({ const handle = await acquireUserBox({
username: workspace.username, username: workspace.username,
workdir: workspace.workdir, workdir: workspace.workdir,
containerHome: workspace.containerHome, containerHome: workspace.containerHome,
}); });
acquired = true; handleHolder.current = handle;
boxName = handle.boxName;
} catch (error) { } catch (error) {
ws.close( ws.close(
1011, 1011,
@@ -95,7 +97,11 @@ const bridge = async (ws: WebSocket, jobId: string) => {
return; return;
} }
if (isReleased()) return; // client disconnected during startup; cleanup ran // Cleanup may have run before the awaited handle existed.
if (isReleased()) {
handleHolder.current.release();
return;
}
// Reattach a persistent tmux session across reconnects when available, else a // Reattach a persistent tmux session across reconnects when available, else a
// plain login shell. `stty` sizes the PTY to the client's viewport up front. // plain login shell. `stty` sizes the PTY to the client's viewport up front.
+92 -24
View File
@@ -1,42 +1,110 @@
import { env } from './env'; import { env } from './env';
import { import {
ensureUserContainer, ensureUserContainer,
listWorkspaceContainerNames,
stopWorkspaceContainer, stopWorkspaceContainer,
userContainerName, userContainerName,
} from './runtime/docker'; } from './runtime/docker';
// Phase 2: one persistent "box" container per user that all of their threads // Phase 2: one persistent "box" container per user. Reference-counted by opaque
// (agent turns + terminal + commands) exec into. Reference-counted so it stays // handles (not a shared integer), serialized per-username with an async mutex,
// up while any thread workspace is active or a terminal is connected, and is // and idle-reaped once no handle is held.
// reaped after an idle period once nothing holds it. export type BoxHandle = { boxName: string; release: () => void };
type Box = { refs: number; idleTimer?: NodeJS.Timeout };
const boxes = new Map<string, Box>();
export const acquireUserBox = async (args: { type Box = {
username: string; name: string;
workdir: string; refs: Set<symbol>;
containerHome: string; idleTimer?: ReturnType<typeof setTimeout>;
}): Promise<string> => { };
const name = await ensureUserContainer(args); const boxes = new Map<string, Box>();
const box = boxes.get(args.username) ?? { refs: 0 }; const locks = new Map<string, Promise<unknown>>();
if (box.idleTimer) {
clearTimeout(box.idleTimer); // Per-username async mutex: chain each operation after the previous one.
box.idleTimer = undefined; const withLock = <T>(username: string, fn: () => Promise<T>): Promise<T> => {
} const previous = locks.get(username) ?? Promise.resolve();
box.refs += 1; const next = previous.then(fn, fn);
boxes.set(args.username, box); locks.set(
return name; username,
next.then(
() => undefined,
() => undefined,
),
);
return next;
}; };
export const releaseUserBox = (username: string) => { const scheduleReapIfIdle = (username: string) => {
const box = boxes.get(username); const box = boxes.get(username);
if (!box) return; if (!box || box.refs.size > 0) return;
box.refs = Math.max(0, box.refs - 1); if (box.idleTimer) clearTimeout(box.idleTimer);
if (box.refs > 0) return;
box.idleTimer = setTimeout(() => { box.idleTimer = setTimeout(() => {
void stopWorkspaceContainer(userContainerName(username)); void stopWorkspaceContainer(userContainerName(username));
boxes.delete(username); boxes.delete(username);
}, env.boxIdleMs); }, env.boxIdleMs);
}; };
const makeHandle = (username: string, box: Box): BoxHandle => {
const token = Symbol('box-ref');
box.refs.add(token);
let released = false;
return {
boxName: box.name,
release: () => {
if (released) return;
released = true;
box.refs.delete(token);
scheduleReapIfIdle(username);
},
};
};
export const acquireUserBox = (args: {
username: string;
workdir: string;
containerHome: string;
}): Promise<BoxHandle> =>
withLock(args.username, async () => {
let box = boxes.get(args.username);
let needsEnsure = false;
if (box?.idleTimer) {
clearTimeout(box.idleTimer);
box.idleTimer = undefined;
}
if (!box) {
box = { name: userContainerName(args.username), refs: new Set() };
boxes.set(args.username, box);
needsEnsure = true;
}
// Register the ref before the slow Docker call so a failed acquire can
// release it without leaking registry state.
const handle = makeHandle(args.username, box);
try {
if (needsEnsure) box.name = await ensureUserContainer(args);
return handle;
} catch (error) {
handle.release();
throw error;
}
});
// Adopt boxes left by a prior worker process so the idle reaper cleans them up.
export const reconcileExistingBoxes = async (): Promise<void> => {
const names = await listWorkspaceContainerNames('spoon-box-');
for (const name of names) {
const username = name.replace(/^spoon-box-/, '');
if (boxes.has(username)) continue;
const box: Box = { name, refs: new Set() };
boxes.set(username, box);
scheduleReapIfIdle(username);
}
};
export const runningBoxUsernames = (): Set<string> => new Set(boxes.keys()); export const runningBoxUsernames = (): Set<string> => new Set(boxes.keys());
export const _resetBoxRegistryForTests = () => {
for (const box of boxes.values()) {
if (box.idleTimer) clearTimeout(box.idleTimer);
}
boxes.clear();
locks.clear();
};
+11 -11
View File
@@ -16,6 +16,7 @@ import { api } from '@spoon/backend/convex/_generated/api.js';
import type { NormalizedAgentEvent } from './agent-events'; import type { NormalizedAgentEvent } from './agent-events';
import type { OpenCodeSession } from './opencode-session'; import type { OpenCodeSession } from './opencode-session';
import type { BoxHandle } from './user-container';
import { normalizeCodexJsonLine } from './agent-events'; import { normalizeCodexJsonLine } from './agent-events';
import { prepareCodexWorkspaceFiles } from './codex-runtime'; import { prepareCodexWorkspaceFiles } from './codex-runtime';
import { env } from './env'; import { env } from './env';
@@ -41,11 +42,7 @@ import {
stopWorkspaceContainer, stopWorkspaceContainer,
streamExecInContainer, streamExecInContainer,
} from './runtime/docker'; } from './runtime/docker';
import { import { acquireUserBox, runningBoxUsernames } from './user-container';
acquireUserBox,
releaseUserBox,
runningBoxUsernames,
} from './user-container';
import { fetchUserEnvironment, materializeUserHome } from './user-environment'; import { fetchUserEnvironment, materializeUserHome } from './user-environment';
type Claim = { type Claim = {
@@ -111,6 +108,7 @@ type ActiveWorkspace = {
repoDir: string; repoDir: string;
// Phase 2: the per-user box container this thread execs into. // Phase 2: the per-user box container this thread execs into.
boxName: string; boxName: string;
boxHandle: BoxHandle;
githubToken: string; githubToken: string;
redact: (value: string) => string; redact: (value: string) => string;
runtimeMode?: 'opencode_server' | 'codex_exec' | 'legacy_cli'; runtimeMode?: 'opencode_server' | 'codex_exec' | 'legacy_cli';
@@ -1323,7 +1321,7 @@ const runClaim = async (claim: Claim) => {
...claim.secrets.map((secret) => secret.value), ...claim.secrets.map((secret) => secret.value),
].filter(Boolean); ].filter(Boolean);
const redact = createRedactor(secretValues); const redact = createRedactor(secretValues);
let acquiredBoxUser: string | undefined; let acquiredBoxHandle: BoxHandle | undefined;
try { try {
if ((claim.job.runtime ?? 'opencode') !== 'opencode') { if ((claim.job.runtime ?? 'opencode') !== 'opencode') {
throw new Error('Legacy OpenAI direct jobs are no longer supported.'); throw new Error('Legacy OpenAI direct jobs are no longer supported.');
@@ -1354,12 +1352,13 @@ const runClaim = async (claim: Claim) => {
// Start (or reuse) the persistent per-user box that this thread — and the // Start (or reuse) the persistent per-user box that this thread — and the
// terminal — exec into. It mounts the home, so the clone below is visible. // terminal — exec into. It mounts the home, so the clone below is visible.
const boxName = await acquireUserBox({ const boxHandle = await acquireUserBox({
username, username,
workdir: homeDir, workdir: homeDir,
containerHome, containerHome,
}); });
acquiredBoxUser = username; acquiredBoxHandle = boxHandle;
const boxName = boxHandle.boxName;
const repoDir = await cloneRepository({ const repoDir = await cloneRepository({
workdir: checkoutParent, workdir: checkoutParent,
@@ -1381,6 +1380,7 @@ const runClaim = async (claim: Claim) => {
containerRepo, containerRepo,
repoDir, repoDir,
boxName, boxName,
boxHandle,
githubToken, githubToken,
redact, redact,
}; };
@@ -1460,7 +1460,7 @@ const runClaim = async (claim: Claim) => {
).catch((stopError: unknown) => { ).catch((stopError: unknown) => {
console.error(stopError); console.error(stopError);
}); });
if (acquiredBoxUser) releaseUserBox(acquiredBoxUser); acquiredBoxHandle?.release();
} }
}; };
@@ -1883,7 +1883,7 @@ export const openWorkspacePullRequest = async (jobId: string) => {
activeWorkspaces.delete(jobId); activeWorkspaces.delete(jobId);
// The persistent per-user home + ~/Code checkouts survive across sessions; // The persistent per-user home + ~/Code checkouts survive across sessions;
// release the box (reaped once no other thread/terminal holds it). // release the box (reaped once no other thread/terminal holds it).
releaseUserBox(workspace.username); workspace.boxHandle.release();
return { return {
pullRequestUrl: pullRequest.html_url, pullRequestUrl: pullRequest.html_url,
pullRequestNumber: pullRequest.number, pullRequestNumber: pullRequest.number,
@@ -1900,7 +1900,7 @@ export const stopWorkspace = async (jobId: string) => {
activeWorkspaces.delete(jobId); activeWorkspaces.delete(jobId);
// The persistent per-user home + ~/Code checkouts survive across sessions; // The persistent per-user home + ~/Code checkouts survive across sessions;
// release the box (reaped once no other thread/terminal holds it). // release the box (reaped once no other thread/terminal holds it).
releaseUserBox(workspace.username); workspace.boxHandle.release();
return { success: true }; return { success: true };
}; };
@@ -0,0 +1,85 @@
import { afterEach, beforeEach, describe, expect, test, vi } from 'vitest';
vi.mock('../../src/env', () => ({ env: { boxIdleMs: 1000 } }));
vi.mock('../../src/runtime/docker', () => ({
ensureUserContainer: vi.fn((args: { username: string }) =>
Promise.resolve(`spoon-box-${args.username}`),
),
stopWorkspaceContainer: vi.fn(() => Promise.resolve()),
userContainerName: (username: string) => `spoon-box-${username}`,
listWorkspaceContainerNames: vi.fn(() => Promise.resolve([])),
}));
const load = async () => await import('../../src/user-container');
describe('box registry', () => {
beforeEach(() => vi.useFakeTimers());
afterEach(async () => {
const module = await load();
module._resetBoxRegistryForTests();
vi.useRealTimers();
vi.resetModules();
vi.clearAllMocks();
});
test('serializes concurrent acquire into a single docker run', async () => {
const module = await load();
const docker = await import('../../src/runtime/docker');
const [a, b] = await Promise.all([
module.acquireUserBox({
username: 'alice',
workdir: '/w',
containerHome: '/home/alice',
}),
module.acquireUserBox({
username: 'alice',
workdir: '/w',
containerHome: '/home/alice',
}),
]);
expect(docker.ensureUserContainer).toHaveBeenCalledTimes(1);
expect(a.boxName).toBe('spoon-box-alice');
a.release();
vi.advanceTimersByTime(2000);
expect(docker.stopWorkspaceContainer).not.toHaveBeenCalled();
b.release();
vi.advanceTimersByTime(2000);
expect(docker.stopWorkspaceContainer).toHaveBeenCalledTimes(1);
});
test('double-release of one handle does not reap a box held elsewhere', async () => {
const module = await load();
const docker = await import('../../src/runtime/docker');
const a = await module.acquireUserBox({
username: 'bob',
workdir: '/w',
containerHome: '/home/bob',
});
const b = await module.acquireUserBox({
username: 'bob',
workdir: '/w',
containerHome: '/home/bob',
});
a.release();
a.release();
vi.advanceTimersByTime(2000);
expect(docker.stopWorkspaceContainer).not.toHaveBeenCalled();
b.release();
vi.advanceTimersByTime(2000);
expect(docker.stopWorkspaceContainer).toHaveBeenCalledTimes(1);
});
test('reconcileExistingBoxes adopts orphans with zero refs and reaps them', async () => {
const module = await load();
const docker = await import('../../src/runtime/docker');
(
docker.listWorkspaceContainerNames as unknown as ReturnType<typeof vi.fn>
).mockResolvedValueOnce(['spoon-box-carol']);
await module.reconcileExistingBoxes();
expect(module.runningBoxUsernames().has('carol')).toBe(true);
vi.advanceTimersByTime(2000);
expect(docker.stopWorkspaceContainer).toHaveBeenCalledWith(
'spoon-box-carol',
);
});
});