DRIVERS-3298 Improve stop-orchestration behavior

.evergreen/orchestration/drivers_orchestration.py

@@ -22,6 +22,8 @@
 from datetime import datetime
 from pathlib import Path
 
+import psutil
+
 # Get global values.
 HERE = Path(__file__).absolute().parent
 EVG_PATH = HERE.parent
@@ -635,36 +637,79 @@ def start(opts):
     LOGGER.info("Starting mongo-orchestration... done.")
 
 
+def shutdown_proc(proc: psutil.Process) -> None:
+    proc.terminate()
+    if proc.wait(1) is None:
+        try:
+            proc.kill()
+        except psutil.NoSuchProcess:
+            pass
+
+
+def shutdown_docker(docker: str, container_id: str) -> None:
+    if "podman" in docker:
+        cmd = f"{docker} rm -f {container_id}"
+    else:
+        cmd = f"{docker} kill {container_id}"
+    run_command(cmd, exit_on_error=False)
+
+
 def stop(opts):
     mo_home = Path(opts.mongo_orchestration_home)
     pid_file = mo_home / "server.pid"
     container_file = mo_home / "container_id.txt"
     docker = get_docker_cmd()
+
+    # First try to shut down using pid file.
     if pid_file.exists():
-        LOGGER.info("Stopping mongo-orchestration...")
-        py_exe = normalize_path(sys.executable)
-        run_command(f"{py_exe} -m mongo_orchestration.server stop")
+        pid = int(pid_file.read_text().strip())
         pid_file.unlink(missing_ok=True)
-        LOGGER.info("Stopping mongo-orchestration... done.")
+        if psutil.pid_exists(pid):
+            LOGGER.info("Stopping mongo-orchestration using pid file...")
+            shutdown_proc(psutil.Process(pid))
+            LOGGER.info("Stopping mongo-orchestration using pid file... done.")
+
+    # Next try using a docker container file.
     if container_file.exists():
-        LOGGER.info("Stopping mongodb_atlas_local...")
-        container_id = container_file.read_text()
-        run_command(f"{docker} kill {container_id}", exit_on_error=False)
+        LOGGER.info("Stopping mongodb_atlas_local using container file...")
+        shutdown_docker(docker, container_file.read_text())
         container_file.unlink()
-        LOGGER.info("Stopping mongodb_atlas_local... done.")
-    elif docker:
-        cmd = f"{docker} ps -a -q -f name=mongodb_atlas_local"
+        LOGGER.info("Stopping mongodb_atlas_local using container file ... done.")
+
+    all_procs = list(psutil.process_iter())
+
+    # Next look for mongo-orchestration by command line arguments.
+    for proc in all_procs:
         try:
-            result = subprocess.check_output(shlex.split(cmd), encoding="utf-8").strip()
-        except Exception:
-            result = None
-        if result:
-            LOGGER.info("Stopping mongodb_atlas_local...")
-            if "podman" in docker:
-                run_command(f"{docker} rm -f {result}")
-            else:
-                run_command(f"{docker} kill {result}")
-            LOGGER.info("Stopping mongodb_atlas_local... done.")
+            cmdline = proc.cmdline()
+        except (psutil.AccessDenied, psutil.NoSuchProcess):
+            continue
+        if "mongo_orchestration.server" in cmdline:
+            LOGGER.info("Stopping mongo-orchestration by process info...")
+            shutdown_proc(proc)
+            LOGGER.info("Stopping mongo-orchestration by process info... done.")
+
+    # Next look for running docker images.
+    if docker:
+        cmd = f"{docker} ps --format '{{.Image}}\t{{.ID}}'"
+        response = subprocess.check_output(cmd, shell=True, encoding="utf-8").strip()
+        for line in response.splitlines():
+            image, container_id = line.split("\t")
+            if image in ["mongodb/mongodb-atlas-local", "mongo"]:
+                LOGGER.info(f"Stopping {image} by image name...")
+                shutdown_docker(docker, container_id)
+                LOGGER.info(f"Stopping {image} by image name... done.")
+
+    # Finally, look for any processes that are named mongod or mongos.
+    for proc in all_procs:
+        try:
+            name = proc.name()
+        except psutil.NoSuchProcess:
+            continue
+        if name in ["mongod", "mongos"]:
+            LOGGER.info(f"Stopping {name} by process name...")
+            shutdown_proc(proc)
+            LOGGER.info(f"Stopping {name} by process name... done.")
 
 
 def main():

.evergreen/orchestration/pyproject.toml

@@ -17,7 +17,10 @@ classifiers = [
   "Programming Language :: Python",
   "Programming Language :: Python :: Implementation :: CPython",
 ]
-dependencies = ["drivers-evergreen-tools @ {root:parent:uri}", "mongo-orchestration @ https://github.com/mongodb/mongo-orchestration/archive/master.tar.gz"]
+dependencies = [
+  "drivers-evergreen-tools @ {root:parent:uri}",
+  "mongo-orchestration @ https://github.com/mongodb/mongo-orchestration/archive/master.tar.gz",
+  "psutil>=6.0"]
 
 [project.scripts]
 drivers-orchestration = "drivers_orchestration:main"