[packager] Optimize archive accesses.

7 files changed, 107 insertions, 30 deletions

diff --git a/python/mozbuild/mozpack/copier.py b/python/mozbuild/mozpack/copier.py
index 386930fe7..15c550e7b 100644
--- a/python/mozbuild/mozpack/copier.py
+++ b/python/mozbuild/mozpack/copier.py
@@ -535,7 +535,7 @@ class Jarrer(FileRegistry, BaseFile):
             dest = Dest(dest)
         assert isinstance(dest, Dest)
 
-        from mozpack.mozjar import JarWriter, JarReader
+        from mozpack.mozjar import JarWriter, JarReader, JAR_BROTLI
         try:
             old_jar = JarReader(fileobj=dest)
         except Exception:
diff --git a/python/mozbuild/mozpack/mozjar.py b/python/mozbuild/mozpack/mozjar.py
index a1ada8594..2010a7f13 100644
--- a/python/mozbuild/mozpack/mozjar.py
+++ b/python/mozbuild/mozpack/mozjar.py
@@ -6,6 +6,7 @@ from __future__ import absolute_import
 
 from io import BytesIO
 import struct
+import subprocess
 import zlib
 import os
 from zipfile import (
@@ -15,9 +16,11 @@ from zipfile import (
 from collections import OrderedDict
 from urlparse import urlparse, ParseResult
 import mozpack.path as mozpath
+from mozbuild.util import memoize
 
 JAR_STORED = ZIP_STORED
 JAR_DEFLATED = ZIP_DEFLATED
+JAR_BROTLI = 0x81
 MAX_WBITS = 15
 
 
@@ -262,13 +265,14 @@ class JarFileReader(object):
         corresponding to the file in the jar archive, data a buffer containing
         the file data.
         '''
-        assert header['compression'] in [JAR_DEFLATED, JAR_STORED]
+        assert header['compression'] in [JAR_DEFLATED, JAR_STORED, JAR_BROTLI]
         self._data = data
         # Copy some local file header fields.
         for name in ['filename', 'compressed_size',
                      'uncompressed_size', 'crc32']:
             setattr(self, name, header[name])
-        self.compressed = header['compression'] == JAR_DEFLATED
+        self.compressed = header['compression'] != JAR_STORED
+        self.compress = header['compression']
 
     def read(self, length=-1):
         '''
@@ -317,7 +321,11 @@ class JarFileReader(object):
         if hasattr(self, '_uncompressed_data'):
             return self._uncompressed_data
         data = self.compressed_data
-        if self.compressed:
+        if self.compress == JAR_STORED:
+            data = data.tobytes()
+        elif self.compress == JAR_BROTLI:
+            data = Brotli.decompress(data.tobytes())
+        elif self.compress == JAR_DEFLATED:
             data = zlib.decompress(data.tobytes(), -MAX_WBITS)
         else:
             data = data.tobytes()
@@ -361,6 +369,13 @@ class JarReader(object):
         del self._data
 
     @property
+    def compression(self):
+        entries = self.entries
+        if not entries:
+            return JAR_STORED
+        return max(f['compression'] for f in entries.itervalues())
+
+    @property
     def entries(self):
         '''
         Return an ordered dict of central directory entries, indexed by
@@ -473,6 +488,8 @@ class JarWriter(object):
             self._data = fileobj
         else:
             self._data = open(file, 'wb')
+        if compress is True:
+            compress = JAR_DEFLATED
         self._compress = compress
         self._compress_level = compress_level
         self._contents = OrderedDict()
@@ -574,12 +591,13 @@ class JarWriter(object):
         '''
         Add a new member to the jar archive, with the given name and the given
         data.
-        The compress option indicates if the given data should be compressed
-        (True), not compressed (False), or compressed according to the default
-        defined when creating the JarWriter (None).
-        When the data should be compressed (True or None with self.compress ==
-        True), it is only really compressed if the compressed size is smaller
-        than the uncompressed size.
+        The compress option indicates how the given data should be compressed
+        (one of JAR_STORED, JAR_DEFLATE or JAR_BROTLI), or compressed according
+        to the default defined when creating the JarWriter (None). True and
+        False are allowed values for backwards compatibility, mapping,
+        respectively, to JAR_DEFLATE and JAR_STORED.
+        When the data should be compressed, it is only really compressed if
+        the compressed size is smaller than the uncompressed size.
         The mode option gives the unix permissions that should be stored
         for the jar entry.
         If a duplicated member is found skip_duplicates will prevent raising
@@ -594,8 +612,12 @@ class JarWriter(object):
             raise JarWriterError("File %s already in JarWriter" % name)
         if compress is None:
             compress = self._compress
-        if (isinstance(data, JarFileReader) and data.compressed == compress) \
-                or (isinstance(data, Deflater) and data.compress == compress):
+        if compress is True:
+            compress = JAR_DEFLATED
+        if compress is False:
+            compress = JAR_STORED
+        if (isinstance(data, (JarFileReader, Deflater)) and \
+                data.compress == compress):
             deflater = data
         else:
             deflater = Deflater(compress, compress_level=self._compress_level)
@@ -619,7 +641,7 @@ class JarWriter(object):
         if deflater.compressed:
             entry['min_version'] = 20  # Version 2.0 supports deflated streams
             entry['general_flag'] = 2  # Max compression
-            entry['compression'] = JAR_DEFLATED
+            entry['compression'] = deflater.compress
         else:
             entry['min_version'] = 10  # Version 1.0 for stored streams
             entry['general_flag'] = 0
@@ -659,14 +681,21 @@ class Deflater(object):
     '''
     def __init__(self, compress=True, compress_level=9):
         '''
-        Initialize a Deflater. The compress argument determines whether to
-        try to compress at all.
+        Initialize a Deflater. The compress argument determines how to
+        compress.
         '''
         self._data = BytesIO()
+        if compress is True:
+            compress = JAR_DEFLATED
+        elif compress is False:
+            compress = JAR_STORED
         self.compress = compress
-        if compress:
-            self._deflater = zlib.compressobj(compress_level, zlib.DEFLATED,
-                                              -MAX_WBITS)
+        if compress in (JAR_DEFLATED, JAR_BROTLI):
+            if compress == JAR_DEFLATED:
+                self._deflater = zlib.compressobj(
+                    compress_level, zlib.DEFLATED, -MAX_WBITS)
+            else:
+                self._deflater = BrotliCompress()
             self._deflated = BytesIO()
         else:
             self._deflater = None
@@ -759,6 +788,46 @@ class Deflater(object):
         return self._data.getvalue()
 
 
+class Brotli(object):
+    @staticmethod
+    @memoize
+    def brotli_tool():
+            from buildconfig import topobjdir, substs
+            return os.path.join(topobjdir, 'dist', 'host', 'bin',
+                               'brotli' + substs.get('BIN_SUFFIX', ''))
+
+    @staticmethod
+    def run_brotli_tool(args, input):
+        proc = subprocess.Popen([Brotli.brotli_tool()] + args,
+                                stdin=subprocess.PIPE,
+                                stdout=subprocess.PIPE)
+        (stdout, _) = proc.communicate(input)
+        ret = proc.wait()
+        if ret != 0:
+            raise Exception("Brotli compression failed")
+        return stdout
+
+    @staticmethod
+    def compress(data):
+        return Brotli.run_brotli_tool(['--lgwin=17'], data)
+
+    @staticmethod
+    def decompress(data):
+        return Brotli.run_brotli_tool(['--decompress'], data)
+
+
+class BrotliCompress(object):
+    def __init__(self):
+        self._buf = BytesIO()
+
+    def compress(self, data):
+        self._buf.write(data)
+        return b''
+
+    def flush(self):
+        return Brotli.compress(self._buf.getvalue())
+
+
 class JarLog(dict):
     '''
     Helper to read the file Gecko generates when setting MOZ_JAR_LOG_FILE.
diff --git a/python/mozbuild/mozpack/packager/l10n.py b/python/mozbuild/mozpack/packager/l10n.py
index 758064f59..e3e05fc89 100644
--- a/python/mozbuild/mozpack/packager/l10n.py
+++ b/python/mozbuild/mozpack/packager/l10n.py
@@ -37,6 +37,7 @@ from mozpack.chrome.manifest import (
     Manifest,
 )
 from mozpack.errors import errors
+from mozpack.mozjar import JAR_DEFLATED
 from mozpack.packager.unpack import UnpackFinder
 from createprecomplete import generate_precomplete
 
@@ -241,16 +242,17 @@ def repack(source, l10n, extra_l10n={}, non_resources=[], non_chrome=set()):
             finders[base] = UnpackFinder(path)
         l10n_finder = ComposedFinder(finders)
     copier = FileCopier()
+    compress = min(app_finder.compressed, JAR_DEFLATED)
     if app_finder.kind == 'flat':
         formatter = FlatFormatter(copier)
     elif app_finder.kind == 'jar':
         formatter = JarFormatter(copier,
                                  optimize=app_finder.optimizedjars,
-                                 compress=app_finder.compressed)
+                                 compress=compress)
     elif app_finder.kind == 'omni':
         formatter = OmniJarFormatter(copier, app_finder.omnijar,
                                      optimize=app_finder.optimizedjars,
-                                     compress=app_finder.compressed,
+                                     compress=compress,
                                      non_resources=non_resources)
 
     with errors.accumulate():
diff --git a/python/mozbuild/mozpack/packager/unpack.py b/python/mozbuild/mozpack/packager/unpack.py
index fa2b474e7..515705c0d 100644
--- a/python/mozbuild/mozpack/packager/unpack.py
+++ b/python/mozbuild/mozpack/packager/unpack.py
@@ -54,7 +54,7 @@ class UnpackFinder(BaseFinder):
         self.omnijar = None
         self.jarlogs = {}
         self.optimizedjars = False
-        self.compressed = True
+        self.compressed = False
 
         jars = set()
 
@@ -146,8 +146,7 @@ class UnpackFinder(BaseFinder):
         jar = JarReader(fileobj=file.open())
         if jar.is_optimized:
             self.optimizedjars = True
-        if not any(f.compressed for f in jar):
-            self.compressed = False
+        self.compressed = max(self.compressed, jar.compression)
         if jar.last_preloaded:
             jarlog = jar.entries.keys()
             self.jarlogs[path] = jarlog[:jarlog.index(jar.last_preloaded) + 1]
diff --git a/toolkit/mozapps/installer/packager.mk b/toolkit/mozapps/installer/packager.mk
index e700b5162..6cae15658 100644
--- a/toolkit/mozapps/installer/packager.mk
+++ b/toolkit/mozapps/installer/packager.mk
@@ -50,7 +50,7 @@ stage-package: $(MOZ_PKG_MANIFEST) $(MOZ_PKG_MANIFEST_DEPS)
 		) \
 		$(if $(JARLOG_DIR),$(addprefix --jarlog ,$(wildcard $(JARLOG_FILE_AB_CD)))) \
 		$(if $(OPTIMIZEJARS),--optimizejars) \
-		$(if $(DISABLE_JAR_COMPRESSION),--disable-compression) \
+		$(addprefix --compress ,$(JAR_COMPRESSION)) \
 		$(addprefix --unify ,$(UNIFY_DIST)) \
 		$(MOZ_PKG_MANIFEST) '$(DIST)' '$(DIST)'/$(STAGEPATH)$(MOZ_PKG_DIR)$(if $(MOZ_PKG_MANIFEST),,$(_BINPATH)) \
 		$(if $(filter omni,$(MOZ_PACKAGER_FORMAT)),$(if $(NON_OMNIJAR_FILES),--non-resource $(NON_OMNIJAR_FILES)))
diff --git a/toolkit/mozapps/installer/packager.py b/toolkit/mozapps/installer/packager.py
index 1a144823c..3b90e97c2 100644
--- a/toolkit/mozapps/installer/packager.py
+++ b/toolkit/mozapps/installer/packager.py
@@ -23,6 +23,7 @@ from mozpack.copier import (
     Jarrer,
 )
 from mozpack.errors import errors
+from mozpack.mozjar import JAR_BROTLI
 from mozpack.unify import UnifiedBuildFinder
 import mozpack.path as mozpath
 import buildconfig
@@ -270,9 +271,9 @@ def main():
                         help='Enable jar optimizations')
     parser.add_argument('--unify', default='',
                         help='Base directory of another build to unify with')
-    parser.add_argument('--disable-compression', action='store_false',
-                        dest='compress', default=True,
-                        help='Disable jar compression')
+    parser.add_argument('--compress', choices=('none', 'deflate', 'brotli'),
+                        default='deflate',
+                        help='Use given jar compression (default: deflate)')
     parser.add_argument('manifest', default=None, nargs='?',
                         help='Manifest file name')
     parser.add_argument('source', help='Source directory')
@@ -290,15 +291,21 @@ def main():
         for name, value in [split_define(d) for d in args.defines]:
             defines[name] = value
 
+    compress = {
+        'none': False,
+        'deflate': True,
+        'brotli': JAR_BROTLI,
+    }[args.compress]
+
     copier = FileCopier()
     if args.format == 'flat':
         formatter = FlatFormatter(copier)
     elif args.format == 'jar':
-        formatter = JarFormatter(copier, compress=args.compress, optimize=args.optimizejars)
+        formatter = JarFormatter(copier, compress=compress, optimize=args.optimizejars)
     elif args.format == 'omni':
         formatter = OmniJarFormatter(copier,
                                      buildconfig.substs['OMNIJAR_NAME'],
-                                     compress=args.compress,
+                                     compress=compress,
                                      optimize=args.optimizejars,
                                      non_resources=args.non_resource)
     else:
diff --git a/toolkit/mozapps/installer/upload-files.mk b/toolkit/mozapps/installer/upload-files.mk
index 82b550fb4..865d2cf6a 100644
--- a/toolkit/mozapps/installer/upload-files.mk
+++ b/toolkit/mozapps/installer/upload-files.mk
@@ -385,7 +385,7 @@ ifneq (android,$(MOZ_WIDGET_TOOLKIT))
   OPTIMIZEJARS = 1
   ifneq (gonk,$(MOZ_WIDGET_TOOLKIT))
     ifdef NIGHTLY_BUILD
-      DISABLE_JAR_COMPRESSION = 1
+      JAR_COMPRESSION ?= none
     endif
   endif
 endif