perf: keep memory use down at low thread counts

.changeset/keep-memory-down.md

@@ -0,0 +1,7 @@
+---
+'@doc-kit/core': patch
+'@doc-kit/generator-react': patch
+'@node-core/doc-kit-legacy': patch
+---
+
+Builds use less memory, most of all with few threads: each worker's young generation is limited to 96MB, a single thread processes one chunk at a time and lets go of the pages' code before rendering them, and Vite's native bundler gets as many threads as doc-kit and returns the memory it frees. On the Node.js API docs, `--threads 1` now builds with Node.js's `--max-old-space-size=512`, and no longer crashes when `legacy-html` is built alongside other targets.

packages/core/src/generators.mjs

@@ -87,16 +87,18 @@ const createGenerator = () => {
           ? createParallelWorker(specifier, generator, pool, configuration)
           : Promise.resolve(null);
 
-        const result = await generate(dependencyInput, await worker);
-
+        // Returned rather than awaited, so this function doesn't hold on to the
+        // input while the generator runs, and the generator can let it go
+        return generate(dependencyInput, await worker);
+      })().then(result => {
         // For streaming generators, "Completed" is logged when the cache
         // finishes collecting, not here when the generator returns
         if (!isAsyncIterable(result)) {
           generatorsLogger.debug(`Completed "${name}"`);
         }
 
         return result;
-      })()
+      })
     );
   };
 

packages/core/src/threading/__tests__/index.test.mjs

@@ -4,6 +4,7 @@ import { fileURLToPath } from 'node:url';
 
 import { LogLevel } from '../../logger/constants.mjs';
 import logger from '../../logger/index.mjs';
+import { WORKER_GEN_SIZE, WORKER_HEAP_SIZE } from '../constants.mjs';
 import createWorkerPool from '../index.mjs';
 
 const reporterSpecifier = fileURLToPath(
@@ -114,4 +115,17 @@ describe('createWorkerPool', () => {
       await pool.destroy();
     }
   });
+
+  it("should limit each worker's old and young generation", async () => {
+    const pool = createWorkerPool(1);
+
+    try {
+      deepStrictEqual(pool.options.resourceLimits, {
+        maxOldGenerationSizeMb: WORKER_HEAP_SIZE,
+        maxYoungGenerationSizeMb: WORKER_GEN_SIZE,
+      });
+    } finally {
+      await pool.destroy();
+    }
+  });
 });

packages/core/src/threading/__tests__/parallel.test.mjs

@@ -177,4 +177,32 @@ describe('createParallelWorker', () => {
 
     await pool.destroy();
   });
+
+  it('should process one chunk at a time, in order, with a single thread', async () => {
+    let running = 0;
+    let mostRunning = 0;
+
+    const generator = {
+      name: 'sequential',
+      async processChunk(items, indices) {
+        mostRunning = Math.max(mostRunning, ++running);
+
+        await new Promise(resolve => setImmediate(resolve));
+
+        running--;
+
+        return indices.map(index => items[index]);
+      },
+    };
+
+    const worker = createParallelWorker('sequential', generator, null, {
+      threads: 1,
+      chunkSize: 2,
+    });
+
+    const chunks = await collectChunks(worker.stream([1, 2, 3, 4, 5], {}));
+
+    strictEqual(mostRunning, 1);
+    deepStrictEqual(chunks, [[1, 2], [3, 4], [5]]);
+  });
 });

packages/core/src/threading/constants.mjs

@@ -1,5 +1,12 @@
 'use strict';
 
+/**
+ * The default `threads` ceiling. Each worker holds a heap of its own, with the
+ * libraries and the pages it is working on, so past a few threads memory, not
+ * CPU, is what runs out. `--threads` raises it.
+ */
+export const DEFAULT_MAX_THREADS = 4;
+
 /**
  * The heap size limit of each worker thread (V8's old space), in MB. V8 lets a
  * heap grow to several times its live data before collecting it, the more the
@@ -8,3 +15,10 @@
  * `--max-old-space-size` replaces it.
  */
 export const WORKER_HEAP_SIZE = 512;
+
+/**
+ * The young generation size limit of each worker thread, in MB. V8 sizes it
+ * from the machine's memory, as if each worker were the only one, so it made up
+ * over half of a worker's heap. Node.js's `--max-semi-space-size` replaces it.
+ */
+export const WORKER_GEN_SIZE = 96;

packages/core/src/threading/index.mjs

@@ -2,7 +2,7 @@ import Piscina from 'piscina';
 
 import logger from '#logger/index.mjs';
 
-import { WORKER_HEAP_SIZE } from './constants.mjs';
+import { WORKER_HEAP_SIZE, WORKER_GEN_SIZE } from './constants.mjs';
 
 const poolLogger = logger.child('WorkerPool');
 
@@ -27,7 +27,10 @@ export default function createWorkerPool(threads) {
     // A worker idle for 250ms ends, so its heap is gone during a long stretch
     // of work on the main thread, such as bundling the site
     idleTimeout: 250,
-    resourceLimits: { maxOldGenerationSizeMb: WORKER_HEAP_SIZE },
+    resourceLimits: {
+      maxOldGenerationSizeMb: WORKER_HEAP_SIZE,
+      maxYoungGenerationSizeMb: WORKER_GEN_SIZE,
+    },
     workerData: { logLevel: logger.getLogLevel() },
   });
 }

packages/core/src/threading/parallel.mjs

@@ -96,19 +96,18 @@ export default function createParallelWorker(
         { generator: name, chunks: chunks.length, chunkSize, threads }
       );
 
-      const runInOneGo = threads <= 1 || items.length <= 2;
+      // On this thread, one chunk at a time, so their work isn't all held at once
+      if (threads <= 1 || items.length <= 2) {
+        for (const indices of chunks) {
+          yield await generator.processChunk(items, indices, extra);
+        }
+
+        return;
+      }
 
       // Submit all tasks to Piscina - each promise resolves to itself for removal
       const pending = new Set(
         chunks.map(indices => {
-          if (runInOneGo) {
-            const promise = generator
-              .processChunk(items, indices, extra)
-              .then(result => ({ promise, result }));
-
-            return promise;
-          }
-
           const promise = pool
             .run(
               createTask(items, indices, extra, configuration, specifier, name)

packages/core/src/utils/configuration/constants.mjs

@@ -1,12 +1,5 @@
 'use strict';
 
-/**
- * The default `threads` ceiling. Each worker holds a heap of its own, with the
- * libraries and the pages it is working on, so past a few threads memory, not
- * CPU, is what runs out. `--threads` raises it.
- */
-export const DEFAULT_MAX_THREADS = 4;
-
 /**
  * The default number of items each worker task processes.
  */

packages/core/src/utils/configuration/index.mjs

@@ -22,7 +22,8 @@ import {
 import { resolveMarkdown } from '#utils/markdown/plugins.mjs';
 import { deepMerge } from '#utils/misc.mjs';
 
-import { DEFAULT_CHUNK_SIZE, DEFAULT_MAX_THREADS } from './constants.mjs';
+import { DEFAULT_MAX_THREADS } from '../../threading/constants.mjs';
+import { DEFAULT_CHUNK_SIZE } from './constants.mjs';
 
 const configExplorer = cosmiconfig('doc-kit');
 

packages/node-legacy/src/legacy-html/__tests__/generate.test.mjs

@@ -0,0 +1,56 @@
+import assert from 'node:assert/strict';
+import { describe, it } from 'node:test';
+
+import { loadGenerator } from '@doc-kit/core/generators/loader.mjs';
+import { setConfig } from '@doc-kit/core/utils/configuration/index.mjs';
+import { loadMarkdownPlugins } from '@doc-kit/core/utils/markdown/plugins.mjs';
+
+import { processChunk } from '../generate.mjs';
+
+await loadMarkdownPlugins(
+  await loadGenerator(import.meta.resolve('../index.mjs'))
+);
+
+const createEntry = (api, name) => {
+  const heading = {
+    type: 'heading',
+    depth: 1,
+    children: [{ type: 'text', value: name }],
+    data: { name, text: name, slug: api, type: 'module' },
+  };
+
+  return {
+    api,
+    path: `/${api}`,
+    heading,
+    stability: null,
+    content: {
+      type: 'root',
+      children: [
+        heading,
+        {
+          type: 'paragraph',
+          children: [{ type: 'text', value: `${name} body` }],
+        },
+      ],
+    },
+  };
+};
+
+describe('legacy-html processChunk', () => {
+  it('does not change its input on the main thread', async () => {
+    await setConfig({ target: ['legacy-html'] });
+
+    const fs = createEntry('fs', 'File system');
+    const before = structuredClone(fs);
+
+    const [page] = await processChunk(
+      [{ head: fs, nodes: [fs], headNodes: [fs] }],
+      [0],
+      '<a class="nav-fs" href="fs.html">File system</a>'
+    );
+
+    assert.match(page.content, /File system body/);
+    assert.deepEqual(fs, before);
+  });
+});

packages/node-legacy/src/legacy-html/generate.mjs

@@ -2,6 +2,7 @@
 
 import { readFile, cp } from 'node:fs/promises';
 import { basename, join } from 'node:path';
+import { isMainThread } from 'node:worker_threads';
 
 import getConfig from '@doc-kit/core/utils/configuration/index.mjs';
 import { writeFile } from '@doc-kit/core/utils/file.mjs';
@@ -33,7 +34,11 @@ export async function processChunk(slicedInput, itemIndices, navigation) {
   const results = [];
 
   for (const idx of itemIndices) {
-    const { head, nodes, headNodes } = slicedInput[idx];
+    // Building the content changes the entries. A worker has its own copy, but
+    // on the main thread they're the metadata other generators read too
+    const { head, nodes, headNodes } = isMainThread
+      ? structuredClone(slicedInput[idx])
+      : slicedInput[idx];
 
     const nav = navigation.replace(
       `class="nav-${head.api}"`,

packages/react/src/html/bundlers/__tests__/index.test.mjs

@@ -0,0 +1,73 @@
+import assert from 'node:assert/strict';
+import { afterEach, describe, it, mock } from 'node:test';
+
+import { setConfig } from '@doc-kit/core/utils/configuration/index.mjs';
+
+mock.module('../vite.mjs', {
+  namedExports: { createViteBundler: () => 'vite' },
+});
+
+const { resolveBundler } = await import('../index.mjs');
+
+const VARIABLES = [
+  'ROLLDOWN_WORKER_THREADS',
+  'RAYON_NUM_THREADS',
+  'MIMALLOC_PURGE_DELAY',
+  'MIMALLOC_ALLOW_THP',
+];
+
+const original = Object.fromEntries(
+  VARIABLES.map(name => [name, process.env[name]])
+);
+
+afterEach(() => {
+  for (const name of VARIABLES) {
+    if (original[name] === undefined) {
+      delete process.env[name];
+    } else {
+      process.env[name] = original[name];
+    }
+  }
+});
+
+describe('resolveBundler', () => {
+  it("sizes the bundler's thread pools and allocator before loading it", async () => {
+    await setConfig({ target: ['html'], threads: 3 });
+
+    for (const name of VARIABLES) {
+      delete process.env[name];
+    }
+
+    assert.equal(await resolveBundler(), 'vite');
+    assert.equal(process.env.ROLLDOWN_WORKER_THREADS, '3');
+    assert.equal(process.env.RAYON_NUM_THREADS, '3');
+    assert.equal(process.env.MIMALLOC_PURGE_DELAY, '0');
+    assert.equal(process.env.MIMALLOC_ALLOW_THP, '0');
+  });
+
+  it('keeps what the environment already sets', async () => {
+    await setConfig({ target: ['html'], threads: 3 });
+
+    process.env.ROLLDOWN_WORKER_THREADS = '8';
+    process.env.MIMALLOC_PURGE_DELAY = '250';
+
+    await resolveBundler();
+
+    assert.equal(process.env.ROLLDOWN_WORKER_THREADS, '8');
+    assert.equal(process.env.MIMALLOC_PURGE_DELAY, '250');
+  });
+
+  it('leaves the environment alone for a custom bundler', async () => {
+    for (const name of VARIABLES) {
+      delete process.env[name];
+    }
+
+    const bundler = {};
+
+    assert.equal(await resolveBundler(bundler), bundler);
+
+    for (const name of VARIABLES) {
+      assert.equal(process.env[name], undefined);
+    }
+  });
+});

packages/react/src/html/bundlers/index.mjs

@@ -1,3 +1,7 @@
+import getConfig from '@doc-kit/core/utils/configuration/index.mjs';
+
+import { BUNDLER_THREAD_VARIABLES } from '../constants.mjs';
+
 /**
  * Returns the configured bundler or lazily creates the default Vite adapter.
  *
@@ -9,6 +13,20 @@ export const resolveBundler = async bundler => {
     return bundler;
   }
 
+  // Vite's native thread pools (rolldown, lightningcss) otherwise start a
+  // thread per CPU, each holding memory of its own
+  const threads = String(getConfig().threads);
+
+  for (const name of BUNDLER_THREAD_VARIABLES) {
+    process.env[name] ??= threads;
+  }
+
+  // rolldown's allocator only returns freed memory on its next call past this
+  // delay, which never comes once bundling ends, so it returns it right away,
+  // and not in huge pages, which it only returns once they're entirely free
+  process.env.MIMALLOC_PURGE_DELAY ??= '0';
+  process.env.MIMALLOC_ALLOW_THP ??= '0';
+
   const { createViteBundler } = await import('./vite.mjs');
   return createViteBundler();
 };

packages/react/src/html/constants.mjs

@@ -92,3 +92,14 @@ export const FONTS = [
   'open-sans-latin-wght-italic.woff2',
   'ibm-plex-mono-latin-400-normal.woff2',
 ];
+
+/**
+ * The environment variables that size the thread pools of Vite's native
+ * bundler, read when they start: rolldown's runtime, and the Rayon pool
+ * rolldown and lightningcss share.
+ * @see https://rolldown.rs/guide/troubleshooting
+ */
+export const BUNDLER_THREAD_VARIABLES = [
+  'ROLLDOWN_WORKER_THREADS',
+  'RAYON_NUM_THREADS',
+];