siteConfig: provide better control over parallelism, disable parallel bundle by default

package.json

@@ -193,7 +193,7 @@
     "rollup": "^4.9.2",
     "rollup-plugin-dts": "^6.1.0",
     "rollup-plugin-esbuild": "^6.1.0",
-    "rpc-magic-proxy": "^1.0.5",
+    "rpc-magic-proxy": "^1.0.6",
     "semver": "^7.5.4",
     "simple-git-hooks": "^2.9.0",
     "sirv": "^2.0.4",

pnpm-lock.yaml

@@ -280,8 +280,8 @@ importers:
         specifier: ^6.1.0
         version: 6.1.0(esbuild@0.19.11)(rollup@4.9.2)(supports-color@9.4.0)
       rpc-magic-proxy:
-        specifier: ^1.0.5
-        version: 1.0.5
+        specifier: ^1.0.6
+        version: 1.0.6
       semver:
         specifier: ^7.5.4
         version: 7.5.4
@@ -4016,8 +4016,8 @@ packages:
       '@rollup/rollup-win32-x64-msvc': 4.9.2
       fsevents: 2.3.3
 
-  /rpc-magic-proxy@1.0.5:
-    resolution: {integrity: sha512-fd9jbcTrde8F1Nb8+dbl17CYbtBOvzxgnZzYBZ8wV+hwEOZDFPiWvl3yfnJjnYXHmMNSadFkLAa/L9lUA4f8eQ==}
+  /rpc-magic-proxy@1.0.6:
+    resolution: {integrity: sha512-0h0zldD09BLpMi/46xuGFsre2h68LVhIIEqjrSq0bAyeP+ph4/BDmoG4ykNFrU7v2Bx4mS3vX/KNZPVlG6JXQQ==}
     dev: true
 
   /rrweb-cssom@0.6.0:

src/node/build/build.ts

@@ -16,10 +16,8 @@ import { bundle } from './bundle'
 import { generateSitemap } from './generateSitemap'
 import { renderPage, type RenderPageContext } from './render'
 import humanizeDuration from 'humanize-duration'
-import { launchWorkers, stopWorkers } from '../worker'
+import { launchWorkers, shouldUseParallel, stopWorkers } from '../worker'
 import { registerWorkload, updateContext } from '../worker'
-import { createMarkdownRenderer } from '../markdown/markdown'
-import type { DefaultTheme } from '../shared'
 
 type RenderFn = (path: string) => Promise<SSGContext>
 
@@ -53,25 +51,7 @@ export async function build(
   const siteConfig = await resolveConfig(root, 'build', 'production')
   const unlinkVue = linkVue()
 
-  if (siteConfig.parallel) {
-    // Dirty fix: md.render() has side effects on env.
-    // When user provides a custom render function, it will be invoked in the
-    // main thread, but md.render will be called in the worker thread. The side
-    // effects on env will be lost. So we need to make _render a curry function,
-    // and use `md` provided in the main thread
-    const config = siteConfig as SiteConfig<DefaultTheme.Config>
-    const search = config.site?.themeConfig?.search
-    if (search?.provider === 'local' && search?.options?._render) {
-      const md = await createMarkdownRenderer(
-        config.srcDir,
-        config.markdown,
-        config.site.base,
-        config.logger
-      )
-      const _render = search.options._render
-      search.options._render = (src, env, _) => _render(src, env, md)
-    }
-
+  if (shouldUseParallel(siteConfig)) {
     await launchWorkers(siteConfig.concurrency, {
       config: siteConfig,
       options: buildOptions
@@ -175,7 +155,7 @@ export async function build(
 
       let task: (page: string) => Promise<void>
 
-      if (siteConfig.parallel) {
+      if (shouldUseParallel(siteConfig, 'render')) {
         const { config, ...additionalContext } = context
         await updateContext({ renderEntry, ...additionalContext })
         task = (page) => dispatchRenderPageWork(page)
@@ -209,9 +189,7 @@ export async function build(
   await generateSitemap(siteConfig)
   await siteConfig.buildEnd?.(siteConfig)
   clearCache()
-
-  if (siteConfig.parallel) await stopWorkers('build finish')
-
+  await stopWorkers('build complete')
   const timeEnd = performance.now()
   const duration = humanizeDuration(timeEnd - timeStart, {
     maxDecimalPoints: 2

src/node/build/bundle.ts

@@ -1,12 +1,13 @@
 import fs from 'fs-extra'
 import path from 'path'
-import { build, type BuildOptions, type Rollup } from 'vite'
+import { build, type BuildOptions, type PluginOption, type Rollup } from 'vite'
 import type { SiteConfig } from '../config'
 import { updateCurrentTask } from '../utils/task'
 import { buildMPAClient } from './buildMPAClient'
-import { registerWorkload } from '../worker'
+import { registerWorkload, shouldUseParallel } from '../worker'
 import resolveViteConfig from './viteConfig'
 import { type WorkerContext } from './build'
+import { createVitePressPlugin } from '../plugin'
 
 const dispatchBundleWorkload = registerWorkload(
   'build:bundle',
@@ -18,18 +19,18 @@ const dispatchBundleWorkload = registerWorkload(
   }
 )
 
-async function bundleWorkload(this: WorkerContext, ssr: boolean) {
-  const pageToHashMap = Object.create(null) as Record<string, string>
-  const clientJSMap = Object.create(null) as Record<string, string>
-  const result = (await build(
+async function bundleWorkload(
+  this: WorkerContext,
+  ssr: boolean,
+  plugins: PluginOption[]
+) {
+  return build(
     await resolveViteConfig(ssr, {
       config: this.config,
       options: this.options,
-      pageToHashMap,
-      clientJSMap
+      plugins
     })
-  )) as Rollup.RollupOutput
-  return { result, pageToHashMap, clientJSMap }
+  ) as Promise<Rollup.RollupOutput>
 }
 
 async function bundleMPA(
@@ -75,27 +76,26 @@ export async function bundle(
   const pageToHashMap = Object.create(null)
   const clientJSMap = Object.create(null)
 
-  const [server, client] = await Promise.all(
-    config.parallel
-      ? [
-          dispatchBundleWorkload(true),
-          config.mpa ? null : dispatchBundleWorkload(false)
-        ]
-      : [
-          bundleWorkload.apply({ config, options }, [true]),
-          config.mpa ? null : bundleWorkload.apply({ config, options }, [false])
-        ]
+  const [serverResult, clientResult] = await Promise.all(
+    [true, false].map(async (ssr) => {
+      if (!ssr && config.mpa) return null
+      const plugins = await createVitePressPlugin(
+        config,
+        ssr,
+        pageToHashMap,
+        clientJSMap
+      )
+      return shouldUseParallel(config, 'bundle')
+        ? dispatchBundleWorkload(ssr, plugins)
+        : bundleWorkload.apply({ config, options }, [ssr, plugins])
+    })
   )
 
-  // Update maps
-  Object.assign(pageToHashMap, server.pageToHashMap, client?.pageToHashMap)
-  Object.assign(clientJSMap, server.clientJSMap, client?.clientJSMap)
-
   return {
     clientResult: config.mpa
-      ? await bundleMPA(config, server.result, clientJSMap)
-      : client?.result!,
-    serverResult: server.result,
+      ? await bundleMPA(config, serverResult!, clientJSMap)
+      : clientResult!,
+    serverResult: serverResult!,
     pageToHashMap
   }
 }

src/node/build/viteConfig.ts

@@ -4,12 +4,12 @@ import {
   normalizePath,
   type BuildOptions,
   type Rollup,
-  type InlineConfig as ViteInlineConfig
+  type InlineConfig as ViteInlineConfig,
+  type PluginOption
 } from 'vite'
 import { APP_PATH } from '../alias'
 import type { SiteConfig } from '../config'
 import { slash } from '../shared'
-import { createVitePressPlugin } from '../plugin'
 import { escapeRegExp, sanitizeFileName } from '../shared'
 
 // https://github.com/vitejs/vite/blob/d2aa0969ee316000d3b957d7e879f001e85e369e/packages/vite/src/node/plugins/splitVendorChunk.ts#L14
@@ -114,13 +114,11 @@ export default async function resolveViteConfig(
   {
     config,
     options,
-    pageToHashMap,
-    clientJSMap
+    plugins
   }: {
     config: SiteConfig
     options: BuildOptions
-    pageToHashMap: Record<string, string>
-    clientJSMap: Record<string, string>
+    plugins: PluginOption[]
   }
 ): Promise<ViteInlineConfig> {
   return {
@@ -128,12 +126,7 @@ export default async function resolveViteConfig(
     cacheDir: config.cacheDir,
     base: config.site.base,
     logLevel: config.vite?.logLevel ?? 'warn',
-    plugins: await createVitePressPlugin(
-      config,
-      ssr,
-      pageToHashMap,
-      clientJSMap
-    ),
+    plugins,
     ssr: {
       noExternal: ['vitepress', '@docsearch/css']
     },
@@ -144,12 +137,8 @@ export default async function resolveViteConfig(
       ssrEmitAssets: config.mpa,
       // minify with esbuild in MPA mode (for CSS)
       minify: ssr
-        ? config.mpa
-          ? 'esbuild'
-          : false
-        : typeof options.minify === 'boolean'
-          ? options.minify
-          : !process.env.DEBUG,
+        ? config.mpa && 'esbuild'
+        : options.minify ?? !process.env.DEBUG,
       outDir: ssr ? config.tempDir : config.outDir,
       cssCodeSplit: false,
       rollupOptions: {

src/node/config.ts

@@ -148,7 +148,7 @@ export async function resolveConfig(
       userConfig.concurrency ?? Math.round(cpus().length / 1.5),
       1 // At least one thread required
     ),
-    parallel: userConfig.parallel ?? true
+    parallel: userConfig.parallel ?? ['render', 'local-search']
   }
 
   // to be shared with content loaders

src/node/plugins/localSearchPlugin.ts

@@ -173,9 +173,12 @@ export async function localSearchPlugin(
         siteConfig.pages.length,
         'indexing local search'
       )
+
+    const parallel = shouldUseParallel(siteConfig, 'local-search')
+
     await pMap(
       siteConfig.pages,
-      (page) => indexFile(page, !!siteConfig.parallel).then(updateProgress),
+      (page) => indexFile(page, parallel).then(updateProgress),
       { concurrency }
     )
 
@@ -293,7 +296,7 @@ async function* splitPageIntoSections(
 }
 
 /*=============================== Worker API ===============================*/
-import { registerWorkload } from '../worker'
+import { registerWorkload, shouldUseParallel } from '../worker'
 import Queue from '../utils/queue'
 
 // Worker proxy (worker thread)

src/node/server.ts

@@ -1,7 +1,7 @@
 import { createServer as createViteServer, type ServerOptions } from 'vite'
 import { resolveConfig } from './config'
 import { createVitePressPlugin } from './plugin'
-import { launchWorkers, stopWorkers } from './worker'
+import { launchWorkers, stopWorkers, shouldUseParallel } from './worker'
 
 export async function createServer(
   root: string = process.cwd(),
@@ -10,7 +10,8 @@ export async function createServer(
 ) {
   const config = await resolveConfig(root)
 
-  if (config.parallel) launchWorkers(config.concurrency, { config: config })
+  if (shouldUseParallel(config))
+    launchWorkers(config.concurrency, { config: config })
 
   if (serverOptions.base) {
     config.site.base = serverOptions.base

src/node/siteConfig.ts

@@ -13,6 +13,7 @@ import type {
   SSGContext,
   SiteData
 } from './shared'
+import type { SupportsParallel } from './worker'
 
 export type RawConfigExports<ThemeConfig = any> =
   | Awaitable<UserConfig<ThemeConfig>>
@@ -165,9 +166,9 @@ export interface UserConfig<ThemeConfig = any>
    * 2. Parallel SSR Rendering
    * 3. Parallel Local Search Indexing (when using default splitter)
    * @experimental
-   * @default true
+   * @default ['render', 'local-search']
    */
-  parallel?: boolean
+  parallel?: boolean | SupportsParallel[]
 
   /**
    * @experimental
@@ -263,5 +264,5 @@ export interface SiteConfig<ThemeConfig = any>
   logger: Logger
   userConfig: UserConfig
   concurrency: number
-  parallel: boolean
+  parallel: boolean | SupportsParallel[]
 }

src/node/worker.ts

@@ -1,8 +1,25 @@
 import { Worker, isMainThread, parentPort, workerData } from 'worker_threads'
 import RpcContext, { deferPromise } from 'rpc-magic-proxy'
 import { task, updateCurrentTask } from './utils/task'
+import c from 'picocolors'
 import Queue from './utils/queue'
 import _debug from 'debug'
+import type { SiteConfig } from 'siteConfig'
+
+export type SupportsParallel = 'bundle' | 'render' | 'local-search'
+
+/**
+ * Checks if the given task should be run in parallel.
+ * If task is omitted, checks if any task should be run in parallel.
+ */
+export function shouldUseParallel(config: SiteConfig, task?: SupportsParallel) {
+  const { parallel = false } = config
+  if (task === undefined)
+    return parallel === true || (Array.isArray(parallel) && parallel.length > 0)
+  if (typeof parallel === 'boolean') return parallel
+  if (Array.isArray(parallel)) return parallel.includes(task)
+  throw new TypeError(`Invalid value for config.parallel: ${parallel}`)
+}
 
 let debug = _debug('vitepress:worker:main')
 const WORKER_MAGIC = 'vitepress:worker'
@@ -66,7 +83,7 @@ export async function launchWorkers(numWorkers: number, context: Object) {
         debug: debug.enabled ? debug : null,
         task,
         updateCurrentTask
-      } as typeof workerMeta,
+      },
       initWorkerHooks,
       getNextTask,
       context
@@ -137,7 +154,7 @@ export function registerWorkload<T extends Object, K extends any[], V>(
 
 // Will keep querying next workload from main thread
 async function workerMainLoop() {
-  const ctx = new RpcContext(parentPort!)
+  const ctx = new RpcContext({ preserveThis: true }).bind(parentPort!)
   const {
     workerMeta: _workerMeta,
     initWorkerHooks,
@@ -174,14 +191,17 @@ async function workerMainLoop() {
       try {
         await init.apply(context)
       } catch (e) {
-        console.error(`worker: failed to init workload "${name}": ${e}`)
+        console.error(c.red(`worker: failed to init workload "${name}":`), e)
       }
       el.init = undefined
     }
     try {
       resolve(await main.apply(context, argv))
     } catch (e) {
-      console.error(`worker:${workerMeta.workerId}: task "${name}" error`, e)
+      console.error(
+        c.red(`worker:${workerMeta.workerId} error running task "${name}":`),
+        e
+      )
       reject(e)
     }
   }