From 9f6bc58da88d1d5c6c3c6e43e670b70fdd3aef29 Mon Sep 17 00:00:00 2001
From: James M Snell <jasnell@gmail.com>
Date: Mon, 22 Feb 2021 16:22:17 -0800
Subject: [PATCH] worker: add setEnvironmentData/getEnvironmentData

These APIs allow arbitrary, cloneable JavaScript values to be set and
passed to all new Worker instances spawned from the current context.
It is similar to `workerData` except that environment data is set
independently of the `new Worker()` constructor, and the the value is
passed automatically to all new Workers.

This is a *partial* fix of https://github.com/nodejs/node/issues/30992
but does not implement a complete fix.

Signed-off-by: James M Snell <jasnell@gmail.com>

PR-URL: https://github.com/nodejs/node/pull/37486
Reviewed-By: Anna Henningsen <anna@addaleax.net>
Reviewed-By: Gireesh Punathil <gpunathi@in.ibm.com>
Reviewed-By: Darshan Sen <raisinten@gmail.com>
---
 doc/api/worker_threads.md                    | 49 ++++++++++++++++++++
 lib/internal/main/worker_thread.js           |  3 ++
 lib/internal/worker.js                       | 25 ++++++++++
 lib/worker_threads.js                        |  4 ++
 test/parallel/test-worker-environmentdata.js | 33 +++++++++++++
 5 files changed, 114 insertions(+)
 create mode 100644 test/parallel/test-worker-environmentdata.js

diff --git a/doc/api/worker_threads.md b/doc/api/worker_threads.md
index 37f19bc135d826..e9e5845d4ca19a 100644
--- a/doc/api/worker_threads.md
+++ b/doc/api/worker_threads.md
@@ -61,6 +61,38 @@ Worker threads inherit non-process-specific options by default. Refer to
 [`Worker constructor options`][] to know how to customize worker thread options,
 specifically `argv` and `execArgv` options.
 
+## `worker.getEnvironmentData(key)`
+<!-- YAML
+added: REPLACEME
+-->
+
+> Stability: 1 - Experimental
+
+* `key` {any} Any arbitrary, cloneable JavaScript value that can be used as a
+  {Map} key.
+* Returns: {any}
+
+Within a worker thread, `worker.getEnvironmentData()` returns a clone
+of data passed to the spawning thread's `worker.setEnvironmentData()`.
+Every new `Worker` receives its own copy of the environment data
+automatically.
+
+```js
+const {
+  Worker,
+  isMainThread,
+  setEnvironmentData,
+  getEnvironmentData,
+} = require('worker_threads');
+
+if (isMainThread) {
+  setEnvironmentData('Hello', 'World!');
+  const worker = new Worker(__filename);
+} else {
+  console.log(getEnvironmentData('Hello'));  // Prints 'World!'.
+}
+```
+
 ## `worker.isMainThread`
 <!-- YAML
 added: v10.5.0
@@ -240,6 +272,23 @@ new Worker('process.env.SET_IN_WORKER = "foo"', { eval: true, env: SHARE_ENV })
   });
 ```
 
+## `worker.setEnvironmentData(key[, value])`
+<!--YAML
+added: REPLACEME
+-->
+
+> Stability: 1 - Experimental
+
+* `key` {any} Any arbitrary, cloneable JavaScript value that can be used as a
+  {Map} key.
+* `value` {any} Any arbitrary, cloneable JavaScript value that will be cloned
+  and passed automatically to all new `Worker` instances. If `value` is passed
+  as `undefined`, any previously set value for the `key` will be deleted.
+
+The `worker.setEnvironmentData()` API sets the content of
+`worker.getEnvironmentData()` in the current thread and all new `Worker`
+instances spawned from the current context.
+
 ## `worker.threadId`
 <!-- YAML
 added: v10.5.0
diff --git a/lib/internal/main/worker_thread.js b/lib/internal/main/worker_thread.js
index e6bc341c33e35d..9db0422dcfa1a8 100644
--- a/lib/internal/main/worker_thread.js
+++ b/lib/internal/main/worker_thread.js
@@ -107,6 +107,7 @@ port.on('message', (message) => {
       filename,
       doEval,
       workerData,
+      environmentData,
       publicPort,
       manifestSrc,
       manifestURL,
@@ -130,6 +131,8 @@ port.on('message', (message) => {
     publicWorker.parentPort = publicPort;
     publicWorker.workerData = workerData;
 
+    require('internal/worker').assignEnvironmentData(environmentData);
+
     // The counter is only passed to the workers created by the main thread, not
     // to workers created by other workers.
     let cachedCwd = '';
diff --git a/lib/internal/worker.js b/lib/internal/worker.js
index 73aa910b0b77ce..2b35791473d1fa 100644
--- a/lib/internal/worker.js
+++ b/lib/internal/worker.js
@@ -13,6 +13,7 @@ const {
   Promise,
   PromiseResolve,
   RegExpPrototypeTest,
+  SafeMap,
   String,
   Symbol,
   SymbolFor,
@@ -85,6 +86,8 @@ let debug = require('internal/util/debuglog').debuglog('worker', (fn) => {
 
 let cwdCounter;
 
+const environmentData = new SafeMap();
+
 if (isMainThread) {
   cwdCounter = new Uint32Array(new SharedArrayBuffer(4));
   const originalChdir = process.chdir;
@@ -94,6 +97,24 @@ if (isMainThread) {
   };
 }
 
+function setEnvironmentData(key, value) {
+  if (value === undefined)
+    environmentData.delete(key);
+  else
+    environmentData.set(key, value);
+}
+
+function getEnvironmentData(key) {
+  return environmentData.get(key);
+}
+
+function assignEnvironmentData(data) {
+  if (data === undefined) return;
+  data.forEach((value, key) => {
+    environmentData.set(key, value);
+  });
+}
+
 class Worker extends EventEmitter {
   constructor(filename, options = {}) {
     super();
@@ -222,6 +243,7 @@ class Worker extends EventEmitter {
       doEval,
       cwdCounter: cwdCounter || workerIo.sharedCwdCounter,
       workerData: options.workerData,
+      environmentData,
       publicPort: port2,
       manifestURL: getOptionValue('--experimental-policy') ?
         require('internal/process/policy').url :
@@ -482,6 +504,9 @@ module.exports = {
   SHARE_ENV,
   resourceLimits:
     !isMainThread ? makeResourceLimits(resourceLimitsRaw) : {},
+  setEnvironmentData,
+  getEnvironmentData,
+  assignEnvironmentData,
   threadId,
   Worker,
 };
diff --git a/lib/worker_threads.js b/lib/worker_threads.js
index 1b9c6ae945e530..969e8f29794cb1 100644
--- a/lib/worker_threads.js
+++ b/lib/worker_threads.js
@@ -4,6 +4,8 @@ const {
   isMainThread,
   SHARE_ENV,
   resourceLimits,
+  setEnvironmentData,
+  getEnvironmentData,
   threadId,
   Worker
 } = require('internal/worker');
@@ -32,4 +34,6 @@ module.exports = {
   Worker,
   parentPort: null,
   workerData: null,
+  setEnvironmentData,
+  getEnvironmentData,
 };
diff --git a/test/parallel/test-worker-environmentdata.js b/test/parallel/test-worker-environmentdata.js
new file mode 100644
index 00000000000000..cac97b746acef8
--- /dev/null
+++ b/test/parallel/test-worker-environmentdata.js
@@ -0,0 +1,33 @@
+'use strict';
+
+require('../common');
+const {
+  Worker,
+  getEnvironmentData,
+  setEnvironmentData,
+  threadId,
+} = require('worker_threads');
+
+const {
+  deepStrictEqual,
+  strictEqual,
+} = require('assert');
+
+if (!process.env.HAS_STARTED_WORKER) {
+  process.env.HAS_STARTED_WORKER = 1;
+  setEnvironmentData('foo', 'bar');
+  setEnvironmentData('hello', { value: 'world' });
+  setEnvironmentData(1, 2);
+  strictEqual(getEnvironmentData(1), 2);
+  setEnvironmentData(1); // Delete it, key won't show up in the worker.
+  new Worker(__filename);
+  setEnvironmentData('hello');  // Delete it. Has no impact on the worker.
+} else {
+  strictEqual(getEnvironmentData('foo'), 'bar');
+  deepStrictEqual(getEnvironmentData('hello'), { value: 'world' });
+  strictEqual(getEnvironmentData(1), undefined);
+
+  // Recurse to make sure the environment data is inherited
+  if (threadId <= 2)
+    new Worker(__filename);
+}