Merge "Store EROFS images uncompressed"

tools/releasetools/add_img_to_target_files.py

@@ -54,6 +54,7 @@ import shutil
 import stat
 import sys
 import uuid
+import tempfile
 import zipfile
 
 import build_image
@@ -104,9 +105,10 @@ class OutputFile(object):
     if self._output_zip:
       self._zip_name = os.path.join(*args)
 
-  def Write(self):
+  def Write(self, compress_type=None):
     if self._output_zip:
-      common.ZipWrite(self._output_zip, self.name, self._zip_name)
+      common.ZipWrite(self._output_zip, self.name,
+                      self._zip_name, compress_type=compress_type)
 
 
 def AddSystem(output_zip, recovery_img=None, boot_img=None):
@@ -139,7 +141,8 @@ def AddSystem(output_zip, recovery_img=None, boot_img=None):
     common.MakeRecoveryPatch(OPTIONS.input_tmp, output_sink, recovery_img,
                              boot_img, info_dict=OPTIONS.info_dict)
 
-  block_list = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", "system.map")
+  block_list = OutputFile(output_zip, OPTIONS.input_tmp,
+                          "IMAGES", "system.map")
   CreateImage(OPTIONS.input_tmp, OPTIONS.info_dict, "system", img,
               block_list=block_list)
   return img.name
@@ -187,7 +190,8 @@ def AddVendor(output_zip, recovery_img=None, boot_img=None):
     common.MakeRecoveryPatch(OPTIONS.input_tmp, output_sink, recovery_img,
                              boot_img, info_dict=OPTIONS.info_dict)
 
-  block_list = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", "vendor.map")
+  block_list = OutputFile(output_zip, OPTIONS.input_tmp,
+                          "IMAGES", "vendor.map")
   CreateImage(OPTIONS.input_tmp, OPTIONS.info_dict, "vendor", img,
               block_list=block_list)
   return img.name
@@ -390,7 +394,8 @@ def AddCustomImages(output_zip, partition_name):
 
   for img_name in OPTIONS.info_dict.get(
           "avb_{}_image_list".format(partition_name)).split():
-    custom_image = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", img_name)
+    custom_image = OutputFile(
+        output_zip, OPTIONS.input_tmp, "IMAGES", img_name)
     if os.path.exists(custom_image.name):
       continue
 
@@ -499,7 +504,9 @@ def AddUserdata(output_zip):
   build_image.BuildImage(user_dir, image_props, img.name)
 
   common.CheckSize(img.name, "userdata.img", OPTIONS.info_dict)
-  img.Write()
+  # Always use compression for useradata image.
+  # As it's likely huge and consist of lots of 0s.
+  img.Write(zipfile.ZIP_DEFLATED)
 
 
 def AddVBMeta(output_zip, partitions, name, needed_partitions):
@@ -820,7 +827,7 @@ def AddImagesToTargetFiles(filename):
     boot_images = OPTIONS.info_dict.get("boot_images")
     if boot_images is None:
       boot_images = "boot.img"
-    for index,b in enumerate(boot_images.split()):
+    for index, b in enumerate(boot_images.split()):
       # common.GetBootableImage() returns the image directly if present.
       boot_image = common.GetBootableImage(
           "IMAGES/" + b, b, OPTIONS.input_tmp, "BOOT")
@@ -841,7 +848,8 @@ def AddImagesToTargetFiles(filename):
     init_boot_image = common.GetBootableImage(
         "IMAGES/init_boot.img", "init_boot.img", OPTIONS.input_tmp, "INIT_BOOT")
     if init_boot_image:
-      partitions['init_boot'] = os.path.join(OPTIONS.input_tmp, "IMAGES", "init_boot.img")
+      partitions['init_boot'] = os.path.join(
+          OPTIONS.input_tmp, "IMAGES", "init_boot.img")
       if not os.path.exists(partitions['init_boot']):
         init_boot_image.WriteToDir(OPTIONS.input_tmp)
         if output_zip:
@@ -1005,6 +1013,35 @@ def AddImagesToTargetFiles(filename):
                           OPTIONS.replace_updated_files_list)
 
 
+def OptimizeCompressedEntries(zipfile_path):
+  """Convert files that do not compress well to uncompressed storage
+
+  EROFS images tend to be compressed already, so compressing them again
+  yields little space savings. Leaving them uncompressed will make
+  downstream tooling's job easier, and save compute time.
+  """
+  if not zipfile.is_zipfile(zipfile_path):
+    return
+  entries_to_store = []
+  with tempfile.TemporaryDirectory() as tmpdir:
+    with zipfile.ZipFile(zipfile_path, "r", allowZip64=True) as zfp:
+      for zinfo in zfp.filelist:
+        if not zinfo.filename.startswith("IMAGES/") and not zinfo.filename.startswith("META"):
+          pass
+        # Don't try to store userdata.img uncompressed, it's usually huge.
+        if zinfo.filename.endswith("userdata.img"):
+          pass
+        if zinfo.compress_size > zinfo.file_size * 0.80 and zinfo.compress_type != zipfile.ZIP_STORED:
+          entries_to_store.append(zinfo)
+          zfp.extract(zinfo, tmpdir)
+    # Remove these entries, then re-add them as ZIP_STORED
+    common.RunAndCheckOutput(
+        ["zip", "-d", zipfile_path] + [entry.filename for entry in entries_to_store])
+    with zipfile.ZipFile(zipfile_path, "a", allowZip64=True) as zfp:
+      for entry in entries_to_store:
+        zfp.write(os.path.join(tmpdir, entry.filename), entry.filename, compress_type=zipfile.ZIP_STORED)
+
+
 def main(argv):
   def option_handler(o, a):
     if o in ("-a", "--add_missing"):
@@ -1036,8 +1073,10 @@ def main(argv):
   common.InitLogging()
 
   AddImagesToTargetFiles(args[0])
+  OptimizeCompressedEntries(args[0])
   logger.info("done.")
 
+
 if __name__ == '__main__':
   try:
     common.CloseInheritedPipes()