You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
608 lines
20 KiB
Python
608 lines
20 KiB
Python
# Copyright 2013 The LUCI Authors. All rights reserved.
|
|
# Use of this source code is governed by the Apache v2.0 license that can be
|
|
# found in the LICENSE file.
|
|
|
|
# This is a direct vendored copy of the original which is located at:
|
|
# https://github.com/luci/luci-py/blob/master/client/utils/subprocess42.py
|
|
# at commit 45eab10a17a55f1978f3bfa02c35457dee1c64e4.
|
|
|
|
"""subprocess42 is the answer to life the universe and everything.
|
|
|
|
It has the particularity of having a Popen implementation that can yield output
|
|
as it is produced while implementing a timeout and NOT requiring the use of
|
|
worker threads.
|
|
|
|
Example:
|
|
Wait for a child process with a timeout, send SIGTERM, wait a grace period
|
|
then send SIGKILL:
|
|
|
|
def wait_terminate_then_kill(proc, timeout, grace):
|
|
try:
|
|
return proc.wait(timeout)
|
|
except subprocess42.TimeoutExpired:
|
|
proc.terminate()
|
|
try:
|
|
return proc.wait(grace)
|
|
except subprocess42.TimeoutExpired:
|
|
proc.kill()
|
|
return proc.wait()
|
|
|
|
|
|
TODO(maruel): Add VOID support like subprocess2.
|
|
"""
|
|
|
|
import contextlib
|
|
import errno
|
|
import os
|
|
import signal
|
|
import threading
|
|
import time
|
|
|
|
import subprocess
|
|
|
|
from subprocess import CalledProcessError, PIPE, STDOUT # pylint: disable=W0611
|
|
from subprocess import list2cmdline
|
|
|
|
|
|
# Default maxsize argument.
|
|
MAX_SIZE = 16384
|
|
|
|
|
|
if subprocess.mswindows:
|
|
import msvcrt # pylint: disable=F0401
|
|
from ctypes import wintypes
|
|
from ctypes import windll
|
|
|
|
|
|
# Which to be received depends on how this process was called and outside the
|
|
# control of this script. See Popen docstring for more details.
|
|
STOP_SIGNALS = (signal.SIGBREAK, signal.SIGTERM)
|
|
|
|
|
|
def ReadFile(handle, desired_bytes):
|
|
"""Calls kernel32.ReadFile()."""
|
|
c_read = wintypes.DWORD()
|
|
buff = wintypes.create_string_buffer(desired_bytes+1)
|
|
windll.kernel32.ReadFile(
|
|
handle, buff, desired_bytes, wintypes.byref(c_read), None)
|
|
# NULL terminate it.
|
|
buff[c_read.value] = '\x00'
|
|
return wintypes.GetLastError(), buff.value
|
|
|
|
def PeekNamedPipe(handle):
|
|
"""Calls kernel32.PeekNamedPipe(). Simplified version."""
|
|
c_avail = wintypes.DWORD()
|
|
c_message = wintypes.DWORD()
|
|
success = windll.kernel32.PeekNamedPipe(
|
|
handle, None, 0, None, wintypes.byref(c_avail),
|
|
wintypes.byref(c_message))
|
|
if not success:
|
|
raise OSError(wintypes.GetLastError())
|
|
return c_avail.value
|
|
|
|
def recv_multi_impl(conns, maxsize, timeout):
|
|
"""Reads from the first available pipe.
|
|
|
|
It will immediately return on a closed connection, independent of timeout.
|
|
|
|
Arguments:
|
|
- maxsize: Maximum number of bytes to return. Defaults to MAX_SIZE.
|
|
- timeout: If None, it is blocking. If 0 or above, will return None if no
|
|
data is available within |timeout| seconds.
|
|
|
|
Returns:
|
|
tuple(int(index), str(data), bool(closed)).
|
|
"""
|
|
assert conns
|
|
assert timeout is None or isinstance(timeout, (int, float)), timeout
|
|
maxsize = max(maxsize or MAX_SIZE, 1)
|
|
|
|
# TODO(maruel): Use WaitForMultipleObjects(). Python creates anonymous pipes
|
|
# for proc.stdout and proc.stderr but they are implemented as named pipes on
|
|
# Windows. Since named pipes are not waitable object, they can't be passed
|
|
# as-is to WFMO(). So this means N times CreateEvent(), N times ReadFile()
|
|
# and finally WFMO(). This requires caching the events handles in the Popen
|
|
# object and remembering the pending ReadFile() calls. This will require
|
|
# some re-architecture to store the relevant event handle and OVERLAPPEDIO
|
|
# object in Popen or the file object.
|
|
start = time.time()
|
|
handles = [
|
|
(i, msvcrt.get_osfhandle(c.fileno())) for i, c in enumerate(conns)
|
|
]
|
|
while True:
|
|
for index, handle in handles:
|
|
try:
|
|
avail = min(PeekNamedPipe(handle), maxsize)
|
|
if avail:
|
|
return index, ReadFile(handle, avail)[1], False
|
|
except OSError:
|
|
# The pipe closed.
|
|
return index, None, True
|
|
|
|
if timeout is not None and (time.time() - start) >= timeout:
|
|
return None, None, False
|
|
# Polling rocks.
|
|
time.sleep(0.001)
|
|
|
|
else:
|
|
import fcntl # pylint: disable=F0401
|
|
import select
|
|
|
|
|
|
# Signals that mean this process should exit quickly.
|
|
STOP_SIGNALS = (signal.SIGINT, signal.SIGTERM)
|
|
|
|
|
|
def recv_multi_impl(conns, maxsize, timeout):
|
|
"""Reads from the first available pipe.
|
|
|
|
It will immediately return on a closed connection, independent of timeout.
|
|
|
|
Arguments:
|
|
- maxsize: Maximum number of bytes to return. Defaults to MAX_SIZE.
|
|
- timeout: If None, it is blocking. If 0 or above, will return None if no
|
|
data is available within |timeout| seconds.
|
|
|
|
Returns:
|
|
tuple(int(index), str(data), bool(closed)).
|
|
"""
|
|
assert conns
|
|
assert timeout is None or isinstance(timeout, (int, float)), timeout
|
|
maxsize = max(maxsize or MAX_SIZE, 1)
|
|
|
|
# select(timeout=0) will block, it has to be a value > 0.
|
|
if timeout == 0:
|
|
timeout = 0.001
|
|
try:
|
|
r, _, _ = select.select(conns, [], [], timeout)
|
|
except select.error:
|
|
r = None
|
|
if not r:
|
|
return None, None, False
|
|
|
|
conn = r[0]
|
|
# Temporarily make it non-blocking.
|
|
# TODO(maruel): This is not very efficient when the caller is doing this in
|
|
# a loop. Add a mechanism to have the caller handle this.
|
|
flags = fcntl.fcntl(conn, fcntl.F_GETFL)
|
|
if not conn.closed:
|
|
# pylint: disable=E1101
|
|
fcntl.fcntl(conn, fcntl.F_SETFL, flags | os.O_NONBLOCK)
|
|
try:
|
|
try:
|
|
data = conn.read(maxsize)
|
|
except IOError as e:
|
|
# On posix, this means the read would block.
|
|
if e.errno == errno.EAGAIN:
|
|
return conns.index(conn), None, False
|
|
raise e
|
|
|
|
if not data:
|
|
# On posix, this means the channel closed.
|
|
return conns.index(conn), None, True
|
|
|
|
return conns.index(conn), data, False
|
|
finally:
|
|
if not conn.closed:
|
|
fcntl.fcntl(conn, fcntl.F_SETFL, flags)
|
|
|
|
|
|
class TimeoutExpired(Exception):
|
|
"""Compatible with python3 subprocess."""
|
|
def __init__(self, cmd, timeout, output=None, stderr=None):
|
|
self.cmd = cmd
|
|
self.timeout = timeout
|
|
self.output = output
|
|
# Non-standard:
|
|
self.stderr = stderr
|
|
super(TimeoutExpired, self).__init__(str(self))
|
|
|
|
def __str__(self):
|
|
return "Command '%s' timed out after %s seconds" % (self.cmd, self.timeout)
|
|
|
|
|
|
class Popen(subprocess.Popen):
|
|
"""Adds timeout support on stdout and stderr.
|
|
|
|
Inspired by
|
|
http://code.activestate.com/recipes/440554-module-to-allow-asynchronous-subprocess-use-on-win/
|
|
|
|
Unlike subprocess, yield_any(), recv_*(), communicate() will close stdout and
|
|
stderr once the child process closes them, after all the data is read.
|
|
|
|
Arguments:
|
|
- detached: If True, the process is created in a new process group. On
|
|
Windows, use CREATE_NEW_PROCESS_GROUP. On posix, use os.setpgid(0, 0).
|
|
|
|
Additional members:
|
|
- start: timestamp when this process started.
|
|
- end: timestamp when this process exited, as seen by this process.
|
|
- detached: If True, the child process was started as a detached process.
|
|
- gid: process group id, if any.
|
|
- duration: time in seconds the process lasted.
|
|
|
|
Additional methods:
|
|
- yield_any(): yields output until the process terminates.
|
|
- recv_any(): reads from stdout and/or stderr with optional timeout.
|
|
- recv_out() & recv_err(): specialized version of recv_any().
|
|
"""
|
|
# subprocess.Popen.__init__() is not threadsafe; there is a race between
|
|
# creating the exec-error pipe for the child and setting it to CLOEXEC during
|
|
# which another thread can fork and cause the pipe to be inherited by its
|
|
# descendents, which will cause the current Popen to hang until all those
|
|
# descendents exit. Protect this with a lock so that only one fork/exec can
|
|
# happen at a time.
|
|
popen_lock = threading.Lock()
|
|
|
|
def __init__(self, args, **kwargs):
|
|
assert 'creationflags' not in kwargs
|
|
assert 'preexec_fn' not in kwargs, 'Use detached=True instead'
|
|
self.start = time.time()
|
|
self.end = None
|
|
self.gid = None
|
|
self.detached = kwargs.pop('detached', False)
|
|
if self.detached:
|
|
if subprocess.mswindows:
|
|
kwargs['creationflags'] = subprocess.CREATE_NEW_PROCESS_GROUP
|
|
else:
|
|
kwargs['preexec_fn'] = lambda: os.setpgid(0, 0)
|
|
with self.popen_lock:
|
|
super(Popen, self).__init__(args, **kwargs)
|
|
self.args = args
|
|
if self.detached and not subprocess.mswindows:
|
|
try:
|
|
self.gid = os.getpgid(self.pid)
|
|
except OSError:
|
|
# sometimes the process can run+finish before we collect its pgid. fun.
|
|
pass
|
|
|
|
def duration(self):
|
|
"""Duration of the child process.
|
|
|
|
It is greater or equal to the actual time the child process ran. It can be
|
|
significantly higher than the real value if neither .wait() nor .poll() was
|
|
used.
|
|
"""
|
|
return (self.end or time.time()) - self.start
|
|
|
|
# pylint: disable=arguments-differ,redefined-builtin
|
|
def communicate(self, input=None, timeout=None):
|
|
"""Implements python3's timeout support.
|
|
|
|
Unlike wait(), timeout=0 is considered the same as None.
|
|
|
|
Raises:
|
|
- TimeoutExpired when more than timeout seconds were spent waiting for the
|
|
process.
|
|
"""
|
|
if not timeout:
|
|
return super(Popen, self).communicate(input=input)
|
|
|
|
assert isinstance(timeout, (int, float)), timeout
|
|
|
|
if self.stdin or self.stdout or self.stderr:
|
|
stdout = '' if self.stdout else None
|
|
stderr = '' if self.stderr else None
|
|
t = None
|
|
if input is not None:
|
|
assert self.stdin, (
|
|
'Can\'t use communicate(input) if not using '
|
|
'Popen(stdin=subprocess42.PIPE')
|
|
# TODO(maruel): Switch back to non-threading.
|
|
def write():
|
|
try:
|
|
self.stdin.write(input)
|
|
except IOError:
|
|
pass
|
|
t = threading.Thread(name='Popen.communicate', target=write)
|
|
t.daemon = True
|
|
t.start()
|
|
|
|
try:
|
|
if self.stdout or self.stderr:
|
|
start = time.time()
|
|
end = start + timeout
|
|
def remaining():
|
|
return max(end - time.time(), 0)
|
|
for pipe, data in self.yield_any(timeout=remaining):
|
|
if pipe is None:
|
|
raise TimeoutExpired(self.args, timeout, stdout, stderr)
|
|
assert pipe in ('stdout', 'stderr'), pipe
|
|
if pipe == 'stdout':
|
|
stdout += data
|
|
else:
|
|
stderr += data
|
|
else:
|
|
# Only stdin is piped.
|
|
self.wait(timeout=timeout)
|
|
finally:
|
|
if t:
|
|
try:
|
|
self.stdin.close()
|
|
except IOError:
|
|
pass
|
|
t.join()
|
|
else:
|
|
# No pipe. The user wanted to use wait().
|
|
self.wait(timeout=timeout)
|
|
return None, None
|
|
|
|
# Indirectly initialize self.end.
|
|
self.wait()
|
|
return stdout, stderr
|
|
|
|
def wait(self, timeout=None): # pylint: disable=arguments-differ
|
|
"""Implements python3's timeout support.
|
|
|
|
Raises:
|
|
- TimeoutExpired when more than timeout seconds were spent waiting for the
|
|
process.
|
|
"""
|
|
assert timeout is None or isinstance(timeout, (int, float)), timeout
|
|
if timeout is None:
|
|
super(Popen, self).wait()
|
|
elif self.returncode is None:
|
|
if subprocess.mswindows:
|
|
WAIT_TIMEOUT = 258
|
|
result = subprocess._subprocess.WaitForSingleObject(
|
|
self._handle, int(timeout * 1000))
|
|
if result == WAIT_TIMEOUT:
|
|
raise TimeoutExpired(self.args, timeout)
|
|
self.returncode = subprocess._subprocess.GetExitCodeProcess(
|
|
self._handle)
|
|
else:
|
|
# If you think the following code is horrible, it's because it is
|
|
# inspired by python3's stdlib.
|
|
end = time.time() + timeout
|
|
delay = 0.001
|
|
while True:
|
|
try:
|
|
pid, sts = subprocess._eintr_retry_call(
|
|
os.waitpid, self.pid, os.WNOHANG)
|
|
except OSError as e:
|
|
if e.errno != errno.ECHILD:
|
|
raise
|
|
pid = self.pid
|
|
sts = 0
|
|
if pid == self.pid:
|
|
# This sets self.returncode.
|
|
self._handle_exitstatus(sts)
|
|
break
|
|
remaining = end - time.time()
|
|
if remaining <= 0:
|
|
raise TimeoutExpired(self.args, timeout)
|
|
delay = min(delay * 2, remaining, .05)
|
|
time.sleep(delay)
|
|
|
|
if not self.end:
|
|
# communicate() uses wait() internally.
|
|
self.end = time.time()
|
|
return self.returncode
|
|
|
|
def poll(self):
|
|
ret = super(Popen, self).poll()
|
|
if ret is not None and not self.end:
|
|
self.end = time.time()
|
|
return ret
|
|
|
|
def yield_any(self, maxsize=None, timeout=None):
|
|
"""Yields output until the process terminates.
|
|
|
|
Unlike wait(), does not raise TimeoutExpired.
|
|
|
|
Yields:
|
|
(pipename, data) where pipename is either 'stdout', 'stderr' or None in
|
|
case of timeout or when the child process closed one of the pipe(s) and
|
|
all pending data on the pipe was read.
|
|
|
|
Arguments:
|
|
- maxsize: See recv_any(). Can be a callable function.
|
|
- timeout: If None, the call is blocking. If set, yields None, None if no
|
|
data is available within |timeout| seconds. It resets itself after
|
|
each yield. Can be a callable function.
|
|
"""
|
|
assert self.stdout or self.stderr
|
|
if timeout is not None:
|
|
# timeout=0 effectively means that the pipe is continuously polled.
|
|
if isinstance(timeout, (int, float)):
|
|
assert timeout >= 0, timeout
|
|
old_timeout = timeout
|
|
timeout = lambda: old_timeout
|
|
else:
|
|
assert callable(timeout), timeout
|
|
|
|
if maxsize is not None and not callable(maxsize):
|
|
assert isinstance(maxsize, (int, float)), maxsize
|
|
|
|
last_yield = time.time()
|
|
while self.poll() is None:
|
|
to = (None if timeout is None
|
|
else max(timeout() - (time.time() - last_yield), 0))
|
|
t, data = self.recv_any(
|
|
maxsize=maxsize() if callable(maxsize) else maxsize, timeout=to)
|
|
if data or to is 0:
|
|
yield t, data
|
|
last_yield = time.time()
|
|
|
|
# Read all remaining output in the pipes.
|
|
# There is 3 cases:
|
|
# - pipes get closed automatically by the calling process before it exits
|
|
# - pipes are closed automated by the OS
|
|
# - pipes are kept open due to grand-children processes outliving the
|
|
# children process.
|
|
while True:
|
|
ms = maxsize
|
|
if callable(maxsize):
|
|
ms = maxsize()
|
|
# timeout=0 is mainly to handle the case where a grand-children process
|
|
# outlives the process started.
|
|
t, data = self.recv_any(maxsize=ms, timeout=0)
|
|
if not data:
|
|
break
|
|
yield t, data
|
|
|
|
def recv_any(self, maxsize=None, timeout=None):
|
|
"""Reads from the first pipe available from stdout and stderr.
|
|
|
|
Unlike wait(), it does not throw TimeoutExpired.
|
|
|
|
Arguments:
|
|
- maxsize: Maximum number of bytes to return. Defaults to MAX_SIZE.
|
|
- timeout: If None, it is blocking. If 0 or above, will return None if no
|
|
data is available within |timeout| seconds.
|
|
|
|
Returns:
|
|
tuple(pipename or None, str(data)). pipename is one of 'stdout' or
|
|
'stderr'.
|
|
"""
|
|
# recv_multi_impl will early exit on a closed connection. Loop accordingly
|
|
# to simplify call sites.
|
|
while True:
|
|
pipes = [
|
|
x for x in ((self.stderr, 'stderr'), (self.stdout, 'stdout')) if x[0]
|
|
]
|
|
# If both stdout and stderr have the exact file handle, they are
|
|
# effectively the same pipe. Deduplicate it since otherwise it confuses
|
|
# recv_multi_impl().
|
|
if len(pipes) == 2 and self.stderr.fileno() == self.stdout.fileno():
|
|
pipes.pop(0)
|
|
|
|
if not pipes:
|
|
return None, None
|
|
start = time.time()
|
|
conns, names = zip(*pipes)
|
|
index, data, closed = recv_multi_impl(conns, maxsize, timeout)
|
|
if index is None:
|
|
return index, data
|
|
if closed:
|
|
self._close(names[index])
|
|
if not data:
|
|
# Loop again. The other pipe may still be open.
|
|
if timeout:
|
|
timeout -= (time.time() - start)
|
|
continue
|
|
|
|
if self.universal_newlines and data:
|
|
data = self._translate_newlines(data)
|
|
return names[index], data
|
|
|
|
def recv_out(self, maxsize=None, timeout=None):
|
|
"""Reads from stdout synchronously with timeout."""
|
|
return self._recv('stdout', maxsize, timeout)
|
|
|
|
def recv_err(self, maxsize=None, timeout=None):
|
|
"""Reads from stderr synchronously with timeout."""
|
|
return self._recv('stderr', maxsize, timeout)
|
|
|
|
def terminate(self):
|
|
"""Tries to do something saner on Windows that the stdlib.
|
|
|
|
Windows:
|
|
self.detached/CREATE_NEW_PROCESS_GROUP determines what can be used:
|
|
- If set, only SIGBREAK can be sent and it is sent to a single process.
|
|
- If not set, in theory only SIGINT can be used and *all processes* in
|
|
the processgroup receive it. In practice, we just kill the process.
|
|
See http://msdn.microsoft.com/library/windows/desktop/ms683155.aspx
|
|
The default on Windows is to call TerminateProcess() always, which is not
|
|
useful.
|
|
|
|
On Posix, always send SIGTERM.
|
|
"""
|
|
try:
|
|
if subprocess.mswindows and self.detached:
|
|
return self.send_signal(signal.CTRL_BREAK_EVENT)
|
|
super(Popen, self).terminate()
|
|
except OSError:
|
|
# The function will throw if the process terminated in-between. Swallow
|
|
# this.
|
|
pass
|
|
|
|
def kill(self):
|
|
"""Kills the process and its children if possible.
|
|
|
|
Swallows exceptions and return True on success.
|
|
"""
|
|
if self.gid:
|
|
try:
|
|
os.killpg(self.gid, signal.SIGKILL)
|
|
except OSError:
|
|
return False
|
|
else:
|
|
try:
|
|
super(Popen, self).kill()
|
|
except OSError:
|
|
return False
|
|
return True
|
|
|
|
def _close(self, which):
|
|
"""Closes either stdout or stderr."""
|
|
getattr(self, which).close()
|
|
setattr(self, which, None)
|
|
|
|
def _recv(self, which, maxsize, timeout):
|
|
"""Reads from one of stdout or stderr synchronously with timeout."""
|
|
conn = getattr(self, which)
|
|
if conn is None:
|
|
return None
|
|
_, data, closed = recv_multi_impl([conn], maxsize, timeout)
|
|
if closed:
|
|
self._close(which)
|
|
if self.universal_newlines and data:
|
|
data = self._translate_newlines(data)
|
|
return data
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def set_signal_handler(signals, handler):
|
|
"""Temporarilly override signals handler.
|
|
|
|
Useful when waiting for a child process to handle signals like SIGTERM, so the
|
|
signal can be propagated to the child process.
|
|
"""
|
|
previous = {s: signal.signal(s, handler) for s in signals}
|
|
try:
|
|
yield
|
|
finally:
|
|
for sig, h in previous.iteritems():
|
|
signal.signal(sig, h)
|
|
|
|
|
|
def call(*args, **kwargs):
|
|
"""Adds support for timeout."""
|
|
timeout = kwargs.pop('timeout', None)
|
|
return Popen(*args, **kwargs).wait(timeout)
|
|
|
|
|
|
def check_call(*args, **kwargs):
|
|
"""Adds support for timeout."""
|
|
retcode = call(*args, **kwargs)
|
|
if retcode:
|
|
raise CalledProcessError(retcode, kwargs.get('args') or args[0])
|
|
return 0
|
|
|
|
|
|
def check_output(*args, **kwargs):
|
|
"""Adds support for timeout."""
|
|
timeout = kwargs.pop('timeout', None)
|
|
if 'stdout' in kwargs:
|
|
raise ValueError('stdout argument not allowed, it will be overridden.')
|
|
process = Popen(stdout=PIPE, *args, **kwargs)
|
|
output, _ = process.communicate(timeout=timeout)
|
|
retcode = process.poll()
|
|
if retcode:
|
|
raise CalledProcessError(retcode, kwargs.get('args') or args[0], output)
|
|
return output
|
|
|
|
|
|
def call_with_timeout(args, timeout, **kwargs):
|
|
"""Runs an executable; kill it in case of timeout."""
|
|
proc = Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, **kwargs)
|
|
try:
|
|
out, err = proc.communicate(timeout=timeout)
|
|
except TimeoutExpired as e:
|
|
out = e.output
|
|
err = e.stderr
|
|
proc.kill()
|
|
proc.wait()
|
|
return out, err, proc.returncode, proc.duration()
|