testing/condprofile/condprof/client.py

   1 # This Source Code Form is subject to the terms of the Mozilla Public
   2 # License, v. 2.0. If a copy of the MPL was not distributed with this
   3 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
   4 #
   5 # This module needs to stay Python 2 and 3 compatible
   6 #
   7 import datetime
   8 import functools
   9 import os
  10 import shutil
  11 import tarfile
  12 import tempfile
  13 import time
  14
  15 from mozprofile.prefs import Preferences
  16
  17 from condprof import progress
  18 from condprof.changelog import Changelog
  19 from condprof.util import (
  20     TASK_CLUSTER,
  21     ArchiveNotFound,
  22     check_exists,
  23     download_file,
  24     logger,
  25 )
  26
  27 TC_SERVICE = "https://firefox-ci-tc.services.mozilla.com"
  28 ROOT_URL = TC_SERVICE + "/api/index"
  29 INDEX_PATH = "gecko.v2.%(repo)s.latest.firefox.condprof-%(platform)s-%(scenario)s"
  30 INDEX_BY_DATE_PATH = "gecko.v2.%(repo)s.pushdate.%(date)s.latest.firefox.condprof-%(platform)s-%(scenario)s"
  31 PUBLIC_DIR = "artifacts/public/condprof"
  32 TC_LINK = ROOT_URL + "/v1/task/" + INDEX_PATH + "/" + PUBLIC_DIR + "/"
  33 TC_LINK_BY_DATE = ROOT_URL + "/v1/task/" + INDEX_BY_DATE_PATH + "/" + PUBLIC_DIR + "/"
  34 ARTIFACT_NAME = "profile%(version)s-%(platform)s-%(scenario)s-%(customization)s.tgz"
  35 CHANGELOG_LINK = (
  36     ROOT_URL + "/v1/task/" + INDEX_PATH + "/" + PUBLIC_DIR + "/changelog.json"
  37 )
  38 ARTIFACTS_SERVICE = "https://taskcluster-artifacts.net"
  39 DIRECT_LINK = ARTIFACTS_SERVICE + "/%(task_id)s/0/public/condprof/"
  40 CONDPROF_CACHE = "~/.condprof-cache"
  41 RETRIES = 3
  42 RETRY_PAUSE = 45
  43
  44
  45 class ServiceUnreachableError(Exception):
  46     pass
  47
  48
  49 class ProfileNotFoundError(Exception):
  50     pass
  51
  52
  53 class RetriesError(Exception):
  54     pass
  55
  56
  57 def _check_service(url):
  58     """Sanity check to see if we can reach the service root url."""
  59
  60     def _check():
  61         exists, _ = check_exists(url, all_types=True)
  62         if not exists:
  63             raise ServiceUnreachableError(url)
  64
  65     try:
  66         return _retries(_check)
  67     except RetriesError:
  68         raise ServiceUnreachableError(url)
  69
  70
  71 def _check_profile(profile_dir):
  72     """Checks for prefs we need to remove or set."""
  73     to_remove = ("gfx.blacklist.", "marionette.")
  74
  75     def _keep_pref(name, value):
  76         for item in to_remove:
  77             if not name.startswith(item):
  78                 continue
  79             logger.info("Removing pref %s: %s" % (name, value))
  80             return False
  81         return True
  82
  83     def _clean_pref_file(name):
  84         js_file = os.path.join(profile_dir, name)
  85         prefs = Preferences.read_prefs(js_file)
  86         cleaned_prefs = dict([pref for pref in prefs if _keep_pref(*pref)])
  87         if name == "prefs.js":
  88             # When we start Firefox, forces startupScanScopes to SCOPE_PROFILE (1)
  89             # otherwise, side loading will be deactivated and the
  90             # Raptor web extension won't be able to run.
  91             cleaned_prefs["extensions.startupScanScopes"] = 1
  92
  93             # adding a marker so we know it's a conditioned profile
  94             cleaned_prefs["profile.conditioned"] = True
  95
  96         with open(js_file, "w") as f:
  97             Preferences.write(f, cleaned_prefs)
  98
  99     _clean_pref_file("prefs.js")
 100     _clean_pref_file("user.js")
 101
 102
 103 def _retries(callable, onerror=None, retries=RETRIES):
 104     _retry_count = 0
 105     pause = RETRY_PAUSE
 106
 107     while _retry_count < retries:
 108         try:
 109             return callable()
 110         except Exception as e:
 111             if onerror is not None:
 112                 onerror(e)
 113             logger.info("Failed, retrying")
 114             _retry_count += 1
 115             time.sleep(pause)
 116             pause *= 1.5
 117
 118     # If we reach that point, it means all attempts failed
 119     if _retry_count >= RETRIES:
 120         logger.error("All attempt failed")
 121     else:
 122         logger.info("Retried %s attempts and failed" % _retry_count)
 123     raise RetriesError()
 124
 125
 126 def get_profile(
 127     target_dir,
 128     platform,
 129     scenario,
 130     customization="default",
 131     task_id=None,
 132     download_cache=True,
 133     repo="mozilla-central",
 134     remote_test_root="/sdcard/test_root/",
 135     version=None,
 136     retries=RETRIES,
 137 ):
 138     """Extract a conditioned profile in the target directory.
 139
 140     If task_id is provided, will grab the profile from that task. when not
 141     provided (default) will grab the latest profile.
 142     """
 143
 144     # XXX assert values
 145     if version:
 146         version = "-v%s" % version
 147     else:
 148         version = ""
 149
 150     # when we bump the Firefox version on trunk, autoland still needs to catch up
 151     # in this case we want to download an older profile- 2 days to account for closures/etc.
 152     oldday = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(days=2)
 153     params = {
 154         "platform": platform,
 155         "scenario": scenario,
 156         "customization": customization,
 157         "task_id": task_id,
 158         "repo": repo,
 159         "version": version,
 160         "date": str(oldday.date()).replace("-", "."),
 161     }
 162     logger.info("Getting conditioned profile with arguments: %s" % params)
 163     filename = ARTIFACT_NAME % params
 164     if task_id is None:
 165         url = TC_LINK % params + filename
 166         _check_service(TC_SERVICE)
 167     else:
 168         url = DIRECT_LINK % params + filename
 169         _check_service(ARTIFACTS_SERVICE)
 170
 171     logger.info("preparing download dir")
 172     if not download_cache:
 173         download_dir = tempfile.mkdtemp()
 174     else:
 175         # using a cache dir in the user home dir
 176         download_dir = os.path.expanduser(CONDPROF_CACHE)
 177         if not os.path.exists(download_dir):
 178             os.makedirs(download_dir)
 179
 180     downloaded_archive = os.path.join(download_dir, filename)
 181     logger.info("Downloaded archive path: %s" % downloaded_archive)
 182
 183     def _get_profile():
 184         logger.info("Getting %s" % url)
 185         try:
 186             archive = download_file(url, target=downloaded_archive)
 187         except ArchiveNotFound:
 188             raise ProfileNotFoundError(url)
 189         try:
 190             with tarfile.open(archive, "r:gz") as tar:
 191                 logger.info("Extracting the tarball content in %s" % target_dir)
 192                 size = len(list(tar))
 193                 with progress.Bar(expected_size=size) as bar:
 194
 195                     def _extract(self, *args, **kw):
 196                         if not TASK_CLUSTER:
 197                             bar.show(bar.last_progress + 1)
 198                         return self.old(*args, **kw)
 199
 200                     tar.old = tar.extract
 201                     tar.extract = functools.partial(_extract, tar)
 202                     tar.extractall(target_dir)
 203         except (OSError, tarfile.ReadError) as e:
 204             logger.info("Failed to extract the tarball")
 205             if download_cache and os.path.exists(archive):
 206                 logger.info("Removing cached file to attempt a new download")
 207                 os.remove(archive)
 208             raise ProfileNotFoundError(str(e))
 209         finally:
 210             if not download_cache:
 211                 shutil.rmtree(download_dir)
 212
 213         _check_profile(target_dir)
 214         logger.info("Success, we have a profile to work with")
 215         return target_dir
 216
 217     def onerror(error):
 218         logger.info("Failed to get the profile.")
 219         if os.path.exists(downloaded_archive):
 220             try:
 221                 os.remove(downloaded_archive)
 222             except Exception:
 223                 logger.error("Could not remove the file")
 224
 225     try:
 226         return _retries(_get_profile, onerror, retries)
 227     except RetriesError:
 228         # look for older profile 2 days previously
 229         filename = ARTIFACT_NAME % params
 230         url = TC_LINK_BY_DATE % params + filename
 231         try:
 232             return _retries(_get_profile, onerror, retries)
 233         except RetriesError:
 234             raise ProfileNotFoundError(url)
 235
 236
 237 def read_changelog(platform, repo="mozilla-central", scenario="settled"):
 238     params = {"platform": platform, "repo": repo, "scenario": scenario}
 239     changelog_url = CHANGELOG_LINK % params
 240     logger.info("Getting %s" % changelog_url)
 241     download_dir = tempfile.mkdtemp()
 242     downloaded_changelog = os.path.join(download_dir, "changelog.json")
 243
 244     def _get_changelog():
 245         try:
 246             download_file(changelog_url, target=downloaded_changelog)
 247         except ArchiveNotFound:
 248             shutil.rmtree(download_dir)
 249             raise ProfileNotFoundError(changelog_url)
 250         return Changelog(download_dir)
 251
 252     try:
 253         return _retries(_get_changelog)
 254     except Exception:
 255         raise ProfileNotFoundError(changelog_url)