refactor: improved worker typings

wyattjoh · wyattjoh · commit 919eaa4c43cf · 2023-10-02T11:00:05.000-06:00
diff --git a/packages/next/src/build/index.ts b/packages/next/src/build/index.ts
@@ -1127,7 +1127,15 @@ export default async function build(
       ) {
         let infoPrinted = false
 
-        return new Worker(staticWorkerPath, {
+        return Worker.create<
+          Pick<
+            typeof import('./worker'),
+            | 'hasCustomGetInitialProps'
+            | 'isPageStatic'
+            | 'getDefinedNamedExports'
+            | 'exportPage'
+          >
+        >(staticWorkerPath, {
           timeout: timeout * 1000,
           onRestart: (method, [arg], attempts) => {
             if (method === 'exportPage') {
@@ -1174,14 +1182,7 @@ export default async function build(
             'getDefinedNamedExports',
             'exportPage',
           ],
-        }) as Worker &
-          Pick<
-            typeof import('./worker'),
-            | 'hasCustomGetInitialProps'
-            | 'isPageStatic'
-            | 'getDefinedNamedExports'
-            | 'exportPage'
-          >
+        })
       }
 
       let CacheHandler: any
@@ -2054,10 +2055,10 @@ export default async function build(
             pages: combinedPages,
             outdir: path.join(distDir, 'export'),
             statusMessage: 'Generating static pages',
-            exportAppPageWorker:
-              appStaticWorkers?.exportPage.bind(appStaticWorkers),
-            exportPageWorker:
-              pagesStaticWorkers.exportPage.bind(pagesStaticWorkers),
+            // The worker already explicitly binds `this` to each of the
+            // exposed methods.
+            exportAppPageWorker: appStaticWorkers?.exportPage,
+            exportPageWorker: pagesStaticWorkers?.exportPage,
             endWorker: async () => {
               await pagesStaticWorkers.end()
               await appStaticWorkers?.end()
@@ -2754,8 +2755,10 @@ export default async function build(
           silent: true,
           threads: config.experimental.cpus,
           outdir: path.join(dir, configOutDir),
-          exportAppPageWorker: appWorker.exportPage.bind(appWorker),
-          exportPageWorker: pagesWorker.exportPage.bind(pagesWorker),
+          // The worker already explicitly binds `this` to each of the
+          // exposed methods.
+          exportAppPageWorker: appWorker?.exportPage,
+          exportPageWorker: pagesWorker?.exportPage,
           endWorker: async () => {
             await pagesWorker.end()
             await appWorker.end()
diff --git a/packages/next/src/export/index.ts b/packages/next/src/export/index.ts
@@ -173,32 +173,35 @@ function setupWorkers(
 
   let infoPrinted = false
 
-  const worker = new Worker(require.resolve('./worker'), {
-    timeout: timeout * 1000,
-    onRestart: (_method, [{ path }], attempts) => {
-      if (attempts >= 3) {
-        throw new ExportError(
-          `Static page generation for ${path} is still timing out after 3 attempts. See more info here https://nextjs.org/docs/messages/static-page-generation-timeout`
-        )
-      }
-      Log.warn(
-        `Restarted static page generation for ${path} because it took more than ${timeout} seconds`
-      )
-      if (!infoPrinted) {
+  const worker = Worker.create<typeof import('./worker')>(
+    require.resolve('./worker'),
+    {
+      timeout: timeout * 1000,
+      onRestart: (_method, [{ path }], attempts) => {
+        if (attempts >= 3) {
+          throw new ExportError(
+            `Static page generation for ${path} is still timing out after 3 attempts. See more info here https://nextjs.org/docs/messages/static-page-generation-timeout`
+          )
+        }
         Log.warn(
-          'See more info here https://nextjs.org/docs/messages/static-page-generation-timeout'
+          `Restarted static page generation for ${path} because it took more than ${timeout} seconds`
         )
-        infoPrinted = true
-      }
-    },
-    maxRetries: 0,
-    numWorkers: threads,
-    enableWorkerThreads: nextConfig.experimental.workerThreads,
-    exposedMethods: ['default'],
-  })
+        if (!infoPrinted) {
+          Log.warn(
+            'See more info here https://nextjs.org/docs/messages/static-page-generation-timeout'
+          )
+          infoPrinted = true
+        }
+      },
+      maxRetries: 0,
+      numWorkers: threads,
+      enableWorkerThreads: nextConfig.experimental.workerThreads,
+      exposedMethods: ['default'],
+    }
+  )
 
   return {
-    pages: (worker as any).default as ExportWorker,
+    pages: worker.default,
     end: async () => {
       await worker.end()
     },
diff --git a/packages/next/src/lib/worker.ts b/packages/next/src/lib/worker.ts
@@ -1,6 +1,10 @@
 import { ChildProcess } from 'child_process'
 import { Worker as JestWorker } from 'next/dist/compiled/jest-worker'
 import { getNodeOptionsWithoutInspect } from '../server/lib/utils'
+
+// We need this as we're using `Promise.withResolvers` which is not available in the node typings
+import '../server/node-environment'
+
 type FarmOptions = ConstructorParameters<typeof JestWorker>[1]
 
 const RESTARTED = Symbol('restarted')
@@ -13,45 +17,55 @@ const cleanupWorkers = (worker: JestWorker) => {
   }
 }
 
-export class Worker {
-  private _worker: JestWorker | undefined
+type Options<T extends object = object> = FarmOptions & {
+  timeout?: number
+  onRestart?: (method: string, args: any[], attempts: number) => void
+  exposedMethods: ReadonlyArray<keyof T>
+  enableWorkerThreads?: boolean
+}
+
+export class Worker<T extends object = object> {
+  private _worker?: JestWorker
 
-  constructor(
+  /**
+   * Creates a new worker with the correct typings associated with the selected
+   * methods.
+   */
+  public static create<T extends object>(
     workerPath: string,
-    options: FarmOptions & {
-      timeout?: number
-      onRestart?: (method: string, args: any[], attempts: number) => void
-      exposedMethods: ReadonlyArray<string>
-      enableWorkerThreads?: boolean
-    }
-  ) {
+    options: Options<T>
+  ): Worker<T> & T {
+    return new Worker(workerPath, options) as Worker<T> & T
+  }
+
+  constructor(workerPath: string, options: Options<T>) {
     let { timeout, onRestart, ...farmOptions } = options
 
     let restartPromise: Promise<typeof RESTARTED>
     let resolveRestartPromise: (arg: typeof RESTARTED) => void
     let activeTasks = 0
 
-    this._worker = undefined
-
     const createWorker = () => {
-      this._worker = new JestWorker(workerPath, {
+      const worker = new JestWorker(workerPath, {
         ...farmOptions,
         forkOptions: {
           ...farmOptions.forkOptions,
           env: {
-            ...((farmOptions.forkOptions?.env || {}) as any),
+            ...farmOptions.forkOptions?.env,
             ...process.env,
-            // we don't pass down NODE_OPTIONS as it can
-            // extra memory usage
+            // We don't pass down NODE_OPTIONS as it can lead to extra memory
+            // usage,
             NODE_OPTIONS: getNodeOptionsWithoutInspect()
               .replace(/--max-old-space-size=[\d]{1,}/, '')
               .trim(),
-          } as any,
+          },
+          stdio: 'inherit',
         },
-      }) as JestWorker
-      restartPromise = new Promise(
-        (resolve) => (resolveRestartPromise = resolve)
-      )
+      })
+
+      const { promise, resolve } = Promise.withResolvers<typeof RESTARTED>()
+      restartPromise = promise
+      resolveRestartPromise = resolve
 
       /**
        * Jest Worker has two worker types, ChildProcessWorker (uses child_process) and NodeThreadWorker (uses worker_threads)
@@ -63,11 +77,14 @@ export class Worker {
        * But this property is not available in NodeThreadWorker, so we need to check if we are using ChildProcessWorker
        */
       if (!farmOptions.enableWorkerThreads) {
-        for (const worker of ((this._worker as any)._workerPool?._workers ||
-          []) as {
-          _child?: ChildProcess
-        }[]) {
-          worker._child?.on('exit', (code, signal) => {
+        const poolWorkers: { _child?: ChildProcess }[] =
+          // @ts-expect-error - we're accessing a private property
+          worker._workerPool?._workers ?? []
+
+        for (const poolWorker of poolWorkers) {
+          if (!poolWorker._child) continue
+
+          poolWorker._child.once('exit', (code, signal) => {
             // log unexpected exit if .end() wasn't called
             if ((code || signal) && this._worker) {
               console.error(
@@ -78,16 +95,22 @@ export class Worker {
         }
       }
 
-      this._worker.getStdout().pipe(process.stdout)
-      this._worker.getStderr().pipe(process.stderr)
+      return worker
     }
-    createWorker()
+
+    // Create the first worker.
+    this._worker = createWorker()
 
     const onHanging = () => {
       const worker = this._worker
       if (!worker) return
+
+      // Grab the current restart promise, and create a new worker.
       const resolve = resolveRestartPromise
-      createWorker()
+      this._worker = createWorker()
+
+      // Once the old worker is ended, resolve the restart promise to signal to
+      // any active tasks that the worker had to be restarted.
       worker.end().then(() => {
         resolve(RESTARTED)
       })
@@ -96,33 +119,63 @@ export class Worker {
     let hangingTimer: NodeJS.Timeout | false = false
 
     const onActivity = () => {
+      // If there was an active hanging timer, clear it.
       if (hangingTimer) clearTimeout(hangingTimer)
-      hangingTimer = activeTasks > 0 && setTimeout(onHanging, timeout)
+
+      // If there are no active tasks, we don't need to start a new hanging
+      // timer.
+      if (activeTasks === 0) return
+
+      hangingTimer = setTimeout(onHanging, timeout)
     }
 
-    for (const method of farmOptions.exposedMethods) {
-      if (method.startsWith('_')) continue
-      ;(this as any)[method] = timeout
-        ? // eslint-disable-next-line no-loop-func
-          async (...args: any[]) => {
-            activeTasks++
-            try {
-              let attempts = 0
-              for (;;) {
-                onActivity()
-                const result = await Promise.race([
-                  (this._worker as any)[method](...args),
-                  restartPromise,
-                ])
-                if (result !== RESTARTED) return result
-                if (onRestart) onRestart(method, args, ++attempts)
-              }
-            } finally {
-              activeTasks--
-              onActivity()
+    const wrapMethodWithTimeout =
+      <M extends (...args: unknown[]) => Promise<unknown> | unknown>(
+        method: M
+      ) =>
+      async (...args: Parameters<M>) => {
+        activeTasks++
+
+        try {
+          let attempts = 0
+          for (;;) {
+            // Mark that we're doing work, we want to ensure that if the worker
+            // halts for any reason, we restart it.
+            onActivity()
+
+            const result = await Promise.race([
+              // Either we'll get the result from the worker, or we'll get the
+              // restart promise to fire.
+              method(...args),
+              restartPromise,
+            ])
+
+            // If the result anything besides `RESTARTED`, we can return it, as
+            // it's the actual result from the worker.
+            if (result !== RESTARTED) {
+              return result
             }
+
+            // Otherwise, we'll need to restart the worker, and try again.
+            if (onRestart) onRestart(method.name, args, ++attempts)
           }
-        : (this._worker as any)[method].bind(this._worker)
+        } finally {
+          activeTasks--
+          onActivity()
+        }
+      }
+
+    for (const name of farmOptions.exposedMethods) {
+      if (name.startsWith('_')) continue
+
+      // @ts-expect-error - we're grabbing a dynamic method on the worker
+      let method = this._worker[name].bind(this._worker)
+      if (timeout) {
+        method = wrapMethodWithTimeout(method)
+      }
+
+      // @ts-expect-error - we're dynamically creating methods
+      this[name] = method
     }
   }
 
