From fa24b4fe54b63e90796e8c7be56ad771c00f439a Mon Sep 17 00:00:00 2001
From: Michal Piechowiak <misiek.piechowiak@gmail.com>
Date: Mon, 18 Oct 2021 12:01:13 +0200
Subject: [PATCH] fix(gatsby): assign correct parentSpans to PQR activities
 (#33568)

---
 packages/gatsby/src/commands/build.ts         |  8 +++---
 .../src/utils/worker/__tests__/queries.ts     |  2 +-
 packages/gatsby/src/utils/worker/pool.ts      | 27 ++++++++++++++-----
 3 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/packages/gatsby/src/commands/build.ts b/packages/gatsby/src/commands/build.ts
index 496b27a721131..bb51593db2a10 100644
--- a/packages/gatsby/src/commands/build.ts
+++ b/packages/gatsby/src/commands/build.ts
@@ -221,7 +221,7 @@ module.exports = async function build(program: IBuildArgs): Promise<void> {
   }
 
   const cacheActivity = report.activityTimer(`Caching Webpack compilations`, {
-    parentSpan: buildActivityTimer.span,
+    parentSpan: buildSpan,
   })
   try {
     cacheActivity.start()
@@ -249,12 +249,14 @@ module.exports = async function build(program: IBuildArgs): Promise<void> {
 
   let waitForWorkerPoolRestart = Promise.resolve()
   if (process.env.GATSBY_EXPERIMENTAL_PARALLEL_QUERY_RUNNING) {
-    await runQueriesInWorkersQueue(workerPool, queryIds)
+    await runQueriesInWorkersQueue(workerPool, queryIds, {
+      parentSpan: buildSpan,
+    })
     // Jobs still might be running even though query running finished
     await waitUntilAllJobsComplete()
     // Restart worker pool before merging state to lower memory pressure while merging state
     waitForWorkerPoolRestart = workerPool.restart()
-    await mergeWorkerState(workerPool)
+    await mergeWorkerState(workerPool, buildSpan)
   } else {
     await runStaticQueries({
       queryIds,
diff --git a/packages/gatsby/src/utils/worker/__tests__/queries.ts b/packages/gatsby/src/utils/worker/__tests__/queries.ts
index 6d33c2eb4541c..961fb82c945cb 100644
--- a/packages/gatsby/src/utils/worker/__tests__/queries.ts
+++ b/packages/gatsby/src/utils/worker/__tests__/queries.ts
@@ -298,7 +298,7 @@ describeWhenLMDB(`worker (queries)`, () => {
     const spy = jest.spyOn(worker.single, `runQueries`)
 
     // @ts-ignore - worker is defined
-    await runQueriesInWorkersQueue(worker, queryIdsBig, 10)
+    await runQueriesInWorkersQueue(worker, queryIdsBig, { chunkSize: 10 })
     const stateFromWorker = await worker.single.getState()
 
     // Called the complete ABC so we can test _a
diff --git a/packages/gatsby/src/utils/worker/pool.ts b/packages/gatsby/src/utils/worker/pool.ts
index af78be727adf8..5fdc6340d72fb 100644
--- a/packages/gatsby/src/utils/worker/pool.ts
+++ b/packages/gatsby/src/utils/worker/pool.ts
@@ -2,6 +2,7 @@ import { WorkerPool } from "gatsby-worker"
 import { chunk } from "lodash"
 import reporter from "gatsby-cli/lib/reporter"
 import { cpuCoreCount } from "gatsby-core-utils"
+import { Span } from "opentracing"
 
 import { IGroupedQueryIds } from "../../services"
 import { initJobsMessagingInMainProcess } from "../jobs/worker-messaging"
@@ -46,16 +47,27 @@ function handleRunQueriesInWorkersQueueError(e: Error): never {
 export async function runQueriesInWorkersQueue(
   pool: GatsbyWorkerPool,
   queryIds: IGroupedQueryIds,
-  chunkSize = queriesChunkSize
+  opts?: {
+    chunkSize?: number
+    parentSpan?: Span
+  }
 ): Promise<void> {
   const activity = reporter.createProgress(
     `run queries in workers`,
-    queryIds.staticQueryIds.length + queryIds.pageQueryIds.length
+    queryIds.staticQueryIds.length + queryIds.pageQueryIds.length,
+    0,
+    { parentSpan: opts?.parentSpan }
   )
   activity.start()
   try {
-    const staticQuerySegments = chunk(queryIds.staticQueryIds, chunkSize)
-    const pageQuerySegments = chunk(queryIds.pageQueryIds, chunkSize)
+    const staticQuerySegments = chunk(
+      queryIds.staticQueryIds,
+      opts?.chunkSize ?? queriesChunkSize
+    )
+    const pageQuerySegments = chunk(
+      queryIds.pageQueryIds,
+      opts?.chunkSize ?? queriesChunkSize
+    )
 
     pool.all.setComponents()
 
@@ -90,8 +102,11 @@ export async function runQueriesInWorkersQueue(
   }
 }
 
-export async function mergeWorkerState(pool: GatsbyWorkerPool): Promise<void> {
-  const activity = reporter.activityTimer(`Merge worker state`)
+export async function mergeWorkerState(
+  pool: GatsbyWorkerPool,
+  parentSpan?: Span
+): Promise<void> {
+  const activity = reporter.activityTimer(`Merge worker state`, { parentSpan })
   activity.start()
 
   for (const { workerId } of pool.getWorkerInfo()) {