From 3a9e633c71038211fe4bffadccd4f7cf2f7f7da4 Mon Sep 17 00:00:00 2001
From: Anna Henningsen <anna.henningsen@mongodb.com>
Date: Thu, 18 Feb 2021 15:43:31 +0100
Subject: [PATCH] fix(cli-repl): defer exit event until evalutation done

This otherwise messes with cases where the input stream ends while
async evaluation is still in progress (e.g. when piping into mongosh).
---
 packages/cli-repl/src/async-repl.spec.ts | 20 ++++++++++++++++++++
 packages/cli-repl/src/async-repl.ts      | 20 ++++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/packages/cli-repl/src/async-repl.spec.ts b/packages/cli-repl/src/async-repl.spec.ts
index 6facb75d0a..e19f231f33 100644
--- a/packages/cli-repl/src/async-repl.spec.ts
+++ b/packages/cli-repl/src/async-repl.spec.ts
@@ -6,6 +6,7 @@ import { once } from 'events';
 import chai, { expect } from 'chai';
 import sinon from 'ts-sinon';
 import sinonChai from 'sinon-chai';
+import { tick } from '../test/repl-helpers';
 chai.use(sinonChai);
 
 const delay = promisify(setTimeout);
@@ -123,6 +124,25 @@ describe('AsyncRepl', () => {
     expect(foundUid).to.be.true;
   });
 
+  it('delays the "exit" event until after asynchronous evaluation is finished', async() => {
+    const { input, repl } = createDefaultAsyncRepl();
+    let exited = false;
+    repl.on('exit', () => { exited = true; });
+
+    let resolve;
+    repl.context.asyncFn = () => new Promise((res) => { resolve = res; });
+
+    input.end('asyncFn()\n');
+    expect(exited).to.be.false;
+
+    await tick();
+    resolve();
+    expect(exited).to.be.false;
+
+    await tick();
+    expect(exited).to.be.true;
+  });
+
   describe('allows handling exceptions from e.g. the writer function', () => {
     it('for succesful completions', async() => {
       const error = new Error('throwme');
diff --git a/packages/cli-repl/src/async-repl.ts b/packages/cli-repl/src/async-repl.ts
index 159c9a01e7..c6100a40ac 100644
--- a/packages/cli-repl/src/async-repl.ts
+++ b/packages/cli-repl/src/async-repl.ts
@@ -49,8 +49,13 @@ export function start(opts: AsyncREPLOptions): REPLServer {
     repl.emit(evalStart, { input } as EvalStartEvent);
 
     try {
+      let exitEventPending = false;
+      const exitListener = () => { exitEventPending = true; };
+      let previousExitListeners: any[] = [];
+
       let sigintListener: (() => void) | undefined = undefined;
       let previousSigintListeners: any[] = [];
+
       try {
         result = await new Promise((resolve, reject) => {
           if (breakEvalOnSigint) {
@@ -68,6 +73,13 @@ export function start(opts: AsyncREPLOptions): REPLServer {
             repl.once('SIGINT', sigintListener);
           }
 
+          // The REPL may become over-eager and emit 'exit' events while our
+          // evaluation is still in progress (because it doesn't expect async
+          // evaluation). If that happens, defer the event until later.
+          previousExitListeners = repl.rawListeners('exit');
+          repl.removeAllListeners('exit');
+          repl.once('exit', exitListener);
+
           const evalResult = asyncEval(originalEval, input, context, filename);
 
           if (sigintListener !== undefined) {
@@ -84,6 +96,14 @@ export function start(opts: AsyncREPLOptions): REPLServer {
         for (const listener of previousSigintListeners) {
           repl.on('SIGINT', listener);
         }
+
+        repl.removeListener('exit', exitListener);
+        for (const listener of previousExitListeners) {
+          repl.on('exit', listener);
+        }
+        if (exitEventPending) {
+          process.nextTick(() => repl.emit('exit'));
+        }
       }
     } catch (err) {
       try {