-
Notifications
You must be signed in to change notification settings - Fork 0
/
subproc.py
314 lines (256 loc) · 11.5 KB
/
subproc.py
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
"""Common tasks for managing child processes.
To have child processes actually be managed by this module, you should use
the Popen() here rather than subprocess.Popen() directly.
Some parts do not yet work fully on windows (sending/trapping signals).
"""
import atexit
import inspect
import os
import signal
import subprocess
import sys
import time
mswindows = (sys.platform == "win32")
if mswindows:
# Options to control this module's behaviour
"""
By default on windows, python calls TerminateProcess to terminate a child
process, but semantically this is equivalent to SIGKILL which bypasses any
final cleanup. Set USE_CTRLBREAK_AS_SIGTERM=1 in the environment to send a
CTRL-BREAK console event instead, which acts more like SIGTERM.
"""
_use_ctrlbreak_as_sigterm = bool(os.getenv("USE_CTRLBREAK_AS_SIGTERM", 1))
"""
Set KILL_CHILDREN_ON_DEATH=1 in the environment to automatically kill all
descendents when this process dies.
"""
# TODO(infinity0): write a test for this, similar to test_killall_kill
# Note: setting this to True defeats the point of some of the tests, so
# keep the default value as False. Perhaps we could make that work better.
_kill_children_on_death = bool(os.getenv("KILL_CHILDREN_ON_DEATH", 0))
from ctypes import byref, windll, WinError
from ctypes.wintypes import DWORD
import win32api, win32job
_CREATE_BREAKAWAY_FROM_JOB = 0x01000000 # see _kill_children_on_death
_SYNCHRONIZE = 0x00100000 # generically useful
_PROCESS_SET_QUOTA = 0x0100 # required for AssignProcessToJobObject
_PROCESS_TERMINATE = 0x0001 # required for AssignProcessToJobObject
_PROCESS_QUERY_INFORMATION = 0x0400 # required for GetExitCodeProcess
_STILL_ACTIVE = 259 # GetExitCodeProcess returns this for still-running process
_CHILD_PROCS = []
# TODO(infinity0): add functionality to detect when any child dies, and
# offer different response strategies for them (e.g. restart the child? or die
# and kill the other children too).
SINK = object()
# get default args from subprocess.Popen to use in subproc.Popen
a = inspect.getargspec(subprocess.Popen.__init__)
_Popen_defaults = zip(a.args[-len(a.defaults):],a.defaults); del a
if mswindows:
# required for os.kill() to work
_Popen_creationflags = subprocess.CREATE_NEW_PROCESS_GROUP
if _kill_children_on_death:
_chJob = win32job.CreateJobObject(None, "")
if not _chJob:
raise WinError()
chJeli = win32job.QueryInformationJobObject(
_chJob, win32job.JobObjectExtendedLimitInformation)
# JOB_OBJECT_LIMIT_BREAKAWAY_OK allows children to assign grandchildren
# to their own jobs
chJeli['BasicLimitInformation']['LimitFlags'] |= (
win32job.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE |
win32job.JOB_OBJECT_LIMIT_BREAKAWAY_OK)
if win32job.SetInformationJobObject(
_chJob, win32job.JobObjectExtendedLimitInformation, chJeli) == 0:
raise WinError()
del chJeli
# If we already belong to a JobObject, our children are auto-assigned
# to that and AssignProcessToJobObject(ch, _chJob) fails. This flag
# prevents this auto-assignment, as long as the parent JobObject has
# JOB_OBJECT_LIMIT_BREAKAWAY_OK set on it as well.
_Popen_creationflags |= _CREATE_BREAKAWAY_FROM_JOB
tmp = dict(_Popen_defaults)
tmp['creationflags'] |= _Popen_creationflags
_Popen_defaults = tmp.items()
del tmp, _Popen_creationflags
class Popen(subprocess.Popen):
"""Wrapper for subprocess.Popen that tracks every child process.
See the subprocess module for documentation.
On windows, you are recommended to use the creationflagsmerge param so as
not to interfere with the required flags that we set in this module.
Additionally, you may use subproc.SINK as the value for either of the
stdout, stderr arguments to tell subprocess to discard anything written
to those channels.
"""
def __init__(self, *args, **kwargs):
kwargs = dict(_Popen_defaults + kwargs.items())
if 'creationflagsmerge' in kwargs:
kwargs['creationflags'] = (
kwargs.get('creationflags', 0) | kwargs['creationflagsmerge'])
del kwargs['creationflagsmerge']
for f in ['stdout', 'stderr']:
if kwargs[f] is SINK:
kwargs[f] = create_sink()
# super() does some magic that makes **kwargs not work, so just call
# our super-constructor directly
subprocess.Popen.__init__(self, *args, **kwargs)
_CHILD_PROCS.append(self)
if mswindows and _kill_children_on_death:
handle = windll.kernel32.OpenProcess(
_SYNCHRONIZE | _PROCESS_SET_QUOTA | _PROCESS_TERMINATE, 0, self.pid)
if win32job.AssignProcessToJobObject(_chJob, handle) == 0:
raise WinError()
if mswindows and _use_ctrlbreak_as_sigterm:
def send_signal(self, sig):
if sig == signal.SIGTERM:
self.terminate()
else:
subprocess.Popen.send_signal(self, sig)
def terminate(self):
os.kill(self.pid, signal.CTRL_BREAK_EVENT)
# TODO(infinity0): perhaps replace Popen.std* with wrapped file objects
# that don't buffer readlines() et. al. Currently one must avoid these and
# use while/readline(); see man page for "python -u" for more details.
def create_sink():
return open(os.devnull, "w", 0)
if mswindows:
# adapted from http://www.madebuild.org/blog/?p=30
def proc_is_alive(pid):
"""Check if a pid is still running."""
handle = windll.kernel32.OpenProcess(
_SYNCHRONIZE | _PROCESS_QUERY_INFORMATION, 0, pid)
if handle == 0:
return False
# If the process exited recently, a pid may still exist for the handle.
# So, check if we can get the exit code.
exit_code = DWORD()
rval = windll.kernel32.GetExitCodeProcess(handle, byref(exit_code))
windll.kernel32.CloseHandle(handle)
if rval == 0: # GetExitCodeProcess failure
raise WinError()
return exit_code.value == _STILL_ACTIVE
else:
# adapted from http://stackoverflow.com/questions/568271/check-if-pid-is-not-in-use-in-python
import errno
def proc_is_alive(pid):
"""Check if a pid is still running."""
try:
os.kill(pid, 0)
except OSError as e:
if e.errno == errno.EPERM:
return True
if e.errno == errno.ESRCH:
return False
raise # something else went wrong
else:
return True
class SignalHandlers(object):
def __init__(self):
self.handlers = {}
self.received = 0
def attach_override_unix(self, signum):
if signal.signal(signum, self.handle) != self.handle:
self.handlers.clear()
def handle(self, signum=0, sframe=None):
self.received += 1
# code snippet adapted from atexit._run_exitfuncs
exc_info = None
for i in xrange(self.received).__reversed__():
for handler in self.handlers.get(i, []).__reversed__():
try:
handler(signum, sframe)
except SystemExit:
exc_info = sys.exc_info()
except:
import traceback
print >> sys.stderr, "Error in SignalHandler.handle:"
traceback.print_exc()
exc_info = sys.exc_info()
if exc_info is not None:
raise exc_info[0], exc_info[1], exc_info[2]
def register(self, handler, ignoreNum):
self.handlers.setdefault(ignoreNum, []).append(handler)
_SIGINT_HANDLERS = SignalHandlers()
def trap_sigint(handler, ignoreNum=0):
"""Register a handler for an INT signal (Unix).
Note: this currently has no effect on windows.
Successive handlers are cumulative. On Unix, they override any previous
handlers registered with signal.signal().
Args:
handler: a signal handler; see signal.signal() for details
ignoreNum: number of signals to ignore before activating the handler,
which will be run on all subsequent signals.
"""
handlers = _SIGINT_HANDLERS
handlers.attach_override_unix(signal.SIGINT)
handlers.register(handler, ignoreNum)
_SIGTERM_HANDLERS = SignalHandlers()
def trap_sigterm(handler, ignoreNum=0):
"""Register a handler for a TERM signal (Unix) or CTRL-BREAK console event
(Windows).
Successive handlers are cumulative. On Unix, they override any previous
handlers registered with signal.signal(). On Windows, they *do not*
override previous handlers registered with win32api.SetConsoleCtrlHandler().
Args:
handler: a signal handler; see signal.signal() for details. For
cross-platform portability, it should accept None as a valid value
for the sframe (second) parameter.
ignoreNum: number of signals to ignore before activating the handler,
which will be run on all subsequent signals.
"""
handlers = _SIGTERM_HANDLERS
if not (mswindows and _use_ctrlbreak_as_sigterm):
handlers.attach_override_unix(signal.SIGTERM)
handlers.register(handler, ignoreNum)
if mswindows and _use_ctrlbreak_as_sigterm:
def _HandlerRoutine(signum):
if signum != signal.CTRL_BREAK_EVENT: return False
return _SIGTERM_HANDLERS.handle(signum)
win32api.SetConsoleCtrlHandler(_HandlerRoutine, True)
_isTerminating = False
def killall(cleanup=lambda:None, wait_s=16):
"""Attempt to gracefully terminate all child processes.
All children are told to terminate gracefully. A waiting period is then
applied, after which all children are killed forcefully. If all children
terminate before this waiting period is over, the function exits early.
Args:
cleanup: Run after all children are dead. For example, if your program
does not automatically terminate after this, you can use this
to signal that it should exit. In particular, Twisted
applications ought to use this to call reactor.stop().
wait_s: Time in seconds to wait before trying to kill children.
"""
# TODO(infinity0): log this somewhere, maybe
global _isTerminating, _CHILD_PROCS
if _isTerminating: return
_isTerminating = True
# terminate all
for proc in _CHILD_PROCS:
if proc.poll() is None:
proc.terminate()
# wait and make sure they're dead
for i in xrange(wait_s):
_CHILD_PROCS = [proc for proc in _CHILD_PROCS
if proc.poll() is None]
if not _CHILD_PROCS: break
time.sleep(1)
# if still existing, kill them
for proc in _CHILD_PROCS:
if proc.poll() is None:
proc.kill()
time.sleep(0.5)
# reap any zombies
for proc in _CHILD_PROCS:
proc.poll()
cleanup()
def auto_killall(ignoreNumSigInts=0, *args, **kwargs):
"""Automatically terminate all child processes on exit.
Args:
ignoreNumSigInts: this number of INT signals will be ignored before
attempting termination. This will be attempted unconditionally in
all other cases, such as on normal exit, or on a TERM signal.
*args, **kwargs: See killall().
"""
killall_handler = lambda signum, sframe: killall(*args, **kwargs)
trap_sigint(killall_handler, ignoreNumSigInts)
trap_sigterm(killall_handler)
atexit.register(killall, *args, **kwargs)