Add ThreadPool.abort() to stop processing early.

utils/threading_utils.py

 # Copyright 2013 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.
 
 """Classes and functions related to threading."""
 
 import functools
 import inspect
 import logging
 import os
@@ skipping 247 matching lines @@
       self._is_closed = True
     for _ in range(len(self._workers)):
       # Enqueueing None causes the worker to stop.
       self.tasks.put(None)
     for t in self._workers:
       t.join()
     logging.debug(
       'Thread pool \'%s\' closed: spawned %d threads total',
       self._prefix, len(self._workers))
 
+  def abort(self):
+    """Empties the queue.
+
+    To be used when the pool should stop early, like when Ctrl-C was detected.
+
+    Returns:
+      Number of tasks cancelled.
+    """
+    index = 0
+    while True:
+      try:
+        self.tasks.get_nowait()
+        self.tasks.task_done()
+        index += 1
+      except Queue.Empty:
+        return index
+
   def __enter__(self):
     """Enables 'with' statement."""
     return self
 
   def __exit__(self, _exc_type, _exc_value, _traceback):
     """Enables 'with' statement."""
     self.close()
 
 
 class AutoRetryThreadPool(ThreadPool):
@@ skipping 84 matching lines @@
     self.start = time.time()
     self.size = size
     # Setting it to True forces a print on the first print_update() call.
     self.value_changed = True
     self.use_cr_only = True
     self.unfinished_commands = set()
 
     # To be used in all threads.
     self.queued_lines = Queue.Queue()
 
-  def update_item(self, name, index=False, size=False, raw=False):
+  def update_item(self, name, index=0, size=0, raw=False):

[Marc-Antoine Ruel (Google), 2013/10/02 14:08:31]

Eventually I want it to support an arbitrary number of column. At least I have a
use case for 3 numbers. That's out of scope for this CL.

     """Queue information to print out.
 
     Arguments:
-      index: index should be incremented.
-      size: total size should be incremented.
+      index: increment to add to index. usually 0 or 1.
+      size: increment to add to size. usually 0 or 1.
       raw: if True, prints the data without the header.
     """
+    assert isinstance(name, str)
+    assert isinstance(index, int)
+    assert isinstance(size, int)
+    assert isinstance(raw, bool)
     self.queued_lines.put((name, index, size, raw))
 
   def print_update(self):
     """Prints the current status."""
     # Flush all the logging output so it doesn't appear within this output.
     for handler in logging.root.handlers:
       handler.flush()
 
     got_one = False
     while True:
       try:
         name, index, size, raw = self.queued_lines.get_nowait()
       except Queue.Empty:
         break
 
-      if size:
-        self.size += 1
-        self.value_changed = True
-      if index:
-        self.index += 1
-        self.value_changed = True
+      self.size += size
+      self.index += index
+      self.value_changed = bool(size or index)
       if not name:
         # Even if raw=True, there's nothing to print.
         continue
 
       got_one = True
       if raw:
         # Prints the data as-is.
         self.last_printed_line = ''
         sys.stdout.write('\n%s\n' % name.strip('\n'))
       else:
@@ skipping 69 matching lines @@
     # between the number of output and the number of input task.
     with self.all_tasks_done:
       self.all_tasks_done.notify_all()
 
   def join(self):
     """Calls print_update() whenever possible."""
     self.progress.print_update()
     with self.all_tasks_done:
       while self.unfinished_tasks:
         self.progress.print_update()
-        self.all_tasks_done.wait(60.)
+        # Use a short wait timeout so updates are printed in a timely manner.
+        self.all_tasks_done.wait(0.1)

[Vadim Sh., 2013/10/02 18:19:48]

I'm unhappy about polling. I think print_update can be rewritten to use blocking
get from self.queued_lines. Then updates will be printed instantly and
performance will also be better.

For instance in currently implementation all log handlers will be flushed each
100ms, it's not great.

But it's probably for some future CL.

[Marc-Antoine Ruel (Google), 2013/10/02 19:41:26]

I agree it's far from awesome but I prefer the Progress to stay "interactive".
So I added a TODO.

       self.progress.print_update()
 
 
 class ThreadPoolWithProgress(ThreadPool):
   QUEUE_CLASS = QueueWithProgress
 
   def __init__(self, progress, *args, **kwargs):
     super(ThreadPoolWithProgress, self).__init__(*args, **kwargs)
     self.tasks.set_progress(progress)
 
@@ skipping 183 matching lines @@
     # Multiprocessing
     import multiprocessing
     return multiprocessing.cpu_count()
   except:  # pylint: disable=W0702
     try:
       # Mac OS 10.6
       return int(os.sysconf('SC_NPROCESSORS_ONLN'))  # pylint: disable=E1101
     except:
       # Some of the windows builders seem to get here.
       return 4