Merge pull request storybookjs#35305 from storybookjs/norbert/oom-docgen-fix

DocgenServer: Fix OOM by recycling the shared TS program

Author: ndelangen

code/renderers/react/src/componentManifest/componentMeta/ComponentMetaManager.test.ts

@@ -1,7 +1,9 @@
 import * as fs from 'node:fs';
 import * as path from 'node:path';
 
-import { afterEach, describe, expect, it } from 'vitest';
+import { once } from 'storybook/internal/node-logger';
+
+import { afterEach, describe, expect, it, vi } from 'vitest';
 
 import { dedent } from 'ts-dedent';
 import ts from 'typescript';
@@ -340,6 +342,140 @@ describe('multi-project management', () => {
     expect(project).toBeDefined();
   });
 
+  it(
+    'recycles the shared program when heap usage crosses the threshold',
+    { timeout: 30_000 },
+    () => {
+      tempDir = createTempDir();
+
+      const files = writeFiles(tempDir, {
+        'tsconfig.json': tsconfigJSON(),
+        'Button.tsx': dedent`
+        import React from 'react';
+        export const Button = (_props: { label: string }) => <button />;
+      `,
+        'Button.stories.tsx': dedent`
+        import { Button } from './Button';
+        export default { component: Button };
+      `,
+      });
+
+      // ratio 0 → threshold 0 → heapUsed is always ≥ 0 → recycle fires after every batchExtract.
+      manager = new ComponentMetaManager(ts, 0);
+      const componentPath = path.join(tempDir, 'Button.tsx');
+      const before = manager.getProjectForFile(componentPath);
+
+      const entries: StoryRef[] = [
+        {
+          storyPath: files['Button.stories.tsx'],
+          component: {
+            componentName: 'Button',
+            importName: 'Button',
+            path: componentPath,
+            isPackage: false,
+          },
+        },
+      ];
+      manager.batchExtract(entries);
+
+      // The disposed program was cleared, so the next lookup rebuilds a fresh instance.
+      const after = manager.getProjectForFile(componentPath);
+      expect(after).not.toBe(before);
+    }
+  );
+
+  it(
+    'warns once, with guidance to raise the memory limit, when it recycles under pressure',
+    { timeout: 30_000 },
+    () => {
+      tempDir = createTempDir();
+
+      const files = writeFiles(tempDir, {
+        'tsconfig.json': tsconfigJSON(),
+        'Button.tsx': dedent`
+        import React from 'react';
+        export const Button = (_props: { label: string }) => <button />;
+      `,
+        'Button.stories.tsx': dedent`
+        import { Button } from './Button';
+        export default { component: Button };
+      `,
+      });
+
+      // Spy on the one-time `once.warn` channel: routing the warning through it is what makes it
+      // surface once per process (the dedupe itself is node-logger's responsibility, tested there).
+      const onceWarnSpy = vi.spyOn(once, 'warn').mockImplementation(() => undefined);
+      try {
+        // ratio 0 → threshold 0 → recycle fires on every batchExtract.
+        manager = new ComponentMetaManager(ts, 0);
+        const componentPath = path.join(tempDir, 'Button.tsx');
+        manager.getProjectForFile(componentPath);
+
+        const entries: StoryRef[] = [
+          {
+            storyPath: files['Button.stories.tsx'],
+            component: {
+              componentName: 'Button',
+              importName: 'Button',
+              path: componentPath,
+              isPackage: false,
+            },
+          },
+        ];
+
+        manager.batchExtract(entries);
+
+        expect(onceWarnSpy).toHaveBeenCalled();
+        expect(onceWarnSpy.mock.calls[0][0]).toContain('--max-old-space-size');
+      } finally {
+        onceWarnSpy.mockRestore();
+        once.clear();
+      }
+    }
+  );
+
+  it(
+    'keeps the shared program while heap usage stays below the threshold',
+    { timeout: 30_000 },
+    () => {
+      tempDir = createTempDir();
+
+      const files = writeFiles(tempDir, {
+        'tsconfig.json': tsconfigJSON(),
+        'Button.tsx': dedent`
+        import React from 'react';
+        export const Button = (_props: { label: string }) => <button />;
+      `,
+        'Button.stories.tsx': dedent`
+        import { Button } from './Button';
+        export default { component: Button };
+      `,
+      });
+
+      // ratio Infinity → threshold Infinity → heapUsed is always below it → recycle never fires.
+      manager = new ComponentMetaManager(ts, Number.POSITIVE_INFINITY);
+      const componentPath = path.join(tempDir, 'Button.tsx');
+      const before = manager.getProjectForFile(componentPath);
+
+      const entries: StoryRef[] = [
+        {
+          storyPath: files['Button.stories.tsx'],
+          component: {
+            componentName: 'Button',
+            importName: 'Button',
+            path: componentPath,
+            isPackage: false,
+          },
+        },
+      ];
+      manager.batchExtract(entries);
+
+      // No recycle → the same program instance is reused across extractions.
+      const after = manager.getProjectForFile(componentPath);
+      expect(after).toBe(before);
+    }
+  );
+
   it('extracts via Path 1 (importId + JSX in story file)', { timeout: 30_000 }, () => {
     tempDir = createTempDir();
 

code/renderers/react/src/componentManifest/componentMeta/ComponentMetaManager.ts

@@ -15,11 +15,13 @@
  * Plus our own file watching layer (fs.watch + debounce) since we're a standalone Node.js process,
  * not running inside an IDE.
  */
-import { logger } from 'storybook/internal/node-logger';
+import { logger, once } from 'storybook/internal/node-logger';
 
 import { type FSWatcher, existsSync, watch } from 'fs';
 import * as path from 'path';
+import { dedent } from 'ts-dedent';
 import type ts from 'typescript';
+import * as v8 from 'v8';
 
 import type { StoryRef } from '../getComponentImports.ts';
 import { groupByToMap } from '../utils.ts';
@@ -36,6 +38,12 @@ const DEFAULT_INFERRED_OPTIONS: ts.CompilerOptions = {
   skipLibCheck: true,
 };
 
+/**
+ * Recycle (dispose + lazily rebuild) the shared TS program(s) once heap usage crosses this fraction
+ * of the V8 heap limit.
+ */
+const RECYCLE_HEAP_PRESSURE_RATIO = 0.7;
+
 export class ComponentMetaManager {
   // Adapted from:
   // https://github.com/volarjs/volar.js/blob/882cd56d46a13d272f34e451f495d3d62251969a/packages/language-server/lib/project/typescriptProject.ts#L34-L37
@@ -56,7 +64,23 @@ export class ComponentMetaManager {
     [number | undefined, ts.IScriptSnapshot | undefined]
   >();
 
-  constructor(private typescript: typeof ts) {}
+  // Heap-pressure recycle: bytes of `heapUsed` past which the TS program(s) are disposed and rebuilt.
+  private readonly heapRecycleThresholdBytes: number;
+
+  /**
+   * @param recycleHeapPressureRatio Fraction of the V8 heap limit at which the shared program(s) are
+   *   recycled (see {@link RECYCLE_HEAP_PRESSURE_RATIO}). Exposed for tuning and for the memory
+   *   regression gate, which passes `Infinity` to disable recycling and assert the OOM still happens
+   *   without the fix. Defaults to {@link RECYCLE_HEAP_PRESSURE_RATIO}.
+   */
+  constructor(
+    private typescript: typeof ts,
+    recycleHeapPressureRatio: number = RECYCLE_HEAP_PRESSURE_RATIO
+  ) {
+    this.heapRecycleThresholdBytes = Math.floor(
+      v8.getHeapStatistics().heap_size_limit * recycleHeapPressureRatio
+    );
+  }
 
   // ---------------------------------------------------------------------------
   // Adapted from:
@@ -95,6 +119,39 @@ export class ComponentMetaManager {
         logger.debug(`[reactComponentMeta] Batch extraction failed: ${err}`);
       }
     }
+
+    this.recycleProjectsIfHeapPressured();
+  }
+
+  /**
+   * Reclaim the shared program(s)' resident type-resolution cache when heap usage approaches the V8
+   * limit, preventing the next extraction's spike from OOMing the dev server.
+   */
+  private recycleProjectsIfHeapPressured(): void {
+    if (this.configProjects.size === 0 && !this.inferredProject) {
+      return;
+    }
+    if (process.memoryUsage().heapUsed < this.heapRecycleThresholdBytes) {
+      return;
+    }
+
+    // We crossed the heap-pressure threshold. Recycling reclaims the type cache and usually prevents
+    // the crash, but a single extraction can still exceed the cap — warn once so the user can raise
+    // their memory limit rather than silently paying repeated rebuilds (or eventually crashing).
+    const heapLimitMb = Math.round(v8.getHeapStatistics().heap_size_limit / (1024 * 1024));
+    once.warn(dedent`
+      Storybook's experimental docgen server is nearing the Node.js memory limit (~${heapLimitMb} MB) while extracting component types, and recycled its TypeScript program to avoid an out-of-memory crash. This can briefly slow down the docs and Controls panels.
+
+      If this happens often, raise Node's memory limit before starting Storybook, for example:
+        NODE_OPTIONS="--max-old-space-size=${heapLimitMb * 2}"
+    `);
+
+    for (const project of this.configProjects.values()) {
+      project.dispose();
+    }
+    this.inferredProject?.dispose();
+    this.configProjects.clear();
+    this.inferredProject = undefined;
   }
 
   // ---------------------------------------------------------------------------