sandwich: Save browser cache as zip archive preserving all timestamps

tools/android/loading/run_sandwich.py

 #! /usr/bin/env python
 # Copyright 2016 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.
 
 """Instructs Chrome to load series of web pages and reports results.
 
 When running Chrome is sandwiched between preprocessed disk caches and
 WepPageReplay serving all connections.
 
 TODO(pasko): implement cache preparation and WPR.
 """
 
 import argparse
+import json
 import logging
 import os
+import shutil
 import sys
+import tempfile
 import time
+import zipfile
 
 _SRC_DIR = os.path.abspath(os.path.join(
     os.path.dirname(__file__), '..', '..', '..'))
 
 sys.path.append(os.path.join(_SRC_DIR, 'third_party', 'catapult', 'devil'))
 from devil.android import device_utils
 
 sys.path.append(os.path.join(_SRC_DIR, 'build', 'android'))
 from pylib import constants
 import devil_chromium
@@ skipping 64 matching lines @@
       json.dump({'traceEvents': events['events'], 'metadata': {}}, f, indent=2)
   except IOError:
     logging.warning('Could not save a trace: %s' % filename)
     # Swallow the exception.
 
 
 def _UpdateTimestampFromAdbStat(filename, stat):
   os.utime(filename, (stat.st_time, stat.st_time))
 
 
-def _SaveBrowserCache(device, output_directory):
+def _PullBrowserCache(device):
   """Pulls the browser cache from the device and saves it locally.
 
   Cache is saved with the same file structure as on the device. Timestamps are
   important to preserve because indexing and eviction depends on them.
 
-  Args:
-    output_directory: name of the directory for saving cache.
+  Returns:
+    Temporary directory containing all the browser cache.
   """
-  save_target = os.path.join(output_directory, _CACHE_DIRECTORY_NAME)
-  try:
-    os.makedirs(save_target)
-  except IOError:
-    logging.warning('Could not create directory: %s' % save_target)
-    raise
-
+  save_target = tempfile.mkdtemp(suffix='.cache')
   cache_directory = '/data/data/' + _CHROME_PACKAGE + '/cache/Cache'
   for filename, stat in device.adb.Ls(cache_directory):
     if filename == '..':
       continue
     if filename == '.':
       cache_directory_stat = stat
       continue
     original_file = os.path.join(cache_directory, filename)
     saved_file = os.path.join(save_target, filename)
     device.adb.Pull(original_file, saved_file)
     _UpdateTimestampFromAdbStat(saved_file, stat)
     if filename == _INDEX_DIRECTORY_NAME:
       # The directory containing the index was pulled recursively, update the
       # timestamps for known files. They are ignored by cache backend, but may
       # be useful for debugging.
       index_dir_stat = stat
       saved_index_dir = os.path.join(save_target, _INDEX_DIRECTORY_NAME)
       saved_index_file = os.path.join(saved_index_dir, _REAL_INDEX_FILE_NAME)
       for sub_file, sub_stat in device.adb.Ls(original_file):
         if sub_file == _REAL_INDEX_FILE_NAME:
           _UpdateTimestampFromAdbStat(saved_index_file, sub_stat)
           break
       _UpdateTimestampFromAdbStat(saved_index_dir, index_dir_stat)
 
   # Store the cache directory modification time. It is important to update it
   # after all files in it have been written. The timestamp is compared with
   # the contents of the index file when freshness is determined.
   _UpdateTimestampFromAdbStat(save_target, cache_directory_stat)
+  return save_target
+
+
+def _ZipDirectoryContent(root_directory_path, archive_dest_path):
+  """Zip a directory's content recursively with all the directories'
+  timestamps preserved.
+
+  Args:
+    root_directory_path: The directory's path to archive.
+    archive_dest_path: Archive destination's path.
+  """
+  with zipfile.ZipFile(archive_dest_path, 'w') as zip_output:
+    timestamps = {}
+    for directory_path, dirnames, filenames in os.walk(root_directory_path):
+      for dirname in dirnames:
+        subdirectory_path = os.path.join(directory_path, dirname)
+        subdirectory_relative_path = os.path.relpath(subdirectory_path,
+                                                     root_directory_path)
+        subdirectory_stats = os.stat(subdirectory_path)
+        timestamps[subdirectory_relative_path] = {
+            'atime': subdirectory_stats.st_atime,
+            'mtime': subdirectory_stats.st_mtime}
+      for filename in filenames:
+        file_path = os.path.join(directory_path, filename)
+        file_archive_name = os.path.join('content',
+            os.path.relpath(file_path, root_directory_path))
+        file_stats = os.stat(file_path)
+        timestamps[file_archive_name[8:]] = {
+            'atime': file_stats.st_atime,
+            'mtime': file_stats.st_mtime}
+        zip_output.write(file_path, arcname=file_archive_name)
+    zip_output.writestr('timestamps.json',
+                        json.dumps(timestamps, indent=2))
+
+
+def _UnzipDirectoryContent(archive_path, directory_dest_path):
+  """Unzip a directory's content recursively with all the directories'
+  timestamps preserved.
+
+  Args:
+    archive_path: Archive's path to unzip.
+    directory_dest_path: Directory destination path.
+  """
+  if not os.path.exists(directory_dest_path):
+    os.makedirs(directory_dest_path)
+
+  with zipfile.ZipFile(archive_path) as zip_input:
+    timestamps = None
+    for file_archive_name in zip_input.namelist():
+      if file_archive_name == 'timestamps.json':
+        timestamps = json.loads(zip_input.read(file_archive_name))
+      elif file_archive_name.startswith('content/'):
+        file_relative_path = file_archive_name[8:]
+        file_output_path = os.path.join(directory_dest_path, file_relative_path)
+        file_parent_directory_path = os.path.dirname(file_output_path)
+        if not os.path.exists(file_parent_directory_path):
+          os.makedirs(file_parent_directory_path)
+        with open(file_output_path, 'w') as f:
+          f.write(zip_input.read(file_archive_name))
+
+    assert timestamps
+    for relative_path, stats in timestamps.iteritems():
+      output_path = os.path.join(directory_dest_path, relative_path)
+      if not os.path.exists(output_path):
+        os.makedirs(output_path)
+      os.utime(output_path, (stats['atime'], stats['mtime']))
 
 
 def main():
   logging.basicConfig(level=logging.INFO)
   devil_chromium.Initialize()
 
   parser = argparse.ArgumentParser()
   parser.add_argument('--job', required=True,
                       help='JSON file with job description.')
   parser.add_argument('--output', required=True,
@@ skipping 39 matching lines @@
           pages_loaded += 1
           _SaveChromeTrace(tracing_track.ToJsonDict(), args.output,
               str(pages_loaded))
 
   if args.save_cache:
     # Move Chrome to background to allow it to flush the index.
     device.adb.Shell('am start com.google.android.launcher')
     time.sleep(_TIME_TO_DEVICE_IDLE_SECONDS)
     device.KillAll(_CHROME_PACKAGE, quiet=True)
     time.sleep(_TIME_TO_DEVICE_IDLE_SECONDS)
-    _SaveBrowserCache(device, args.output)
+
+    cache_directory_path = _PullBrowserCache(device)
+    _ZipDirectoryContent(cache_directory_path,
+                         os.path.join(args.output, 'cache.zip'))
+    shutil.rmtree(cache_directory_path)
 
 
 if __name__ == '__main__':
   sys.exit(main())