From 7eab9a5f6772d7ebc1a18f784ef0139ef038b97f Mon Sep 17 00:00:00 2001
From: Quentin Barbe <forty@everteam.org>
Date: Sat, 22 Jan 2022 15:20:19 +0100
Subject: [PATCH] feat(parallel): assign each worker a worker-id (#4813)

This uses a new feature from the workerpool module, which allows
overriding process parameters for each process, which is why the
workerpool module has been updated to the latest version (6.2.0).
---
 docs/index.md                                  |  6 ++++++
 lib/nodejs/buffered-worker-pool.js             | 18 +++++++++++++++++-
 package.json                                   |  2 +-
 .../fixtures/parallel/testworkerid1.mjs        |  7 +++++++
 .../fixtures/parallel/testworkerid2.mjs        |  7 +++++++
 .../fixtures/parallel/testworkerid3.mjs        |  9 +++++++++
 test/integration/parallel.spec.js              | 12 ++++++++++++
 7 files changed, 59 insertions(+), 2 deletions(-)
 create mode 100644 test/integration/fixtures/parallel/testworkerid1.mjs
 create mode 100644 test/integration/fixtures/parallel/testworkerid2.mjs
 create mode 100644 test/integration/fixtures/parallel/testworkerid3.mjs

diff --git a/docs/index.md b/docs/index.md
index 3abf41cafd..be27aa9a3b 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -1379,6 +1379,12 @@ Free-tier cloud CI services may not provide a suitable multi-core container or V
 
 It's unlikely (but not impossible) to see a performance gain from a [job count](#-jobs-count-j-count) _greater than_ the number of available CPU cores. That said, _play around with the job count_--there's no one-size-fits all, and the unique characteristics of your tests will determine the optimal number of jobs; it may even be that fewer is faster!
 
+### Parallel Mode Worker IDs
+
+> _New in v9.2.0_
+
+Each process launched by parallel mode is assigned a unique id, from 0 for the first process to be launched, to N-1 for the Nth process. This worker id may be accessed in tests via the environment variable `MOCHA_WORKER_ID`. It can be used for example to assign a different database, service port, etc for each test process.
+
 ## Root Hook Plugins
 
 > _New in v8.0.0_
diff --git a/lib/nodejs/buffered-worker-pool.js b/lib/nodejs/buffered-worker-pool.js
index 9b0d4516e5..fe4abf945e 100644
--- a/lib/nodejs/buffered-worker-pool.js
+++ b/lib/nodejs/buffered-worker-pool.js
@@ -75,7 +75,23 @@ class BufferedWorkerPool {
       process.execArgv.join(' ')
     );
 
-    this.options = {...WORKER_POOL_DEFAULT_OPTS, opts, maxWorkers};
+    let counter = 0;
+    const onCreateWorker = ({forkOpts}) => {
+      return {
+        forkOpts: {
+          ...forkOpts,
+          // adds an incremental id to all workers, which can be useful to allocate resources for each process
+          env: {...process.env, MOCHA_WORKER_ID: counter++}
+        }
+      };
+    };
+
+    this.options = {
+      ...WORKER_POOL_DEFAULT_OPTS,
+      ...opts,
+      maxWorkers,
+      onCreateWorker
+    };
     this._pool = workerpool.pool(WORKER_PATH, this.options);
   }
 
diff --git a/package.json b/package.json
index 4273a88ae4..fd52de9110 100644
--- a/package.json
+++ b/package.json
@@ -72,7 +72,7 @@
     "strip-json-comments": "3.1.1",
     "supports-color": "8.1.1",
     "which": "2.0.2",
-    "workerpool": "6.1.5",
+    "workerpool": "6.2.0",
     "yargs": "16.2.0",
     "yargs-parser": "20.2.4",
     "yargs-unparser": "2.0.0"
diff --git a/test/integration/fixtures/parallel/testworkerid1.mjs b/test/integration/fixtures/parallel/testworkerid1.mjs
new file mode 100644
index 0000000000..2e1f55e6a3
--- /dev/null
+++ b/test/integration/fixtures/parallel/testworkerid1.mjs
@@ -0,0 +1,7 @@
+import assert from 'assert';
+
+describe('test1', () => {
+  it('should always run on worker with id 0', () => {
+    assert.ok(process.env.MOCHA_WORKER_ID === '0');
+  });
+});
diff --git a/test/integration/fixtures/parallel/testworkerid2.mjs b/test/integration/fixtures/parallel/testworkerid2.mjs
new file mode 100644
index 0000000000..c9bf710ca0
--- /dev/null
+++ b/test/integration/fixtures/parallel/testworkerid2.mjs
@@ -0,0 +1,7 @@
+import assert from 'assert';
+
+describe('test2', () => {
+  it('should always run on worker with id 1', () => {
+    assert.ok(process.env.MOCHA_WORKER_ID === '1');
+  });
+});
diff --git a/test/integration/fixtures/parallel/testworkerid3.mjs b/test/integration/fixtures/parallel/testworkerid3.mjs
new file mode 100644
index 0000000000..40a1ab43c5
--- /dev/null
+++ b/test/integration/fixtures/parallel/testworkerid3.mjs
@@ -0,0 +1,9 @@
+import assert from 'assert';
+
+describe('test3', () => {
+  it('should run on worker with either id 0 or 1', () => {
+    assert.ok(
+      process.env.MOCHA_WORKER_ID === '0' || process.env.MOCHA_WORKER_ID === '1'
+    );
+  });
+});
diff --git a/test/integration/parallel.spec.js b/test/integration/parallel.spec.js
index b019307414..3cdecfcf18 100644
--- a/test/integration/parallel.spec.js
+++ b/test/integration/parallel.spec.js
@@ -18,4 +18,16 @@ describe('parallel run', () => {
     assert.strictEqual(result.stats.failures, 1);
     assert.strictEqual(result.stats.passes, 2);
   });
+
+  it('should correctly set worker ids for each process', async () => {
+    const result = await runMochaJSONAsync('parallel/testworkerid3.mjs', [
+      '--parallel',
+      '--jobs',
+      '2',
+      require.resolve('./fixtures/parallel/testworkerid1.mjs'),
+      require.resolve('./fixtures/parallel/testworkerid2.mjs')
+    ]);
+    assert.strictEqual(result.stats.failures, 0);
+    assert.strictEqual(result.stats.passes, 3);
+  });
 });