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)