SSG thread pool

packages/docusaurus/package.json

@@ -69,6 +69,7 @@
     "semver": "^7.5.4",
     "serve-handler": "^6.1.6",
     "shelljs": "^0.8.5",
+    "tinypool": "^1.0.2",
     "tslib": "^2.6.0",
     "update-notifier": "^6.0.2",
     "webpack": "^5.95.0",

packages/docusaurus/src/ssg/ssg.ts

@@ -152,13 +152,15 @@ Troubleshooting guide: https://github.com/facebook/docusaurus/discussions/10580
   }
 }
 
+export type GenerateStaticFilesResult = {collectedData: SiteCollectedData};
+
 export async function generateStaticFiles({
   pathnames,
   params,
 }: {
   pathnames: string[];
   params: SSGParams;
-}): Promise<{collectedData: SiteCollectedData}> {
+}): Promise<GenerateStaticFilesResult> {
   const [renderer, htmlMinifier, ssgTemplate] = await Promise.all([
     PerfLogger.async('Load App renderer', () =>
       loadAppRenderer({
@@ -302,3 +304,10 @@ It might also require to wrap your client code in ${logger.code(
 
   return parts.join('\n');
 }
+
+export default async function worker(arg: {
+  pathnames: string[];
+  params: SSGParams;
+}): Promise<{collectedData: SiteCollectedData}> {
+  return generateStaticFiles(arg);
+}

packages/docusaurus/src/ssg/ssgExecutor.ts

@@ -5,13 +5,113 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+import * as path from 'path';
+import os from 'os';
+import _ from 'lodash';
 import {PerfLogger} from '@docusaurus/logger';
-import {createSSGParams} from './ssgParams';
 import {generateStaticFiles} from './ssg';
+import {createSSGParams} from './ssgParams';
 import {renderHashRouterTemplate} from './ssgTemplate';
 import {generateHashRouterEntrypoint} from './ssgUtils';
 import type {Props, RouterType} from '@docusaurus/types';
 import type {SiteCollectedData} from '../common';
+import type {GenerateStaticFilesResult} from './ssg';
+import type {SSGParams} from './ssgParams';
+
+function mergeResults(
+  results: GenerateStaticFilesResult[],
+): GenerateStaticFilesResult {
+  return {
+    collectedData: Object.assign({}, ...results.map((r) => r.collectedData)),
+  };
+}
+
+type SSGExecutor = {
+  run: () => Promise<GenerateStaticFilesResult>;
+  destroy: () => Promise<void>;
+};
+
+type CreateSSGExecutor = (params: {
+  params: SSGParams;
+  pathnames: string[];
+}) => Promise<SSGExecutor>;
+
+const createSimpleSSGExecutor: CreateSSGExecutor = async ({
+  params,
+  pathnames,
+}) => {
+  return {
+    run: () => {
+      return PerfLogger.async('Generate static files', () =>
+        generateStaticFiles({
+          pathnames,
+          params,
+        }),
+      );
+    },
+
+    destroy: async () => {
+      // nothing to do
+    },
+  };
+};
+
+const createPooledSSGExecutor: CreateSSGExecutor = async ({
+  params,
+  pathnames,
+}) => {
+  // TODO make this configurable
+  // Sensible default that gives the best improvement so far:
+  const numberOfThreads = os.cpus().length / 2;
+
+  const pathnamesChunks = _.chunk(
+    pathnames,
+    Math.ceil(pathnames.length / numberOfThreads),
+  );
+
+  const pool = await PerfLogger.async(
+    `Create SSG pool with ${numberOfThreads} threads`,
+    async () => {
+      const Tinypool = await import('tinypool').then((m) => m.default);
+      return new Tinypool({
+        filename: path.resolve(__dirname, './ssg.js'),
+        minThreads: numberOfThreads,
+        maxThreads: numberOfThreads,
+        concurrentTasksPerWorker: 1,
+        runtime: 'worker_threads',
+      });
+    },
+  );
+
+  return {
+    run: async () => {
+      const results = await PerfLogger.async(
+        'Generate static files - pooled',
+        async () => {
+          return Promise.all(
+            pathnamesChunks.map((chunk, chunkIndex) => {
+              return PerfLogger.async(
+                `Generate static files for chunk=${chunkIndex} with ${chunk.length} pathnames`,
+                () => {
+                  return pool.run({
+                    pathnames: chunk,
+                    params,
+                  }) as Promise<GenerateStaticFilesResult>;
+                },
+              );
+            }),
+          );
+        },
+      );
+
+      return mergeResults(results);
+    },
+
+    destroy: async () => {
+      await pool.destroy();
+    },
+  };
+};
 
 // TODO Docusaurus v4 - introduce SSG worker threads
 export async function executeSSG({
@@ -39,12 +139,12 @@ export async function executeSSG({
     return {collectedData: {}};
   }
 
-  const ssgResult = await PerfLogger.async('Generate static files', () =>
-    generateStaticFiles({
-      pathnames: props.routesPaths,
-      params,
-    }),
-  );
+  const createExecutor =
+    process.env.DOCUSAURUS_DISABLE_SSG_POOL === 'true'
+      ? createSimpleSSGExecutor
+      : createPooledSSGExecutor;
+
+  const executor = await createExecutor({params, pathnames: props.routesPaths});
 
-  return ssgResult;
+  return executor.run();
 }

yarn.lock

@@ -17312,6 +17312,11 @@ tinyglobby@^0.2.9:
     fdir "^6.4.0"
     picomatch "^4.0.2"
 
+tinypool@^1.0.2:
+  version "1.0.2"
+  resolved "https://registry.yarnpkg.com/tinypool/-/tinypool-1.0.2.tgz#706193cc532f4c100f66aa00b01c42173d9051b2"
+  integrity sha512-al6n+QEANGFOMf/dmUMsuS5/r9B06uwlyNjZZql/zv8J7ybHCgoihBNORZCY2mzUuAnomQa2JdhyHKzZxPCrFA==
+
 tmp-promise@^3.0.3:
   version "3.0.3"
   resolved "https://registry.yarnpkg.com/tmp-promise/-/tmp-promise-3.0.3.tgz#60a1a1cc98c988674fcbfd23b6e3367bdeac4ce7"