fix(gatsby): wait for worker jobs to complete (#35513) (#35522)

Co-authored-by: Matt Kane <[email protected]>

Author: gatsbybot

packages/gatsby/src/commands/build.ts

@@ -62,6 +62,7 @@ import {
 } from "../utils/page-mode"
 import { validateEngines } from "../utils/validate-engines"
 import { constructConfigObject } from "../utils/gatsby-cloud-config"
+import { waitUntilWorkerJobsAreComplete } from "../utils/jobs/worker-messaging"
 
 module.exports = async function build(
   program: IBuildArgs,
@@ -295,7 +296,10 @@ module.exports = async function build(
       parentSpan: buildSpan,
     })
     // Jobs still might be running even though query running finished
-    await waitUntilAllJobsComplete()
+    await Promise.all([
+      waitUntilAllJobsComplete(),
+      waitUntilWorkerJobsAreComplete(),
+    ])
     // Restart worker pool before merging state to lower memory pressure while merging state
     waitForWorkerPoolRestart = workerPool.restart()
     await mergeWorkerState(workerPool, buildSpan)

packages/gatsby/src/utils/jobs/worker-messaging.ts

@@ -11,11 +11,18 @@ import { internalActions } from "../../redux/actions"
 import { GatsbyWorkerPool } from "../worker/types"
 import { isWorker, getMessenger } from "../worker/messaging"
 
+let hasActiveWorkerJobs: pDefer.DeferredPromise<void> | null = null
+let activeWorkerJobs = 0
+
 export function initJobsMessagingInMainProcess(
   workerPool: GatsbyWorkerPool
 ): void {
   workerPool.onMessage((msg, workerId) => {
     if (msg.type === MESSAGE_TYPES.JOB_CREATED) {
+      if (activeWorkerJobs === 0) {
+        hasActiveWorkerJobs = pDefer<void>()
+      }
+      activeWorkerJobs++
       store
         .dispatch(internalActions.createJobV2FromInternalJob(msg.payload))
         .then(result => {
@@ -42,10 +49,19 @@ export function initJobsMessagingInMainProcess(
             workerId
           )
         })
+        .finally(() => {
+          if (--activeWorkerJobs === 0) {
+            hasActiveWorkerJobs?.resolve()
+            hasActiveWorkerJobs = null
+          }
+        })
     }
   })
 }
 
+export const waitUntilWorkerJobsAreComplete = (): Promise<void> =>
+  hasActiveWorkerJobs ? hasActiveWorkerJobs.promise : Promise.resolve()
+
 /**
  * This map is ONLY used in worker. It's purpose is to keep track of promises returned to plugins
  * when creating jobs (in worker context), so that we can resolve or reject those once main process