1 files changed, 27 insertions, 2 deletions
diff --git a/catapult/devil/devil/android/tools/system_app.py b/catapult/devil/devil/android/tools/system_app.py
index 62bf5a59..50d85595 100755
--- a/catapult/devil/devil/android/tools/system_app.py
+++ b/catapult/devil/devil/android/tools/system_app.py
@@ -22,6 +22,7 @@ from devil.android import decorators
 from devil.android import device_errors
 from devil.android import device_temp_file
 from devil.android.sdk import version_codes
+from devil.android.sdk import adb_wrapper
 from devil.android.tools import script_common
 from devil.utils import cmd_helper
 from devil.utils import parallelizer
@@ -128,7 +129,15 @@ _ENABLE_MODIFICATION_PROP = 'devil.modify_sys_apps'
 
 
 def _ShouldRetryModification(exc):
-  return not isinstance(exc, device_errors.CommandTimeoutError)
+  try:
+    if isinstance(exc, device_errors.CommandTimeoutError):
+      logger.info('Restarting the adb server')
+      adb_wrapper.RestartServer()
+    return True
+  except Exception: # pylint: disable=broad-except
+    logger.exception(('Caught an exception when deciding'
+                      ' to retry system modification'))
+    return False
 
 
 # timeout and retries are both required by the decorator, but neither
@@ -172,7 +181,7 @@ def _SetUpSystemAppModification(device, timeout=None, retries=None):
       # Point the user to documentation, since there's a good chance they can
       # workaround this on an emulator.
       docs_url = ('https://chromium.googlesource.com/chromium/src/+/'
-                  'master/docs/android_emulator.md#writable-system-partition')
+                  'HEAD/docs/android_emulator.md#writable-system-partition')
       logger.error(
           'Did you start the emulator with "-writable-system?"\n'
           'See %s\n', docs_url)
@@ -187,6 +196,22 @@ def _TearDownSystemAppModification(device,
                                    timeout=None,
                                    retries=None):
   try:
+    # The function may be re-entered after the the device loses root
+    # privilege. For instance if the adb server is restarted before
+    # re-entering the function then the device may lose root privilege.
+    # Therefore we need to do a sanity check for root privilege
+    # on the device and then re-enable root privilege if the device
+    # does not have it.
+    if not device.HasRoot():
+      logger.warning('Need to re-enable root.')
+      device.EnableRoot()
+
+      if not device.HasRoot():
+        raise device_errors.CommandFailedError(
+          ('Failed to tear down modification of '
+           'system apps on non-rooted device.'),
+          str(device))
+
     device.SetProp(_ENABLE_MODIFICATION_PROP, '0')
     device.Reboot()
     device.WaitUntilFullyBooted()