diff options
Diffstat (limited to 'site_scons/site_tools/command_output.py')
| -rw-r--r-- | site_scons/site_tools/command_output.py | 130 |
1 files changed, 117 insertions, 13 deletions
diff --git a/site_scons/site_tools/command_output.py b/site_scons/site_tools/command_output.py index ce702b7..795f3eb 100644 --- a/site_scons/site_tools/command_output.py +++ b/site_scons/site_tools/command_output.py @@ -32,11 +32,69 @@ import os -import SCons.Script +import signal import subprocess +import sys +import threading +import time +import SCons.Script + + +# TODO(rspangler): Move KillProcessTree() and RunCommand() into their own +# module. + + +def KillProcessTree(pid): + """Kills the process and all of its child processes. + + Args: + pid: process to kill. + + Raises: + OSError: Unsupported OS. + """ + if sys.platform in ('win32', 'cygwin'): + # Use Windows' taskkill utility + killproc_path = '%s;%s\\system32;%s\\system32\\wbem' % ( + (os.environ['SYSTEMROOT'],) * 3) + killproc_cmd = 'taskkill /F /T /PID %d' % pid + killproc_task = subprocess.Popen(killproc_cmd, shell=True, + stdout=subprocess.PIPE, + env={'PATH':killproc_path}) + killproc_task.communicate() + + elif sys.platform in ('linux', 'linux2', 'darwin'): + # Use ps to get a list of processes + ps_task = subprocess.Popen(['/bin/ps', 'x', '-o', 'pid,ppid'], stdout=subprocess.PIPE) + ps_out = ps_task.communicate()[0] + + # Parse out a dict of pid->ppid + ppid = {} + for ps_line in ps_out.split('\n'): + w = ps_line.strip().split() + if len(w) < 2: + continue # Not enough words in this line to be a process list + try: + ppid[int(w[0])] = int(w[1]) + except ValueError: + pass # Header or footer + + # For each process, kill it if it or any of its parents is our child + for p in ppid: + p2 = p + while p2: + if p2 == pid: + os.kill(p, signal.SIGKILL) + break + p2 = ppid.get(p2) -def RunCommand(cmdargs, cwdir=None, env=None, echo_output=True): + else: + raise OSError('Unsupported OS for KillProcessTree()') + + +def RunCommand(cmdargs, cwdir=None, env=None, echo_output=True, timeout=None, + timeout_errorlevel=14): """Runs an external command. Args: @@ -45,28 +103,62 @@ def RunCommand(cmdargs, cwdir=None, env=None, echo_output=True): cwdir: Working directory for the command, if not None. env: Environment variables dict, if not None. echo_output: If True, output will be echoed to stdout. + timeout: If not None, timeout for command in seconds. If command times + out, it will be killed and timeout_errorlevel will be returned. + timeout_errorlevel: The value to return if the command times out. Returns: The integer errorlevel from the command. The combined stdout and stderr as a string. """ - # Force unicode string in the environment to strings. if env: env = dict([(k, str(v)) for k, v in env.items()]) + start_time = time.time() child = subprocess.Popen(cmdargs, cwd=cwdir, env=env, shell=True, + universal_newlines=True, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) child_out = [] child_retcode = None - # Need to poll the child process, since the stdout pipe can fill. + def _ReadThread(): + """Thread worker function to read output from child process. + + Necessary since there is no cross-platform way of doing non-blocking + reads of the output pipe. + """ + read_run = True + while read_run: + # Need to have a delay of 1 cycle between child completing and + # thread exit, to pick up the final output from the child. + if child_retcode is not None: + read_run = False + new_out = child.stdout.read() + if new_out: + if echo_output: + print new_out, + child_out.append(new_out) + + read_thread = threading.Thread(target=_ReadThread) + read_thread.start() + + # Wait for child to exit or timeout while child_retcode is None: + time.sleep(1) # So we don't poll too frequently child_retcode = child.poll() - new_out = child.stdout.read() - if echo_output: - print new_out, - child_out.append(new_out) + if timeout and child_retcode is None: + elapsed = time.time() - start_time + if elapsed > timeout: + print '*** RunCommand() timeout:', cmdargs + KillProcessTree(child.pid) + child_retcode = timeout_errorlevel + + # Wait for worker thread to pick up final output and die + read_thread.join(5) + if read_thread.isAlive(): + print '*** Error: RunCommand() read thread did not exit.' + sys.exit(1) if echo_output: print # end last line of output @@ -100,8 +192,13 @@ def CommandOutputBuilder(target, source, env): env.AppendENVPath('LD_LIBRARY_PATH', cwdir) else: cwdir = None + cmdecho = env.get('COMMAND_OUTPUT_ECHO', True) + timeout = env.get('COMMAND_OUTPUT_TIMEOUT') + timeout_errorlevel = env.get('COMMAND_OUTPUT_TIMEOUT_ERRORLEVEL') - retcode, output = RunCommand(cmdline, cwdir=cwdir, env=env['ENV']) + retcode, output = RunCommand(cmdline, cwdir=cwdir, env=env['ENV'], + echo_output=cmdecho, timeout=timeout, + timeout_errorlevel=timeout_errorlevel) # Save command line output output_file = open(str(target[0]), 'w') @@ -116,10 +213,17 @@ def generate(env): """SCons entry point for this tool.""" # Add the builder and tell it which build environment variables we use. - action = SCons.Script.Action(CommandOutputBuilder, varlist=[ - 'COMMAND_OUTPUT_CMDLINE', - 'COMMAND_OUTPUT_RUN_DIR', - ]) + action = SCons.Script.Action( + CommandOutputBuilder, + 'Output "$COMMAND_OUTPUT_CMDLINE" to $TARGET', + varlist=[ + 'COMMAND_OUTPUT_CMDLINE', + 'COMMAND_OUTPUT_RUN_DIR', + 'COMMAND_OUTPUT_TIMEOUT', + 'COMMAND_OUTPUT_TIMEOUT_ERRORLEVEL', + # We use COMMAND_OUTPUT_ECHO also, but that doesn't change the + # command being run or its output. + ], ) builder = SCons.Script.Builder(action = action) env.Append(BUILDERS={'CommandOutput': builder}) |