11import { ChildProcess } from 'child_process'
22import { Worker as JestWorker } from 'next/dist/compiled/jest-worker'
33import { getNodeOptionsWithoutInspect } from '../server/lib/utils'
4+
5+ // We need this as we're using `Promise.withResolvers` which is not available in the node typings
6+ import '../server/node-environment'
7+
48type FarmOptions = ConstructorParameters < typeof JestWorker > [ 1 ]
59
610const RESTARTED = Symbol ( 'restarted' )
@@ -13,45 +17,55 @@ const cleanupWorkers = (worker: JestWorker) => {
1317 }
1418}
1519
16- export class Worker {
17- private _worker : JestWorker | undefined
20+ type Options < T extends object = object > = FarmOptions & {
21+ timeout ?: number
22+ onRestart ?: ( method : string , args : any [ ] , attempts : number ) => void
23+ exposedMethods : ReadonlyArray < keyof T >
24+ enableWorkerThreads ?: boolean
25+ }
26+
27+ export class Worker < T extends object = object > {
28+ private _worker ?: JestWorker
1829
19- constructor (
30+ /**
31+ * Creates a new worker with the correct typings associated with the selected
32+ * methods.
33+ */
34+ public static create < T extends object > (
2035 workerPath : string ,
21- options : FarmOptions & {
22- timeout ?: number
23- onRestart ?: ( method : string , args : any [ ] , attempts : number ) => void
24- exposedMethods : ReadonlyArray < string >
25- enableWorkerThreads ?: boolean
26- }
27- ) {
36+ options : Options < T >
37+ ) : Worker < T > & T {
38+ return new Worker ( workerPath , options ) as Worker < T > & T
39+ }
40+
41+ constructor ( workerPath : string , options : Options < T > ) {
2842 let { timeout, onRestart, ...farmOptions } = options
2943
3044 let restartPromise : Promise < typeof RESTARTED >
3145 let resolveRestartPromise : ( arg : typeof RESTARTED ) => void
3246 let activeTasks = 0
3347
34- this . _worker = undefined
35-
3648 const createWorker = ( ) => {
37- this . _worker = new JestWorker ( workerPath , {
49+ const worker = new JestWorker ( workerPath , {
3850 ...farmOptions ,
3951 forkOptions : {
4052 ...farmOptions . forkOptions ,
4153 env : {
42- ...( ( farmOptions . forkOptions ?. env || { } ) as any ) ,
54+ ...farmOptions . forkOptions ?. env ,
4355 ...process . env ,
44- // we don't pass down NODE_OPTIONS as it can
45- // extra memory usage
56+ // We don't pass down NODE_OPTIONS as it can lead to extra memory
57+ // usage,
4658 NODE_OPTIONS : getNodeOptionsWithoutInspect ( )
4759 . replace ( / - - m a x - o l d - s p a c e - s i z e = [ \d ] { 1 , } / , '' )
4860 . trim ( ) ,
49- } as any ,
61+ } ,
62+ stdio : 'inherit' ,
5063 } ,
51- } ) as JestWorker
52- restartPromise = new Promise (
53- ( resolve ) => ( resolveRestartPromise = resolve )
54- )
64+ } )
65+
66+ const { promise, resolve } = Promise . withResolvers < typeof RESTARTED > ( )
67+ restartPromise = promise
68+ resolveRestartPromise = resolve
5569
5670 /**
5771 * Jest Worker has two worker types, ChildProcessWorker (uses child_process) and NodeThreadWorker (uses worker_threads)
@@ -63,11 +77,14 @@ export class Worker {
6377 * But this property is not available in NodeThreadWorker, so we need to check if we are using ChildProcessWorker
6478 */
6579 if ( ! farmOptions . enableWorkerThreads ) {
66- for ( const worker of ( ( this . _worker as any ) . _workerPool ?. _workers ||
67- [ ] ) as {
68- _child ?: ChildProcess
69- } [ ] ) {
70- worker . _child ?. on ( 'exit' , ( code , signal ) => {
80+ const poolWorkers : { _child ?: ChildProcess } [ ] =
81+ // @ts -expect-error - we're accessing a private property
82+ worker . _workerPool ?. _workers ?? [ ]
83+
84+ for ( const poolWorker of poolWorkers ) {
85+ if ( ! poolWorker . _child ) continue
86+
87+ poolWorker . _child . once ( 'exit' , ( code , signal ) => {
7188 // log unexpected exit if .end() wasn't called
7289 if ( ( code || signal ) && this . _worker ) {
7390 console . error (
@@ -78,16 +95,22 @@ export class Worker {
7895 }
7996 }
8097
81- this . _worker . getStdout ( ) . pipe ( process . stdout )
82- this . _worker . getStderr ( ) . pipe ( process . stderr )
98+ return worker
8399 }
84- createWorker ( )
100+
101+ // Create the first worker.
102+ this . _worker = createWorker ( )
85103
86104 const onHanging = ( ) => {
87105 const worker = this . _worker
88106 if ( ! worker ) return
107+
108+ // Grab the current restart promise, and create a new worker.
89109 const resolve = resolveRestartPromise
90- createWorker ( )
110+ this . _worker = createWorker ( )
111+
112+ // Once the old worker is ended, resolve the restart promise to signal to
113+ // any active tasks that the worker had to be restarted.
91114 worker . end ( ) . then ( ( ) => {
92115 resolve ( RESTARTED )
93116 } )
@@ -96,33 +119,63 @@ export class Worker {
96119 let hangingTimer : NodeJS . Timeout | false = false
97120
98121 const onActivity = ( ) => {
122+ // If there was an active hanging timer, clear it.
99123 if ( hangingTimer ) clearTimeout ( hangingTimer )
100- hangingTimer = activeTasks > 0 && setTimeout ( onHanging , timeout )
124+
125+ // If there are no active tasks, we don't need to start a new hanging
126+ // timer.
127+ if ( activeTasks === 0 ) return
128+
129+ hangingTimer = setTimeout ( onHanging , timeout )
101130 }
102131
103- for ( const method of farmOptions . exposedMethods ) {
104- if ( method . startsWith ( '_' ) ) continue
105- ; ( this as any ) [ method ] = timeout
106- ? // eslint-disable-next-line no-loop-func
107- async ( ...args : any [ ] ) => {
108- activeTasks ++
109- try {
110- let attempts = 0
111- for ( ; ; ) {
112- onActivity ( )
113- const result = await Promise . race ( [
114- ( this . _worker as any ) [ method ] ( ...args ) ,
115- restartPromise ,
116- ] )
117- if ( result !== RESTARTED ) return result
118- if ( onRestart ) onRestart ( method , args , ++ attempts )
119- }
120- } finally {
121- activeTasks --
122- onActivity ( )
132+ const wrapMethodWithTimeout =
133+ < M extends ( ...args : unknown [ ] ) => Promise < unknown > | unknown > (
134+ method : M
135+ ) =>
136+ async ( ...args : Parameters < M > ) => {
137+ activeTasks ++
138+
139+ try {
140+ let attempts = 0
141+ for ( ; ; ) {
142+ // Mark that we're doing work, we want to ensure that if the worker
143+ // halts for any reason, we restart it.
144+ onActivity ( )
145+
146+ const result = await Promise . race ( [
147+ // Either we'll get the result from the worker, or we'll get the
148+ // restart promise to fire.
149+ method ( ...args ) ,
150+ restartPromise ,
151+ ] )
152+
153+ // If the result anything besides `RESTARTED`, we can return it, as
154+ // it's the actual result from the worker.
155+ if ( result !== RESTARTED ) {
156+ return result
123157 }
158+
159+ // Otherwise, we'll need to restart the worker, and try again.
160+ if ( onRestart ) onRestart ( method . name , args , ++ attempts )
124161 }
125- : ( this . _worker as any ) [ method ] . bind ( this . _worker )
162+ } finally {
163+ activeTasks --
164+ onActivity ( )
165+ }
166+ }
167+
168+ for ( const name of farmOptions . exposedMethods ) {
169+ if ( name . startsWith ( '_' ) ) continue
170+
171+ // @ts -expect-error - we're grabbing a dynamic method on the worker
172+ let method = this . _worker [ name ] . bind ( this . _worker )
173+ if ( timeout ) {
174+ method = wrapMethodWithTimeout ( method )
175+ }
176+
177+ // @ts -expect-error - we're dynamically creating methods
178+ this [ name ] = method
126179 }
127180 }
128181
0 commit comments