Bug 1637381: Add support for extracting zst files to mozbuild tooltool r=glandium
authorChris AtLee <catlee@mozilla.com>
Thu, 21 May 2020 13:28:33 +0000
changeset 531428 4aae9dc67310a0d7cc8eae3c415205b9930c153d
parent 531427 5d1fbed42ab7dcd32dc2c9cbbfd89c25796b3480
child 531429 aab5979b57e2654a22d0041a330343b8ff1665c7
push id37439
push userbtara@mozilla.com
push dateThu, 21 May 2020 21:49:34 +0000
treeherdermozilla-central@92c11f0bf14b [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewersglandium
bugs1637381
milestone78.0a1
first release with
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
last release without
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
Bug 1637381: Add support for extracting zst files to mozbuild tooltool r=glandium Differential Revision: https://phabricator.services.mozilla.com/D75203
python/mozbuild/mozbuild/action/tooltool.py
python/mozbuild/mozbuild/artifact_commands.py
python/mozbuild/mozbuild/base.py
taskcluster/mach_commands.py
--- a/python/mozbuild/mozbuild/action/tooltool.py
+++ b/python/mozbuild/mozbuild/action/tooltool.py
@@ -883,17 +883,17 @@ def clean_path(dirname):
         shutil.rmtree(dirname)
 
 
 CHECKSUM_SUFFIX = ".checksum"
 
 
 def unpack_file(filename):
     """Untar `filename`, assuming it is uncompressed or compressed with bzip2,
-    xz, gzip, or unzip a zip file. The file is assumed to contain a single
+    xz, gzip, zst, or unzip a zip file. The file is assumed to contain a single
     directory with a name matching the base of the given filename.
     Xz support is handled by shelling out to 'tar'."""
     if os.path.isfile(filename) and tarfile.is_tarfile(filename):
         tar_file, zip_ext = os.path.splitext(filename)
         base_file, tar_ext = os.path.splitext(tar_file)
         clean_path(base_file)
         log.info('untarring "%s"' % filename)
         tar = tarfile.open(filename)
@@ -909,16 +909,26 @@ def unpack_file(filename):
         if process.returncode != 0:
             return False
         fileobj = BytesIO()
         fileobj.write(stdout)
         fileobj.seek(0)
         tar = tarfile.open(fileobj=fileobj, mode='r|')
         tar.extractall()
         tar.close()
+    elif os.path.isfile(filename) and filename.endswith('.tar.zst'):
+        import zstandard
+        base_file = filename.replace('.tar.zst', '')
+        clean_path(base_file)
+        log.info('untarring "%s"' % filename)
+        dctx = zstandard.ZstdDecompressor()
+        with dctx.stream_reader(open(filename, "rb")) as fileobj:
+            tar = tarfile.open(fileobj=fileobj, mode='r|')
+            tar.extractall()
+            tar.close()
     elif os.path.isfile(filename) and zipfile.is_zipfile(filename):
         base_file = filename.replace('.zip', '')
         clean_path(base_file)
         log.info('unzipping "%s"' % filename)
         z = zipfile.ZipFile(filename)
         z.extractall()
         z.close()
     else:
--- a/python/mozbuild/mozbuild/artifact_commands.py
+++ b/python/mozbuild/mozbuild/artifact_commands.py
@@ -321,24 +321,33 @@ class PackageFrontend(MachCommandBase):
                     b = 'toolchain-{}'.format(b)
 
                 task = toolchains.get(aliases.get(b, b))
                 if not task:
                     self.log(logging.ERROR, 'artifact', {'build': user_value},
                              'Could not find a toolchain build named `{build}`')
                     return 1
 
+                artifact_name = task.attributes.get('toolchain-artifact')
+                self.log(logging.DEBUG, 'artifact',
+                         {'name': artifact_name,
+                          'index': task.optimization.get('index-search')},
+                         'Searching for {name} in {index}')
                 task_id = IndexSearch().should_replace_task(
                     task, {}, task.optimization.get('index-search', []))
-                artifact_name = task.attributes.get('toolchain-artifact')
                 if task_id in (True, False) or not artifact_name:
                     self.log(logging.ERROR, 'artifact', {'build': user_value},
                              _COULD_NOT_FIND_ARTIFACTS_TEMPLATE)
                     return 1
 
+                self.log(logging.DEBUG, 'artifact',
+                         {'name': artifact_name,
+                          'task_id': task_id},
+                         'Found {name} in {task_id}')
+
                 record = ArtifactRecord(task_id, artifact_name)
                 records[record.filename] = record
 
         # Handle the list of files of the form path@task-id on the command
         # line. Each of those give a path to an artifact to download.
         for f in files:
             if '@' not in f:
                 self.log(logging.ERROR, 'artifact', {},
@@ -423,17 +432,26 @@ class PackageFrontend(MachCommandBase):
                         data = fh.read(1024 * 1024)
                         if not data:
                             break
                         h.update(data)
                 artifacts[record.url] = {
                     'sha256': h.hexdigest(),
                 }
             if record.unpack and not no_unpack:
-                unpack_file(local)
+                # Try to unpack the file. If we get an exception importing
+                # zstandard when calling unpack_file, we can try installing
+                # zstandard locally and trying again
+                try:
+                    unpack_file(local)
+                except ImportError as e:
+                    if e.name != "zstandard":
+                        raise
+                    self._ensure_zstd()
+                    unpack_file(local)
                 os.unlink(local)
 
         if not downloaded:
             self.log(logging.ERROR, 'artifact', {}, 'Nothing to download')
             if files:
                 return 1
 
         if artifacts:
--- a/python/mozbuild/mozbuild/base.py
+++ b/python/mozbuild/mozbuild/base.py
@@ -854,16 +854,23 @@ class MozbuildObject(ProcessExecutionMix
         return pipenv
 
     def activate_pipenv(self, pipfile=None, populate=False, python=None):
         if pipfile is not None and not os.path.exists(pipfile):
             raise Exception('Pipfile not found: %s.' % pipfile)
         self.ensure_pipenv()
         self.virtualenv_manager.activate_pipenv(pipfile, populate, python)
 
+    def _ensure_zstd(self):
+        try:
+            import zstandard  # noqa: F401
+        except (ImportError, AttributeError):
+            self._activate_virtualenv()
+            self.virtualenv_manager.install_pip_package('zstandard>=0.9.0,<=0.13.0')
+
 
 class MachCommandBase(MozbuildObject):
     """Base class for mach command providers that wish to be MozbuildObjects.
 
     This provides a level of indirection so MozbuildObject can be refactored
     without having to change everything that inherits from it.
     """
 
--- a/taskcluster/mach_commands.py
+++ b/taskcluster/mach_commands.py
@@ -429,23 +429,16 @@ class MachCommands(MachCommandBase):
             print(json.dumps(actions, sort_keys=True, indent=2, separators=(',', ': ')))
         except Exception:
             traceback.print_exc()
             sys.exit(1)
 
 
 @CommandProvider
 class TaskClusterImagesProvider(MachCommandBase):
-    def _ensure_zstd(self):
-        try:
-            import zstandard  # noqa: F401
-        except (ImportError, AttributeError):
-            self._activate_virtualenv()
-            self.virtualenv_manager.install_pip_package('zstandard==0.9.0')
-
     @Command('taskcluster-load-image', category="ci",
              description="Load a pre-built Docker image. Note that you need to "
                          "have docker installed and running for this to work.")
     @CommandArgument('--task-id',
                      help="Load the image at public/image.tar.zst in this task, "
                           "rather than searching the index")
     @CommandArgument('-t', '--tag',
                      help="tag that the image should be loaded as. If not "