Directly hook sys.stdout for thread annotated output.

gclient_utils.py

 # Copyright 2009 Google Inc.  All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #      http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
@@ skipping 279 matching lines @@
         filter_fn(line)
     kwargs['filter_fn'] = filter_msg
     kwargs['call_filter_on_first_line'] = True
   # Obviously.
   kwargs['print_stdout'] = True
   return CheckCallAndFilter(args, **kwargs)
 
 
 def SoftClone(obj):
   """Clones an object. copy.copy() doesn't work on 'file' objects."""
-  class NewObject(object): pass
-  new_obj = NewObject()
+  if obj.__class__.__name__ == 'SoftCloned':
+    return obj
+  class SoftCloned(object): pass
+  new_obj = SoftCloned()
   for member in dir(obj):
     if member.startswith('_'):
       continue
     setattr(new_obj, member, getattr(obj, member))
   return new_obj
 
 
 def MakeFileAutoFlush(fileobj, delay=10):
   """Creates a file object clone to automatically flush after N seconds."""
   if hasattr(fileobj, 'last_flushed_at'):
     # Already patched. Just update delay.
     fileobj.delay = delay
     return fileobj
 
   new_fileobj = SoftClone(fileobj)
-  new_fileobj.lock = threading.Lock()
+  if not hasattr(new_fileobj, 'lock'):
+    new_fileobj.lock = threading.Lock()
   new_fileobj.last_flushed_at = time.time()
   new_fileobj.delay = delay
-  new_fileobj.old_auto_flush_write = fileobj.write
+  new_fileobj.old_auto_flush_write = new_fileobj.write
   # Silence pylint.
   new_fileobj.flush = fileobj.flush
 
   def auto_flush_write(out):
     new_fileobj.old_auto_flush_write(out)
     should_flush = False
     new_fileobj.lock.acquire()
     try:
       if (new_fileobj.delay and
           (time.time() - new_fileobj.last_flushed_at) > new_fileobj.delay):
         should_flush = True
         new_fileobj.last_flushed_at = time.time()
     finally:
       new_fileobj.lock.release()
     if should_flush:
       new_fileobj.flush()
 
   new_fileobj.write = auto_flush_write
   return new_fileobj
 
 
-class StdoutAnnotated(object):
-  """Prepends every line with a string."""
-  def __init__(self, prepend, stdout):
-    self.prepend = prepend
-    self.buf = ''
-    self.stdout = stdout
+def MakeFileAnnotated(fileobj):
+  """Creates a file object clone to automatically prepends every line in worker
+  threads with a NN> prefix."""
+  if hasattr(fileobj, 'output_buffers'):
+    # Already patched.
+    return fileobj
 
-  def write(self, out):
-    self.buf += out
-    while '\n' in self.buf:
-      line, self.buf = self.buf.split('\n', 1)
-      self.stdout.write(self.prepend + line + '\n')
+  new_fileobj = SoftClone(fileobj)
+  if not hasattr(new_fileobj, 'lock'):
+    new_fileobj.lock = threading.Lock()
+  new_fileobj.output_buffers = {}
+  new_fileobj.old_annotated_write = new_fileobj.write
 
-  def flush(self):
-    pass
+  def annotated_write(out):
+    index = getattr(threading.currentThread(), 'index', None)
+    if index is None:
+      # Undexed threads aren't buffered.
+      new_fileobj.old_annotated_write(out)
+      return
 
-  def full_flush(self):
-    if self.buf:
-      self.stdout.write(self.prepend + self.buf)
-      self.stdout.flush()
-      self.buf = ''
+    new_fileobj.lock.acquire()
+    try:
+      # Use a dummy array to hold the string so the code can be lockless.
+      # Strings are immutable, requiring to keep a lock for the whole dictionary
+      # otherwise. Using an array is faster than using a dummy object.
+      if not index in new_fileobj.output_buffers:
+        obj = new_fileobj.output_buffers[index] = ['']
+      else:
+        obj = new_fileobj.output_buffers[index]
+    finally:
+      new_fileobj.lock.release()
+
+    # Continue lockless.
+    obj[0] += out
+    while '\n' in obj[0]:
+      line, remaining = obj[0].split('\n', 1)
+      new_fileobj.old_annotated_write('%d>%s\n' % (index, line))
+      obj[0] = remaining
+
+  def full_flush():
+    """Flush buffered output."""
+    orphans = []
+    new_fileobj.lock.acquire()
+    try:
+      # Detect threads no longer existing.
+      indexes = (getattr(t, 'index', None) for t in threading.enumerate())
+      indexed = filter(None, indexes)
+      for index in new_fileobj.output_buffers:
+        if not index in indexes:
+          orphans.append((index, new_fileobj.output_buffers[index][0]))
+      for orphan in orphans:
+        del new_fileobj.output_buffers[orphan[0]]
+    finally:
+      new_fileobj.lock.release()
+
+    # Don't keep the lock while writting. Will append \n when it shouldn't.
+    for orphan in orphans:
+      new_fileobj.old_annotated_write('%d>%s\n' % (orphan[0], orphan[1]))
+
+  new_fileobj.write = annotated_write
+  new_fileobj.full_flush = full_flush
+  return new_fileobj
 
 
 def CheckCallAndFilter(args, stdout=None, filter_fn=None,
                        print_stdout=None, call_filter_on_first_line=False,
                        **kwargs):
   """Runs a command and calls back a filter function if needed.
 
   Accepts all subprocess.Popen() parameters plus:
     print_stdout: If True, the command's stdout is forwarded to stdout.
     filter_fn: A function taking a single string argument called with each line
@@ skipping 248 matching lines @@
         if self.progress:
           self.progress.update(1)
         assert not t.item.name in self.ran
         if not t.item.name in self.ran:
           self.ran.append(t.item.name)
 
   def _run_one_task(self, task_item, args, kwargs):
     if self.jobs > 1:
       # Start the thread.
       index = len(self.ran) + len(self.running) + 1
-      # Copy 'options' and add annotated stdout.
+      # Copy 'options'.
       task_kwargs = kwargs.copy()
       task_kwargs['options'] = copy.copy(task_kwargs['options'])
-      task_kwargs['options'].stdout = StdoutAnnotated(
-          '%d>' % index, task_kwargs['options'].stdout)
-      new_thread = self._Worker(task_item, args, task_kwargs)
+      new_thread = self._Worker(task_item, index, args, task_kwargs)
       self.running.append(new_thread)
       new_thread.start()
     else:
       # Run the 'thread' inside the main thread. Don't try to catch any
       # exception.
       task_item.run(*args, **kwargs)
       self.ran.append(task_item.name)
       if self.progress:
         self.progress.update(1)
 
   class _Worker(threading.Thread):
     """One thread to execute one WorkItem."""
-    def __init__(self, item, args, kwargs):
+    def __init__(self, item, index, args, kwargs):
       threading.Thread.__init__(self, name=item.name or 'Worker')
       logging.info(item.name)
       self.item = item
+      self.index = index
       self.args = args
       self.kwargs = kwargs
 
     def run(self):
       """Runs in its own thread."""
       logging.debug('running(%s)' % self.item.name)
       work_queue = self.kwargs['work_queue']
       try:
         self.item.run(*self.args, **self.kwargs)
       except Exception:
         # Catch exception location.
         logging.info('Caught exception in thread %s' % self.item.name)
         logging.info(str(sys.exc_info()))
         work_queue.exceptions.put(sys.exc_info())
       logging.info('Task %s done' % self.item.name)
 
       work_queue.ready_cond.acquire()
       try:
         work_queue.ready_cond.notifyAll()
       finally:
         work_queue.ready_cond.release()