run_isolated.py: install CIPD packages

client/cipd.py

Index: client/cipd.py
diff --git a/client/cipd.py b/client/cipd.py
new file mode 100644
index 0000000000000000000000000000000000000000..6fa93dbfcb3871bb03a4d0f4a159d6ef1f1c8b5b
--- /dev/null
+++ b/client/cipd.py
@@ -0,0 +1,356 @@
+# Copyright 2012 The LUCI Authors. All rights reserved.

[M-A Ruel, 2016/06/06 23:34:50]

2012?

[nodir, 2016/06/07 18:46:35]

Done.

+# Use of this source code is governed under the Apache License, Version 2.0
+# that can be found in the LICENSE file.
+
+"""Fetches CIPD client and installs packages."""
+
+__version__ = '0.1'
+
+import hashlib
+import logging
+import optparse
+import platform
+import sys
+import tempfile
+import urllib
+
+from utils import fs
+from utils import net
+from utils import subprocess42
+from utils import tools
+import isolated_format
+import isolateserver
+
+
+class Error(Exception):
+  """Raised on CIPD errors."""
+
+
+def add_cipd_options(parser):
+  group = optparse.OptionGroup(parser, 'CIPD')
+  group.add_option(
+      '--cipd-server',
+      help=(

[M-A Ruel, 2016/06/06 23:34:50]

() are not needed for str concatenation

[nodir, 2016/06/07 18:46:35]

Done.

+          'URL of the CIPD server. Only relevant with --cipd-package. '
+          'Default: "%default"'),
+      default='https://chrome-infra-packages.appspot.com')
+  group.add_option(
+      '--cipd-client-package',
+      help=(
+          'Package of CIPD client. See --cipd-package for format. '
+          'Only relevant with --cipd-package. '
+          'Default: "%default"'),
+      default='infra/tools/cipd/${platform}:latest')
+  group.add_option(
+      '--cipd-package',
+      help=(
+          'CIPD package to install. '
+          'Format: "<package name template>:<version>". '
+          'Package name template is a CIPD package name with optional '
+          '${platform} and/or ${os_ver} parameters. '
+          '${platform} will be expanded to "<os>-<architecture>" and '
+          '${os_ver} will be expanded to OS version name. '
+          'This option can be specified more than once.'),
+      action='append')
+  group.add_option(
+      '--cipd-cache',
+      help=(
+          'CIPD cache directory, separate from isolate cache. '
+          'Only relevant with --cipd-package. '
+          'Default: "%default".'))
+  parser.add_option_group(group)
+
+
+def validate_cipd_options(parser, options):
+  """Calls parser.error on first found error among cipd options."""
+  if not options.cipd_package:
+    return
+  for p in options.cipd_package:
+    try:
+      parse_package(p)
+    except ValueError as ex:
+      parser.error('Invalid cipd package %r: %s' % (p, ex))
+
+  if not options.cipd_server:
+    parser.error('--cipd-package requires non-empty --cipd-server')
+
+  if not options.cipd_client_package:
+    parser.error('--cipd-package requires non-empty --cipd-client-package')
+  try:
+    parse_package(options.cipd_client_package)
+  except ValueError as ex:
+    parser.error(
+        'Invalid cipd client package %r: %s' %
+        (options.cipd_client_package, ex))
+
+
+class CipdClient(object):
+  """Installs packages."""
+
+  def __init__(self, binary_path, service_url=None):
+    """Initializes CipdClient.
+
+    Args:
+      binary_path (str): path to the CIPD client binary.
+      service_url (str): if not None, URL of the CIPD backend that overrides
+        the default one.
+    """
+    self.binary_path = binary_path
+    self.service_url = service_url
+
+  def ensure(self, site_root, packages, cache_dir=None, timeout=None):
+    """Ensures that packages installed in |site_root| equals |packages| set.
+
+    Blocking call.
+
+    Args:
+      site_root (str): where to install packages.
+      packages (str): list of packages to install, parsable by parse_pacakge().
+      cache_dir (str): if set, cache dir for cipd binary own cache.
+        Typically contains packages and tags.
+      timeout (int): if not None, timeout in seconds for this function to run.
+
+    Raises:
+      Error if could not install packages or timed out.
+    """
+    timeouter = tools.Timeouter(timeout)
+    logging.info('Installing packages %r into %s', packages, site_root)
+    with tempfile.NamedTemporaryFile(prefix='cipd-ensure-list-') as list_file:

[Vadim Sh., 2016/06/06 21:32:13]

keep it under cache_dir or site_root, not in global /tmp

[nodir, 2016/06/07 18:46:35]

Done.

+      for p in packages:
+        pkg, version = parse_package(p)
+        pkg = render_package_name_template(pkg)
+        list_file.write('%s %s\n' % (pkg, version))
+      list_file.flush()

[Vadim Sh., 2016/06/06 21:32:13]

there's а chance this won't work on Windows due to file locking. Consider
properly closing the file before passing to the cipd.

[M-A Ruel, 2016/06/06 23:34:50]

Vadim is right, this is broken on Windows.

[nodir, 2016/06/07 18:46:35]

Done.

+
+      cmd = [
+        self.binary_path, 'ensure',
+        '-root', site_root,
+        '-list', list_file.name,
+        '-verbose',  # this is safe because cipd-ensure does not print a lot
+      ]
+      if cache_dir:
+        cmd += ['-cache-dir', cache_dir]
+      if self.service_url:
+        cmd += ['-service-url', self.service_url]
+
+      logging.debug('Running %r', cmd)
+      process = subprocess42.Popen(
+          cmd, stdout=subprocess42.PIPE, stderr=subprocess42.PIPE)
+      output = []
+      for pipe_name, line in process.yield_any_line(timeout=timeouter.left):
+        if pipe_name is None:

[M-A Ruel, 2016/06/06 23:34:51]

It could also mean that one pipe closed, not necessarily that it timed out.

[nodir, 2016/06/07 18:46:34]

Done

+          raise Error(
+              'Could not install packages; took more than %d seconds' % timeout)

[Vadim Sh., 2016/06/06 21:32:13]

We should probably add this to the cipd client too :( It should be relatively
simple, internal API supports context.Context already.

[nodir, 2016/06/07 18:46:35]

I started doing that, but it turned out much more work than doing it here,
primarily because cipd code does not check for ctx.Err(). Please consider adding
timeout when you have time.

+        output.append(line)
+        if pipe_name == 'stderr':
+          logging.debug('cipd client: %s', line)
+        else:
+          logging.info('cipd client: %s', line)
+
+      exit_code = process.wait(timeout=timeouter.left())
+      if exit_code != 0:
+        raise Error(
+            'Could not install packages; exit code %d\noutput:%s' % (
+            exit_code, '\n'.join(output)))
+
+
+def get_platform():
+  """Returns ${platform} parameter value.
+
+  Borrowed from
+  https://chromium.googlesource.com/infra/infra/+/aaf9586/build/build.py#204
+  """
+  # linux, mac or windows.
+  platform_variant = {
+    'darwin': 'mac',
+    'linux2': 'linux',
+    'win32': 'windows',
+  }.get(sys.platform)
+  if not platform_variant:
+    raise Error('Unknown OS: %s' % sys.platform)
+
+  # amd64, 386, etc.
+  platform_arch = {
+    'amd64': 'amd64',
+    'i386': '386',
+    'i686': '386',
+    'x86': '386',
+    'x86_64': 'amd64',
+  }.get(platform.machine().lower())

[M-A Ruel, 2016/06/06 23:34:50]

https://github.com/luci/luci-py/blob/master/appengine/swarming/swarming_bot/a...

this works on ARM too.

[nodir, 2016/06/07 18:46:34]

that function cannot be imported here because this code is in client/ and that
one is in swarming_bot.
I've adapted that one line

+  if not platform_arch:
+    raise Error('Unknown machine arch: %s' % platform.machine())
+  return '%s-%s' % (platform_variant, platform_arch)
+
+
+def get_os_ver():
+  """Returns ${os_ver} parameter value.
+
+  Examples: 'ubuntu14_04' or 'mac10_9' or 'win6_1'.
+
+  Borrowed from
+  https://chromium.googlesource.com/infra/infra/+/aaf9586/build/build.py#204
+  """
+  if sys.platform == 'darwin':
+    # platform.mac_ver()[0] is '10.9.5'.
+    dist = platform.mac_ver()[0].split('.')
+    return 'mac%s_%s' % (dist[0], dist[1])
+
+  if sys.platform == 'linux2':
+    # platform.linux_distribution() is ('Ubuntu', '14.04', ...).
+    dist = platform.linux_distribution()
+    return '%s%s' % (dist[0].lower(), dist[1].replace('.', '_'))
+
+  if sys.platform == 'win32':
+    # platform.version() is '6.1.7601'.
+    dist = platform.version().split('.')
+    return 'win%s_%s' % (dist[0], dist[1])
+  raise Error('Unknown OS: %s' % sys.platform)
+
+
+def render_package_name_template(template):
+  """Expands template variables in a CIPD package name template."""
+  result = template

[M-A Ruel, 2016/06/06 23:34:50]

this line is not useful.

[nodir, 2016/06/07 18:46:34]

Done.

+  result = result.replace('${platform}', get_platform())
+  result = result.replace('${os_ver}', get_os_ver())
+  return result

[M-A Ruel, 2016/06/06 23:34:50]

this is not needed either

[nodir, 2016/06/07 18:46:35]

Done

+
+
+def parse_package(package):
+  """Parses a package in --cipd-package format.
+
+  Returns:
+    (package_name_template, version) tuple.
+
+  Raises:
+    ValueError if package name or version is not specified.
+  """
+  if not package:
+    raise ValueError('package is not specified')
+  parts = package.split(':', 1)
+  if len(parts) != 2:
+    raise ValueError('version is not specified')
+  return tuple(parts)
+
+
+def _check_response(res, fmt, *args):
+  """Raises Error if response is bad."""
+  if res and res.get('status') == 'SUCCESS':

[M-A Ruel, 2016/06/06 23:34:50]

I'd replace the whole function with:

if not res:
  raise Error('%s: no response' % (fmt & args))
if res.get('status') != 'SUCCESS':
  raise Error('%s: %s' % (
      fmt & args,
      res.get('error_message') or 'status is %s' % res.get('status')))

[nodir, 2016/06/07 18:46:35]

Done.

+    return
+
+  if not res:
+    reason = 'no response'
+  else:
+    reason = res.get('error_message') or 'status is %s' % res.get('status')
+  raise Error(
+      '%s: %s' % (fmt & args, reason))
+
+
+def resolve_version(cipd_server, package_name, version, timeout=None):
+  """Resolves a package instance version (e.g. a tag) to an instance id."""
+  url = '%s/_ah/api/repo/v1/instance/resolve?%s' % (
+      cipd_server,
+      urllib.urlencode({
+        'package_name': package_name,
+        'version': version,
+      }))
+  res = net.url_read_json(url, timeout=timeout)
+  _check_response(res, 'Could not resolve version %s:%s', package_name, version)
+  instance_id = res.get('instance_id')
+  if not instance_id:
+    raise Error('Invalid resolveVersion response: no instance id')
+  return instance_id
+
+
+def fetch_client_binary(cipd_server, package_name, instance_id, timeout=None):
+  """Returns a net.HttpResponse with CIPD client binary contents.
+
+  Raises:
+    Error if cannot fetch the client.
+  """
+  # Fetch the URL of the binary from CIPD backend.
+  package_name = render_package_name_template(package_name)
+  url = '%s/_ah/api/repo/v1/client?%s' % (cipd_server, urllib.urlencode({
+    'package_name': package_name,
+    'instance_id': instance_id,
+  }))
+  res = net.url_read_json(url, timeout=timeout)
+  _check_response(
+      res, 'Could not fetch CIPD client %s:%s',package_name, instance_id)
+  fetch_url = res.get('client_binary', {}).get('fetch_url')
+  if not fetch_url:
+    raise Error('Invalid fetchClientBinary response: no fetch_url')
+  logging.debug('CIPD client binary fetch URL: %s', fetch_url)
+
+  # Now fetch the actual binary from Google Storage.
+  res = net.url_open(fetch_url, stream=False)

[M-A Ruel, 2016/06/06 23:34:50]

that assumes the client binary fits in memory, we're generally fine but it could
be stretching it on a lower end rPi.

[nodir, 2016/06/07 18:46:34]

I reworked fetch_client_binary to (url_open() | disk_cache.write) with retries
and exponential back-off and added more exception handling

+  if res is None:
+    raise Error(
+        'Could not fetch CIPD client %s:%s: no response' %
+        (package_name, instance_id))
+  return res
+
+
+def get_client(
+      service_url, package_name, version, client_cache, version_cache=None,
+      timeout=None):
+  """Creates a CipdClient. A blocking call.
+
+  The returned client points to a file in client_cache, thus client_cache
+  should not be mutated while CipdClient is used, otherwise the binary may be
+  deleted.
+
+  Args:
+      service_url (str): URL of the CIPD backend.
+      package_name (str): package name template of the CIPD client.
+      version (str): version of CIPD client package.
+      client_cache (isolateserver.DiskCache): cache for client binaries.
+      version_cache (isolateserver.LocalCache): if not None, cache for
+        {client version -> instance id} mapping.
+      timeout (int): if not None, timeout in seconds for this function.
+
+  Returns:
+    CipdClient.
+
+  Raises:
+    Error if CIPD client version cannot be resolved or client cannot be fetched.
+  """
+  timeouter = tools.Timeouter(timeout)
+  assert isinstance(client_cache, isolateserver.DiskCache)
+
+  package_name = render_package_name_template(package_name)
+
+  # Resolve version to instance id.
+  # Is it an instance id already? They look like HEX SHA1.
+  if isolated_format.is_valid_hash(version, hashlib.sha1):
+    instance_id = version
+  else:
+    instance_id = None
+    version_digest = None
+    if version_cache:
+      # Convert |version| to a string that may be used as a filename in disk
+      # cache by hashing it.
+      sha1 = hashlib.sha1()
+      sha1.update(version)
+      version_digest = sha1.hexdigest()

[M-A Ruel, 2016/06/06 23:34:50]

version_digest = hashlib.sha1(version).hexdigest()

[nodir, 2016/06/07 18:46:35]

Done.

+      try:
+        instance_id = version_cache.read(version_digest)
+      except isolateserver.CacheMiss:
+        pass
+    if not instance_id:
+      instance_id = resolve_version(
+          service_url, package_name, version, timeout=timeouter.left())
+      if version_cache:
+        version_cache.write(version_digest, instance_id)
+
+  # Get the path to the client binary inside client cache.
+  try:
+    binary_path = client_cache.item_path(instance_id)
+  except isolateserver.CacheMiss:
+    logging.info('Fetching CIPD client %s:%s', package_name, instance_id)
+    res = fetch_client_binary(
+        service_url, package_name, instance_id, timeout=timeouter.left())
+    client_cache.write(instance_id, res.iter_content(64 * 1024))

[M-A Ruel, 2016/06/06 23:34:50]

iter_content() doesn't matter here since stream=False is used.

[nodir, 2016/06/07 18:46:34]

changed to streaming

+    binary_path = client_cache.item_path(instance_id)
+    fs.chmod(binary_path, 0711)  # -rwx--x--x
+    logging.info('Fetched CIPD client into %s', binary_path)
+  return CipdClient(binary_path)