feat/multithread-render: it works!

package.json

@@ -101,6 +101,7 @@
     "mark.js": "8.11.1",
     "minisearch": "^6.3.0",
     "mrmime": "^2.0.0",
+    "rpc-magic-proxy": "0.0.0-beta.0",
     "shikiji": "^0.9.12",
     "shikiji-transformers": "^0.9.12",
     "vite": "^5.0.10",

pnpm-lock.yaml

@@ -44,6 +44,9 @@ importers:
       mrmime:
         specifier: ^2.0.0
         version: 2.0.0
+      rpc-magic-proxy:
+        specifier: 0.0.0-beta.0
+        version: 0.0.0-beta.0
       shikiji:
         specifier: ^0.9.12
         version: 0.9.12
@@ -3913,6 +3916,10 @@ packages:
       '@rollup/rollup-win32-x64-msvc': 4.9.1
       fsevents: 2.3.3
 
+  /rpc-magic-proxy@0.0.0-beta.0:
+    resolution: {integrity: sha512-1UFsu4fpV/OGL8c1knrnx613zzZY9XSbqLUd5pSmOfh0ErAz1R8Dve5S65otFAtNtTO56SWFHbutHcw1lUkkcg==}
+    dev: false
+
   /run-parallel@1.2.0:
     resolution: {integrity: sha512-5l4VyZR86LZ/lDxZTR6jqL8AFE2S0IFLMP26AbjsLVADxHdhB/c0GUsH+y39UfCi3dzz8OlQuPmnaJOMoDHQBA==}
     dependencies:

rollup.config.ts

@@ -46,12 +46,28 @@ const plugins = [
   json()
 ]
 
+const node_root = r('src/node/')
+
 const esmBuild: RollupOptions = {
   input: [r('src/node/index.ts'), r('src/node/cli.ts')],
   output: {
     format: 'esm',
     entryFileNames: `[name].js`,
-    chunkFileNames: 'serve-[hash].js',
+    chunkFileNames(chunk) {
+      console.log('chunkFileNames =>', chunk.name)
+      return `${chunk.name}-[hash].js`
+    },
+    manualChunks(id) {
+      // All workers will be chunked into a single file
+      if (!id.startsWith(node_root)) return
+      id = id.slice(node_root.length).replace(/^\//, '')
+      const match = /^.*\-worker(?=(\.(js|ts))?$)/i.exec(id)
+      if (match) {
+        const [name] = match
+        console.log('manualChunks worker =>', name)
+        return name
+      }
+    },
     dir: r('dist/node'),
     sourcemap: DEV
   },

src/node/build/build.ts

@@ -14,7 +14,7 @@ import { deserializeFunctions, serializeFunctions } from '../utils/fnSerialize'
 import { task } from '../utils/task'
 import { bundle } from './bundle'
 import { generateSitemap } from './generateSitemap'
-import { renderPage } from './render'
+import { renderPage, type RenderPageContext } from './render'
 import humanizeDuration from 'humanize-duration'
 
 export async function build(
@@ -52,10 +52,9 @@ export async function build(
       return
     }
 
-    const entryPath = path.join(siteConfig.tempDir, 'app.js')
-    const { render } = await import(pathToFileURL(entryPath).toString())
-
     await task('rendering pages', async (updateProgress) => {
+      const entryPath = path.join(siteConfig.tempDir, 'app.js')
+
       const appChunk =
         clientResult &&
         (clientResult.output.find(
@@ -109,27 +108,34 @@ export async function build(
         }
       }
 
+      const context: RenderPageContext = {
+        config: siteConfig,
+        result: clientResult,
+        appChunk,
+        cssChunk,
+        assets,
+        pageToHashMap,
+        metadataScript,
+        additionalHeadTags
+      }
+
       const pages = ['404.md', ...siteConfig.pages]
-      let count_done = 0
-      await pMap(
-        pages,
-        async (page) => {
-          await renderPage(
-            render,
-            siteConfig,
-            siteConfig.rewrites.map[page] || page,
-            clientResult,
-            appChunk,
-            cssChunk,
-            assets,
-            pageToHashMap,
-            metadataScript,
-            additionalHeadTags
-          )
-          updateProgress(++count_done, pages.length)
-        },
-        { concurrency: siteConfig.buildConcurrency }
-      )
+
+      if (siteConfig.multithreadRender) {
+        const { default: cluster } = await import('./render-worker')
+        await cluster(entryPath, context, pages, updateProgress)
+      } else {
+        let count_done = 0
+        const { render } = await import(pathToFileURL(entryPath).toString())
+        await pMap(
+          pages,
+          async (page) => {
+            await renderPage(render, page, context)
+            updateProgress(++count_done, pages.length)
+          },
+          { concurrency: siteConfig.buildConcurrency }
+        )
+      }
     })
 
     // emit page hash map for the case where a user session is open

src/node/build/render-worker.ts

@@ -0,0 +1,77 @@
+import { Worker, workerData, isMainThread, parentPort } from 'worker_threads'
+import { type UpdateHandle } from '../utils/task'
+import { type RenderPageContext } from './render'
+
+type TaskAllocator<T> = () => Promise<T | undefined>
+
+import RpcContext from 'rpc-magic-proxy'
+
+export default async function cluster(
+  entryPath: string,
+  context: RenderPageContext,
+  pages: string[],
+  update: UpdateHandle
+) {
+  const concurrency = context.config.buildConcurrency || 1
+  const num_tasks = pages.length
+  let progress = -concurrency
+
+  const pageAlloc: TaskAllocator<string> = async () => {
+    progress++
+    if (progress >= 0) update(progress, num_tasks)
+    return pages.shift()
+  }
+
+  const tasks = []
+
+  for (let _ = 0; _ < concurrency; _++) {
+    const ctx = new RpcContext()
+    const workerData = await ctx.serialize({
+      entryPath,
+      pageAlloc,
+      context,
+      workload: 'render'
+    })
+    const worker = new Worker(new URL(import.meta.url), { workerData })
+    ctx.bind(worker)
+    tasks.push(
+      new Promise((res, rej) =>
+        worker.once('exit', (code) => {
+          if (code === 0) res(code)
+          else rej()
+        })
+      )
+    )
+  }
+
+  await Promise.all(tasks)
+}
+
+async function renderWorker() {
+  const ctx = new RpcContext(parentPort!)
+  try {
+    const {
+      entryPath,
+      pageAlloc,
+      context
+    }: {
+      entryPath: string
+      pageAlloc: TaskAllocator<string>
+      context: RenderPageContext
+    } = ctx.deserialize(workerData)
+    const { pathToFileURL } = await import('url')
+    const { renderPage } = await import('./render')
+    const { render } = await import(pathToFileURL(entryPath).toString())
+    while (true) {
+      const page = await pageAlloc()
+      if (!page) break
+      await renderPage(render, page, context)
+    }
+  } catch (e) {
+    console.error(e)
+  } finally {
+    ctx.reset()
+  }
+}
+
+if (!isMainThread && workerData?.workload === 'render') renderWorker()

src/node/build/render.ts

@@ -19,18 +19,34 @@ import {
 } from '../shared'
 import { version } from '../../../package.json'
 
+export interface RenderPageContext {
+  config: SiteConfig
+  result: Rollup.RollupOutput | null
+  appChunk: Rollup.OutputChunk | null
+  cssChunk: Rollup.OutputAsset | null
+  assets: string[]
+  pageToHashMap: Record<string, string>
+  metadataScript: { html: string; inHead: boolean }
+  additionalHeadTags: HeadConfig[]
+}
+
 export async function renderPage(
   render: (path: string) => Promise<SSGContext>,
-  config: SiteConfig,
-  page: string, // foo.md
-  result: Rollup.RollupOutput | null,
-  appChunk: Rollup.OutputChunk | null,
-  cssChunk: Rollup.OutputAsset | null,
-  assets: string[],
-  pageToHashMap: Record<string, string>,
-  metadataScript: { html: string; inHead: boolean },
-  additionalHeadTags: HeadConfig[]
+  pageNameRaw: string,
+  renderContext: RenderPageContext
 ) {
+  const {
+    config,
+    result,
+    appChunk,
+    cssChunk,
+    assets,
+    pageToHashMap,
+    metadataScript,
+    additionalHeadTags
+  } = renderContext
+
+  const page = config.rewrites.inv[pageNameRaw] || pageNameRaw
   const routePath = `/${page.replace(/\.md$/, '')}`
   const siteData = resolveSiteDataByRoute(config.site, routePath)
 

src/node/siteConfig.ts

@@ -156,6 +156,13 @@ export interface UserConfig<ThemeConfig = any>
    */
   buildConcurrency?: number
 
+  /**
+   * This option allows you to enable or disable the multithread render.
+   * @experimental
+   * @default false
+   */
+  multithreadRender?: boolean
+
   /**
    * @experimental
    *
@@ -250,4 +257,5 @@ export interface SiteConfig<ThemeConfig = any>
   logger: Logger
   userConfig: UserConfig
   buildConcurrency: number
+  multithreadRender: boolean
 }