test/util/killableprocess.py

   1 # killableprocess - subprocesses which can be reliably killed
   2 #
   3 # Parts of this module are copied from the subprocess.py file contained
   4 # in the Python distribution.
   5 #
   6 # Copyright (c) 2003-2004 by Peter Astrand <astrand@lysator.liu.se>
   7 #
   8 # Additions and modifications written by Benjamin Smedberg
   9 # <benjamin@smedbergs.us> are Copyright (c) 2006 by the Mozilla Foundation
  10 # <http://www.mozilla.org/>
  11 #
  12 # By obtaining, using, and/or copying this software and/or its
  13 # associated documentation, you agree that you have read, understood,
  14 # and will comply with the following terms and conditions:
  15 #
  16 # Permission to use, copy, modify, and distribute this software and
  17 # its associated documentation for any purpose and without fee is
  18 # hereby granted, provided that the above copyright notice appears in
  19 # all copies, and that both that copyright notice and this permission
  20 # notice appear in supporting documentation, and that the name of the
  21 # author not be used in advertising or publicity pertaining to
  22 # distribution of the software without specific, written prior
  23 # permission.
  24 #
  25 # THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE,
  26 # INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS.
  27 # IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, INDIRECT OR
  28 # CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS
  29 # OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
  30 # NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION
  31 # WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
  32
  33 r"""killableprocess - Subprocesses which can be reliably killed
  34
  35 This module is a subclass of the builtin "subprocess" module. It allows
  36 processes that launch subprocesses to be reliably killed on Windows (via the Popen.kill() method.
  37
  38 It also adds a timeout argument to Wait() for a limited period of time before
  39 forcefully killing the process.
  40
  41 Note: On Windows, this module requires Windows 2000 or higher (no support for
  42 Windows 95, 98, or NT 4.0). It also requires ctypes, which is bundled with
  43 Python 2.5+ or available from http://python.net/crew/theller/ctypes/
  44 """
  45
  46 import subprocess
  47 import sys
  48 import os
  49 import time
  50 import types
  51
  52 try:
  53     from subprocess import CalledProcessError
  54 except ImportError:
  55     # Python 2.4 doesn't implement CalledProcessError
  56     class CalledProcessError(Exception):
  57         """This exception is raised when a process run by check_call() returns
  58         a non-zero exit status. The exit status will be stored in the
  59         returncode attribute."""
  60         def __init__(self, returncode, cmd):
  61             self.returncode = returncode
  62             self.cmd = cmd
  63         def __str__(self):
  64             return "Command '%s' returned non-zero exit status %d" % (self.cmd, self.returncode)
  65
  66 mswindows = (sys.platform == "win32")
  67
  68 if mswindows:
  69     import winprocess
  70 else:
  71     import signal
  72
  73 def call(*args, **kwargs):
  74     waitargs = {}
  75     if "timeout" in kwargs:
  76         waitargs["timeout"] = kwargs.pop("timeout")
  77
  78     return Popen(*args, **kwargs).wait(**waitargs)
  79
  80 def check_call(*args, **kwargs):
  81     """Call a program with an optional timeout. If the program has a non-zero
  82     exit status, raises a CalledProcessError."""
  83
  84     retcode = call(*args, **kwargs)
  85     if retcode:
  86         cmd = kwargs.get("args")
  87         if cmd is None:
  88             cmd = args[0]
  89         raise CalledProcessError(retcode, cmd)
  90
  91 if not mswindows:
  92     def DoNothing(*args):
  93         pass
  94
  95 class Popen(subprocess.Popen):
  96     if not mswindows:
  97         # Override __init__ to set a preexec_fn
  98         def __init__(self, *args, **kwargs):
  99             if len(args) >= 7:
 100                 raise Exception("Arguments preexec_fn and after must be passed by keyword.")
 101
 102             real_preexec_fn = kwargs.pop("preexec_fn", None)
 103             def setpgid_preexec_fn():
 104                 os.setpgid(0, 0)
 105                 if real_preexec_fn:
 106                     apply(real_preexec_fn)
 107
 108             kwargs['preexec_fn'] = setpgid_preexec_fn
 109
 110             subprocess.Popen.__init__(self, *args, **kwargs)
 111
 112     if mswindows:
 113         def _execute_child(self, args, executable, preexec_fn, close_fds,
 114                            cwd, env, universal_newlines, startupinfo,
 115                            creationflags, shell,
 116                            p2cread, p2cwrite,
 117                            c2pread, c2pwrite,
 118                            errread, errwrite):
 119             if not isinstance(args, types.StringTypes):
 120                 args = subprocess.list2cmdline(args)
 121
 122             if startupinfo is None:
 123                 startupinfo = winprocess.STARTUPINFO()
 124
 125             if None not in (p2cread, c2pwrite, errwrite):
 126                 startupinfo.dwFlags |= winprocess.STARTF_USESTDHANDLES
 127
 128                 startupinfo.hStdInput = int(p2cread)
 129                 startupinfo.hStdOutput = int(c2pwrite)
 130                 startupinfo.hStdError = int(errwrite)
 131             if shell:
 132                 startupinfo.dwFlags |= winprocess.STARTF_USESHOWWINDOW
 133 #                startupinfo.wShowWindow = winprocess.SW_HIDE
 134                 comspec = os.environ.get("COMSPEC", "cmd.exe")
 135                 args = comspec + " /c " + args
 136
 137             # We create a new job for this process, so that we can kill
 138             # the process and any sub-processes
 139             self._job = winprocess.CreateJobObject()
 140
 141             creationflags |= winprocess.CREATE_SUSPENDED
 142             creationflags |= winprocess.CREATE_UNICODE_ENVIRONMENT
 143
 144             hp, ht, pid, tid = winprocess.CreateProcess(
 145                 executable, args,
 146                 None, None, # No special security
 147                 1, # Must inherit handles!
 148                 creationflags,
 149                 winprocess.EnvironmentBlock(env),
 150                 cwd, startupinfo)
 151
 152             self._child_created = True
 153             self._handle = hp
 154             self._thread = ht
 155             self.pid = pid
 156
 157             winprocess.AssignProcessToJobObject(self._job, hp)
 158             winprocess.ResumeThread(ht)
 159
 160             if p2cread is not None:
 161                 p2cread.Close()
 162             if c2pwrite is not None:
 163                 c2pwrite.Close()
 164             if errwrite is not None:
 165                 errwrite.Close()
 166
 167     def kill(self, group=True):
 168         """Kill the process. If group=True, all sub-processes will also be killed."""
 169         if mswindows:
 170             if group:
 171                 winprocess.TerminateJobObject(self._job, 127)
 172             else:
 173                 winprocess.TerminateProcess(self._handle, 127)
 174             self.returncode = 127
 175         else:
 176             if group:
 177                 os.killpg(self.pid, signal.SIGKILL)
 178             else:
 179                 os.kill(self.pid, signal.SIGKILL)
 180             self.returncode = -9
 181
 182     def wait(self, timeout=-1, group=True):
 183         """Wait for the process to terminate. Returns returncode attribute.
 184         If timeout seconds are reached and the process has not terminated,
 185         it will be forcefully killed. If timeout is -1, wait will not
 186         time out."""
 187
 188         if self.returncode is not None:
 189             return self.returncode
 190
 191         if mswindows:
 192             if timeout != -1:
 193                 timeout = timeout * 1000
 194             rc = winprocess.WaitForSingleObject(self._handle, timeout)
 195             if rc == winprocess.WAIT_TIMEOUT:
 196                 self.kill(group)
 197             else:
 198                 self.returncode = winprocess.GetExitCodeProcess(self._handle)
 199         else:
 200             if timeout == -1:
 201                 subprocess.Popen.wait(self)
 202                 return self.returncode
 203
 204             starttime = time.time()
 205
 206             # Make sure there is a signal handler for SIGCHLD installed
 207             oldsignal = signal.signal(signal.SIGCHLD, DoNothing)
 208
 209             while time.time() < starttime + timeout - 0.01:
 210                 pid, sts = os.waitpid(self.pid, os.WNOHANG)
 211                 if pid != 0:
 212                     self._handle_exitstatus(sts)
 213                     signal.signal(signal.SIGCHLD, oldsignal)
 214                     return self.returncode
 215
 216                 # time.sleep is interrupted by signals (good!)
 217                 newtimeout = timeout - time.time() + starttime
 218                 time.sleep(newtimeout)
 219
 220             self.kill(group)
 221             signal.signal(signal.SIGCHLD, oldsignal)
 222             subprocess.Popen.wait(self)
 223
 224         return self.returncode