From 8213c4fa2256c36c97f137f9b6256cf59c9bdc6d Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Sun, 2 Mar 2025 16:48:01 +0100
Subject: [PATCH] make wait script more robust

---
 acceptance/bin/wait_pid.py | 59 +++++++++++++++++++++++---------------
 1 file changed, 36 insertions(+), 23 deletions(-)

diff --git a/acceptance/bin/wait_pid.py b/acceptance/bin/wait_pid.py
index def958fae..989d4414d 100755
--- a/acceptance/bin/wait_pid.py
+++ b/acceptance/bin/wait_pid.py
@@ -5,45 +5,58 @@ import time
 import platform
 import subprocess
 
+def is_finished_windows(pid):
+    assert int(pid) > 0
+
+    # Run task list but filter the list by pid..
+    output = subprocess.check_output(f"tasklist /fi \"PID eq {pid}\"", text=True)
+
+    # If tasklist does not find any process, that means the process we are
+    # waiting for has terminated.
+    unaliveMsg = "No tasks are running which match the specified criteria."
+    if unaliveMsg in output:
+        return True
+
+    return False
+
+def is_finished_unix(pid):
+    assert int(pid) > 0
+
+    # Send signal 0 to the process to check if it exists.
+    # From the docs:
+    #    If sig is 0, then no signal is sent, but existence and permission
+    #    checks are still performed; this can be used to check for the
+    #    existence of a process ID or process group ID that the caller is
+    #    permitted to signal.
+    # ref: https://man7.org/linux/man-pages/man2/kill.2.html
+    try:
+        os.kill(pid, 0)
+    except OSError:
+        return True
+
+    return False
+
 
 def wait_pid(pid):
     max_attempts = 600  # 600 * 0.1 seconds = 1 minute
     sleep_time = 0.1
 
     for i in range(max_attempts):
-        # Check if we are on Windows or a Unix system.
         if platform.system().lower() == "windows":
-            # Windows approach: use tasklist to check for the existence of the process.
-            try:
-                # Get the output of 'tasklist'
-                output = subprocess.check_output(["tasklist"], text=True)
-            except subprocess.CalledProcessError:
-                print("[wait_pid] Error retrieving tasklist. Assuming process has ended.")
-                return 0
-
-            # if the PID is not found in the list then assume the process ended.
-            if str(pid) not in output:
+            if is_finished_windows(pid):
                 print("[wait_pid] process has ended")
                 return 0
+
         else:
-            # Linux/macOS approach: using os.kill with signal 0 to check if the process exists.
-            try:
-                os.kill(pid, 0)
-            except OSError:
+            if is_finished_unix(pid):
                 print("[wait_pid] process has ended")
                 return 0
 
         time.sleep(sleep_time)
 
+
     print(f"Timeout: Process {pid} did not end within 1 minute")
     return 1
 
-
-try:
-    pid = int(sys.argv[1])
-except ValueError:
-    print("Error: PID must be an integer.")
-    sys.exit(1)
-
-exit_code = wait_pid(pid)
+exit_code = wait_pid(int(sys.argv[1]))
 sys.exit(exit_code)