Parallelize SSG (#650)

packages/zudoku/package.json

@@ -212,6 +212,7 @@
     "object-hash": "3.0.0",
     "openapi-types": "12.1.3",
     "picocolors": "1.1.1",
+    "piscina": "5.0.0-alpha.1",
     "postcss": "8.5.1",
     "posthog-node": "4.4.1",
     "prism-react-renderer": "2.4.1",

packages/zudoku/src/app/entry.server.tsx

@@ -13,7 +13,7 @@ import "virtual:zudoku-theme.css";
 import "vite/modulepreload-polyfill";
 import { BootstrapStatic, ServerError } from "zudoku/components";
 import type { ZudokuConfig } from "../config/config.js";
-import type { FileWritingResponse } from "../vite/prerender.js";
+import type { FileWritingResponse } from "../vite/prerender/FileWritingResponse.js";
 import "./main.css";
 import { getRoutesByConfig } from "./main.js";
 

packages/zudoku/src/vite/build.ts

@@ -6,7 +6,7 @@ import { joinUrl } from "../lib/util/joinUrl.js";
 import { getViteConfig, loadZudokuConfig } from "./config.js";
 import { getBuildHtml } from "./html.js";
 import { writeOutput } from "./output.js";
-import { prerender } from "./prerender.js";
+import { prerender } from "./prerender/prerender.js";
 
 const DIST_DIR = "dist";
 

packages/zudoku/src/vite/plugin-config.ts

@@ -19,7 +19,7 @@ const viteConfigPlugin = (getConfig: () => ZudokuPluginOptions): Plugin => {
     async transform(code, id) {
       if (id !== getConfig().__meta.path) return;
 
-      return code.replaceAll(
+      const replacedCode = code.replaceAll(
         /process\.env\.([a-z_][a-z0-9_]*)/gi,
         (_, envVar) => {
           if (!envVar.startsWith(viteConfig.envPrefix)) {
@@ -40,6 +40,8 @@ const viteConfigPlugin = (getConfig: () => ZudokuPluginOptions): Plugin => {
           return value;
         },
       );
+
+      return { code: replacedCode, map: null };
     },
   };
 };

packages/zudoku/src/vite/prerender/FileWritingResponse.ts

@@ -0,0 +1,46 @@
+import fs from "node:fs/promises";
+import path from "path";
+
+export class FileWritingResponse {
+  private buffer = "";
+  private dontSave = false;
+  private resolve = () => {};
+  private resolved = new Promise<void>((res) => (this.resolve = res));
+
+  set() {}
+  status(status: number) {
+    if (status >= 300) {
+      this.dontSave = true;
+    }
+  }
+  on() {}
+
+  constructor(private readonly fileName: string) {}
+
+  redirect() {
+    this.buffer = "redirected";
+    this.dontSave = true;
+    this.resolve();
+  }
+
+  send = async (chunk: string) => {
+    this.write(chunk);
+    await this.end();
+  };
+
+  write(chunk: string, _encoding?: string) {
+    this.buffer += chunk;
+  }
+
+  async end(chunk = "") {
+    if (!this.dontSave) {
+      await fs.mkdir(path.dirname(this.fileName), { recursive: true });
+      await fs.writeFile(this.fileName, this.buffer + chunk);
+    }
+    this.resolve();
+  }
+
+  isSent() {
+    return this.resolved;
+  }
+}

packages/zudoku/src/vite/prerender/prerender.ts

@@ -0,0 +1,103 @@
+import path from "node:path";
+import { pathToFileURL } from "node:url";
+import PiscinaImport from "piscina";
+import type { getRoutesByConfig } from "../../app/main.js";
+import { type ZudokuConfig } from "../../config/validators/validate.js";
+import { generateSitemap } from "../sitemap.js";
+import { type WorkerData } from "./worker.js";
+
+const Piscina = PiscinaImport as unknown as typeof PiscinaImport.default;
+
+const routesToPaths = (routes: ReturnType<typeof getRoutesByConfig>) => {
+  const paths: string[] = [];
+  const addPaths = (routes: ReturnType<typeof getRoutesByConfig>) => {
+    for (const route of routes) {
+      // skip catch-all routes
+      if (route.path?.includes("*")) {
+        continue;
+      }
+
+      if (route.path) {
+        paths.push(route.path.startsWith("/") ? route.path : `/${route.path}`);
+      }
+      if (route.children) {
+        addPaths(route.children);
+      }
+    }
+  };
+  addPaths(routes);
+  return paths;
+};
+
+export const prerender = async ({
+  html,
+  dir,
+  basePath = "",
+  serverConfigFilename,
+}: {
+  html: string;
+  dir: string;
+  basePath?: string;
+  serverConfigFilename: string;
+}) => {
+  // eslint-disable-next-line no-console
+  console.log("Prerendering...");
+  const distDir = path.join(dir, "dist", basePath);
+  const config: ZudokuConfig = await import(
+    pathToFileURL(path.join(distDir, "server", serverConfigFilename)).href
+  ).then((m) => m.default);
+
+  const module = await import(
+    pathToFileURL(path.join(distDir, "server/entry.server.js")).href
+  );
+  const getRoutes = module.getRoutesByConfig as typeof getRoutesByConfig;
+
+  const routes = getRoutes(config);
+  const paths = routesToPaths(routes);
+
+  const pool = new Piscina({
+    filename: new URL("./worker.js", import.meta.url).href,
+  });
+
+  const start = performance.now();
+  let completedCount = 0;
+  const writtenFiles = await Promise.all(
+    paths.map(async (urlPath) => {
+      const filename = urlPath === "/" ? "/index.html" : `${urlPath}.html`;
+      const outputPath = path.join(distDir, filename);
+      const url = `http://localhost${config.basePath ?? ""}${urlPath}`;
+      const serverPath = path.join(distDir, "server");
+
+      await pool.run({
+        template: html,
+        outputPath,
+        url,
+        serverConfigPath: path.join(serverPath, serverConfigFilename),
+        entryServerPath: path.join(serverPath, "entry.server.js"),
+      } satisfies WorkerData);
+
+      completedCount++;
+
+      if (process.stdout.isTTY) {
+        process.stdout.write(
+          `\rWritten ${completedCount}/${paths.length} pages`,
+        );
+      }
+      return outputPath;
+    }),
+  );
+
+  const seconds = ((performance.now() - start) / 1000).toFixed(1);
+  process.stdout.write(`\rWritten ${paths.length} pages in ${seconds} seconds`);
+
+  await pool.destroy();
+
+  await generateSitemap({
+    basePath: config.basePath,
+    outputUrls: paths,
+    config: config.sitemap,
+    baseOutputDir: distDir,
+  });
+
+  return writtenFiles;
+};

packages/zudoku/src/vite/prerender/worker.ts

@@ -0,0 +1,47 @@
+import { pathToFileURL } from "node:url";
+import type { render } from "../../app/entry.server.js";
+import type { ZudokuConfig } from "../../config/validators/validate.js";
+import { FileWritingResponse } from "./FileWritingResponse.js";
+
+export type WorkerData = {
+  template: string;
+  outputPath: string;
+  url: string;
+  serverConfigPath: string;
+  entryServerPath: string;
+};
+
+let initialized = false;
+let renderFn: typeof render;
+let config: ZudokuConfig;
+
+const initialize = async ({
+  serverConfigPath,
+  entryServerPath,
+}: WorkerData) => {
+  if (initialized) return;
+
+  const [module, configModule] = await Promise.all([
+    import(pathToFileURL(entryServerPath).href),
+    import(pathToFileURL(serverConfigPath).href),
+  ]);
+
+  renderFn = module.render;
+  config = configModule.default;
+  initialized = true;
+};
+
+const renderPage = async (data: WorkerData): Promise<string> => {
+  await initialize(data);
+
+  const { url, template, outputPath } = data;
+  const request = new Request(url);
+  const response = new FileWritingResponse(outputPath);
+
+  await renderFn({ template, request, response, config });
+  await response.isSent();
+
+  return data.outputPath;
+};
+
+export default renderPage;