Issue 116: provide a 'timeout' parameter for Process.wait() and define a new TimeoutExpired exception.

giampaolo · giampaolo · commit 886710dafd1b · 2011-02-18T00:17:55.000Z
diff --git a/psutil/__init__.py b/psutil/__init__.py
@@ -12,7 +12,7 @@
 
 __all__ = [
     # exceptions
-    "Error", "NoSuchProcess", "AccessDenied",
+    "Error", "NoSuchProcess", "AccessDenied", "TimeoutExpired",
     # constants
     "NUM_CPUS", "TOTAL_PHYMEM", "BOOT_TIME",
     "version_info", "__version__",
@@ -39,7 +39,7 @@
 except ImportError:
     pwd = None
 
-from psutil.error import Error, NoSuchProcess, AccessDenied
+from psutil.error import Error, NoSuchProcess, AccessDenied, TimeoutExpired
 from psutil._compat import property
 from psutil._common import (STATUS_RUNNING, STATUS_IDLE, STATUS_SLEEPING,
                             STATUS_DISK_SLEEP, STATUS_STOPPED,
@@ -468,11 +468,11 @@ def kill(self):
         else:
             self._platform_impl.kill_process()
 
-    def wait(self):
+    def wait(self, timeout=None):
         """Wait for process to terminate and, if process is a children
         of the current one also return its exit code, else None.
         """
-        return self._platform_impl.process_wait()
+        return self._platform_impl.process_wait(timeout)
 
 
 class Popen(Process):
diff --git a/psutil/_psbsd.py b/psutil/_psbsd.py
@@ -9,7 +9,7 @@
 import _psutil_bsd
 import _psutil_posix
 import _psposix
-from psutil.error import AccessDenied, NoSuchProcess
+from psutil.error import AccessDenied, NoSuchProcess, TimeoutExpired
 from psutil._compat import namedtuple
 from psutil._common import *
 
@@ -173,8 +173,11 @@ def get_connections(self):
         return lsof.get_process_connections()
 
     @wrap_exceptions
-    def process_wait(self):
-        return _psposix.wait_pid(self.pid)
+    def process_wait(self, timeout=None):
+        try:
+            return _psposix.wait_pid(self.pid, timeout)
+        except TimeoutExpired:
+            raise TimeoutExpired(self.pid, self._process_name)
 
     @wrap_exceptions
     def get_process_nice(self):
diff --git a/psutil/_pslinux.py b/psutil/_pslinux.py
@@ -13,7 +13,7 @@
 import _psutil_posix
 import _psutil_linux
 from psutil import _psposix
-from psutil.error import AccessDenied, NoSuchProcess
+from psutil.error import AccessDenied, NoSuchProcess, TimeoutExpired
 from psutil._compat import namedtuple
 from psutil._common import *
 
@@ -294,8 +294,11 @@ def get_cpu_times(self):
         return ntuple_cputimes(utime, stime)
 
     @wrap_exceptions
-    def process_wait(self):
-        return _psposix.wait_pid(self.pid)
+    def process_wait(self, timeout=None):
+        try:
+            return _psposix.wait_pid(self.pid, timeout)
+        except TimeoutExpired:
+            raise TimeoutExpired(self.pid, self._process_name)
 
     @wrap_exceptions
     def get_process_create_time(self):
diff --git a/psutil/_psosx.py b/psutil/_psosx.py
@@ -9,7 +9,7 @@
 import _psutil_osx
 import _psutil_posix
 import _psposix
-from psutil.error import AccessDenied, NoSuchProcess
+from psutil.error import AccessDenied, NoSuchProcess, TimeoutExpired
 from psutil._compat import namedtuple
 from psutil._common import *
 
@@ -162,8 +162,11 @@ def get_connections(self):
         return lsof.get_process_connections()
 
     @wrap_exceptions
-    def process_wait(self):
-        return _psposix.wait_pid(self.pid)
+    def process_wait(self, timeout=None):
+        try:
+            return _psposix.wait_pid(self.pid, timeout)
+        except TimeoutExpired:
+            raise TimeoutExpired(self.pid, self._process_name)
 
     @wrap_exceptions
     def get_process_nice(self):
diff --git a/psutil/_psposix.py b/psutil/_psposix.py
@@ -15,7 +15,7 @@
 import warnings
 import time
 
-from psutil.error import AccessDenied, NoSuchProcess
+from psutil.error import AccessDenied, NoSuchProcess, TimeoutExpired
 from psutil._compat import namedtuple
 
 
@@ -30,29 +30,50 @@ def pid_exists(pid):
     else:
         return True
 
-def wait_pid(pid):
+def wait_pid(pid, timeout=None):
     """Wait for process with pid 'pid' to terminate and return its
     exit status code as an integer.
 
     If pid is not a children of os.getpid() (current process) just
     waits until the process disappears and return None.
 
     If pid does not exist at all return None immediately.
+
+    Raise TimeoutExpired on timeout expired.
     """
-    while True:
+    def check_timeout():
+        if timeout:
+            if time.time() >= stop_at:
+                raise TimeoutExpired
+            time.sleep(0.001)
+
+    if timeout:
+        waitcall = lambda: os.waitpid(pid, os.WNOHANG)
+        stop_at = time.time() + timeout
+    else:
+        waitcall = lambda: os.waitpid(pid, 0)
+
+    while 1:
         try:
-            pid, status = os.waitpid(pid, 0)
+            retpid, status = waitcall()
         except OSError, err:
             if err.errno == errno.EINTR:
+                check_timeout()
                 continue
             elif err.errno == errno.ECHILD:
                 # not a child of os.getpid(): poll until pid has
                 # disappeared and return None instead
-                while pid_exists(pid):
-                    time.sleep(0.05)
-                return
-            raise
+                while 1:
+                    if pid_exists(pid):
+                        check_timeout()
+                    else:
+                        return
+            else:
+                raise
         else:
+            if retpid == 0:
+                check_timeout()
+                continue
             # process exited due to a signal; return the integer of
             # that signal
             if os.WIFSIGNALED(status):
diff --git a/psutil/error.py b/psutil/error.py
@@ -39,13 +39,32 @@ def __init__(self, pid=None, name=None, msg=None):
         self.name = name
         self.msg = msg
         if msg is None:
-            if (pid is not None) and (name is not None):
+            if pid and name:
                 self.msg = "(pid=%s, name=%s)" % (pid, repr(name))
-            elif (pid is not None):
+            elif pid:
                 self.msg = "(pid=%s)" % self.pid
             else:
                 self.msg = ""
 
     def __str__(self):
         return self.msg
 
+
+class TimeoutExpired(Error):
+    """Raised on Process.wait(timeout) if timeout expires and process
+    is still alive.
+    """
+
+    def __init__(self, pid=None, name=None):
+        self.pid = pid
+        self.name = name
+        if pid and name:
+            self.msg = "(pid=%s, name=%s)" % (pid, repr(name))
+        elif pid:
+            self.msg = "(pid=%s)" % self.pid
+        else:
+            self.msg = ""
+
+    def __str__(self):
+        return self.msg
+
diff --git a/test/test_psutil.py b/test/test_psutil.py
@@ -229,14 +229,14 @@ def test_send_signal(self):
         self.assertFalse(psutil.pid_exists(test_pid) and name == PYTHON)
 
     def test_wait(self):
+        # check exit code signal
         sproc = get_test_subprocess()
         p = psutil.Process(sproc.pid)
         p.kill()
         code = p.wait()
         if os.name == 'posix':
             self.assertEqual(code, signal.SIGKILL)
         else:
-            # windows
             self.assertEqual(code, 0)
         self.assertFalse(p.is_running())
 
@@ -247,10 +247,10 @@ def test_wait(self):
         if os.name == 'posix':
             self.assertEqual(code, signal.SIGTERM)
         else:
-            # windows
             self.assertEqual(code, 0)
         self.assertFalse(p.is_running())
 
+        # check sys.exit() code
         code = "import time, sys; time.sleep(0.01); sys.exit(5);"
         sproc = get_test_subprocess([PYTHON, "-c", code])
         p = psutil.Process(sproc.pid)
@@ -266,6 +266,28 @@ def test_wait(self):
         self.assertEqual(p.wait(), 5)
         self.assertTrue(p.wait() in (5, None))
 
+        # test timeout
+        sproc = get_test_subprocess()
+        p = psutil.Process(sproc.pid)
+        p.name
+        self.assertRaises(psutil.TimeoutExpired, p.wait, 0.01)
+
+    def test_wait_non_children(self):
+        # test wait() against processes which are not our children
+        code = "import sys;"
+        code += "from subprocess import Popen, PIPE;"
+        code += "cmd = ['%s', '-c', 'import time; time.sleep(10)'];" %PYTHON
+        code += "sp = Popen(cmd, stdout=PIPE);"
+        code += "sys.stdout.write(str(sp.pid));"
+        sproc = get_test_subprocess([PYTHON, "-c", code], stdout=subprocess.PIPE)
+
+        grandson_pid = int(sproc.stdout.read())
+        grandson_proc = psutil.Process(grandson_pid)
+        self.assertRaises(psutil.TimeoutExpired, grandson_proc.wait, 0.01)
+        grandson_proc.kill()
+        ret = grandson_proc.wait()
+        self.assertEqual(ret, None)
+
     def test_TOTAL_PHYMEM(self):
         x = psutil.TOTAL_PHYMEM
         self.assertTrue(isinstance(x, (int, long)))
