cache: Use 'pixz' when extracting .xz tarballs if it's installed.
When untar'ing an xz-compressed tarball, we can speed things up
signficiantly by using pixz, which utilizes more CPU cores rather than
just a single core.
So this uses pixz if it's installed to extract .xz tarballs, and
will use up to half of the system's cores.
BUG=chromium:1034063
TEST=unittest
TEST=simplechrome'd locally
Change-Id: If6f4a3680a450c9d886a3f6bf5f1b7e9f0566e96
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/2383130
Tested-by: Ben Pastene <bpastene@chromium.org>
Commit-Queue: Mike Frysinger <vapier@chromium.org>
Reviewed-by: Mike Frysinger <vapier@chromium.org>
diff --git a/lib/cache.py b/lib/cache.py
index 51f459d..780c240 100644
--- a/lib/cache.py
+++ b/lib/cache.py
@@ -322,7 +322,11 @@
def Untar(path, cwd, sudo=False):
"""Untar a tarball."""
functor = cros_build_lib.sudo_run if sudo else cros_build_lib.run
- functor(['tar', '-xpf', path], cwd=cwd, debug_level=logging.DEBUG, quiet=True)
+ comp = cros_build_lib.CompressionExtToType(path)
+ cmd = ['tar']
+ if comp != cros_build_lib.COMP_NONE:
+ cmd += ['-I', cros_build_lib.FindCompressor(comp)]
+ functor(cmd + ['-xpf', path], cwd=cwd, debug_level=logging.DEBUG, quiet=True)
class TarballCache(RemoteCache):
diff --git a/lib/cache_unittest.py b/lib/cache_unittest.py
index 926cf67..63cb40d 100644
--- a/lib/cache_unittest.py
+++ b/lib/cache_unittest.py
@@ -12,9 +12,10 @@
import mock
-from chromite.lib import gs_unittest
-from chromite.lib import cros_test_lib
from chromite.lib import cache
+from chromite.lib import cros_build_lib
+from chromite.lib import cros_test_lib
+from chromite.lib import gs_unittest
from chromite.lib import osutils
from chromite.lib import partial_mock
from chromite.lib import retry_util
@@ -304,3 +305,24 @@
testAssign = CacheTestCase._testAssign
testAssignData = CacheTestCase._testAssignData
testRemove = CacheTestCase._testRemove
+
+
+class UntarTest(cros_test_lib.RunCommandTestCase):
+ """Tests cache.Untar()."""
+
+ @mock.patch('chromite.lib.cros_build_lib.CompressionExtToType')
+ def testNoneCompression(self, mock_compression_type):
+ """Tests Untar with an uncompressed tarball."""
+ mock_compression_type.return_value = cros_build_lib.COMP_NONE
+ cache.Untar('/some/tarball.tar.gz', '/')
+ self.assertCommandContains(['tar', '-xpf', '/some/tarball.tar.gz'])
+
+ @mock.patch('chromite.lib.cros_build_lib.CompressionExtToType')
+ @mock.patch('chromite.lib.cros_build_lib.FindCompressor')
+ def testCompression(self, mock_find_compressor, mock_compression_type):
+ """Tests Untar with a compressed tarball."""
+ mock_compression_type.return_value = 'some-compression'
+ mock_find_compressor.return_value = '/bin/custom/xz'
+ cache.Untar('/some/tarball.tar.xz', '/')
+ self.assertCommandContains(
+ ['tar', '-I', '/bin/custom/xz', '-xpf', '/some/tarball.tar.xz'])
diff --git a/scripts/xz_auto.py b/scripts/xz_auto.py
index 4d2fb4a..c2b299b 100644
--- a/scripts/xz_auto.py
+++ b/scripts/xz_auto.py
@@ -5,14 +5,59 @@
"""Run xz from PATH with a thread for each core in the system."""
+from __future__ import division
from __future__ import print_function
+import multiprocessing
import os
import sys
+from chromite.lib import commandline
+from chromite.lib import osutils
+from chromite.utils import memoize
+
assert sys.version_info >= (3, 6), 'This module requires Python 3.6+'
+@memoize.Memoize
+def HasPixz():
+ """Returns path to pixz if it's on PATH or None otherwise."""
+ return osutils.Which('pixz')
+
+
+@memoize.Memoize
+def GetJobCount():
+ """Returns half of the total number of the machine's CPUs as a string.
+
+ Returns half rather than all of them to avoid starving out other parallel
+ processes on the same machine.
+ """
+ return str(int(max(1, multiprocessing.cpu_count() / 2)))
+
+
+def GetDecompressCommand():
+ """Returns decompression command."""
+ if HasPixz():
+ return ['pixz', '-d', '-p', GetJobCount()]
+ return ['xz', '-d']
+
+
+def GetParser():
+ """Return a command line parser."""
+ parser = commandline.ArgumentParser(description=__doc__)
+ parser.add_argument(
+ '-d', '--decompress', '--uncompress',
+ help='Decompress rather than compress.',
+ action='store_true')
+ return parser
+
+
def main(argv):
+ parser = GetParser()
+ known_args, argv = parser.parse_known_args()
+ # xz doesn't support multi-threaded decompression, so try using pixz for that.
+ if known_args.decompress:
+ args = GetDecompressCommand()
+ os.execvp(args[0], args + argv)
os.execvp('xz', ['xz', '-T0'] + argv)