Cache device apk checksums in device_utils.py

build/android/pylib/device/device_utils.py

Index: build/android/pylib/device/device_utils.py
diff --git a/build/android/pylib/device/device_utils.py b/build/android/pylib/device/device_utils.py
index f201ef36b8bf04fa20881109d11c6225f8322bfb..04f3022bf9d2a8623ce89e912e96871cc3ec06c8 100644
--- a/build/android/pylib/device/device_utils.py
+++ b/build/android/pylib/device/device_utils.py
@@ -176,9 +176,26 @@ class DeviceUtils(object):
     self._default_retries = default_retries
     self._cache = {}
     self._client_caches = {}
+    self._package_to_device_apk_checksums_cache = {}
+    self._package_to_device_apk_paths_cache = {}
     assert hasattr(self, decorators.DEFAULT_TIMEOUT_ATTR)
     assert hasattr(self, decorators.DEFAULT_RETRIES_ATTR)
 
+    # We need to flush the cache on uninstall, so wrap self.adb.Uninstall
+    # in order to ensure we detect calls to it.
+    old_adb_uninstall = self.adb.Uninstall
+    def WrappedAdbUninstall(package_name, *args, **kwargs):
+      try:
+        old_adb_uninstall(package_name, *args, **kwargs)
+      except:
+        # Clear cache since we can't be sure of the state.
+        self._package_to_device_apk_paths_cache.pop(package_name, 0)
+        self._package_to_device_apk_checksums_cache.pop(package_name, 0)
+        raise
+      self._package_to_device_apk_paths_cache[package_name] = []

[jbudorick, 2015/08/06 19:28:22]

I would still prefer to use self._cache and do some additional cache management
for this key.

[agrieve, 2015/08/11 14:07:46]

Done.

+      self._package_to_device_apk_checksums_cache[package_name] = set()
+    self.adb.Uninstall = WrappedAdbUninstall
+
   def __eq__(self, other):
     """Checks whether |other| refers to the same device as |self|.
 
@@ -352,6 +369,12 @@ class DeviceUtils(object):
     Returns:
       List of paths to the apks on the device for the given package.
     """
+    return self._GetApplicationPathsInternal(package)
+
+  def _GetApplicationPathsInternal(self, package, skip_cache=False):
+    cached_result = self._package_to_device_apk_paths_cache.get(package)
+    if cached_result is not None and not skip_cache:
+      return list(cached_result)
     # 'pm path' is liable to incorrectly exit with a nonzero number starting
     # in Lollipop.
     # TODO(jbudorick): Check if this is fixed as new Android versions are
@@ -366,6 +389,7 @@ class DeviceUtils(object):
         raise device_errors.CommandFailedError(
             'pm path returned: %r' % '\n'.join(output), str(self))
       apks.append(line[len('package:'):])
+    self._package_to_device_apk_paths_cache[package] = list(apks)
     return apks
 
   @decorators.WithTimeoutAndRetriesFromInstance()
@@ -418,7 +442,7 @@ class DeviceUtils(object):
 
     def pm_ready():
       try:
-        return self.GetApplicationPaths('android')
+        return self._GetApplicationPathsInternal('android', skip_cache=True)
       except device_errors.CommandFailedError:
         return False
 
@@ -488,21 +512,30 @@ class DeviceUtils(object):
       DeviceUnreachableError on missing device.
     """
     package_name = apk_helper.GetPackageName(apk_path)
-    device_paths = self.GetApplicationPaths(package_name)
+    device_paths = self._GetApplicationPathsInternal(package_name)
     if device_paths:
       if len(device_paths) > 1:
         logging.warning(
             'Installing single APK (%s) when split APKs (%s) are currently '
             'installed.', apk_path, ' '.join(device_paths))
-      (files_to_push, _) = self._GetChangedAndStaleFiles(
-          apk_path, device_paths[0])
-      should_install = bool(files_to_push)
+      apks_to_install, host_checksums = (
+          self._ComputeStaleApks(package_name, [apk_path]))
+      should_install = bool(apks_to_install)
       if should_install and not reinstall:
         self.adb.Uninstall(package_name)
     else:
       should_install = True
+      host_checksums = None
     if should_install:
-      self.adb.Install(apk_path, reinstall=reinstall)
+      # We won't know the resulting device apk names.
+      self._package_to_device_apk_paths_cache.pop(package_name, 0)
+      try:
+        self.adb.Install(apk_path, reinstall=reinstall)
+        self._package_to_device_apk_checksums_cache[package_name] = (
+            host_checksums)
+      except:
+        self._package_to_device_apk_paths_cache.pop(package_name, 0)
+        raise
 
   @decorators.WithTimeoutAndRetriesDefaults(
       INSTALL_DEFAULT_TIMEOUT,
@@ -531,14 +564,12 @@ class DeviceUtils(object):
     all_apks = [base_apk] + split_select.SelectSplits(
         self, base_apk, split_apks)
     package_name = apk_helper.GetPackageName(base_apk)
-    device_apk_paths = self.GetApplicationPaths(package_name)
+    device_apk_paths = self._GetApplicationPathsInternal(package_name)
 
     if device_apk_paths:
       partial_install_package = package_name
-      device_checksums = md5sum.CalculateDeviceMd5Sums(device_apk_paths, self)
-      host_checksums = md5sum.CalculateHostMd5Sums(all_apks)
-      apks_to_install = [k for (k, v) in host_checksums.iteritems()
-                         if v not in device_checksums.values()]
+      apks_to_install, host_checksums = (
+          self._ComputeStaleApks(package_name, all_apks))
       if apks_to_install and not reinstall:
         self.adb.Uninstall(package_name)
         partial_install_package = None
@@ -546,9 +577,19 @@ class DeviceUtils(object):
     else:
       partial_install_package = None
       apks_to_install = all_apks
+      host_checksums = None
     if apks_to_install:
-      self.adb.InstallMultiple(
-          apks_to_install, partial=partial_install_package, reinstall=reinstall)
+      # We won't know the resulting device apk names.
+      self._package_to_device_apk_paths_cache.pop(package_name, 0)
+      try:
+        self.adb.InstallMultiple(
+            apks_to_install, partial=partial_install_package,
+            reinstall=reinstall)
+        self._package_to_device_apk_checksums_cache[package_name] = (
+            host_checksums)
+      except:
+        self._package_to_device_apk_paths_cache.pop(package_name, 0)

[jbudorick, 2015/08/06 19:28:22]

Do you need to do this? If we hit an exception in InstallMultiple, we won't have
put the package_name entry back in yet.

[agrieve, 2015/08/11 14:07:46]

Nice catch. Gone.

+        raise
 
   def _CheckSdkLevel(self, required_sdk_level):
     """Raises an exception if the device does not have the required SDK level.
@@ -893,7 +934,7 @@ class DeviceUtils(object):
     # may never return.
     if ((self.build_version_sdk >=
          constants.ANDROID_SDK_VERSION_CODES.JELLY_BEAN_MR2)
-        or self.GetApplicationPaths(package)):
+        or self._GetApplicationPathsInternal(package)):
       self.RunShellCommand(['pm', 'clear', package], check_return=True)
 
   @decorators.WithTimeoutAndRetriesFromInstance()
@@ -1017,6 +1058,22 @@ class DeviceUtils(object):
       to_delete = device_checksums.keys()
       return (to_push, to_delete)
 
+  def _ComputeDeviceChecksumsForApks(self, package_name):
+    ret = self._package_to_device_apk_checksums_cache.get(package_name)
+    if ret is None:
+      device_paths = self._GetApplicationPathsInternal(package_name)
+      file_to_checksums = md5sum.CalculateDeviceMd5Sums(device_paths, self)
+      ret = set(file_to_checksums.values())
+      self._package_to_device_apk_checksums_cache[package_name] = ret
+    return ret
+
+  def _ComputeStaleApks(self, package_name, host_apk_paths):

[jbudorick, 2015/08/06 19:28:22]

I'm curious about the component time savings of caching the application path &
of caching the checksum. I don't particularly like replicating the
GetChangedFiles logic, and if this isn't saving very much time I'd rather not
add it.

Alternatively, maybe we should be doing some amount of caching like this more
broadly.

[agrieve, 2015/08/11 14:07:46]

When I first started down this path, I tried to add caching to
_GetChangedAndStaleFiles but hit a wall. 

The fundamental difference is that there's no way to (quickly) know the
on-device file name of an apk given the host apk. That's why in InstallSplitApk,
it checks whether _any_ of the on-device .apk checksums match a host checksum to
know whether or not the file is stale.

This approach was already used by InstallSplitApk, but now is a bit more
formalized in these functions (with caching)

+    host_checksums = md5sum.CalculateHostMd5Sums(host_apk_paths)
+    device_checksums = self._ComputeDeviceChecksumsForApks(package_name)
+    stale_apks = [k for (k, v) in host_checksums.iteritems()
+                  if v not in device_checksums]
+    return stale_apks, set(host_checksums.values())
+
   def _PushFilesImpl(self, host_device_tuples, files):
     size = sum(host_utils.GetRecursiveDiskUsage(h) for h, _ in files)
     file_count = len(files)
@@ -1712,6 +1769,8 @@ class DeviceUtils(object):
     for client in self._client_caches:
       self._client_caches[client].clear()
     self._cache.clear()
+    self._package_to_device_apk_paths_cache.clear()
+    self._package_to_device_apk_checksums_cache.clear()
 
   @classmethod
   def parallel(cls, devices=None, async=False):