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

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

 # Copyright 2014 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.
 
 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
 from devil.android import device_utils
 from devil.android import logcat_monitor
 from devil.utils import file_utils
 from devil.utils import parallelizer
 from pylib import constants
 from pylib.base import environment
 
 
 def _DeviceCachePath(device):
   file_name = 'device_cache_%s.json' % device.adb.GetDeviceSerial()
   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):
     super(LocalDeviceEnvironment, self).__init__()
     self._blacklist = (device_blacklist.Blacklist(args.blacklist_file)
                        if args.blacklist_file
                        else None)
     self._device_serial = args.test_device
     self._devices_lock = threading.Lock()
     self._devices = []
@@ skipping 22 matching lines @@
             'Read device list %s from target devices file.', str(device_arg))
     elif self._device_serial:
       device_arg = self._device_serial
 
     self._devices = device_utils.DeviceUtils.HealthyDevices(
         self._blacklist, enable_device_files_cache=self._enable_device_cache,
         default_retries=self._max_tries - 1, device_arg=device_arg)
     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)
           with open(cache_path) as f:
             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(),
                        datetime.datetime.utcnow().strftime('%Y%m%dT%H%M%S')))
         monitor = logcat_monitor.LogcatMonitor(
             d.adb, clear=True, output_file=logcat_file)
         self._logcat_monitors.append(monitor)
         monitor.Start()
 
+    self.parallel_devices.pMap(prepare_device)
+
   @property
   def blacklist(self):
     return self._blacklist
 
   @property
   def concurrent_adb(self):
     return self._concurrent_adb
 
   @property
   def devices(self):
@@ skipping 12 matching lines @@
   @property
   def skip_clear_data(self):
     return self._skip_clear_data
 
   @property
   def tool(self):
     return self._tool_name
 
   #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,
           [m.output_file for m in self._logcat_monitors])
       shutil.rmtree(self._logcat_output_dir)
 
   def BlacklistDevice(self, device, reason='local_device_failure'):
     device_serial = device.adb.GetDeviceSerial()
     if self._blacklist:
       self._blacklist.Extend([device_serial], reason=reason)
     with self._devices_lock:
       self._devices = [d for d in self._devices if str(d) != device_serial]