[Android] Parallelize provision_devices.py.

build/android/provision_devices.py

Index: build/android/provision_devices.py
diff --git a/build/android/provision_devices.py b/build/android/provision_devices.py
index 35cfb814c6fc14fd499fc8cb9d8a8455a0864636..07d972c1e04d0c22cfc1ef47b01555b106c26d78 100755
--- a/build/android/provision_devices.py
+++ b/build/android/provision_devices.py
@@ -24,6 +24,7 @@ from pylib import device_settings
 from pylib.device import device_blacklist
 from pylib.device import device_errors
 from pylib.device import device_utils
+from pylib.utils import run_tests_helper
 
 sys.path.append(os.path.join(constants.DIR_SOURCE_ROOT,
                              'third_party', 'android_testrunner'))
@@ -34,8 +35,7 @@ def KillHostHeartbeat():
   stdout, _ = ps.communicate()
   matches = re.findall('\\n.*host_heartbeat.*', stdout)
   for match in matches:
-    print 'An instance of host heart beart running... will kill'
-    pid = re.findall('(\d+)', match)[1]
+    pid = re.findall('(\d+)', match)[0]
     subprocess.call(['kill', str(pid)])
 
 
@@ -48,37 +48,34 @@ def LaunchHostHeartbeat():
                                  'build/android/host_heartbeat.py')])
 
 
-def PushAndLaunchAdbReboot(devices, target):
+def PushAndLaunchAdbReboot(device, target):
   """Pushes and launches the adb_reboot binary on the device.
 
   Arguments:
-    devices: The list of serial numbers of the device to which the
-             adb_reboot binary should be pushed.
-    target : The build target (example, Debug or Release) which helps in
-             locating the adb_reboot binary.
+    device: The DeviceUtils instance for the device to which the adb_reboot
+            binary should be pushed.
+    target: The build target (example, Debug or Release) which helps in
+            locating the adb_reboot binary.
   """
-  for device_serial in devices:
-    print 'Will push and launch adb_reboot on %s' % device_serial
-    device = device_utils.DeviceUtils(device_serial)
-    # Kill if adb_reboot is already running.
-    try:
-      # Don't try to kill adb_reboot more than once. We don't expect it to be
-      # running at all.
-      device.KillAll('adb_reboot', blocking=True, timeout=2, retries=0)
-    except device_errors.CommandFailedError:
-      # We can safely ignore the exception because we don't expect adb_reboot
-      # to be running.
-      pass
-    # Push adb_reboot
-    print '  Pushing adb_reboot ...'
-    adb_reboot = os.path.join(constants.DIR_SOURCE_ROOT,
-                              'out/%s/adb_reboot' % target)
-    device.PushChangedFiles(adb_reboot, '/data/local/tmp/')
-    # Launch adb_reboot
-    print '  Launching adb_reboot ...'
-    device.old_interface.GetAndroidToolStatusAndOutput(
-        '/data/local/tmp/adb_reboot')
-  LaunchHostHeartbeat()
+  print 'Will push and launch adb_reboot on %s' % str(device)
+  # Kill if adb_reboot is already running.
+  try:
+    # Don't try to kill adb_reboot more than once. We don't expect it to be
+    # running at all.
+    device.KillAll('adb_reboot', blocking=True, timeout=2, retries=0)
+  except device_errors.CommandFailedError:
+    # We can safely ignore the exception because we don't expect adb_reboot
+    # to be running.
+    pass
+  # Push adb_reboot
+  print '  Pushing adb_reboot ...'
+  adb_reboot = os.path.join(constants.DIR_SOURCE_ROOT,
+                            'out/%s/adb_reboot' % target)
+  device.PushChangedFiles(adb_reboot, '/data/local/tmp/')
+  # Launch adb_reboot
+  print '  Launching adb_reboot ...'
+  device.old_interface.GetAndroidToolStatusAndOutput(
+      '/data/local/tmp/adb_reboot')
 
 
 def _ConfigureLocalProperties(device, is_perf):
@@ -137,125 +134,94 @@ def WipeDeviceData(device):
                            as_root=True)
 
 
-def WipeDevicesIfPossible(devices):
-  devices_to_reboot = []
-  for device_serial in devices:
-    device = device_utils.DeviceUtils(device_serial)
-    if not device.old_interface.EnableAdbRoot():
-      continue
+def WipeDeviceIfPossible(device):
+  try:
+    device.EnableRoot()
     WipeDeviceData(device)
-    devices_to_reboot.append(device)
-
-  if devices_to_reboot:
-    try:
-      device_utils.DeviceUtils.parallel(devices_to_reboot).Reboot(True)
-    except errors.DeviceUnresponsiveError:
-      pass
-    for device_serial in devices_to_reboot:
-      device.WaitUntilFullyBooted(timeout=90)
-
-
-def ProvisionDevice(device_serial, is_perf, disable_location):
-  device = device_utils.DeviceUtils(device_serial)
-  device.old_interface.EnableAdbRoot()
-  _ConfigureLocalProperties(device, is_perf)
-  device_settings_map = device_settings.DETERMINISTIC_DEVICE_SETTINGS
-  if disable_location:
-    device_settings_map.update(device_settings.DISABLE_LOCATION_SETTING)
-  else:
-    device_settings_map.update(device_settings.ENABLE_LOCATION_SETTING)
-  device_settings.ConfigureContentSettingsDict(device, device_settings_map)
-  device_settings.SetLockScreenSettings(device)
-  if is_perf:
-    # TODO(tonyg): We eventually want network on. However, currently radios
-    # can cause perfbots to drain faster than they charge.
-    device_settings.ConfigureContentSettingsDict(
-      device, device_settings.NETWORK_DISABLED_SETTINGS)
-    # Some perf bots run benchmarks with USB charging disabled which leads
-    # to gradual draining of the battery. We must wait for a full charge
-    # before starting a run in order to keep the devices online.
-    try:
-      battery_info = device.old_interface.GetBatteryInfo()
-    except Exception as e:
-      battery_info = {}
-      logging.error('Unable to obtain battery info for %s, %s',
-                    device_serial, e)
-
-    while int(battery_info.get('level', 100)) < 95:
-      if not device.old_interface.IsDeviceCharging():
-        if device.old_interface.CanControlUsbCharging():
-          device.old_interface.EnableUsbCharging()
-        else:
-          logging.error('Device is not charging')
-          break
-      logging.info('Waiting for device to charge. Current level=%s',
+    device.Reboot(True, timeout=180, retries=0)
+  except (errors.DeviceUnresponsiveError, device_errors.CommandFailedError):
+    pass
+
+
+def ProvisionDevice(device, options, is_perf):
+  try:
+    if not options.skip_wipe:
+      WipeDeviceIfPossible(device)
+    device.EnableRoot()
+    _ConfigureLocalProperties(device, is_perf)
+    device_settings.ConfigureContentSettings(
+        device, device_settings.DETERMINISTIC_DEVICE_SETTINGS)
+    if options.disable_location:
+      device_settings.ConfigureContentSettings(
+          device, device_settings.DISABLE_LOCATION_SETTINGS)
+    else:
+      device_settings.ConfigureContentSettings(
+          device, device_settings.ENABLE_LOCATION_SETTINGS)
+    device_settings.SetLockScreenSettings(device)
+    if is_perf:
+      # TODO(tonyg): We eventually want network on. However, currently radios
+      # can cause perfbots to drain faster than they charge.
+      device_settings.ConfigureContentSettings(
+          device, device_settings.NETWORK_DISABLED_SETTINGS)
+      # Some perf bots run benchmarks with USB charging disabled which leads
+      # to gradual draining of the battery. We must wait for a full charge
+      # before starting a run in order to keep the devices online.
+      try:
+        battery_info = device.old_interface.GetBatteryInfo()
+      except Exception as e:
+        battery_info = {}
+        logging.error('Unable to obtain battery info for %s, %s',
+                      str(device), e)
+
+      while int(battery_info.get('level', 100)) < 95:
+        if not device.old_interface.IsDeviceCharging():
+          if device.old_interface.CanControlUsbCharging():
+            device.old_interface.EnableUsbCharging()
+          else:
+            logging.error('Device is not charging')
+            break
+        logging.info('Waiting for device to charge. Current level=%s',
                      battery_info.get('level', 0))
-      time.sleep(60)
-      battery_info = device.old_interface.GetBatteryInfo()
-  device.RunShellCommand('date -u %f' % time.time(), as_root=True)
+        time.sleep(60)
+        battery_info = device.old_interface.GetBatteryInfo()
+    device.RunShellCommand('date -u %f' % time.time(), as_root=True)
+    device.Reboot(True, timeout=180, retries=0)
+    props = device.RunShellCommand('getprop')
+    for prop in props:
+      logging.info('  %s' % prop)
+    if options.auto_reconnect:
+      PushAndLaunchAdbReboot(device, options.target)
+  except (errors.WaitForResponseTimedOutError,
+          device_errors.CommandTimeoutError):
+    logging.info('Timed out waiting for device %s. Adding to blacklist.',
+                 str(device))
+    # Device black list is reset by bb_device_status_check.py per build.
+    device_blacklist.ExtendBlacklist([str(device)])
+  except (device_errors.CommandFailedError):
+    logging.info('Failed to provision device %s. Adding to blacklist.',
+                 str(device))
+    device_blacklist.ExtendBlacklist([str(device)])
 
 
 def ProvisionDevices(options):
   is_perf = 'perf' in os.environ.get('BUILDBOT_BUILDERNAME', '').lower()
-  # TODO(jbudorick): Parallelize provisioning of all attached devices after
-  # switching from AndroidCommands.
   if options.device is not None:
     devices = [options.device]
   else:
     devices = android_commands.GetAttachedDevices()
 
-  # Wipe devices (unless --skip-wipe was specified)
-  if not options.skip_wipe:
-    WipeDevicesIfPossible(devices)
-
-  bad_devices = []
-  # Provision devices
-  for device_serial in devices:
-    try:
-      ProvisionDevice(device_serial, is_perf, options.disable_location)
-    except errors.WaitForResponseTimedOutError:
-      logging.info('Timed out waiting for device %s. Adding to blacklist.',
-                   device_serial)
-      bad_devices.append(device_serial)
-      # Device black list is reset by bb_device_status_check.py per build.
-      device_blacklist.ExtendBlacklist([device_serial])
-  devices = [device for device in devices if device not in bad_devices]
-
-  # If there are no good devices
-  if not devices:
-    raise device_errors.NoDevicesError
-
-  try:
-    device_utils.DeviceUtils.parallel(devices).Reboot(True)
-  except errors.DeviceUnresponsiveError:
-    pass
-
-  bad_devices = []
-  for device_serial in devices:
-    device = device_utils.DeviceUtils(device_serial)
-    try:
-      device.WaitUntilFullyBooted(timeout=90)
-      (_, prop) = device.old_interface.GetShellCommandStatusAndOutput('getprop')
-      for p in prop:
-        print p
-    except errors.WaitForResponseTimedOutError:
-      logging.info('Timed out waiting for device %s. Adding to blacklist.',
-                   device_serial)
-      bad_devices.append(device_serial)
-      # Device black list is reset by bb_device_status_check.py per build.
-      device_blacklist.ExtendBlacklist([device_serial])
-  devices = [device for device in devices if device not in bad_devices]
-
-  # If there are no good devices
-  if not devices:
-    raise device_errors.NoDevicesError
-
-  if options.auto_reconnect:
-    PushAndLaunchAdbReboot(devices, options.target)
+  parallel_devices = device_utils.DeviceUtils.parallel(devices)
+  parallel_devices.pMap(ProvisionDevice, options, is_perf)
+  LaunchHostHeartbeat()
+  blacklist = device_blacklist.ReadBlacklist()
+  return 1 if all(d in blacklist for d in devices) else 0
 
 
 def main(argv):
-  logging.basicConfig(level=logging.INFO)
+  custom_handler = logging.StreamHandler(sys.stdout)
+  custom_handler.setFormatter(run_tests_helper.CustomFormatter())
+  logging.getLogger().addHandler(custom_handler)
+  logging.getLogger().setLevel(logging.INFO)
 
   parser = optparse.OptionParser()
   parser.add_option('--skip-wipe', action='store_true', default=False,
@@ -275,7 +241,7 @@ def main(argv):
     print >> sys.stderr, 'Unused args %s' % args
     return 1
 
-  ProvisionDevices(options)
+  return ProvisionDevices(options)
 
 
 if __name__ == '__main__':