mirror of https://github.com/vuejs/vitepress
Yuxuan Zhang
2 years ago
commit
046b6530b4
@ -0,0 +1,36 @@
|
|||||||
|
// Asynchronous queue with a close method
|
||||||
|
export default class Queue<T> {
|
||||||
|
private queue: Array<T> = []
|
||||||
|
private pending: Array<(data: T | null) => void> = []
|
||||||
|
#closed: boolean = false
|
||||||
|
|
||||||
|
get closed() {
|
||||||
|
return this.#closed
|
||||||
|
}
|
||||||
|
|
||||||
|
async *items() {
|
||||||
|
while (true) {
|
||||||
|
const item = await this.dequeue()
|
||||||
|
if (item === null) break
|
||||||
|
yield item
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
enqueue(data: T) {
|
||||||
|
if (this.closed)
|
||||||
|
throw new Error(`Failed to enqueue ${data}, queue already closed`)
|
||||||
|
if (this.pending.length) this.pending.shift()!(data)
|
||||||
|
else this.queue.push(data)
|
||||||
|
}
|
||||||
|
|
||||||
|
async dequeue(): Promise<T | null> {
|
||||||
|
if (this.closed) return null
|
||||||
|
if (this.queue.length) return this.queue.shift()!
|
||||||
|
return new Promise((res) => this.pending.push(res))
|
||||||
|
}
|
||||||
|
|
||||||
|
close() {
|
||||||
|
this.#closed = true
|
||||||
|
for (const res of this.pending) res(null)
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -0,0 +1,114 @@
|
|||||||
|
import { Worker, isMainThread, parentPort, workerData } from 'worker_threads'
|
||||||
|
import RpcContext from 'rpc-magic-proxy'
|
||||||
|
import Queue from './utils/queue'
|
||||||
|
|
||||||
|
const WORKER_MAGIC = '::vitepress::build-worker::'
|
||||||
|
|
||||||
|
interface WorkerTask {
|
||||||
|
name: string
|
||||||
|
argv: any[]
|
||||||
|
resolve: (retVal: any) => void
|
||||||
|
reject: (error?: any) => void
|
||||||
|
}
|
||||||
|
|
||||||
|
/*=============================== Main Thread ===============================*/
|
||||||
|
|
||||||
|
// Owned by main thread, will be distributed to workers
|
||||||
|
const taskQueue = new Queue<WorkerTask>()
|
||||||
|
|
||||||
|
// This function will be exposed to workers via magic proxy
|
||||||
|
function getNextTask() {
|
||||||
|
return taskQueue.dequeue()
|
||||||
|
}
|
||||||
|
|
||||||
|
export function dispatchWork(name: string, ...argv: any[]): Promise<any> {
|
||||||
|
return new Promise((resolve, reject) =>
|
||||||
|
taskQueue.enqueue({ name, argv, resolve, reject })
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
const workers: Worker[] = []
|
||||||
|
|
||||||
|
export async function launchWorkers(numWorkers: number, context: Object) {
|
||||||
|
const ctx = new RpcContext()
|
||||||
|
const workerData = await ctx.serialize({
|
||||||
|
[WORKER_MAGIC]: '',
|
||||||
|
getNextTask,
|
||||||
|
context
|
||||||
|
})
|
||||||
|
for (let i = 0; i < numWorkers; i++) {
|
||||||
|
const worker = new Worker(new URL(import.meta.url), { workerData })
|
||||||
|
ctx.bind(worker)
|
||||||
|
workers.push(worker)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export function updateContext(context: Object) {
|
||||||
|
for (const worker of workers) {
|
||||||
|
worker.postMessage({
|
||||||
|
[WORKER_MAGIC]: 'update/context',
|
||||||
|
context
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Wait for workers to drain the taskQueue and exit.
|
||||||
|
export function waitWorkers() {
|
||||||
|
const allClosed = workers.map(
|
||||||
|
(w) => new Promise((res) => w.once('exit', res))
|
||||||
|
)
|
||||||
|
taskQueue.close()
|
||||||
|
return Promise.all(allClosed)
|
||||||
|
}
|
||||||
|
|
||||||
|
/*============================== Worker Thread ==============================*/
|
||||||
|
|
||||||
|
const registry: Map<string, { main: Function; init?: Function }> = new Map()
|
||||||
|
|
||||||
|
export function registerWorkload(
|
||||||
|
name: string,
|
||||||
|
main: (...argv: any[]) => any,
|
||||||
|
init?: () => void
|
||||||
|
) {
|
||||||
|
// Only register workload in worker threads
|
||||||
|
if (isMainThread) return
|
||||||
|
if (registry.has(name)) {
|
||||||
|
throw new Error(`Workload "${name}" already registered.`)
|
||||||
|
}
|
||||||
|
registry.set(name, { main, init })
|
||||||
|
}
|
||||||
|
|
||||||
|
// Will keep querying next workload from main thread
|
||||||
|
async function workerMain() {
|
||||||
|
const ctx = new RpcContext(parentPort!)
|
||||||
|
const {
|
||||||
|
getNextTask,
|
||||||
|
context
|
||||||
|
}: {
|
||||||
|
getNextTask: () => Promise<WorkerTask | null>
|
||||||
|
context: Object
|
||||||
|
} = ctx.deserialize(workerData)
|
||||||
|
|
||||||
|
parentPort!.on('message', (msg) => {
|
||||||
|
if (msg?.[WORKER_MAGIC] === 'update/context') {
|
||||||
|
Object.assign(context, msg.context)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
while (true) {
|
||||||
|
const task = await getNextTask()
|
||||||
|
if (task === null) break
|
||||||
|
const { name, argv, resolve, reject } = task
|
||||||
|
if (!registry.has(name)) throw new Error(`No task "${name}" registered.`)
|
||||||
|
const { main, init } = registry.get(name)!
|
||||||
|
if (init) {
|
||||||
|
init.apply(context)
|
||||||
|
delete registry.get(name)!.init
|
||||||
|
}
|
||||||
|
await (main.apply(context, argv) as Promise<any>).then(resolve, reject)
|
||||||
|
}
|
||||||
|
ctx.reset()
|
||||||
|
process.exit(0)
|
||||||
|
}
|
||||||
|
|
||||||
|
if (!isMainThread && WORKER_MAGIC in workerData) workerMain()
|
||||||
Loading…
Reference in new issue