summaryrefslogtreecommitdiffstats
path: root/addon-sdk/source/python-lib/mozrunner/killableprocess.py
blob: 21eac696521d9f800a0a179d3298c7cf56376879 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
# killableprocess - subprocesses which can be reliably killed
#
# Parts of this module are copied from the subprocess.py file contained
# in the Python distribution.
#
# Copyright (c) 2003-2004 by Peter Astrand <astrand@lysator.liu.se>
#
# Additions and modifications written by Benjamin Smedberg
# <benjamin@smedbergs.us> are Copyright (c) 2006 by the Mozilla Foundation
# <http://www.mozilla.org/>
#
# More Modifications
# Copyright (c) 2006-2007 by Mike Taylor <bear@code-bear.com>
# Copyright (c) 2007-2008 by Mikeal Rogers <mikeal@mozilla.com>
#
# By obtaining, using, and/or copying this software and/or its
# associated documentation, you agree that you have read, understood,
# and will comply with the following terms and conditions:
#
# Permission to use, copy, modify, and distribute this software and
# its associated documentation for any purpose and without fee is
# hereby granted, provided that the above copyright notice appears in
# all copies, and that both that copyright notice and this permission
# notice appear in supporting documentation, and that the name of the
# author not be used in advertising or publicity pertaining to
# distribution of the software without specific, written prior
# permission.
#
# THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE,
# INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS.
# IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, INDIRECT OR
# CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS
# OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
# NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION
# WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

"""killableprocess - Subprocesses which can be reliably killed

This module is a subclass of the builtin "subprocess" module. It allows
processes that launch subprocesses to be reliably killed on Windows (via the Popen.kill() method.

It also adds a timeout argument to Wait() for a limited period of time before
forcefully killing the process.

Note: On Windows, this module requires Windows 2000 or higher (no support for
Windows 95, 98, or NT 4.0). It also requires ctypes, which is bundled with
Python 2.5+ or available from http://python.net/crew/theller/ctypes/
"""

import subprocess
import sys
import os
import time
import datetime
import types
import exceptions

try:
    from subprocess import CalledProcessError
except ImportError:
    # Python 2.4 doesn't implement CalledProcessError
    class CalledProcessError(Exception):
        """This exception is raised when a process run by check_call() returns
        a non-zero exit status. The exit status will be stored in the
        returncode attribute."""
        def __init__(self, returncode, cmd):
            self.returncode = returncode
            self.cmd = cmd
        def __str__(self):
            return "Command '%s' returned non-zero exit status %d" % (self.cmd, self.returncode)

mswindows = (sys.platform == "win32")

if mswindows:
    import winprocess
else:
    import signal

# This is normally defined in win32con, but we don't want
# to incur the huge tree of dependencies (pywin32 and friends)
# just to get one constant.  So here's our hack
STILL_ACTIVE = 259

def call(*args, **kwargs):
    waitargs = {}
    if "timeout" in kwargs:
        waitargs["timeout"] = kwargs.pop("timeout")

    return Popen(*args, **kwargs).wait(**waitargs)

def check_call(*args, **kwargs):
    """Call a program with an optional timeout. If the program has a non-zero
    exit status, raises a CalledProcessError."""

    retcode = call(*args, **kwargs)
    if retcode:
        cmd = kwargs.get("args")
        if cmd is None:
            cmd = args[0]
        raise CalledProcessError(retcode, cmd)

if not mswindows:
    def DoNothing(*args):
        pass

class Popen(subprocess.Popen):
    kill_called = False
    if mswindows:
        def _execute_child(self, *args_tuple):
            # workaround for bug 958609
            if sys.hexversion < 0x02070600: # prior to 2.7.6
                (args, executable, preexec_fn, close_fds,
                    cwd, env, universal_newlines, startupinfo,
                    creationflags, shell,
                    p2cread, p2cwrite,
                    c2pread, c2pwrite,
                    errread, errwrite) = args_tuple
                to_close = set()
            else: # 2.7.6 and later
                (args, executable, preexec_fn, close_fds,
                    cwd, env, universal_newlines, startupinfo,
                    creationflags, shell, to_close,
                    p2cread, p2cwrite,
                    c2pread, c2pwrite,
                    errread, errwrite) = args_tuple

            if not isinstance(args, types.StringTypes):
                args = subprocess.list2cmdline(args)

            # Always or in the create new process group
            creationflags |= winprocess.CREATE_NEW_PROCESS_GROUP

            if startupinfo is None:
                startupinfo = winprocess.STARTUPINFO()

            if None not in (p2cread, c2pwrite, errwrite):
                startupinfo.dwFlags |= winprocess.STARTF_USESTDHANDLES

                startupinfo.hStdInput = int(p2cread)
                startupinfo.hStdOutput = int(c2pwrite)
                startupinfo.hStdError = int(errwrite)
            if shell:
                startupinfo.dwFlags |= winprocess.STARTF_USESHOWWINDOW
                startupinfo.wShowWindow = winprocess.SW_HIDE
                comspec = os.environ.get("COMSPEC", "cmd.exe")
                args = comspec + " /c " + args

            # determine if we can create create a job
            canCreateJob = winprocess.CanCreateJobObject()

            # set process creation flags
            creationflags |= winprocess.CREATE_SUSPENDED
            creationflags |= winprocess.CREATE_UNICODE_ENVIRONMENT
            if canCreateJob:
                # Uncomment this line below to discover very useful things about your environment
                #print "++++ killableprocess: releng twistd patch not applied, we can create job objects"
                creationflags |= winprocess.CREATE_BREAKAWAY_FROM_JOB

            # create the process
            hp, ht, pid, tid = winprocess.CreateProcess(
                executable, args,
                None, None, # No special security
                1, # Must inherit handles!
                creationflags,
                winprocess.EnvironmentBlock(env),
                cwd, startupinfo)
            self._child_created = True
            self._handle = hp
            self._thread = ht
            self.pid = pid
            self.tid = tid

            if canCreateJob:
                # We create a new job for this process, so that we can kill
                # the process and any sub-processes
                self._job = winprocess.CreateJobObject()
                winprocess.AssignProcessToJobObject(self._job, int(hp))
            else:
                self._job = None

            winprocess.ResumeThread(int(ht))
            ht.Close()

            if p2cread is not None:
                p2cread.Close()
            if c2pwrite is not None:
                c2pwrite.Close()
            if errwrite is not None:
                errwrite.Close()
            time.sleep(.1)

    def kill(self, group=True):
        """Kill the process. If group=True, all sub-processes will also be killed."""
        self.kill_called = True

        if mswindows:
            if group and self._job:
                winprocess.TerminateJobObject(self._job, 127)
            else:
                winprocess.TerminateProcess(self._handle, 127)
            self.returncode = 127
        else:
            if group:
                try:
                    os.killpg(self.pid, signal.SIGKILL)
                except: pass
            else:
                os.kill(self.pid, signal.SIGKILL)
            self.returncode = -9

    def wait(self, timeout=None, group=True):
        """Wait for the process to terminate. Returns returncode attribute.
        If timeout seconds are reached and the process has not terminated,
        it will be forcefully killed. If timeout is -1, wait will not
        time out."""
        if timeout is not None:
            # timeout is now in milliseconds
            timeout = timeout * 1000

        starttime = datetime.datetime.utcnow()

        if mswindows:
            if timeout is None:
                timeout = -1
            rc = winprocess.WaitForSingleObject(self._handle, timeout)

            if (rc == winprocess.WAIT_OBJECT_0 or
                rc == winprocess.WAIT_ABANDONED or
                rc == winprocess.WAIT_FAILED):
                # Object has either signaled, or the API call has failed.  In 
                # both cases we want to give the OS the benefit of the doubt
                # and supply a little time before we start shooting processes
                # with an M-16.

                # Returns 1 if running, 0 if not, -1 if timed out
                def check():
                    now = datetime.datetime.utcnow()
                    diff = now - starttime
                    if (diff.seconds * 1000000 + diff.microseconds) < (timeout * 1000):    # (1000*1000)
                        if self._job:
                            if (winprocess.QueryInformationJobObject(self._job, 8)['BasicInfo']['ActiveProcesses'] > 0):
                                # Job Object is still containing active processes
                                return 1
                        else:
                            # No job, we use GetExitCodeProcess, which will tell us if the process is still active
                            self.returncode = winprocess.GetExitCodeProcess(self._handle)
                            if (self.returncode == STILL_ACTIVE):
                                # Process still active, continue waiting
                                return 1
                        # Process not active, return 0
                        return 0
                    else:
                        # Timed out, return -1
                        return -1

                notdone = check()
                while notdone == 1:
                    time.sleep(.5)
                    notdone = check()

                if notdone == -1:
                    # Then check timed out, we have a hung process, attempt
                    # last ditch kill with explosives
                    self.kill(group)
                                
            else:
                # In this case waitforsingleobject timed out.  We have to
                # take the process behind the woodshed and shoot it.
                self.kill(group)

        else:
            if sys.platform in ('linux2', 'sunos5', 'solaris') \
                    or sys.platform.startswith('freebsd'):
                def group_wait(timeout):
                    try:
                        os.waitpid(self.pid, 0)
                    except OSError, e:
                        pass # If wait has already been called on this pid, bad things happen
                    return self.returncode
            elif sys.platform == 'darwin':
                def group_wait(timeout):
                    try:
                        count = 0
                        if timeout is None and self.kill_called:
                            timeout = 10 # Have to set some kind of timeout or else this could go on forever
                        if timeout is None:
                            while 1:
                                os.killpg(self.pid, signal.SIG_DFL)
                        while ((count * 2) <= timeout):
                            os.killpg(self.pid, signal.SIG_DFL)
                            # count is increased by 500ms for every 0.5s of sleep
                            time.sleep(.5); count += 500
                    except exceptions.OSError:
                        return self.returncode
                        
            if timeout is None:
                if group is True:
                    return group_wait(timeout)
                else:
                    subprocess.Popen.wait(self)
                    return self.returncode

            returncode = False

            now = datetime.datetime.utcnow()
            diff = now - starttime
            while (diff.seconds * 1000 * 1000 + diff.microseconds) < (timeout * 1000) and ( returncode is False ):
                if group is True:
                    return group_wait(timeout)
                else:
                    if subprocess.poll() is not None:
                        returncode = self.returncode
                time.sleep(.5)
                now = datetime.datetime.utcnow()
                diff = now - starttime
            return self.returncode

        return self.returncode
    # We get random maxint errors from subprocesses __del__
    __del__ = lambda self: None        
        
def setpgid_preexec_fn():
    os.setpgid(0, 0)
   
def runCommand(cmd, **kwargs):
    if sys.platform != "win32":
        return Popen(cmd, preexec_fn=setpgid_preexec_fn, **kwargs)
    else:
        return Popen(cmd, **kwargs)