SERVER-48068 Fix handling of program registry and wait_for_pid in mongo shell

(cherry picked from commit 1af88ecaa952995d44e545708633196b2cd2bbb2)

3 files changed, 131 insertions, 69 deletions

diff --git a/jstests/noPassthrough/shell_parallel_wait_for_pid.js b/jstests/noPassthrough/shell_parallel_wait_for_pid.js
new file mode 100644
index 00000000000..a7afaac8506
--- /dev/null
+++ b/jstests/noPassthrough/shell_parallel_wait_for_pid.js
@@ -0,0 +1,26 @@
+// @tags: [
+//   uses_parallel_shell,
+// ]
+
+// SERVER-48068: check that runningChildPids doesn't unregister a pid multiple
+// times from the registry
+// Verify that an invariant failure doesn't occur in the program registry
+try {
+    var cleanup = startParallelShell("MongoRunner.runningChildPids();", undefined, true);
+    var cleanup2 = startParallelShell("MongoRunner.runningChildPids();", undefined, true);
+    sleep(5000);
+
+    try {
+        MongoRunner.runningChildPids();
+        throw new Error('Simulating assert.soon() failure');
+    } finally {
+        cleanup();
+        cleanup2();
+    }
+
+} catch (e) {
+    assert.eq(e instanceof Error, true);
+    assert.eq(e.message, 'Simulating assert.soon() failure');
+}
+
+print("shell_parallel_wait_for_pid.js SUCCESS");
diff --git a/src/mongo/shell/shell_utils_launcher.cpp b/src/mongo/shell/shell_utils_launcher.cpp
index d0cc331d10f..d24abf2874c 100644
--- a/src/mongo/shell/shell_utils_launcher.cpp
+++ b/src/mongo/shell/shell_utils_launcher.cpp
@@ -194,6 +194,92 @@ void ProgramRegistry::registerReaderThread(ProcessId pid, stdx::thread reader) {
     _outputReaderThreads.emplace(pid, std::move(reader));
 }
 
+bool ProgramRegistry::waitForPid(const ProcessId pid, const bool block, int* const exit_code) {
+    stdx::lock_guard<stdx::recursive_mutex> lk(_mutex);
+    // unregistered pids are dead
+    if (!this->isPidRegistered(pid)) {
+        if (exit_code) {
+            const auto code = _pidToExitCode.find(pid);
+            if (code != _pidToExitCode.end()) {
+                *exit_code = code->second;
+            } else {
+                // If you hit this invariant, you're waiting on a PID that was
+                // never a child of this process.
+                MONGO_UNREACHABLE;
+            }
+        }
+        return true;
+    }
+#ifdef _WIN32
+    HANDLE h = getHandleForPid(pid);
+
+    // wait until the process object is signaled before getting its
+    // exit code. do this even when block is false to ensure that all
+    // file handles open in the process have been closed.
+
+    DWORD ret = WaitForSingleObject(h, (block ? INFINITE : 0));
+    if (ret == WAIT_TIMEOUT) {
+        return false;
+    } else if (ret != WAIT_OBJECT_0) {
+        const auto ewd = errnoWithDescription();
+        LOGV2_INFO(
+            22811, "ProgramRegistry::waitForPid: WaitForSingleObject failed", "error"_attr = ewd);
+    }
+
+    DWORD tmp;
+    if (GetExitCodeProcess(h, &tmp)) {
+        if (tmp == STILL_ACTIVE) {
+            uassert(
+                ErrorCodes::UnknownError, "Process is STILL_ACTIVE even after blocking", !block);
+            return false;
+        }
+        CloseHandle(h);
+        eraseHandleForPid(pid);
+        if (exit_code)
+            *exit_code = tmp;
+        _pidToExitCode[pid] = tmp;
+
+        unregisterProgram(pid);
+        return true;
+    } else {
+        const auto ewd = errnoWithDescription();
+        LOGV2_INFO(22812, "GetExitCodeProcess failed", "error"_attr = ewd);
+        return false;
+    }
+#else
+    int status;
+    int ret;
+    do {
+        errno = 0;
+        ret = waitpid(pid.toNative(), &status, (block ? 0 : WNOHANG));
+    } while (ret == -1 && errno == EINTR);
+    if (ret) {
+        int code;
+        if (WIFEXITED(status)) {
+            code = WEXITSTATUS(status);
+        } else if (WIFSIGNALED(status)) {
+            code = -WTERMSIG(status);
+        } else {
+            MONGO_UNREACHABLE;
+        }
+        _pidToExitCode[pid] = code;
+        if (exit_code) {
+            *exit_code = code;
+        }
+    }
+    if (ret) {
+        unregisterProgram(pid);
+    } else if (block) {
+        uasserted(ErrorCodes::UnknownError, "Process did not exit after blocking");
+    }
+    return ret == pid.toNative();
+#endif
+}
+
+bool ProgramRegistry::isPidDead(const ProcessId pid, int* const exit_code) {
+    return this->waitForPid(pid, false, exit_code);
+}
+
 void ProgramRegistry::getRegisteredPorts(vector<int>& ports) {
     stdx::lock_guard<stdx::recursive_mutex> lk(_mutex);
     for (const auto& portPid : _portToPidMap) {
@@ -702,68 +788,6 @@ void ProgramRunner::launchProcess(int child_stdout) {
 #endif
 }
 
-// returns true if process exited
-// If this function returns true, it will always call `registry.unregisterProgram(pid);`
-// If block is true, this will throw if it cannot wait for the processes to exit.
-bool wait_for_pid(ProcessId pid, bool block = true, int* exit_code = nullptr) {
-#ifdef _WIN32
-    HANDLE h = registry.getHandleForPid(pid);
-
-    // wait until the process object is signaled before getting its
-    // exit code. do this even when block is false to ensure that all
-    // file handles open in the process have been closed.
-
-    DWORD ret = WaitForSingleObject(h, (block ? INFINITE : 0));
-    if (ret == WAIT_TIMEOUT) {
-        return false;
-    } else if (ret != WAIT_OBJECT_0) {
-        const auto ewd = errnoWithDescription();
-        LOGV2_INFO(22811, "wait_for_pid: WaitForSingleObject failed", "error"_attr = ewd);
-    }
-
-    DWORD tmp;
-    if (GetExitCodeProcess(h, &tmp)) {
-        if (tmp == STILL_ACTIVE) {
-            uassert(
-                ErrorCodes::UnknownError, "Process is STILL_ACTIVE even after blocking", !block);
-            return false;
-        }
-        CloseHandle(h);
-        registry.eraseHandleForPid(pid);
-        if (exit_code)
-            *exit_code = tmp;
-
-        registry.unregisterProgram(pid);
-        return true;
-    } else {
-        const auto ewd = errnoWithDescription();
-        LOGV2_INFO(22812, "GetExitCodeProcess failed", "error"_attr = ewd);
-        return false;
-    }
-#else
-    int tmp;
-    int ret;
-    do {
-        ret = waitpid(pid.toNative(), &tmp, (block ? 0 : WNOHANG));
-    } while (ret == -1 && errno == EINTR);
-    if (ret && exit_code) {
-        if (WIFEXITED(tmp)) {
-            *exit_code = WEXITSTATUS(tmp);
-        } else if (WIFSIGNALED(tmp)) {
-            *exit_code = -WTERMSIG(tmp);
-        } else {
-            MONGO_UNREACHABLE;
-        }
-    }
-    if (ret) {
-        registry.unregisterProgram(pid);
-    } else if (block) {
-        uasserted(ErrorCodes::UnknownError, "Process did not exit after blocking");
-    }
-    return ret == pid.toNative();
-#endif
-}
-
 // Output up to BSONObjMaxUserSize characters of the most recent log output in order to
 // avoid hitting the 16MB size limit of a BSONObject.
 BSONObj RawMongoProgramOutput(const BSONObj& args, void* data) {
@@ -783,7 +807,7 @@ BSONObj ClearRawMongoProgramOutput(const BSONObj& args, void* data) {
 BSONObj CheckProgram(const BSONObj& args, void* data) {
     ProcessId pid = ProcessId::fromNative(singleArg(args).numberInt());
     int exit_code = -123456;  // sentinel value
-    bool isDead = wait_for_pid(pid, false, &exit_code);
+    bool isDead = registry.isPidDead(pid, &exit_code);
     if (!isDead) {
         return BSON("" << BSON("alive" << true));
     }
@@ -793,7 +817,7 @@ BSONObj CheckProgram(const BSONObj& args, void* data) {
 BSONObj WaitProgram(const BSONObj& a, void* data) {
     ProcessId pid = ProcessId::fromNative(singleArg(a).numberInt());
     int exit_code = -123456;  // sentinel value
-    wait_for_pid(pid, true, &exit_code);
+    registry.waitForPid(pid, true, &exit_code);
     return BSON(string("") << exit_code);
 }
 
@@ -810,7 +834,7 @@ BSONObj WaitMongoProgram(const BSONObj& a, void* data) {
         return BSON(string("") << 0);
     }
     pid = registry.pidForPort(port);
-    wait_for_pid(pid, true, &exit_code);
+    registry.waitForPid(pid, true, &exit_code);
     return BSON(string("") << exit_code);
 }
 
@@ -856,7 +880,7 @@ BSONObj RunProgram(const BSONObj& a, void* data, bool isMongo) {
     stdx::thread t(r);
     registry.registerReaderThread(r.pid(), std::move(t));
     int exit_code = -123456;  // sentinel value
-    wait_for_pid(r.pid(), true, &exit_code);
+    registry.waitForPid(r.pid(), true, &exit_code);
     return BSON(string("") << exit_code);
 }
 
@@ -1053,7 +1077,7 @@ int killDb(int port, ProcessId _pid, int signal, const BSONObj& opt, bool waitPi
     int exitCode = EXIT_FAILURE;
     try {
         LOGV2_INFO(22819, "Waiting for process to terminate.", "pid"_attr = pid);
-        wait_for_pid(pid, true, &exitCode);
+        registry.waitForPid(pid, true, &exitCode);
     } catch (...) {
         LOGV2_WARNING(22828, "Process failed to terminate.", "pid"_attr = pid);
         return EXIT_FAILURE;
@@ -1171,8 +1195,7 @@ std::vector<ProcessId> getRunningMongoChildProcessIds() {
                  registeredPids.end(),
                  std::back_inserter(outPids),
                  [](const ProcessId& pid) {
-                     const bool block = false;
-                     bool isDead = wait_for_pid(pid, block);
+                     bool isDead = registry.isPidDead(pid);
                      return !isDead;
                  });
     return outPids;
diff --git a/src/mongo/shell/shell_utils_launcher.h b/src/mongo/shell/shell_utils_launcher.h
index 9464100aa9b..7a06dd44f71 100644
--- a/src/mongo/shell/shell_utils_launcher.h
+++ b/src/mongo/shell/shell_utils_launcher.h
@@ -93,12 +93,25 @@ public:
     void unregisterProgram(ProcessId pid);
 
     bool isPidRegistered(ProcessId pid) const;
+    /** platform-agnostic wrapper around waitpid that automatically cleans up
+     * the program registry
+     * @param pid the processid
+     * @param block if true, block the thread until the child has exited
+     * @param exit_code[out] if set, and an exit code is available, the code will be stored here
+     * @return true if the process has exited, false otherwise */
+    bool waitForPid(const ProcessId pid, const bool block, int* const exit_code = nullptr);
+    /** check if a child process is alive. Never blocks
+     * @param pid the processid
+     * @param exit_code[out] if set, and an exit code is available, the code will be stored here
+     * @return true if the process has exited, false otherwise */
+    bool isPidDead(const ProcessId pids, int* const exit_code = nullptr);
     void getRegisteredPorts(std::vector<int>& ports);
     void getRegisteredPids(std::vector<ProcessId>& pids);
 
 private:
     stdx::unordered_set<ProcessId> _registeredPids;
     stdx::unordered_map<int, ProcessId> _portToPidMap;
+    stdx::unordered_map<ProcessId, int> _pidToExitCode;
     stdx::unordered_map<ProcessId, stdx::thread> _outputReaderThreads;
     mutable stdx::recursive_mutex _mutex;