fix(worker): make user-box lifecycle concurrency-safe
This commit is contained in:
@@ -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();
|
||||||
|
|||||||
@@ -362,6 +362,7 @@ export const ensureUserContainer = async (args: {
|
|||||||
[
|
[
|
||||||
'run',
|
'run',
|
||||||
'-d',
|
'-d',
|
||||||
|
'--init',
|
||||||
'--name',
|
'--name',
|
||||||
name,
|
name,
|
||||||
'--memory',
|
'--memory',
|
||||||
|
|||||||
@@ -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.
|
||||||
|
|||||||
@@ -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();
|
||||||
|
};
|
||||||
|
|||||||
@@ -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',
|
||||||
|
);
|
||||||
|
});
|
||||||
|
});
|
||||||
Reference in New Issue
Block a user