import { ChildProcess , fork } from "child_process" ; import { EventEmitter } from "events" ; import { Readable } from "stream" ; import { Worker } from "worker_threads" ; import { WorkerMessage , WorkerExecResult } from "./worker" ; export interface WorkerExecResultWithOutput extends WorkerExecResult { stdout : string ; stderr : string ; } export class WorkerPool { private readonly threadWorkerPool : Worker [ ] = [ ] ; private readonly processWorkerPool : ChildProcess [ ] = [ ] ; constructor ( public readonly useWorkerThreads , private readonly memoryUsageLimit ) { } private getThreadWorker ( ) { let worker = this . threadWorkerPool . pop ( ) ; if ( ! worker ) { worker = new Worker ( `${ __dirname }/worker.js` , { stdout : true , stderr : true } ) ; } return worker ; } private getProcessWorker ( ) { let worker = this . processWorkerPool . pop ( ) ; if ( ! worker ) { worker = fork ( `${ __dirname }/worker.js` , this . memoryUsageLimit !== - 1 ? [ "--memoryUsage" ] : undefined , { silent : true } ) ; } return worker ; } public async runOnWorker ( workerName , command , cwd ) { const workerMessage = { workerName , command , cwd } ; const cleanup = [ ] ; const installTemporaryListener = ( object , event , handler ) => { object . on ( event , handler ) ; cleanup . push ( ( ) => object . off ( event , handler ) ) ; } ; const setupWorker = ( worker , res ) => { let stdout = "" ; let stderr = "" ; if ( worker . stdout ) { installTemporaryListener ( worker . stdout , "data" , ( chunk ) => { stdout += chunk ; } ) ; } if ( worker . stderr ) { installTemporaryListener ( worker . stderr , "data" , ( chunk ) => { stderr += chunk ; } ) ; } worker . once ( "message" , ( result ) => { res ( { ... result , stdout , stderr } ) ; } ) ; } ; try { if ( this . useWorkerThreads ) { const worker = this . getThreadWorker ( ) ; const res = await new Promise < WorkerExecResultWithOutput > ( ( res , rej ) => { setupWorker ( worker , res ) ; worker . postMessage ( workerMessage ) ; } ) ; this . threadWorkerPool . push ( worker ) ; return res ; } else { const worker = this . getProcessWorker ( ) ; const res = await new Promise < WorkerExecResultWithOutput > ( ( res , rej ) => { const setupErrorListener = ( event ) => { installTemporaryListener ( worker , event , ( ) => { rej ( new Error ( `Worker ${ event }` ) ) ; } ) ; } ; setupWorker ( worker , res ) ; setupErrorListener ( "close" ) ; setupErrorListener ( "disconnect" ) ; setupErrorListener ( "error" ) ; setupErrorListener ( "exit" ) ; worker . send ( workerMessage ) ; } ) ; if ( this . memoryUsageLimit >= 0 && ( res . memoryUsage ?. rss ?? 0 ) > this . memoryUsageLimit ) { worker . kill ( ) ; } else { this . processWorkerPool . push ( worker ) ; } return res ; } } finally { cleanup . forEach ( value => value ( ) ) ; } } public reset ( ) { this . threadWorkerPool . forEach ( ( worker ) => { worker . terminate ( ) ; } ) ; this . threadWorkerPool . length = 0 ; this . processWorkerPool . forEach ( ( worker ) => { worker . kill ( ) ; } ) ; this . processWorkerPool . length = 0 ; } }