diff --git a/src/node/build/build.ts b/src/node/build/build.ts index e4fcf925..b8c88e96 100644 --- a/src/node/build/build.ts +++ b/src/node/build/build.ts @@ -16,7 +16,7 @@ import { bundle } from './bundle' import { generateSitemap } from './generateSitemap' import { renderPage, type RenderPageContext } from './render' import humanizeDuration from 'humanize-duration' -import { launchWorkers, waitWorkers } from '../worker' +import { launchWorkers, stopWorkers } from '../worker' import { registerWorkload, updateContext } from '../worker' import { createMarkdownRenderer } from '../markdown/markdown' import type { DefaultTheme } from '../shared' @@ -210,7 +210,7 @@ export async function build( await siteConfig.buildEnd?.(siteConfig) clearCache() - if (siteConfig.parallel) await waitWorkers('build finish') + if (siteConfig.parallel) await stopWorkers('build finish') const timeEnd = performance.now() const duration = humanizeDuration(timeEnd - timeStart, { diff --git a/src/node/server.ts b/src/node/server.ts index d1792edd..ef4d819b 100644 --- a/src/node/server.ts +++ b/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 } from './worker' +import { launchWorkers, stopWorkers } from './worker' export async function createServer( root: string = process.cwd(), @@ -25,5 +25,9 @@ export async function createServer( server: serverOptions, customLogger: config.logger, configFile: config.vite?.configFile - }) + }).then((server) => + Object.assign({}, server, { + close: () => server.close().then(() => stopWorkers('server.close()')) + }) + ) } diff --git a/src/node/worker.ts b/src/node/worker.ts index d30da076..39fff6e5 100644 --- a/src/node/worker.ts +++ b/src/node/worker.ts @@ -28,22 +28,10 @@ interface WorkerTask { // Owned by main thread, will be distributed to workers const taskQueue = new Queue() -// This function will be exposed to workers via magic proxy -async function getNextTask() { - const task = await taskQueue.dequeue() - if (task !== null) { - debug('[proxy] got task', task.name, '(', debugArgv(...task.argv), ')') - } - debugger - return task -} - function dispatchWork(name: string, ...argv: any[]): Promise { if (workerMeta) { - 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 }) ) @@ -62,6 +50,7 @@ const workers: Array = [] export async function launchWorkers(numWorkers: number, context: Object) { const allInitialized: Array> = [] const ctx = new RpcContext() + const getNextTask = () => taskQueue.dequeue() for (let i = 0; i < numWorkers; i++) { const workerId = (i + 1).toString().padStart(2, '0') const { promise, resolve } = deferPromise() @@ -98,12 +87,24 @@ export function updateContext(context: Object) { } // Wait for workers to drain the taskQueue and exit. -export function waitWorkers() { +export async function stopWorkers(reason: string = 'exit') { const allClosed = workers.map( (w) => new Promise((res) => w.once('exit', res)) ) taskQueue.close() - return Promise.all(allClosed) + debug('waiting for workers, exiting because', reason) + const success = await Promise.any([ + Promise.all(allClosed).then(() => true), + new Promise((res) => setTimeout(() => res(false), 1000)) + ]) + if (!success) { + console.warn('forcefully terminating workers') + for (const w of workers) { + try { + w.terminate() + } catch (e) {} + } + } } /*============================== Worker Thread ==============================*/ @@ -184,7 +185,6 @@ async function workerMainLoop() { reject(e) } } - ctx.reset() }