From 4ead65fe69c3fa083433ff3c4bc46ab734b69010 Mon Sep 17 00:00:00 2001
From: Brady Holt <brady.holt@gmail.com>
Date: Mon, 2 Oct 2023 11:41:47 -0500
Subject: [PATCH] Wait for streams to drain before exiting (#26)

---
 src/index.ts | 34 ++++++++++++++++++++++++++++++----
 1 file changed, 30 insertions(+), 4 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index de202f9..dde0ef5 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -31,7 +31,7 @@ const ECHO_RED_FORMAT = "\x1b[31m%s\x1b[0m";
 const ECHO_BLUE_FORMAT = "\x1b[34m%s\x1b[0m";
 
 global.dirname = path.dirname;
-global.exit = process.exit;
+
 /**
  * Echos error message to stdout and then exits with the specified exit code (defaults to 1)
  * @param error The error message string or Error object to print
@@ -39,7 +39,7 @@ global.exit = process.exit;
  */
 const _error = (error: string | Error, exitCode: number = 1) => {
   console.error(ECHO_RED_FORMAT, error); // Will print to stderr
-  exit(exitCode);
+  _exit(exitCode);
 };
 global.error = _error;
 
@@ -57,7 +57,7 @@ const _printUsageAndExit = (exitCode = 1, additionalMessage?: string): void => {
     }
   }
 
-  exit(exitCode);
+  _exit(exitCode);
 };
 
 const _usage = (message: string, printAndExitIfHelpArgumentSpecified = true) => {
@@ -200,6 +200,32 @@ const _stdin = () => {
 };
 global.stdin = _stdin;
 
+const _exit = (exitCode: number = 0) => {
+  // Ensure all streams are drained before exiting (code pulled from https://github.com/cowboy/node-exit)
+  let streams = [process.stdout, process.stderr];
+  let drainCount = 0;
+  function tryToExit() {
+    if (drainCount === streams.length) {
+      process.exit(exitCode);
+    }
+  }
+  streams.forEach(function (stream) {
+    // Count drained streams now, but monitor non-drained streams.
+    if (stream.writableLength === 0) {
+      drainCount++;
+    } else {
+      stream.write("", "utf-8", function () {
+        drainCount++;
+        tryToExit();
+      });
+    }
+    // Prevent further writing to stream
+    stream.write = <any>function () {};
+  });
+  tryToExit();
+};
+global.exit = _exit;
+
 // Echoing
 /**
  * Prints content to stdout with a trailing newline. Multiple arguments can be passed, with the first used as the primary message and all additional used as substitution values
@@ -768,7 +794,7 @@ declare global {
   var __filename: string;
   var __dirname: string;
   var dirname: typeof path.dirname;
-  var exit: typeof process.exit;
+  var exit: typeof _exit;
   var error: typeof _error;
   var echo: typeof _echo;
   var printf: typeof _printf;