Don't buffer stdout/stderr in command_wrapper.py

command_wrapper/bin/command_wrapper.py

 #!/usr/bin/python
 #
-# Copyright (c) 2011 The Chromium Authors. All rights reserved.
+# Copyright (c) 2012 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.
 
 """Wrapper that does auto-retry and stats logging for command invocation.
 
 Various command line tools in use: gsutil, curl have spurious failure.
 This wrapper will track stats to an AppEngine based service to
 help track down the cause of failures, as well as add retry logic.
 """
 
 
 import optparse
 import os
 import platform
 import socket
 import subprocess
 import sys
 import threading
 import time
 import urllib
 import uuid
 
 
 LOG_TIMEOUT = 10
+ON_POSIX = 'posix' in sys.builtin_module_names
 
 
 def LogCommand(options, command_id,
                attempt, cmd, returncode, stdout, stderr, runtime):
   """Log a command invocation and result to a central location.
 
   Arguments:
     options: parsed options
     command_id: unique id for this command (shared by all retries)
     attempt: which try numbered from 0
@@ skipping 36 matching lines @@
 def RunWithTimeout(timeout, func, *args, **kwargs):
   wrapper = { 'result': None }
   def CallFunc():
     wrapper['result'] = func(*args, **kwargs)
   th = threading.Thread(target=CallFunc)
   th.start()
   th.join(timeout)
   return wrapper['result']
 
 
+def Tee(fd, string_buffer, forward_fd):
+  """Read characters from fd and both append them to a buffer and write them to
+  forward_fd."""
+  for char in iter(lambda: fd.read(1), ''):
+    string_buffer += char
+    forward_fd.write(char)
+  fd.close()
+
+
 def main(argv):
   parser = optparse.OptionParser()
   parser.add_option('-r', '--retries', dest='retries',
                     type='int', default=10,
                     help='number of times to retry on failure')
   parser.add_option('-u', '--logurl', dest='logurl',
                     default='https://command-wrapper.appspot.com/log',
                     help='URL to log invocations/failures to')
   (options, args) = parser.parse_args(args=argv[1:])
 
@@ skipping 12 matching lines @@
 
   # Log that we're even starting.
   RunWithTimeout(LOG_TIMEOUT, LogCommand,
                  options, command_id, -1, cmd, -1, '', '', 0)
 
   # Try up to a certain number of times.
   for r in range(options.retries):
     tm = time.time()
     p = subprocess.Popen(cmd, shell=True,
                          stdout=subprocess.PIPE,
-                         stderr=subprocess.PIPE)
-    (p_stdout, p_stderr) = p.communicate()
-    sys.stdout.write(p_stdout)
-    sys.stderr.write(p_stderr)
+                         stderr=subprocess.PIPE,
+                         close_fds=ON_POSIX)
+    p_stdout = ''
+    t_stdout = threading.Thread(target=Tee,
+                                args=(p.stdout, p_stdout, sys.stdout))
+    t_stdout.start()
+
+    p_stderr = ''
+    t_stderr = threading.Thread(target=Tee,
+                                args=(p.stderr, p_stderr, sys.stderr))
+    t_stderr.start()
+
+    p.wait()
+
+    t_stdout.join()
+    t_stderr.join()
+
+
     runtime = time.time() - tm
     accept = RunWithTimeout(LOG_TIMEOUT, LogCommand,
                             options, command_id, r, cmd,
                             p.returncode, p_stdout, p_stderr, runtime)
     if accept:
       return p.returncode
     if p.returncode == 0:
       return 0
     print 'Command %s failed with retcode %d, try %d.' % (
         ' '.join(args), p.returncode, r + 1)
   print 'Command %s failed %d retries, giving up.' % (
       ' '.join(args), options.retries)
 
   return p.returncode
 
 
 if __name__ == '__main__':
   sys.exit(main(sys.argv))