everything is working !!!

pull/3386/head
Yuxuan Zhang 2 years ago
parent 3479538d5d
commit d79465ab10
No known key found for this signature in database
GPG Key ID: 6910B04F3351EF7D

@ -17,20 +17,26 @@ import { generateSitemap } from './generateSitemap'
import { renderPage, type RenderPageContext } from './render' import { renderPage, type RenderPageContext } from './render'
import humanizeDuration from 'humanize-duration' import humanizeDuration from 'humanize-duration'
import { launchWorkers, waitWorkers } from '../worker' import { launchWorkers, waitWorkers } from '../worker'
import { registerWorkload, updateContext, dispatchWork } from '../worker' import { registerWorkload, updateContext } from '../worker'
type RenderFn = (path: string) => Promise<SSGContext> type RenderFn = (path: string) => Promise<SSGContext>
// Worker: workload functions will be called with `this` context
export interface WorkerContext {
config: SiteConfig
options: BuildOptions
}
// Worker proxy (worker thread) // Worker proxy (worker thread)
registerWorkload( const dispatchRenderPageWork = registerWorkload(
'build::render-page', 'build:render-page',
function workload( function (page: string) {
this: RenderPageContext & { render: RenderFn },
page: string
) {
return renderPage(this.render, page, this) return renderPage(this.render, page, this)
}, },
async function init(this: { renderEntry: string; render: RenderFn }) { async function init(
this: WorkerContext &
RenderPageContext & { render: RenderFn; renderEntry: string }
) {
this.render = (await import(this.renderEntry)).render as RenderFn this.render = (await import(this.renderEntry)).render as RenderFn
} }
) )
@ -46,7 +52,10 @@ export async function build(
const unlinkVue = linkVue() const unlinkVue = linkVue()
if (siteConfig.parallel) if (siteConfig.parallel)
launchWorkers(siteConfig.concurrency, { config: siteConfig }) launchWorkers(siteConfig.concurrency, {
config: siteConfig,
options: buildOptions
})
if (buildOptions.base) { if (buildOptions.base) {
siteConfig.site.base = buildOptions.base siteConfig.site.base = buildOptions.base
@ -64,9 +73,9 @@ export async function build(
} }
try { try {
const { clientResult, serverResult, pageToHashMap } = await bundle( const { clientResult, serverResult, pageToHashMap } = await task(
siteConfig, 'building client + server bundles',
buildOptions () => bundle(siteConfig, buildOptions)
) )
if (process.env.BUNDLE_ONLY) { if (process.env.BUNDLE_ONLY) {
@ -148,7 +157,7 @@ export async function build(
if (siteConfig.parallel) { if (siteConfig.parallel) {
const { config, ...additionalContext } = context const { config, ...additionalContext } = context
await updateContext({ renderEntry, ...additionalContext }) await updateContext({ renderEntry, ...additionalContext })
task = (page) => dispatchWork('build::render-page', page) task = (page) => dispatchRenderPageWork(page)
} else { } else {
const { render } = await import(renderEntry) const { render } = await import(renderEntry)
task = (page) => renderPage(render, page, context) task = (page) => renderPage(render, page, context)

@ -1,38 +1,59 @@
import fs from 'fs-extra' import fs from 'fs-extra'
import path from 'path' import path from 'path'
import { fileURLToPath } from 'url' import { build, type BuildOptions, type Rollup } from 'vite'
import {
build,
normalizePath,
type BuildOptions,
type Rollup,
type InlineConfig as ViteInlineConfig
} from 'vite'
import { APP_PATH } from '../alias'
import type { SiteConfig } from '../config' import type { SiteConfig } from '../config'
import { createVitePressPlugin } from '../plugin' import { updateCurrentTask } from '../utils/task'
import { escapeRegExp, sanitizeFileName, slash } from '../shared'
import { task } from '../utils/task'
import { buildMPAClient } from './buildMPAClient' import { buildMPAClient } from './buildMPAClient'
import { registerWorkload } from '../worker'
import resolveViteConfig from './viteConfig'
import { type WorkerContext } from './build'
const dispatchBundleWorkload = registerWorkload('build:bundle', bundleWorkload)
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(
await resolveViteConfig(ssr, {
config: this.config,
options: this.options,
pageToHashMap,
clientJSMap
})
)) as Rollup.RollupOutput
return { result, pageToHashMap, clientJSMap }
}
// https://github.com/vitejs/vite/blob/d2aa0969ee316000d3b957d7e879f001e85e369e/packages/vite/src/node/plugins/splitVendorChunk.ts#L14 async function bundleMPA(
const CSS_LANGS_RE = config: SiteConfig,
/\.(css|less|sass|scss|styl|stylus|pcss|postcss|sss)(?:$|\?)/ serverResult: Rollup.RollupOutput,
clientJSMap: Record<string, string>
const clientDir = normalizePath( ) {
path.resolve(path.dirname(fileURLToPath(import.meta.url)), '../client') updateCurrentTask(0, 1, 'bundling MPA')
) // in MPA mode, we need to copy over the non-js asset files from the
// server build since there is no client-side build.
// these deps are also being used in the client code (outside of the theme) await Promise.all(
// exclude them from the theme chunk so there is no circular dependency serverResult.output.map(async (chunk) => {
const excludedModules = [ if (!chunk.fileName.endsWith('.js')) {
'/@siteData', const tempPath = path.resolve(config.tempDir, chunk.fileName)
'node_modules/@vueuse/core/', const outPath = path.resolve(config.outDir, chunk.fileName)
'node_modules/@vueuse/shared/', await fs.copy(tempPath, outPath)
'node_modules/vue/', }
'node_modules/vue-demi/', })
clientDir )
] // also copy over public dir
const publicDir = path.resolve(config.srcDir, 'public')
if (fs.existsSync(publicDir)) {
await fs.copy(publicDir, config.outDir)
}
updateCurrentTask()
// build <script client> bundle
if (Object.keys(clientJSMap).length) {
return buildMPAClient(clientJSMap, config)
} else {
return null
}
}
// bundles the VitePress app for both client AND server. // bundles the VitePress app for both client AND server.
export async function bundle( export async function bundle(
@ -46,216 +67,27 @@ export async function bundle(
const pageToHashMap = Object.create(null) const pageToHashMap = Object.create(null)
const clientJSMap = Object.create(null) const clientJSMap = Object.create(null)
// define custom rollup input const [server, client] = await Promise.all(
// this is a multi-entry build - every page is considered an entry chunk config.parallel
// the loading is done via filename conversion rules so that the ? [
// metadata doesn't need to be included in the main chunk. dispatchBundleWorkload(true),
const input: Record<string, string> = {} config.mpa ? null : dispatchBundleWorkload(false)
config.pages.forEach((file) => { ]
// page filename conversion : [
// foo/bar.md -> foo_bar.md bundleWorkload.apply({ config, options }, [true]),
const alias = config.rewrites.map[file] || file config.mpa ? null : bundleWorkload.apply({ config, options }, [false])
input[slash(alias).replace(/\//g, '_')] = path.resolve(config.srcDir, file) ]
})
const themeEntryRE = new RegExp(
`^${escapeRegExp(
path.resolve(config.themeDir, 'index.js').replace(/\\/g, '/')
).slice(0, -2)}m?(j|t)s`
) )
// resolve options to pass to vite // Update maps
const { rollupOptions } = options Object.assign(pageToHashMap, server.pageToHashMap, client?.pageToHashMap)
Object.assign(clientJSMap, server.clientJSMap, client?.clientJSMap)
const resolveViteConfig = async ( return {
ssr: boolean clientResult: config.mpa
): Promise<ViteInlineConfig> => ({ ? await bundleMPA(config, server.result, clientJSMap)
root: config.srcDir, : client?.result!,
cacheDir: config.cacheDir, serverResult: server.result,
base: config.site.base, pageToHashMap
logLevel: config.vite?.logLevel ?? 'warn',
plugins: await createVitePressPlugin(
config,
ssr,
pageToHashMap,
clientJSMap
),
ssr: {
noExternal: ['vitepress', '@docsearch/css']
},
build: {
...options,
emptyOutDir: true,
ssr,
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,
outDir: ssr ? config.tempDir : config.outDir,
cssCodeSplit: false,
rollupOptions: {
...rollupOptions,
input: {
...input,
// use different entry based on ssr or not
app: path.resolve(APP_PATH, ssr ? 'ssr.js' : 'index.js')
},
// important so that each page chunk and the index export things for each
// other
preserveEntrySignatures: 'allow-extension',
output: {
sanitizeFileName,
...rollupOptions?.output,
assetFileNames: `${config.assetsDir}/[name].[hash].[ext]`,
...(ssr
? {
entryFileNames: '[name].js',
chunkFileNames: '[name].[hash].js'
}
: {
entryFileNames: `${config.assetsDir}/[name].[hash].js`,
chunkFileNames(chunk) {
// avoid ads chunk being intercepted by adblock
return /(?:Carbon|BuySell)Ads/.test(chunk.name)
? `${config.assetsDir}/chunks/ui-custom.[hash].js`
: `${config.assetsDir}/chunks/[name].[hash].js`
},
manualChunks(id, ctx) {
// move known framework code into a stable chunk so that
// custom theme changes do not invalidate hash for all pages
if (id.startsWith('\0vite')) {
return 'framework'
}
if (id.includes('plugin-vue:export-helper')) {
return 'framework'
}
if (
id.includes(`${clientDir}/app`) &&
id !== `${clientDir}/app/index.js`
) {
return 'framework'
}
if (
isEagerChunk(id, ctx.getModuleInfo) &&
/@vue\/(runtime|shared|reactivity)/.test(id)
) {
return 'framework'
}
if (
(id.startsWith(`${clientDir}/theme-default`) ||
!excludedModules.some((i) => id.includes(i))) &&
staticImportedByEntry(
id,
ctx.getModuleInfo,
cacheTheme,
themeEntryRE
)
) {
return 'theme'
}
}
})
}
}
},
configFile: config.vite?.configFile
})
const serverResult = await task(
'building server bundle',
() => resolveViteConfig(true).then(build) as Promise<Rollup.RollupOutput>
)
const clientResult = !config.mpa
? await task(
'building client bundle',
() =>
resolveViteConfig(false).then(build) as Promise<Rollup.RollupOutput>
)
: await task('building client bundle (MPA)', async () => {
// in MPA mode, we need to copy over the non-js asset files from the
// server build since there is no client-side build.
await Promise.all(
serverResult.output.map(async (chunk) => {
if (!chunk.fileName.endsWith('.js')) {
const tempPath = path.resolve(config.tempDir, chunk.fileName)
const outPath = path.resolve(config.outDir, chunk.fileName)
await fs.copy(tempPath, outPath)
}
})
)
// also copy over public dir
const publicDir = path.resolve(config.srcDir, 'public')
if (fs.existsSync(publicDir)) {
await fs.copy(publicDir, config.outDir)
}
// build <script client> bundle
if (Object.keys(clientJSMap).length) {
return buildMPAClient(clientJSMap, config)
} else {
return null
}
})
return { clientResult, serverResult, pageToHashMap }
}
const cache = new Map<string, boolean>()
const cacheTheme = new Map<string, boolean>()
/**
* Check if a module is statically imported by at least one entry.
*/
function isEagerChunk(id: string, getModuleInfo: Rollup.GetModuleInfo) {
if (
id.includes('node_modules') &&
!CSS_LANGS_RE.test(id) &&
staticImportedByEntry(id, getModuleInfo, cache)
) {
return true
} }
} }
function staticImportedByEntry(
id: string,
getModuleInfo: Rollup.GetModuleInfo,
cache: Map<string, boolean>,
entryRE: RegExp | null = null,
importStack: string[] = []
): boolean {
if (cache.has(id)) {
return !!cache.get(id)
}
if (importStack.includes(id)) {
// circular deps!
cache.set(id, false)
return false
}
const mod = getModuleInfo(id)
if (!mod) {
cache.set(id, false)
return false
}
if (entryRE ? entryRE.test(id) : mod.isEntry) {
cache.set(id, true)
return true
}
const someImporterIs = mod.importers.some((importer: string) =>
staticImportedByEntry(
importer,
getModuleInfo,
cache,
entryRE,
importStack.concat(id)
)
)
cache.set(id, someImporterIs)
return someImporterIs
}

@ -0,0 +1,223 @@
import path from 'path'
import { fileURLToPath } from 'url'
import {
normalizePath,
type BuildOptions,
type Rollup,
type InlineConfig as ViteInlineConfig
} 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
const CSS_LANGS_RE =
/\.(css|less|sass|scss|styl|stylus|pcss|postcss|sss)(?:$|\?)/
const clientDir = normalizePath(
path.resolve(path.dirname(fileURLToPath(import.meta.url)), '../client')
)
// these deps are also being used in the client code (outside of the theme)
// exclude them from the theme chunk so there is no circular dependency
const excludedModules = [
'/@siteData',
'node_modules/@vueuse/core/',
'node_modules/@vueuse/shared/',
'node_modules/vue/',
'node_modules/vue-demi/',
clientDir
]
const themeEntryRE = (themeDir: string) =>
new RegExp(
`^${escapeRegExp(
path.resolve(themeDir, 'index.js').replace(/\\/g, '/')
).slice(0, -2)}m?(j|t)s`
)
const cache = new Map<string, boolean>()
const cacheTheme = new Map<string, boolean>()
/**
* Check if a module is statically imported by at least one entry.
*/
function isEagerChunk(id: string, getModuleInfo: Rollup.GetModuleInfo) {
if (
id.includes('node_modules') &&
!CSS_LANGS_RE.test(id) &&
staticImportedByEntry(id, getModuleInfo, cache)
) {
return true
}
}
function staticImportedByEntry(
id: string,
getModuleInfo: Rollup.GetModuleInfo,
cache: Map<string, boolean>,
entryRE: RegExp | null = null,
importStack: string[] = []
): boolean {
if (cache.has(id)) {
return !!cache.get(id)
}
if (importStack.includes(id)) {
// circular deps!
cache.set(id, false)
return false
}
const mod = getModuleInfo(id)
if (!mod) {
cache.set(id, false)
return false
}
if (entryRE ? entryRE.test(id) : mod.isEntry) {
cache.set(id, true)
return true
}
const someImporterIs = mod.importers.some((importer: string) =>
staticImportedByEntry(
importer,
getModuleInfo,
cache,
entryRE,
importStack.concat(id)
)
)
cache.set(id, someImporterIs)
return someImporterIs
}
// define custom rollup input
// this is a multi-entry build - every page is considered an entry chunk
// the loading is done via filename conversion rules so that the
// metadata doesn't need to be included in the main chunk.
const resolveInput = (config: SiteConfig) =>
Object.fromEntries(
config.pages.map((file) => {
// page filename conversion
// foo/bar.md -> foo_bar.md
const alias = config.rewrites.map[file] || file
return [
slash(alias).replace(/\//g, '_'),
path.resolve(config.srcDir, file)
]
})
) as Record<string, string>
export default async function resolveViteConfig(
ssr: boolean,
{
config,
options,
pageToHashMap,
clientJSMap
}: {
config: SiteConfig
options: BuildOptions
pageToHashMap: Record<string, string>
clientJSMap: Record<string, string>
}
): Promise<ViteInlineConfig> {
return {
root: config.srcDir,
cacheDir: config.cacheDir,
base: config.site.base,
logLevel: config.vite?.logLevel ?? 'warn',
plugins: await createVitePressPlugin(
config,
ssr,
pageToHashMap,
clientJSMap
),
ssr: {
noExternal: ['vitepress', '@docsearch/css']
},
build: {
...options,
emptyOutDir: true,
ssr,
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,
outDir: ssr ? config.tempDir : config.outDir,
cssCodeSplit: false,
rollupOptions: {
...options.rollupOptions,
input: {
...resolveInput(config),
// use different entry based on ssr or not
app: path.resolve(APP_PATH, ssr ? 'ssr.js' : 'index.js')
},
// important so that each page chunk and the index export things for each
// other
preserveEntrySignatures: 'allow-extension',
output: {
sanitizeFileName,
...options.rollupOptions?.output,
assetFileNames: `${config.assetsDir}/[name].[hash].[ext]`,
...(ssr
? {
entryFileNames: '[name].js',
chunkFileNames: '[name].[hash].js'
}
: {
entryFileNames: `${config.assetsDir}/[name].[hash].js`,
chunkFileNames(chunk) {
// avoid ads chunk being intercepted by adblock
return /(?:Carbon|BuySell)Ads/.test(chunk.name)
? `${config.assetsDir}/chunks/ui-custom.[hash].js`
: `${config.assetsDir}/chunks/[name].[hash].js`
},
manualChunks(id, ctx) {
// move known framework code into a stable chunk so that
// custom theme changes do not invalidate hash for all pages
if (id.startsWith('\0vite')) {
return 'framework'
}
if (id.includes('plugin-vue:export-helper')) {
return 'framework'
}
if (
id.includes(`${clientDir}/app`) &&
id !== `${clientDir}/app/index.js`
) {
return 'framework'
}
if (
isEagerChunk(id, ctx.getModuleInfo) &&
/@vue\/(runtime|shared|reactivity)/.test(id)
) {
return 'framework'
}
if (
(id.startsWith(`${clientDir}/theme-default`) ||
!excludedModules.some((i) => id.includes(i))) &&
staticImportedByEntry(
id,
ctx.getModuleInfo,
cacheTheme,
themeEntryRE(config.themeDir)
)
) {
return 'theme'
}
}
})
}
}
},
configFile: config.vite?.configFile
}
}

@ -72,10 +72,17 @@ export async function createMarkdownToVueRenderFn(
const cacheKey = JSON.stringify({ src, file: fileOrig }) const cacheKey = JSON.stringify({ src, file: fileOrig })
if (isBuild || options.cache !== false) { if (isBuild || options.cache !== false) {
const metrics = {
lookUpTime: performance.now(),
keyLength: cacheKey.length
}
const cached = cache.get(cacheKey) const cached = cache.get(cacheKey)
metrics.lookUpTime = performance.now() - metrics.lookUpTime
if (cached) { if (cached) {
debug(`[cache hit] ${relativePath}`) debug(`[cache hit] ${relativePath} ${JSON.stringify(metrics)}`)
return cached return cached
} else {
debug(`[cache miss] ${relativePath} ${JSON.stringify(metrics)}`)
} }
} }

@ -10,7 +10,8 @@ import {
resolveSiteDataByRoute, resolveSiteDataByRoute,
slash, slash,
type DefaultTheme, type DefaultTheme,
type MarkdownEnv type MarkdownEnv,
type Awaitable
} from '../shared' } from '../shared'
import { processIncludes } from '../utils/processIncludes' import { processIncludes } from '../utils/processIncludes'
import { updateCurrentTask, clearLine } from '../utils/task' import { updateCurrentTask, clearLine } from '../utils/task'
@ -292,30 +293,17 @@ async function* splitPageIntoSections(
} }
/*=============================== Worker API ===============================*/ /*=============================== Worker API ===============================*/
import { registerWorkload, dispatchWork } from '../worker' import { registerWorkload } from '../worker'
import Queue from '../utils/queue' import Queue from '../utils/queue'
// Worker proxy (main thread)
function parallelSplitter(html: string, fileId: string) {
const queue = new Queue<PageSplitSection>()
dispatchWork(
'local-search::split',
html,
fileId,
queue.enqueue.bind(queue),
queue.close.bind(queue)
)
return queue.items()
}
// Worker proxy (worker thread) // Worker proxy (worker thread)
registerWorkload( const dispatchPageSplitWork = registerWorkload(
'local-search::split', 'local-search:split',
async ( async (
html: string, html: string,
fileId: string, fileId: string,
_yield: (section: PageSplitSection) => Promise<void>, _yield: (section: PageSplitSection) => Awaitable<void>,
_end: () => Promise<void> _end: () => Awaitable<void>
) => { ) => {
for await (const section of splitPageIntoSections(html, fileId)) { for await (const section of splitPageIntoSections(html, fileId)) {
await _yield(section) await _yield(section)
@ -323,3 +311,15 @@ registerWorkload(
await _end() await _end()
} }
) )
// Worker proxy (main thread)
function parallelSplitter(html: string, fileId: string) {
const queue = new Queue<PageSplitSection>()
dispatchPageSplitWork(
html,
fileId,
queue.enqueue.bind(queue),
queue.close.bind(queue)
)
return queue.items()
}

@ -1,6 +1,7 @@
import ora from 'ora' import ora from 'ora'
import humanizeDuration from 'humanize-duration' import humanizeDuration from 'humanize-duration'
import c from 'picocolors' import c from 'picocolors'
import { workerMeta } from '../worker'
export const okMark = c.green('✓') export const okMark = c.green('✓')
export const failMark = c.red('✖') export const failMark = c.red('✖')
@ -15,7 +16,8 @@ export type UpdateHandle = (
let updateHandle: UpdateHandle | null = null let updateHandle: UpdateHandle | null = null
export const updateCurrentTask: UpdateHandle = (...args) => { export const updateCurrentTask: UpdateHandle = (...args) => {
if (updateHandle) updateHandle(...args) if (workerMeta) workerMeta.updateCurrentTask(...args)
else if (updateHandle) updateHandle(...args)
else if (!process.stderr.isTTY) { else if (!process.stderr.isTTY) {
return return
} else if (args.length === 0) { } else if (args.length === 0) {
@ -32,6 +34,14 @@ export async function task<T>(
taskName: string, taskName: string,
task: (update: UpdateHandle) => Promise<T> task: (update: UpdateHandle) => Promise<T>
): Promise<T> { ): Promise<T> {
if (workerMeta) {
let retVal: T
await workerMeta.task(taskName, async (handle: UpdateHandle) => {
retVal = await task(handle)
})
return retVal!
}
const spinner = ora({ discardStdin: false }) const spinner = ora({ discardStdin: false })
spinner.start(taskName + '...') spinner.start(taskName + '...')

@ -1,9 +1,23 @@
import { Worker, isMainThread, parentPort, workerData } from 'worker_threads' import { Worker, isMainThread, parentPort, workerData } from 'worker_threads'
import RpcContext from 'rpc-magic-proxy' import RpcContext, { deferPromise } from 'rpc-magic-proxy'
import { task, updateCurrentTask } from './utils/task'
import Queue from './utils/queue' import Queue from './utils/queue'
import _debug from 'debug'
const WORKER_MAGIC = '::vitepress::build-worker::'
let debug = _debug('vitepress:worker:main')
const WORKER_MAGIC = 'vitepress:worker'
function debugArgv(...argv: any[]) {
if (!debug.enabled) return ''
return argv
.map((v) => {
const t = typeof v
if (v?.length !== undefined) return `${t}[${v.length}]`
else return t
})
.join(', ')
}
/*=============================== Main Thread ===============================*/
interface WorkerTask { interface WorkerTask {
name: string name: string
argv: any[] argv: any[]
@ -11,60 +25,66 @@ interface WorkerTask {
reject: (error?: any) => void reject: (error?: any) => void
} }
interface WorkerHooks {
// Update worker's context
context: (ctx: Object | null) => void
}
function deferPromise<T>(): {
promise: Promise<T>
resolve: (val: T) => void
reject: (error?: any) => void
} {
let resolve: (val: T) => void
let reject: (error?: any) => void
const promise = new Promise<T>((res, rej) => {
resolve = res
reject = rej
})
return { promise, resolve: resolve!, reject: reject! }
}
/*=============================== Main Thread ===============================*/
// Owned by main thread, will be distributed to workers // Owned by main thread, will be distributed to workers
const taskQueue = new Queue<WorkerTask>() const taskQueue = new Queue<WorkerTask>()
// This function will be exposed to workers via magic proxy // This function will be exposed to workers via magic proxy
function getNextTask() { async function getNextTask() {
return taskQueue.dequeue() const task = await taskQueue.dequeue()
if (task !== null) {
debug('[proxy] got task', task.name, '(', debugArgv(...task.argv), ')')
}
debugger
return task
} }
export function dispatchWork(name: string, ...argv: any[]): Promise<any> { function dispatchWork(name: string, ...argv: any[]): Promise<any> {
return new Promise((resolve, reject) => if (workerMeta) {
taskQueue.enqueue({ name, argv, resolve, reject }) debug('dispatch', name, '(', debugArgv(...argv), ')')
) return workerMeta.dispatchWork(name, ...argv)
} else {
debug('dispatch', name, '(', debugArgv(...argv), ')')
return new Promise((resolve, reject) =>
taskQueue.enqueue({ name, argv, resolve, reject })
)
}
} }
const workers: Array<Worker & { hooks: WorkerHooks }> = [] type WorkerWithHooks = Worker & {
hooks: {
// Update worker's context
updateContext: (ctx: Object | null) => void
}
}
const workers: Array<WorkerWithHooks> = []
export async function launchWorkers(numWorkers: number, context: Object) { export async function launchWorkers(numWorkers: number, context: Object) {
const allInitialized: Array<Promise<void>> = [] const allInitialized: Array<Promise<void>> = []
const ctx = new RpcContext() const ctx = new RpcContext()
for (let i = 0; i < numWorkers; i++) { for (let i = 0; i < numWorkers; i++) {
const { promise, resolve } = deferPromise<void>() const workerId = (i + 1).toString().padStart(2, '0')
const initWorkerHooks = (hooks: WorkerHooks) => { const { promise, resolve } = deferPromise()
const initWorkerHooks = (hooks: WorkerWithHooks['hooks']) => {
worker.hooks = hooks worker.hooks = hooks
resolve() resolve()
} }
const debug = _debug(`vitepress:worker:${workerId.padEnd(4)}`)
const payload = await ctx.serialize({ const payload = await ctx.serialize({
workerMeta: {
workerId,
dispatchWork,
debug: debug.enabled ? debug : null,
task,
updateCurrentTask
} as typeof workerMeta,
initWorkerHooks, initWorkerHooks,
getNextTask, getNextTask,
context context
}) })
const worker = new Worker(new URL(import.meta.url), { const worker = new Worker(new URL(import.meta.url), {
workerData: { [WORKER_MAGIC]: payload } workerData: { [WORKER_MAGIC]: payload }
}) as Worker & { hooks: WorkerHooks } }) as WorkerWithHooks
ctx.bind(worker) ctx.bind(worker)
workers.push(worker) workers.push(worker)
allInitialized.push(promise) allInitialized.push(promise)
@ -74,7 +94,7 @@ export async function launchWorkers(numWorkers: number, context: Object) {
} }
export function updateContext(context: Object) { export function updateContext(context: Object) {
return Promise.all(workers.map(({ hooks }) => hooks.context(context))) return Promise.all(workers.map(({ hooks }) => hooks.updateContext(context)))
} }
// Wait for workers to drain the taskQueue and exit. // Wait for workers to drain the taskQueue and exit.
@ -88,37 +108,54 @@ export function waitWorkers() {
/*============================== Worker Thread ==============================*/ /*============================== Worker Thread ==============================*/
export let workerMeta: {
workerId: string
dispatchWork: typeof dispatchWork
debug: typeof debug
task: typeof task
updateCurrentTask: typeof updateCurrentTask
} | null = null
const registry: Map<string, { main: Function; init?: Function }> = new Map() const registry: Map<string, { main: Function; init?: Function }> = new Map()
export function registerWorkload( export function registerWorkload<T extends Object, K extends any[], V>(
name: string, name: string,
main: (...argv: any[]) => any, main: (this: T, ...args: K) => V,
init?: () => void init?: (this: T, ...args: void[]) => void
) { ) {
// Only register workload in worker threads if (!isMainThread) {
if (isMainThread) return // Only register workload in worker threads
if (registry.has(name)) { if (registry.has(name)) {
throw new Error(`Workload "${name}" already registered.`) throw new Error(`Workload "${name}" already registered.`)
}
registry.set(name, { main, init })
} }
registry.set(name, { main, init }) return (...args: Parameters<typeof main>) =>
dispatchWork(name, ...args) as Promise<Awaited<ReturnType<typeof main>>>
} }
// Will keep querying next workload from main thread // Will keep querying next workload from main thread
async function workerMain() { async function workerMainLoop() {
const ctx = new RpcContext(parentPort!) const ctx = new RpcContext(parentPort!)
const { const {
workerMeta: _workerMeta,
initWorkerHooks, initWorkerHooks,
getNextTask, getNextTask,
context context
}: { }: {
workerMeta: typeof workerMeta
getNextTask: () => Promise<WorkerTask | null> getNextTask: () => Promise<WorkerTask | null>
initWorkerHooks: (hooks: Object) => Promise<void> initWorkerHooks: (hooks: Object) => Promise<void>
context: Object context: Object
} = ctx.deserialize(workerData[WORKER_MAGIC]) } = ctx.deserialize(workerData[WORKER_MAGIC])
// Set up magic proxy to main thread dispatchWork
workerMeta = _workerMeta!
if (workerMeta.debug) debug = workerMeta.debug
else debug = (() => {}) as any as typeof debug
// Upon worker initialization, report back the hooks that main thread can use // Upon worker initialization, report back the hooks that main thread can use
// to reach this worker. // to reach this worker.
await initWorkerHooks({ await initWorkerHooks({
context(ctx: Object | null) { updateContext(ctx: Object | null) {
if (ctx === null) for (const k in context) delete (context as any)[k] if (ctx === null) for (const k in context) delete (context as any)[k]
else Object.assign(context, ctx) else Object.assign(context, ctx)
} }
@ -128,6 +165,7 @@ async function workerMain() {
const task = await getNextTask() const task = await getNextTask()
if (task === null) break if (task === null) break
const { name, argv, resolve, reject } = task const { name, argv, resolve, reject } = task
debug('got task', name, '(', debugArgv(...argv), ')')
if (!registry.has(name)) throw new Error(`No task "${name}" registered.`) if (!registry.has(name)) throw new Error(`No task "${name}" registered.`)
const el = registry.get(name)! const el = registry.get(name)!
const { main, init } = el const { main, init } = el
@ -142,7 +180,7 @@ async function workerMain() {
try { try {
resolve(await main.apply(context, argv)) resolve(await main.apply(context, argv))
} catch (e) { } catch (e) {
console.error(`worker: task "${name}" error`, e) console.error(`worker:${workerMeta.workerId}: task "${name}" error`, e)
reject(e) reject(e)
} }
} }
@ -150,4 +188,4 @@ async function workerMain() {
ctx.reset() ctx.reset()
} }
if (!isMainThread && WORKER_MAGIC in workerData) workerMain() if (!isMainThread && WORKER_MAGIC in workerData) workerMainLoop()

Loading…
Cancel
Save