Add callback support for stdout and stderr.

subprocess2.py

 # coding=utf8
 # Copyright (c) 2011 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 """Collection of subprocess wrapper functions.
 
 In theory you shouldn't need anything else in subprocess, or this module failed.
 """
 
 from __future__ import with_statement
+import cStringIO
 import errno
 import logging
 import os
+import Queue
 import subprocess
 import sys
-import tempfile
 import time
 import threading
 
 # Constants forwarded from subprocess.
 PIPE = subprocess.PIPE
 STDOUT = subprocess.STDOUT
 # Sends stdout or stderr to os.devnull.
 VOID = object()
 # Error code when a process was killed because it timed out.
 TIMED_OUT = -2001
@@ skipping 142 matching lines @@
   logging.debug(tmp_str)
 
   def fix(stream):
     if kwargs.get(stream) in (VOID, os.devnull):
       # Replaces VOID with handle to /dev/null.
       # Create a temporary file to workaround python's deadlock.
       # http://docs.python.org/library/subprocess.html#subprocess.Popen.wait
       # When the pipe fills up, it will deadlock this process. Using a real file
       # works around that issue.
       kwargs[stream] = open(os.devnull, 'w')
+    if callable(kwargs.get(stream)):
+      # Callable stdout/stderr should be used only with call() wrappers.
+      kwargs[stream] = PIPE
 
   fix('stdout')
   fix('stderr')
 
   try:
     return subprocess.Popen(args, **kwargs)
   except OSError, e:
     if e.errno == errno.EAGAIN and sys.platform == 'cygwin':
       # Convert fork() emulation failure into a CygwinRebaseError().
       raise CygwinRebaseError(
           e.errno,
           args,
           kwargs.get('cwd'),
           None,
           'Visit '
           'http://code.google.com/p/chromium/wiki/CygwinDllRemappingFailure to '
           'learn how to fix this error; you need to rebase your cygwin dlls')
     # Popen() can throw OSError when cwd or args[0] doesn't exist. Let it go
     # through
     raise
 
 
+def _queue_pipe_read(pipe, name, done, dest):
+  """Queue characters read from a pipe into a queue.
+
+  Left outside the _tee_threads function to not introduce a function closure
+  to speed up variable lookup.
+  """
+  while not done.isSet():
+    data = pipe.read(1)
+    if not data:
+      break
+    dest.put((name, data))

[Dirk Pranke, 2011/11/04 01:15:11]

This timeout doesn't really work if the pipes are blocking, right? I don't think
the thread will ever wake up if nothing is written.

[M-A Ruel, 2011/11/04 01:41:48]

If the process is killed, the pipes are closed, which unblocks the read() calls.

+  dest.put(name)
+
+
+def _tee_threads(proc, timeout, start, stdin, args, kwargs):
+  """Does I/O for a process's pipes using thread.
+
+  It's the simplest and slowest implementation. Expect very slow behavior.
+
+  If there is a callback and it doesn't keep up with the calls, the timeout
+  effectiveness will be delayed accordingly.
+  """
+  # TODO(maruel): Implement a select based implementation on POSIX and a Windows
+  # one using WaitForMultipleObjects().
+  #
+  # Queue of either of <threadname> when done or (<threadname>, data).
+  # In theory we would like to limit to ~64kb items to not cause large memory
+  # usage when the callback blocks. It is not done because it slows down
+  # processing on OSX10.6 by a factor of 2x, making it even slower than Windows!
+  # Revisit this decision if it becomes a problem, e.g. crash because of
+  # memory exhaustion.
+  queue = Queue.Queue()
+  done = threading.Event()
+
+  def write_stdin():
+    stdin_io = cStringIO.StringIO(stdin)
+    while not done.isSet():
+      data = stdin_io.read(1024)
+      if data:
+        proc.stdin.write(data)
+      else:
+        proc.stdin.close()
+        break
+    queue.put('stdin')
+
+  def timeout_fn():
+    done.wait(timeout)
+    # No need to close the pipes since killing should be sufficient.
+    queue.put('timeout')
+
+  # Starts up to 4 threads:
+  # Read stdout
+  # Read stderr
+  # Write stdin
+  # Timeout
+  threads = {}
+  if timeout is not None:
+    threads['timeout'] = threading.Thread(target=timeout_fn)
+  if callable(kwargs.get('stdout')):
+    threads['stdout'] = threading.Thread(
+      target=_queue_pipe_read, args=(proc.stdout, 'stdout', done, queue))
+  if callable(kwargs.get('stderr')):
+    threads['stderr'] = threading.Thread(
+      target=_queue_pipe_read,
+      args=(proc.stderr, 'stderr', done, queue))
+  if isinstance(stdin, str):
+    threads['stdin'] = threading.Thread(target=write_stdin)
+  for t in threads.itervalues():
+    t.daemon = True
+    t.start()
+
+  timed_out = False
+  try:
+    while proc.returncode is None:
+      assert threads
+      proc.poll()
+      item = queue.get()
+      if isinstance(item, str):
+        threads[item].join()
+        del threads[item]
+        if item == 'timeout' and not timed_out and proc.poll() is None:
+          logging.debug('Timed out: killing')
+          proc.kill()
+          timed_out = True
+        if not threads:
+          # We won't be waken up anymore. Need to busy loop.
+          break
+      else:
+        kwargs[item[0]](item[1])
+  finally:
+    # Stop the threads.
+    done.set()
+    # Join threads
+    for thread in threads.itervalues():
+      thread.join()
+
+  # Flush the queue.
+  try:
+    while True:
+      item = queue.get(False)
+      if isinstance(item, str):
+        if item == 'timeout':
+          # TODO(maruel): Does it make sense at that point?
+          if not timed_out and proc.poll() is None:
+            logging.debug('Timed out: killing')
+            proc.kill()
+            timed_out = True
+      else:
+        kwargs[item[0]](item[1])
+  except Queue.Empty:
+    pass
+
+  # Get the remainder.
+  if callable(kwargs.get('stdout')):
+    data = proc.stdout.read()
+    while data:
+      kwargs['stdout'](data)
+      data = proc.stdout.read()
+  if callable(kwargs.get('stderr')):
+    data = proc.stderr.read()
+    while data:
+      kwargs['stderr'](data)
+      data = proc.stderr.read()
+
+  if proc.returncode is None:
+    # Usually happens when killed with timeout but not listening to pipes.
+    proc.wait()
+
+  if timed_out:
+    return TIMED_OUT
+
+  return proc.returncode
+
+
 def communicate(args, timeout=None, **kwargs):
   """Wraps subprocess.Popen().communicate().
 
   Returns ((stdout, stderr), returncode).
 
   - The process will be killed after |timeout| seconds and returncode set to
     TIMED_OUT.
   - Automatically passes stdin content as input so do not specify stdin=PIPE.
   """
+  if timeout and kwargs.get('shell'):
+    raise TypeError(
+        'Using timeout and shell simultaneously will cause a process leak '
+        'since the shell will be killed instead of the child process.')
+
   stdin = kwargs.pop('stdin', None)
   if stdin is not None:
     if stdin is VOID:
       kwargs['stdin'] = open(os.devnull, 'r')
       stdin = None
     else:
       assert isinstance(stdin, basestring)
       # When stdin is passed as an argument, use it as the actual input data and
       # set the Popen() parameter accordingly.
       kwargs['stdin'] = PIPE
 
-  if not timeout:
+  start = time.time()
+  proc = Popen(args, **kwargs)
+  need_buffering = (timeout or
+      callable(kwargs.get('stdout')) or callable(kwargs.get('stderr')))
+
+  if not need_buffering:
     # Normal workflow.
-    proc = Popen(args, **kwargs)
-    if stdin is not None:
+    if stdin not in (None, VOID):
       return proc.communicate(stdin), proc.returncode
     else:
       return proc.communicate(), proc.returncode
 
-  # Create a temporary file to workaround python's deadlock.
+  stdout = None
+  stderr = None
+  # Convert to a lambda to workaround python's deadlock.
   # http://docs.python.org/library/subprocess.html#subprocess.Popen.wait
-  # When the pipe fills up, it will deadlock this process. Using a real file
-  # works around that issue.
-  with tempfile.TemporaryFile() as buff:
-    start = time.time()
-    kwargs['stdout'] = buff
-    proc = Popen(args, **kwargs)
-    if stdin is not None:
-      proc.stdin.write(stdin)
-    while proc.returncode is None:
-      proc.poll()
-      if timeout and (time.time() - start) > timeout:
-        proc.kill()
-        proc.wait()
-        # It's -9 on linux and 1 on Windows. Standardize to TIMED_OUT.
-        proc.returncode = TIMED_OUT
-      time.sleep(0.001)
-    # Now that the process died, reset the cursor and read the file.
-    buff.seek(0)
-    out = [buff.read(), None]
-  return out, proc.returncode
+  # When the pipe fills up, it will deadlock this process. Using a thread
+  # works around that issue. No need for thread safe function since the call
+  # backs are guaranteed to be called from the main thread.
+  if kwargs.get('stdout') == PIPE:
+    stdout = []
+    kwargs['stdout'] = stdout.append
+  if kwargs.get('stderr') == PIPE:
+    stderr = []
+    kwargs['stderr'] = stderr.append
+  returncode = _tee_threads(proc, timeout, start, stdin, args, kwargs)
+  if not stdout is None:
+    stdout = ''.join(stdout)
+  if not stderr is None:
+    stderr = ''.join(stderr)
+  return (stdout, stderr), returncode
 
 
 def call(args, **kwargs):
   """Emulates subprocess.call().
 
   Automatically convert stdout=PIPE or stderr=PIPE to VOID.
   In no case they can be returned since no code path raises
   subprocess2.CalledProcessError.
   """
   if kwargs.get('stdout') == PIPE:
@@ skipping 40 matching lines @@
 
   Captures stdout of a process call and returns stdout only.
 
   - Throws if return code is not 0.
   - Works even prior to python 2.7.
   - Blocks stdin by default if not specified since no output will be visible.
   - As per doc, "The stdout argument is not allowed as it is used internally."
   """
   kwargs.setdefault('stdin', VOID)
   return check_call_out(args, stdout=PIPE, **kwargs)[0]