[Android] Blacklist devices on failures during environment set up + tear down.

build/android/pylib/local/device/local_device_environment.py

Index: build/android/pylib/local/device/local_device_environment.py
diff --git a/build/android/pylib/local/device/local_device_environment.py b/build/android/pylib/local/device/local_device_environment.py
index f8cd53048f097c66f444d10b995e153b5c1ddd66..431614377efa68bd4cffecc6158f13d0d6673965 100644
--- a/build/android/pylib/local/device/local_device_environment.py
+++ b/build/android/pylib/local/device/local_device_environment.py
@@ -3,12 +3,14 @@
 # found in the LICENSE file.
 
 import datetime
+import functools
 import logging
 import os
 import shutil
 import tempfile
 import threading
 
+from devil import base_error
 from devil.android import device_blacklist
 from devil.android import device_errors
 from devil.android import device_list
@@ -25,6 +27,50 @@ def _DeviceCachePath(device):
   return os.path.join(constants.GetOutDirectory(), file_name)
 
 
+def handle_shard_failures(f):
+  """A decorator that handles device failures for per-device functions.
+
+  Args:
+    f: the function being decorated. The function must take at least one
+      argument, and that argument must be the device.
+  """
+  return handle_shard_failures_with(None)(f)
+
+
+# TODO(jbudorick): Refactor this to work as a decorator or context manager.
+def handle_shard_failures_with(on_failure):
+  """A decorator that handles device failures for per-device functions.
+
+  This calls on_failure in the event of a failure.
+
+  Args:
+    f: the function being decorated. The function must take at least one
+      argument, and that argument must be the device.
+    on_failure: A binary function to call on failure.
+  """
+  def decorator(f):
+    @functools.wraps(f)
+    def wrapper(dev, *args, **kwargs):
+      try:
+        return f(dev, *args, **kwargs)
+      except device_errors.CommandTimeoutError:
+        logging.exception('Shard timed out: %s(%s)', f.__name__, str(dev))
+      except device_errors.DeviceUnreachableError:
+        logging.exception('Shard died: %s(%s)', f.__name__, str(dev))
+      except base_error.BaseError:
+        logging.exception('Shard failed: %s(%s)', f.__name__, str(dev))
+      except SystemExit:
+        logging.exception('Shard killed: %s(%s)', f.__name__, str(dev))
+        raise
+      if on_failure:
+        on_failure(dev, f.__name__)
+      return None
+
+    return wrapper
+
+  return decorator
+
+
 class LocalDeviceEnvironment(environment.Environment):
 
   def __init__(self, args, _error_func):
@@ -67,8 +113,12 @@ class LocalDeviceEnvironment(environment.Environment):
     if not self._devices:
       raise device_errors.NoDevicesError
 
-    if self._enable_device_cache:
-      for d in self._devices:
+    if self._logcat_output_file:
+      self._logcat_output_dir = tempfile.mkdtemp()
+
+    @handle_shard_failures_with(on_failure=self.BlacklistDevice)
+    def prepare_device(d):
+      if self._enable_device_cache:
         cache_path = _DeviceCachePath(d)
         if os.path.exists(cache_path):
           logging.info('Using device cache: %s', cache_path)
@@ -76,10 +126,8 @@ class LocalDeviceEnvironment(environment.Environment):
             d.LoadCacheData(f.read())
           # Delete cached file so that any exceptions cause it to be cleared.
           os.unlink(cache_path)
-    if self._logcat_output_file:
-      self._logcat_output_dir = tempfile.mkdtemp()
-    if self._logcat_output_dir:
-      for d in self._devices:
+
+      if self._logcat_output_dir:
         logcat_file = os.path.join(
             self._logcat_output_dir,
             '%s_%s' % (d.adb.GetDeviceSerial(),
@@ -89,6 +137,8 @@ class LocalDeviceEnvironment(environment.Environment):
         self._logcat_monitors.append(monitor)
         monitor.Start()
 
+    self.parallel_devices.pMap(prepare_device)
+
   @property
   def blacklist(self):
     return self._blacklist
@@ -121,17 +171,27 @@ class LocalDeviceEnvironment(environment.Environment):
 
   #override
   def TearDown(self):
-    # Write the cache even when not using it so that it will be ready the first
-    # time that it is enabled. Writing it every time is also necessary so that
-    # an invalid cache can be flushed just by disabling it for one run.
-    for d in self._devices:
+    @handle_shard_failures_with(on_failure=self.BlacklistDevice)
+    def tear_down_device(d):
+      # Write the cache even when not using it so that it will be ready the
+      # first time that it is enabled. Writing it every time is also necessary
+      # so that an invalid cache can be flushed just by disabling it for one
+      # run.
       cache_path = _DeviceCachePath(d)
       with open(cache_path, 'w') as f:
         f.write(d.DumpCacheData())
         logging.info('Wrote device cache: %s', cache_path)
+
+    self.parallel_devices.pMap(tear_down_device)
+
     for m in self._logcat_monitors:
-      m.Stop()
-      m.Close()
+      try:
+        m.Stop()
+        m.Close()
+      except base_error.BaseError:
+        logging.exception('Failed to stop logcat monitor for %s',
+                          m.adb.GetDeviceSerial())
+
     if self._logcat_output_file:
       file_utils.MergeFiles(
           self._logcat_output_file,