fix(async-rewriter2): fix interaction with domains MONGOSH-662 (#751)

The fix I previously inserted for the “Node.js prints a warning
if we intentionally disregard a rejected Promise from the generated
inner async function” problem turned out to be problematic, because
it creates an infinite loop in the domain handler system.

This problem goes away with Node.js v14.15.2 (due to
nodejs/node@0ad4f70),
but since we also support Node.js-v12.x environments
(and we would want to be resilient against changes to the
platform internals anyway), it makes sense to aim for a better
solution here.

This patch makes sure that the generated inner async function only
returns a meaningful value if its return value is actually used,
i.e. if we’re encountering an async-evaluation case where the
generated outer function actually just forwards the returned
Promise.

Author: addaleax

packages/async-rewriter2/README.md

@@ -114,18 +114,25 @@ made for readability).
 
   const _asynchronousReturnValue = (async () => {
     try {
-      // All return statements are decorated with `return _synchronousReturnValue = ...`
-      return _synchronousReturnValue = (
-        // Most expressions are wrapped in ('original source', _ex = ..., _isp(_ex) ? await _ex : _ex)
-        _ex = ('db.test.find()',
-          _ex = ('db.test',
-            _ex = ('db',
-              _ex = db, _isp(_ex) ? await _ex : _ex
-            ).test, _isp(_ex) ? await _ex : _ex
-          ).find(), _isp(_ex) ? await _ex : _ex
-        ).toArray()
-        , _isp(_ex) ? await _ex : _ex
-      );
+      // All return statements are decorated with
+      // `return (_synchronousReturnValue = ..., _functionState === 'async' ? _synchronousReturnValue : null)`
+      // The function state check is here that, if we are returning synchronously,
+      // we know that we are going to discard the value of `_asynchronousReturnValue`,
+      // which is not what we want if the return value happens to be a rejected
+      // Promise (because Node.js print a warning in that case).
+      return (
+        _synchronousReturnValue = (
+          // Most expressions are wrapped in ('original source', _ex = ..., _isp(_ex) ? await _ex : _ex)
+          _ex = ('db.test.find()',
+            _ex = ('db.test',
+              _ex = ('db',
+                _ex = db, _isp(_ex) ? await _ex : _ex
+              ).test, _isp(_ex) ? await _ex : _ex
+            ).find(), _isp(_ex) ? await _ex : _ex
+          ).toArray()
+          , _isp(_ex) ? await _ex : _ex
+        ),
+        _functionState === 'async' ? _synchronousReturnValue : null);
     } catch (err) {
       err = _demangleError(err);
       if (_functionState === "sync") {
@@ -144,12 +151,6 @@ made for readability).
     }
   })();
 
-  if (_functionState !== 'sync') {
-    // Add a .catch with a no-op function, because if we're here, then that
-    // means that we'll discard the async return value even if it results in a
-    // rejected Promise (and Node.js would otherwise warn about this).
-    _asynchronousReturnValue.catch(() => {});
-  }
   if (_functionState === "returned") {
     return _synchronousReturnValue;
   } else if (_functionState === "threw") {

packages/async-rewriter2/package.json

@@ -4,6 +4,7 @@
   "description": "MongoDB Shell Async Rewriter Package",
   "main": "./lib/index.js",
   "scripts": {
+    "pretest": "npm run compile-ts",
     "test": "mocha -r \"../../scripts/import-expansions.js\" --timeout 60000 --colors -r ts-node/register \"./{src,lib}/**/*.spec.ts\"",
     "test-ci": "mocha -r \"../../scripts/import-expansions.js\" --timeout 60000 -r ts-node/register \"./{src,lib}/**/*.spec.ts\"",
     "lint": "eslint \"**/*.{js,ts,tsx}\"",

packages/async-rewriter2/src/async-writer-babel.spec.ts

@@ -1,9 +1,13 @@
 import AsyncWriter from './';
+import childProcess from 'child_process';
+import path from 'path';
+import { promisify } from 'util';
 import vm from 'vm';
 import sinon from 'ts-sinon';
 import chai, { expect } from 'chai';
 import sinonChai from 'sinon-chai';
 chai.use(sinonChai);
+const execFile = promisify(childProcess.execFile);
 
 describe('AsyncWriter', () => {
   let implicitlyAsyncFn: sinon.SinonStub;
@@ -49,6 +53,10 @@ describe('AsyncWriter', () => {
     };
   });
 
+  before(() => {
+    process.on('unhandledRejection', err => { throw err; });
+  });
+
   context('basic testing', () => {
     it('evaluates plain literal expressions', () => {
       expect(runTranspiledCode('42')).to.equal(42);
@@ -65,6 +73,15 @@ describe('AsyncWriter', () => {
       expect(runTranspiledCode('Promise.resolve([]).constructor').name).to.equal('Promise');
       expect(runTranspiledCode('Promise.resolve([]).constructor.name')).to.equal('Promise');
     });
+
+    it('works fine when immediately receiving a rejected Promise', async() => {
+      try {
+        await runTranspiledCode('Promise.reject(42)');
+        expect.fail('missed exception');
+      } catch (err) {
+        expect(err).to.equal(42);
+      }
+    });
   });
 
   context('scoping', () => {
@@ -594,4 +611,14 @@ describe('AsyncWriter', () => {
         .to.throw('abcdefghijklm ... uvwxyz is not a function');
     });
   });
+
+  context('domain support', () => {
+    it('works fine when run inside a Node.js domain context', async() => {
+      await execFile(process.execPath, [
+        path.resolve(__dirname, '..', 'test', 'fixtures', 'with-domain.js')
+      ], {
+        timeout: 5_000
+      });
+    });
+  });
 });

packages/async-rewriter2/src/async-writer-babel.ts

@@ -284,17 +284,13 @@ export const makeMaybeAsyncFunctionPlugin = ({ types: t }: { types: typeof Babel
     }
   `);
 
-  // Function.prototype is a good no-op function in situations where
-  // function literals receive special treatment :)
   const wrapperFunctionTemplate = babel.template.statements(`
     let FUNCTION_STATE_IDENTIFIER = "sync",
         SYNC_RETURN_VALUE_IDENTIFIER,
         EXPRESSION_HOLDER_IDENTIFIER;
 
     const ASYNC_RETURN_VALUE_IDENTIFIER = (ASYNC_TRY_CATCH_WRAPPER)();
 
-    if (FUNCTION_STATE_IDENTIFIER !== "sync")
-      ASYNC_RETURN_VALUE_IDENTIFIER.catch(Function.prototype);
     if (FUNCTION_STATE_IDENTIFIER === "returned")
       return SYNC_RETURN_VALUE_IDENTIFIER;
     else if (FUNCTION_STATE_IDENTIFIER === "threw")
@@ -335,6 +331,11 @@ export const makeMaybeAsyncFunctionPlugin = ({ types: t }: { types: typeof Babel
     }
   `, { placeholderPattern: false, placeholderWhitelist: new Set(['DE_IDENTIFIER']) });
 
+  const returnValueWrapperTemplate = babel.template.expression(`(
+    SYNC_RETURN_VALUE_IDENTIFIER = NODE,
+    FUNCTION_STATE_IDENTIFIER === 'async' ? SYNC_RETURN_VALUE_IDENTIFIER : null
+  )`);
+
   return {
     pre(file: babel.types.File) {
       this.file = file;
@@ -502,9 +503,14 @@ export const makeMaybeAsyncFunctionPlugin = ({ types: t }: { types: typeof Babel
             identifierGroup = path.getFunctionParent().getFunctionParent().getData(identifierGroupKey);
             if (path.parentPath.isReturnStatement() && !path.node[isGeneratedHelper]) {
               // If this is inside a return statement that we have not already handled,
-              // we replace the `return ...` with `return synchronousReturnValue = ...`.
+              // we replace the `return ...` with
+              // `return (_synchronousReturnValue = ..., _functionState === 'async' ? _synchronousReturnValue : null)`.
               path.replaceWith(Object.assign(
-                t.assignmentExpression('=', identifierGroup.synchronousReturnValue, path.node),
+                returnValueWrapperTemplate({
+                  SYNC_RETURN_VALUE_IDENTIFIER: identifierGroup.synchronousReturnValue,
+                  FUNCTION_STATE_IDENTIFIER: identifierGroup.functionState,
+                  NODE: path.node
+                }),
                 { [isGeneratedHelper]: true }
               ));
               return;
@@ -579,7 +585,9 @@ export const makeMaybeAsyncFunctionPlugin = ({ types: t }: { types: typeof Babel
           const { expressionHolder, isSyntheticPromise } = identifierGroup;
           const originalSource = t.stringLiteral(
             '\ufeff' + limitStringLength(
-              (this.file as any).code.slice(path.node.start, path.node.end), 24) +
+              path.node.start !== undefined ?
+                (this.file as any).code.slice(path.node.start, path.node.end) :
+                '<unknown>', 24) +
             '\ufeff');
           path.replaceWith(Object.assign(
             awaitSyntheticPromiseTemplate({

packages/async-rewriter2/test/fixtures/with-domain.js

@@ -0,0 +1,18 @@
+// eslint-disable-next-line strict
+'use strict';
+const { default: AsyncWriter } = require('../../');
+const domain = require('domain');
+const vm = require('vm');
+
+// This script should not enter an infinite loop. In the past, it did on
+// Node.js < 14.15.2, because domains use Array.prototype.every when being
+// entered, for which we provide a polyfill, which could lead to an async call
+// because of the async wrapping, which in turn would lead to the domain being
+// entered, and so on.
+
+const d = domain.create();
+const aw = new AsyncWriter();
+d.run(() => {
+  vm.runInThisContext(aw.runtimeSupportCode());
+  return vm.runInThisContext('(async() => 42)()');
+});