From da3773eb400658c1ffa334b64a8d4eebe00c681e Mon Sep 17 00:00:00 2001
From: Anna Henningsen <anna.henningsen@mongodb.com>
Date: Mon, 5 Dec 2022 14:25:30 +0100
Subject: [PATCH] fix(cli-repl): flush output before exiting MONGOSH-1300

---
 packages/cli-repl/src/cli-repl.ts             | 16 ++++++++++++-
 .../cli-repl/src/tls-certificate-selector.ts  |  2 ++
 packages/cli-repl/test/e2e.spec.ts            | 19 +++++++++++++++
 packages/cli-repl/test/test-shell.ts          | 23 +++++++++++--------
 4 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/packages/cli-repl/src/cli-repl.ts b/packages/cli-repl/src/cli-repl.ts
index 21feceb997..2c323e5859 100644
--- a/packages/cli-repl/src/cli-repl.ts
+++ b/packages/cli-repl/src/cli-repl.ts
@@ -9,7 +9,7 @@ import { Editor } from '@mongosh/editor';
 import { redactSensitiveData } from '@mongosh/history';
 import Analytics from 'analytics-node';
 import askpassword from 'askpassword';
-import { EventEmitter } from 'events';
+import { EventEmitter, once } from 'events';
 import yaml from 'js-yaml';
 import ConnectionString from 'mongodb-connection-string-url';
 import semver from 'semver';
@@ -680,6 +680,20 @@ class CliRepl implements MongoshIOProvider {
     if (this.closing) {
       return;
     }
+    if (!this.output.destroyed) {
+      // Wait for output to be fully flushed before exiting.
+      if (this.output.writableEnded) {
+        // .end() has been called but not finished; 'close' will be emitted in that case.
+        // (This should not typically happen in the context of mongosh, but there's also
+        // no reason not to handle this case properly.)
+        try {
+          await once(this.output, 'close');
+        } catch { /* ignore */ }
+      } else {
+        // .end() has not been called; write an empty chunk and wait for it to be fully written.
+        await new Promise(resolve => this.output.write('', resolve));
+      }
+    }
     this.closing = true;
     const analytics = this.segmentAnalytics;
     let flushError: string | null = null;
diff --git a/packages/cli-repl/src/tls-certificate-selector.ts b/packages/cli-repl/src/tls-certificate-selector.ts
index a0ac2ce1b7..f4ec2c3ce4 100644
--- a/packages/cli-repl/src/tls-certificate-selector.ts
+++ b/packages/cli-repl/src/tls-certificate-selector.ts
@@ -29,11 +29,13 @@ export function getTlsCertificateSelector(
 }
 
 declare global {
+  // eslint-disable-next-line camelcase
   const __non_webpack_require__: undefined | typeof require;
 }
 
 function getCertificateExporter(): TlsCertificateExporter | undefined {
   if (process.env.TEST_OS_EXPORT_CERTIFICATE_AND_KEY_PATH) {
+    // eslint-disable-next-line camelcase
     if (typeof __non_webpack_require__ === 'function') {
       return __non_webpack_require__(process.env.TEST_OS_EXPORT_CERTIFICATE_AND_KEY_PATH);
     }
diff --git a/packages/cli-repl/test/e2e.spec.ts b/packages/cli-repl/test/e2e.spec.ts
index de58351e3a..92f005c5fe 100644
--- a/packages/cli-repl/test/e2e.spec.ts
+++ b/packages/cli-repl/test/e2e.spec.ts
@@ -174,6 +174,25 @@ describe('e2e', function() {
         expect(await shell.executeLine('[crypto.getFips()]')).to.include('[ 1 ]');
       }
     });
+    it('prints full output even when that output is buffered', async() => {
+      shell = TestShell.start({
+        args: ['--nodb', '--quiet', '--eval', 'EJSON.stringify([...Array(100_000).keys()].map(i=>({i})),null,2)'],
+        consumeStdio: false
+      });
+      let buffer = '';
+      let hasWaited = false;
+      // Start reading data, wait a bit, then read the rest
+      for await (const chunk of shell.process.stdout.setEncoding('utf8')) {
+        buffer += chunk;
+        if (buffer.includes('"i": 100') && !hasWaited) {
+          // This delay is relevant for reproducing this bug; it gives
+          // the mongosh process time to exit before all output has been printed.
+          await new Promise(resolve => setTimeout(resolve, 1000));
+          hasWaited = true;
+        }
+      }
+      expect(buffer).to.include('"i": 99999');
+    });
   });
 
   describe('set db', () => {
diff --git a/packages/cli-repl/test/test-shell.ts b/packages/cli-repl/test/test-shell.ts
index a1e317cf99..3138159037 100644
--- a/packages/cli-repl/test/test-shell.ts
+++ b/packages/cli-repl/test/test-shell.ts
@@ -27,6 +27,7 @@ export class TestShell {
     removeSigintListeners?: boolean;
     cwd?: string;
     forceTerminal?: boolean;
+    consumeStdio?: boolean;
   } = { args: [] }): TestShell {
     let shellProcess: ChildProcess;
 
@@ -64,7 +65,7 @@ export class TestShell {
       });
     }
 
-    const shell = new TestShell(shellProcess);
+    const shell = new TestShell(shellProcess, options.consumeStdio);
     TestShell._openShells.push(shell);
 
     return shell;
@@ -95,20 +96,22 @@ export class TestShell {
   private _rawOutput: string;
   private _onClose: Promise<number>;
 
-  constructor(shellProcess: ChildProcess) {
+  constructor(shellProcess: ChildProcess, consumeStdio = true) {
     this._process = shellProcess;
     this._output = '';
     this._rawOutput = '';
 
-    shellProcess.stdout.setEncoding('utf8').on('data', (chunk) => {
-      this._output += stripAnsi(chunk);
-      this._rawOutput += chunk;
-    });
+    if (consumeStdio) {
+      shellProcess.stdout.setEncoding('utf8').on('data', (chunk) => {
+        this._output += stripAnsi(chunk);
+        this._rawOutput += chunk;
+      });
 
-    shellProcess.stderr.setEncoding('utf8').on('data', (chunk) => {
-      this._output += stripAnsi(chunk);
-      this._rawOutput += chunk;
-    });
+      shellProcess.stderr.setEncoding('utf8').on('data', (chunk) => {
+        this._output += stripAnsi(chunk);
+        this._rawOutput += chunk;
+      });
+    }
 
     this._onClose = (async() => {
       const [ code ] = await once(shellProcess, 'close');