diagnose_apk_bloat.py: support pre supersize revs.

tools/binary_size/diagnose_apk_bloat.py

 #!/usr/bin/env python
 # Copyright 2017 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.
 
 """Tool for finding the cause of APK bloat.
 
 Run diagnose_apk_bloat.py -h for detailed usage help.
 """
 
 import argparse
 import collections
+from contextlib import contextmanager
 import distutils.spawn
-import itertools
 import json
 import multiprocessing
 import os
 import shutil
 import subprocess
 import sys
 import tempfile
 import zipfile
 
 _COMMIT_COUNT_WARN_THRESHOLD = 15
 _ALLOWED_CONSECUTIVE_FAILURES = 2
 _BUILDER_URL = \
     'https://build.chromium.org/p/chromium.perf/builders/Android%20Builder'
 _CLOUD_OUT_DIR = os.path.join('out', 'Release')
 _SRC_ROOT = os.path.abspath(
     os.path.join(os.path.dirname(__file__), os.pardir, os.pardir))
 _DEFAULT_ARCHIVE_DIR = os.path.join(_SRC_ROOT, 'binary-size-bloat')
 _DEFAULT_OUT_DIR = os.path.join(_SRC_ROOT, 'out', 'diagnose-apk-bloat')
 _DEFAULT_TARGET = 'monochrome_public_apk'
 
 
 _global_restore_checkout_func = None
 
 
-def _RestoreFunc(subrepo):
+def _SetRestoreFunc(subrepo):
   branch = _GitCmd(['rev-parse', '--abbrev-ref', 'HEAD'], subrepo)
-  return lambda: _GitCmd(['checkout', branch], subrepo)
+  global _global_restore_checkout_func
+  _global_restore_checkout_func = lambda: _GitCmd(['checkout', branch], subrepo)
 
 
 class BaseDiff(object):
   """Base class capturing binary size diffs."""
   def __init__(self, name):
     self.name = name
     self.banner = '\n' + '*' * 30 + name + '*' * 30
 
   def AppendResults(self, logfile):
     """Print and write diff results to an open |logfile|."""
@@ skipping 182 matching lines @@
 
 class _BuildArchive(object):
   """Class for managing a directory with build results and build metadata."""
   def __init__(self, rev, base_archive_dir, build, subrepo):
     self.build = build
     self.dir = os.path.join(base_archive_dir, rev)
     metadata_path = os.path.join(self.dir, 'metadata.txt')
     self.rev = rev
     self.metadata = _GenerateMetadata([self], build, metadata_path, subrepo)
 
-  def ArchiveBuildResults(self):
+  def ArchiveBuildResults(self, bs_dir):
     """Save build artifacts necessary for diffing."""
     _Print('Saving build results to: {}', self.dir)
     _EnsureDirsExist(self.dir)
     build = self.build
     self._ArchiveFile(build.main_lib_path)
     lib_name_noext = os.path.splitext(os.path.basename(build.main_lib_path))[0]
     size_path = os.path.join(self.dir, lib_name_noext + '.size')
-    supersize_path = os.path.join(_SRC_ROOT, 'tools/binary_size/supersize')
+    supersize_path = os.path.join(bs_dir, 'supersize')
     tool_prefix = _FindToolPrefix(build.output_directory)
     supersize_cmd = [supersize_path, 'archive', size_path, '--elf-file',
                      build.main_lib_path, '--tool-prefix', tool_prefix,
                      '--output-directory', build.output_directory,
                      '--no-source-paths']
     if build.IsAndroid():
       supersize_cmd += ['--apk-file', build.abs_apk_path]
       self._ArchiveFile(build.abs_apk_path)
 
     _RunCmd(supersize_cmd)
@@ skipping 208 matching lines @@
     _Print('Failure: please ensure working directory is clean.')
     sys.exit()
 
 
 def _Die(s, *args, **kwargs):
   _Print('Failure: ' + s, *args, **kwargs)
   _global_restore_checkout_func()
   sys.exit(1)
 
 
-def _DownloadBuildArtifacts(archive, build, depot_tools_path=None):
+def _DownloadBuildArtifacts(archive, build, bs_dir, depot_tools_path):
   """Download artifacts from arm32 chromium perf builder."""
   if depot_tools_path:
     gsutil_path = os.path.join(depot_tools_path, 'gsutil.py')
   else:
     gsutil_path = distutils.spawn.find_executable('gsutil.py')
 
   if not gsutil_path:
     _Die('gsutil.py not found, please provide path to depot_tools via '
          '--depot-tools-path or add it to your PATH')
 
   download_dir = tempfile.mkdtemp(dir=_SRC_ROOT)
   try:
-    _DownloadAndArchive(gsutil_path, archive, download_dir, build)
+    _DownloadAndArchive(gsutil_path, archive, download_dir, build, bs_dir)
   finally:
     shutil.rmtree(download_dir)
 
 
-def _DownloadAndArchive(gsutil_path, archive, dl_dir, build):
+def _DownloadAndArchive(gsutil_path, archive, dl_dir, build, bs_dir):
   dl_file = 'full-build-linux_%s.zip' % archive.rev
   dl_url = 'gs://chrome-perf/Android Builder/%s' % dl_file
   dl_dst = os.path.join(dl_dir, dl_file)
   _Print('Downloading build artifacts for {}', archive.rev)
   # gsutil writes stdout and stderr to stderr, so pipe stdout and stderr to
   # sys.stdout.
   retcode = subprocess.call([gsutil_path, 'cp', dl_url, dl_dir],
                              stdout=sys.stdout, stderr=subprocess.STDOUT)
   if retcode:
       _Die('unexpected error while downloading {}. It may no longer exist on '
            'the server or it may not have been uploaded yet (check {}). '
            'Otherwise, you may not have the correct access permissions.',
            dl_url, _BUILDER_URL)
 
   # Files needed for supersize and resource_sizes. Paths relative to out dir.
   to_extract = [build.main_lib_name, build.map_file_name, 'args.gn',
                 'build_vars.txt', build.apk_path]
   extract_dir = os.path.join(os.path.splitext(dl_dst)[0], 'unzipped')
   # Storage bucket stores entire output directory including out/Release prefix.
   _Print('Extracting build artifacts')
   with zipfile.ZipFile(dl_dst, 'r') as z:
     _ExtractFiles(to_extract, _CLOUD_OUT_DIR, extract_dir, z)
     dl_out = os.path.join(extract_dir, _CLOUD_OUT_DIR)
     build.output_directory, output_directory = dl_out, build.output_directory
-    archive.ArchiveBuildResults()
+    archive.ArchiveBuildResults(bs_dir)
     build.output_directory = output_directory
 
 
 def _ExtractFiles(to_extract, prefix, dst, z):
   zip_infos = z.infolist()
   assert all(info.filename.startswith(prefix) for info in zip_infos), (
       'Storage bucket folder structure doesn\'t start with %s' % prefix)
   to_extract = [os.path.join(prefix, f) for f in to_extract]
   for f in to_extract:
     z.extract(f, path=dst)
 
 
 def _Print(s, *args, **kwargs):
   print s.format(*args, **kwargs)
 
 
 def _PrintAndWriteToFile(logfile, s):
   """Print |s| to |logfile| and stdout."""
   _Print(s)
   logfile.write('%s\n' % s)
 
 
+@contextmanager
+def _TmpBinarySizeDir():
+  """Recursively copy files to a temp dir and yield the tmp binary_size dir."""
+  # Needs to be at same level of nesting as the real //tools/binary_size
+  # since supersize uses this to find d3 in //third_party.
+  tmp_dir = tempfile.mkdtemp(dir=_SRC_ROOT)
+  try:
+    bs_dir = os.path.join(tmp_dir, 'binary_size')
+    shutil.copytree(os.path.join(_SRC_ROOT, 'tools', 'binary_size'), bs_dir)
+    yield bs_dir
+  finally:
+    shutil.rmtree(tmp_dir)
+
+
 def main():
   parser = argparse.ArgumentParser(
       description='Find the cause of APK size bloat.')
-  parser.add_argument('--archive-dir',
+  parser.add_argument('--archive-directory',
                       default=_DEFAULT_ARCHIVE_DIR,
                       help='Where results are stored.')
   parser.add_argument('rev',
                       help='Find binary size bloat for this commit.')
   parser.add_argument('--reference-rev',
                       help='Older rev to diff against. If not supplied, '
                       'the previous commit to rev will be used.')
   parser.add_argument('--all',
                       action='store_true',
                       help='Build/download all revs from --reference-rev to '
@@ skipping 50 matching lines @@
   args = parser.parse_args()
   build = _BuildHelper(args)
   if build.IsCloud():
     if build.IsLinux():
       parser.error('--cloud only works for android')
     if args.subrepo:
       parser.error('--subrepo doesn\'t work with --cloud')
 
   subrepo = args.subrepo or _SRC_ROOT
   _EnsureDirectoryClean(subrepo)
-  _global_restore_checkout_func = _RestoreFunc(subrepo)
+  _SetRestoreFunc(subrepo)
   revs = _GenerateRevList(args.rev,
                           args.reference_rev or args.rev + '^',
                           args.all,
                           subrepo)
   diffs = []
   if build.IsAndroid():
     diffs +=  [
         ResourceSizesDiff(
             build.apk_name, slow_options=args.include_slow_options)
     ]
-  diff_mngr = _DiffArchiveManager(revs, args.archive_dir, diffs, build, subrepo)
+  diff_mngr = _DiffArchiveManager(
+      revs, args.archive_directory, diffs, build, subrepo)
   consecutive_failures = 0
-  for i, archive in enumerate(diff_mngr.IterArchives()):
-    if archive.Exists():
-      _Print('Found matching metadata for {}, skipping build step.',
-             archive.rev)
-    else:
-      if build.IsCloud():
-        _DownloadBuildArtifacts(archive, build,
-                                depot_tools_path=args.depot_tools_path)
+  with _TmpBinarySizeDir() as bs_dir:
+    for i, archive in enumerate(diff_mngr.IterArchives()):
+      if archive.Exists():
+        _Print('Found matching metadata for {}, skipping build step.',
+               archive.rev)
       else:
-        build_success = _SyncAndBuild(archive, build, subrepo)
-        if not build_success:
-          consecutive_failures += 1
-          if consecutive_failures > _ALLOWED_CONSECUTIVE_FAILURES:
-            _Die('{} builds failed in a row, last failure was {}.',
-                 consecutive_failures, archive.rev)
+        if build.IsCloud():
+          _DownloadBuildArtifacts(archive, build, bs_dir, args.depot_tools_path)
         else:
-          archive.ArchiveBuildResults()
-          consecutive_failures = 0
+          build_success = _SyncAndBuild(archive, build, subrepo)
+          if not build_success:
+            consecutive_failures += 1
+            if consecutive_failures > _ALLOWED_CONSECUTIVE_FAILURES:
+              _Die('{} builds failed in a row, last failure was {}.',
+                   consecutive_failures, archive.rev)
+          else:
+            archive.ArchiveBuildResults(bs_dir)
+            consecutive_failures = 0
 
-    if i != 0:
-      diff_mngr.MaybeDiff(i - 1, i)
+      if i != 0:
+        diff_mngr.MaybeDiff(i - 1, i)
 
   _global_restore_checkout_func()
 
 if __name__ == '__main__':
   sys.exit(main())