Implement accelerated tee support for POSIX.

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 select
 import subprocess
 import sys
 import time
 import threading
 
+if sys.platform != 'win32':
+  import fcntl
+
+
 # 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
 
 # Globals.
 # Set to True if you somehow need to disable this hack.
@@ skipping 165 matching lines @@
           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.
+  """Queues 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))
   dest.put(name)
@@ skipping 112 matching lines @@
   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 _read_pipe(handles, pipe, out_fn):
+  """Reads bytes from a pipe and calls the output callback."""
+  data = pipe.read()
+  if not data:
+    del handles[pipe]
+  else:
+    out_fn(data)
+
+
+def _tee_posix(proc, timeout, start, stdin, args, kwargs):
+  """Polls a process and its pipe using select.select().
+
+  TODO(maruel): Implement a non-polling method for OSes that support it.
+  """
+  handles_r = {}
+  if callable(kwargs.get('stdout')):
+    handles_r[proc.stdout] = lambda: _read_pipe(
+        handles_r, proc.stdout, kwargs['stdout'])
+  if callable(kwargs.get('stderr')):
+    handles_r[proc.stderr] = lambda: _read_pipe(
+        handles_r, proc.stderr, kwargs['stderr'])
+
+  handles_w = {}
+  if isinstance(stdin, str):
+    stdin_io = cStringIO.StringIO(stdin)
+    def write_stdin():
+      data = stdin_io.read(1)
+      if data:
+        proc.stdin.write(data)
+      else:
+        del handles_w[proc.stdin]
+        proc.stdin.close()
+    handles_w[proc.stdin] = write_stdin
+  else:
+    # TODO(maruel): Fix me, it could be VOID.
+    assert stdin is None
+
+  # Make all the file objects of the child process non-blocking file.
+  # TODO(maruel): Test if a pipe is handed to the child process.
+  for pipe in (proc.stdin, proc.stdout, proc.stderr):
+    fileno = pipe and getattr(pipe, 'fileno', lambda: None)()
+    if fileno:
+      # Note: making a pipe non-blocking means the C stdio could act wrong. In
+      # particular, readline() cannot be used. Work around is to use os.read().
+      fl = fcntl.fcntl(fileno, fcntl.F_GETFL)
+      fcntl.fcntl(fileno, fcntl.F_SETFL, fl | os.O_NONBLOCK)
+
+  timed_out = False
+  while handles_r or handles_w or (timeout and proc.poll() is None):
+    period = None
+    if timeout:
+      period = max(0, timeout - (time.time() - start))
+      if not period and not timed_out:
+        proc.kill()
+        timed_out = True
+    if timed_out:
+      period = 0.001
+
+    # It reconstructs objects on each loop, not very efficient.
+    reads, writes, _, = select.select(
+        handles_r.keys(), handles_w.keys(), [], period)
+    for read in reads:
+      handles_r[read]()
+    for write in writes:
+      handles_w[write]()
+
+  # No pipe open anymore and if there was a time out, the child process was
+  # killed already.
+  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'):
@@ skipping 30 matching lines @@
   # http://docs.python.org/library/subprocess.html#subprocess.Popen.wait
   # 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 sys.platform == 'win32':
+    # On cygwin, ctypes._FUNCFLAG_STDCALL, which is used by ctypes.WINFUNCTYPE,
+    # doesn't exist so _tee_win() cannot be used yet.
+    returncode = _tee_threads(proc, timeout, start, stdin, args, kwargs)
+  else:
+    returncode = _tee_posix(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().
 
@@ skipping 45 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]