diff --git a/packages/runtime/src/scr_arena.c b/packages/runtime/src/scr_arena.c index e04a6ed35..73901f2c1 100644 --- a/packages/runtime/src/scr_arena.c +++ b/packages/runtime/src/scr_arena.c @@ -251,6 +251,24 @@ static void scr_arena_free(ScrArena *a) { free(a); } +void scr_arena_thread_exit(void) { + while (scr_arena_pool) { + ScrArena *a = scr_arena_pool; + scr_arena_pool = a->next; + free(a->disposal_base); + free(a); + } + scr_arena_pooled = 0; +#if SCR_SMALL_ALLOC + for (ScrArenaChunk *c = scr_arena_chunk_pool, *next; c; c = next) { + next = c->next; + free(c); + } + scr_arena_chunk_pool = NULL; + scr_arena_chunk_bytes = 0; +#endif +} + void scr_arena_drop(void **list) { ScrArena *a = (ScrArena *)*list; *list = NULL; diff --git a/packages/runtime/src/scr_arena.h b/packages/runtime/src/scr_arena.h index cdfe5a476..4c704eeff 100644 --- a/packages/runtime/src/scr_arena.h +++ b/packages/runtime/src/scr_arena.h @@ -79,6 +79,10 @@ void scr_arena_drop(void **list); void *scr_arena_obj(void *arena, size_t size, void *cyc); /* Appends `obj` to `arena`'s dispose list, which the drop runs. */ void scr_arena_disposal(void *arena, void *obj, ScrArenaDispose dispose); +/* Frees the arenas and chunks the calling thread keeps for reuse. A worker + * thread calls it when its script has finished, after every arena it + * opened has been dropped. */ +void scr_arena_thread_exit(void); /* The leading fields of an arena (scr_arena.c's ScrArena), for the * runtime's own inline allocation. */ diff --git a/packages/runtime/src/scr_worker.c b/packages/runtime/src/scr_worker.c index 59e245ea0..5eed2469d 100644 --- a/packages/runtime/src/scr_worker.c +++ b/packages/runtime/src/scr_worker.c @@ -331,6 +331,9 @@ static void scr_worker_thread_run(ScrWorkerThread *worker) { scr_mailbox_close(worker->inbox); scr_stdin_fed_clear(); /* chunks the worker never read */ scr_context_cleanup(); + /* Frees the group arenas and chunks the thread kept for reuse before the + * parent sees the exit. */ + scr_arena_thread_exit(); if (stats) { uint64_t now = scr_hrtime_ns(); fprintf(stderr, "scriptc worker %llu: started %.2f ms after spawn, script %.2f ms, exit %.2f ms\n", diff --git a/tests/harness/worker-arena-pools.test.ts b/tests/harness/worker-arena-pools.test.ts new file mode 100644 index 000000000..a512d5a7d --- /dev/null +++ b/tests/harness/worker-arena-pools.test.ts @@ -0,0 +1,99 @@ +/* Group arenas (runtime scr_arena.c): each thread keeps up to 8 MB of + * arena chunks and its dropped arenas for reuse, and a worker thread frees + * them when its script finishes. This suite runs workers one after another, + * each building and dropping trees of about 12 MB in a group arena, and + * checks that the resident set stays flat across them: a worker that kept + * its pool would add about 8 MB each. Under SCRIPTC_SAN=1 the runtime + * keeps every worker's memory, so the check does not apply. */ +import { execFile } from "node:child_process"; +import { mkdtemp, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { promisify } from "node:util"; +import { expect, test } from "vitest"; +import { compile } from "@scriptc/compiler"; + +const exec = promisify(execFile); +const sanitize = process.env["SCRIPTC_SAN"] === "1"; + +const PROGRAM = ` +import { isMainThread, parentPort, Worker, workerData } from "node:worker_threads"; + +interface TreeNode { + left: TreeNode | null; + right: TreeNode | null; +} + +function tree(depth: number): TreeNode { + return depth > 0 ? { left: tree(depth - 1), right: tree(depth - 1) } : { left: null, right: null }; +} + +function count(node: TreeNode): number { + if (node.left === null || node.right === null) return 1; + return 1 + count(node.left) + count(node.right); +} + +function work(depth: number): number { + let total = 0; + for (let i = 0; i < 2; i++) total += count(tree(depth)); + return total; +} + +function runWorker(depth: number): Promise { + return new Promise((resolve) => { + let total = 0; + const worker = new Worker(new URL(import.meta.url), { workerData: depth }); + worker.on("message", (message: number) => { + total = message; + }); + worker.on("exit", () => resolve(total)); + }); +} + +async function main(): Promise { + let total = 0; + let settled = 0; + for (let i = 0; i < 24; i++) { + total += await runWorker(18); + if (i === 3) settled = process.memoryUsage.rss(); + } + const growth = Math.round((process.memoryUsage.rss() - settled) / (1024 * 1024)); + console.log(total, growth); +} + +if (isMainThread) { + await main(); +} else { + parentPort!.postMessage(work(workerData as number)); +} +`; + +test.skipIf(process.platform !== "linux" || sanitize)( + "workers that use group arenas free their arena pools when they exit", + { timeout: 300_000 }, + async () => { + const dir = await mkdtemp(join(tmpdir(), "scriptc-worker-arena-pools-")); + try { + const source = join(dir, "pools.ts"); + await writeFile(source, PROGRAM); + const result = await compile(source, { + backend: "llvm", + outDir: dir, + outPath: join(dir, "pools"), + }); + expect(result.ok, JSON.stringify(result.diagnostics)).toBe(true); + if (!result.ok) return; + const { stdout } = await exec(result.binaryPath, [], { + env: { ...process.env, SCR_SA_DISCARD_DELAY_MS: "0" }, + timeout: 120_000, + }); + const [total, growth] = stdout.trim().split(" ").map(Number); + // 24 workers, each with 2 trees of 2^19 - 1 nodes. + expect(total).toBe(24 * 2 * (2 ** 19 - 1)); + // 20 workers that kept their pools would add about 160 MiB. + expect(growth).toBeLessThan(32); + } finally { + await rm(dir, { recursive: true, force: true }); + } + }, +);