From e7340f440b6527b29e553ee9cdfe628b60c14f2c Mon Sep 17 00:00:00 2001
From: Jarred Sumner <jarred@jarredsumner.com>
Date: Tue, 18 Jun 2024 13:59:27 -0700
Subject: [PATCH] Drain microtasks in the Bun Shell

---
 src/shell/interpreter.zig                      |  9 ++++++++-
 src/shell/subproc.zig                          | 10 +++++++++-
 test/js/bun/shell/shell-keepalive-fixture-1.js |  6 ++++++
 test/js/bun/shell/shell-keepalive-fixture-2.js |  6 ++++++
 test/js/bun/shell/shell-keepalive.test.ts      | 11 +++++++++++
 5 files changed, 40 insertions(+), 2 deletions(-)
 create mode 100644 test/js/bun/shell/shell-keepalive-fixture-1.js
 create mode 100644 test/js/bun/shell/shell-keepalive-fixture-2.js
 create mode 100644 test/js/bun/shell/shell-keepalive.test.ts

diff --git a/src/shell/interpreter.zig b/src/shell/interpreter.zig
index 9474882853..08763d39aa 100644
--- a/src/shell/interpreter.zig
+++ b/src/shell/interpreter.zig
@@ -843,6 +843,7 @@ pub const Interpreter = struct {
     async_commands_executing: u32 = 0,
 
     globalThis: *JSC.JSGlobalObject,
+    ref: bun.Async.KeepAlive = .{},
 
     flags: packed struct(u8) {
         done: bool = false,
@@ -1235,7 +1236,7 @@ pub const Interpreter = struct {
         interpreter.this_jsvalue = JSC.Codegen.JSShellInterpreter.toJS(interpreter, globalThis);
         JSC.Codegen.JSShellInterpreter.resolveSetCached(interpreter.this_jsvalue, globalThis, resolve);
         JSC.Codegen.JSShellInterpreter.rejectSetCached(interpreter.this_jsvalue, globalThis, reject);
-
+        interpreter.ref.activate(JSC.VirtualMachine.get().uwsLoop());
         bun.Analytics.Features.shell += 1;
         return interpreter.this_jsvalue;
     }
@@ -1656,6 +1657,8 @@ pub const Interpreter = struct {
         defer decrPendingActivityFlag(&this.has_pending_activity);
 
         if (this.event_loop == .js) {
+            this.ref.disable();
+
             defer this.deinitAfterJSRun();
             this.exit_code = exit_code;
             if (this.this_jsvalue != .zero) {
@@ -1676,6 +1679,8 @@ pub const Interpreter = struct {
         defer decrPendingActivityFlag(&this.has_pending_activity);
 
         if (this.event_loop == .js) {
+            this.ref.disable();
+
             if (this.this_jsvalue != .zero) {
                 if (JSC.Codegen.JSShellInterpreter.rejectGetCached(this.this_jsvalue)) |reject| {
                     reject.call(this.globalThis, &[_]JSValue{ JSValue.jsNumberFromChar(1), this.getBufferedStdout(), this.getBufferedStderr() });
@@ -1692,6 +1697,7 @@ pub const Interpreter = struct {
         for (this.jsobjs) |jsobj| {
             jsobj.unprotect();
         }
+        this.ref.disable();
         this.root_io.deref();
         this.root_shell.deinitImpl(false, false);
         this.this_jsvalue = .zero;
@@ -1705,6 +1711,7 @@ pub const Interpreter = struct {
             this.root_shell._buffered_stdout.owned.deinitWithAllocator(bun.default_allocator);
         }
         this.this_jsvalue = .zero;
+        this.ref.disable();
         this.allocator.destroy(this);
     }
 
diff --git a/src/shell/subproc.zig b/src/shell/subproc.zig
index bcf0fdfd49..e26c0a37ed 100644
--- a/src/shell/subproc.zig
+++ b/src/shell/subproc.zig
@@ -949,7 +949,15 @@ pub const ShellSubprocess = struct {
         if (exit_code) |code| {
             if (this.cmd_parent) |cmd| {
                 if (cmd.exit_code == null) {
-                    cmd.onExit(code);
+                    if (this.event_loop == .js) {
+                        const loop = this.event_loop.js;
+
+                        loop.enter();
+                        defer loop.exit();
+                        cmd.onExit(code);
+                    } else {
+                        cmd.onExit(code);
+                    }
                 }
             }
         }
diff --git a/test/js/bun/shell/shell-keepalive-fixture-1.js b/test/js/bun/shell/shell-keepalive-fixture-1.js
new file mode 100644
index 0000000000..54fcf76ef2
--- /dev/null
+++ b/test/js/bun/shell/shell-keepalive-fixture-1.js
@@ -0,0 +1,6 @@
+process.exitCode = 1;
+
+(async () => {
+  await Bun.$`ls .`;
+  process.exit(0);
+})();
diff --git a/test/js/bun/shell/shell-keepalive-fixture-2.js b/test/js/bun/shell/shell-keepalive-fixture-2.js
new file mode 100644
index 0000000000..4a9497cdbd
--- /dev/null
+++ b/test/js/bun/shell/shell-keepalive-fixture-2.js
@@ -0,0 +1,6 @@
+process.exitCode = 1;
+
+(async () => {
+  await Bun.$`${process.execPath} -e "console.log('hi')"`;
+  process.exit(0);
+})();
diff --git a/test/js/bun/shell/shell-keepalive.test.ts b/test/js/bun/shell/shell-keepalive.test.ts
new file mode 100644
index 0000000000..fb082cb74a
--- /dev/null
+++ b/test/js/bun/shell/shell-keepalive.test.ts
@@ -0,0 +1,11 @@
+import { expect, test } from "bun:test";
+import { join } from "path";
+import "harness";
+
+test("shell should stay alive while a builtin command is in progress", async () => {
+  expect([join(import.meta.dir, "shell-keepalive-fixture-1.js")]).toRun();
+});
+
+test("shell should stay alive while a non-builtin command is in progress", async () => {
+  expect([join(import.meta.dir, "shell-keepalive-fixture-2.js")]).toRun();
+});