127 lines
3.6 KiB
TypeScript
127 lines
3.6 KiB
TypeScript
import { env } from './env';
|
|
import {
|
|
ensureUserContainer,
|
|
listWorkspaceContainerNames,
|
|
stopWorkspaceContainer,
|
|
userContainerName,
|
|
} from './runtime/docker';
|
|
|
|
// Phase 2: one persistent "box" container per user. Reference-counted by opaque
|
|
// handles (not a shared integer), serialized per-username with an async mutex,
|
|
// and idle-reaped once no handle is held.
|
|
export type BoxHandle = { boxName: string; release: () => void };
|
|
|
|
type Box = {
|
|
name: string;
|
|
initialized: boolean;
|
|
refs: Set<symbol>;
|
|
idleTimer?: ReturnType<typeof setTimeout>;
|
|
};
|
|
const boxes = new Map<string, Box>();
|
|
const locks = new Map<string, Promise<unknown>>();
|
|
|
|
// Per-username async mutex: chain each operation after the previous one.
|
|
const withLock = <T>(username: string, fn: () => Promise<T>): Promise<T> => {
|
|
const previous = locks.get(username) ?? Promise.resolve();
|
|
const next = previous.then(fn, fn);
|
|
locks.set(
|
|
username,
|
|
next.then(
|
|
() => undefined,
|
|
() => undefined,
|
|
),
|
|
);
|
|
return next;
|
|
};
|
|
|
|
const scheduleReapIfIdle = (username: string) => {
|
|
const box = boxes.get(username);
|
|
if (!box || box.refs.size > 0) return;
|
|
if (box.idleTimer) clearTimeout(box.idleTimer);
|
|
box.idleTimer = setTimeout(() => {
|
|
void withLock(username, async () => {
|
|
const current = boxes.get(username);
|
|
if (current !== box || current.refs.size > 0) return;
|
|
await stopWorkspaceContainer(current.name);
|
|
if (boxes.get(username) === current && current.refs.size === 0) {
|
|
boxes.delete(username);
|
|
}
|
|
});
|
|
}, 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);
|
|
if (box?.idleTimer) {
|
|
clearTimeout(box.idleTimer);
|
|
box.idleTimer = undefined;
|
|
}
|
|
if (!box) {
|
|
box = {
|
|
name: userContainerName(args.username),
|
|
initialized: false,
|
|
refs: new Set(),
|
|
};
|
|
boxes.set(args.username, box);
|
|
}
|
|
// 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 (!box.initialized) {
|
|
box.name = await ensureUserContainer(args);
|
|
box.initialized = true;
|
|
}
|
|
return handle;
|
|
} catch (error) {
|
|
handle.release();
|
|
if (boxes.get(args.username) === box && box.refs.size === 0) {
|
|
if (box.idleTimer) clearTimeout(box.idleTimer);
|
|
boxes.delete(args.username);
|
|
}
|
|
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, initialized: false, refs: new Set() };
|
|
boxes.set(username, box);
|
|
scheduleReapIfIdle(username);
|
|
}
|
|
};
|
|
|
|
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();
|
|
};
|