parallel: Halt tasks on foreground error with halt_on_error set.

Before this CL, halt_on_error would halt remaining background tasks
whenever one of the background tasks failed. This CL extends the
behaviour to halt on exceptions raised from the foreground process as
well. It exposes the new behaviour to clients through
BackgroundTaskRunner context function.

While there, fix timeout in parallel_unittest so that they work
correctly even if softer signals fail to kill background tasks.

BUG=chromium:396723
TEST=unittests, including a new test that fails without the
     halt_on_error flag.

Change-Id: Ib5105c6658d78e5f2a5107500de72c3716b04e02
Reviewed-on: https://chromium-review.googlesource.com/261969
Reviewed-by: David James <davidjames@chromium.org>
Commit-Queue: Prathmesh Prabhu <pprabhu@chromium.org>
Tested-by: Prathmesh Prabhu <pprabhu@chromium.org>

1 files changed, 9 insertions, 3 deletions

diff --git a/lib/parallel.py b/lib/parallel.py
index dd263c1f8..4db9061ad 100644
--- a/lib/parallel.py
+++ b/lib/parallel.py
@@ -537,9 +537,10 @@ class _BackgroundTask(multiprocessing.Process):
       except BaseException:
         foreground_except = sys.exc_info()
       finally:
-        # Wait for each step to complete.
         errors = []
-        while bg_tasks:
+        skip_bg_wait = halt_on_error and foreground_except is not None
+        # Wait for each step to complete.
+        while not skip_bg_wait and bg_tasks:
           task = bg_tasks.popleft()
           task_errors = task.Wait()
           if task_errors:
@@ -717,11 +718,16 @@ def BackgroundTaskRunner(task, *args, **kwargs):
     processes: Number of processes to launch.
     onexit: Function to run in each background process after all inputs are
       processed.
+    halt_on_error: After the first exception occurs, halt any running steps, and
+      squelch any further output, including any exceptions that might occur.
+      Halts on exceptions in any of the background processes, or in the
+      foreground process using the BackgroundTaskRunner.
   """
 
   queue = kwargs.pop('queue', None)
   processes = kwargs.pop('processes', None)
   onexit = kwargs.pop('onexit', None)
+  halt_on_error = kwargs.pop('halt_on_error', False)
 
   with cros_build_lib.ContextManagerStack() as stack:
     if queue is None:
@@ -735,7 +741,7 @@ def BackgroundTaskRunner(task, *args, **kwargs):
                               onexit=onexit, task_args=args,
                               task_kwargs=kwargs)
     steps = [child] * processes
-    with _BackgroundTask.ParallelTasks(steps):
+    with _BackgroundTask.ParallelTasks(steps, halt_on_error=halt_on_error):
       try:
         yield queue
       finally: