Add LogDog bootstrapping to `annotated_run.py`.

scripts/slave/launcher.py

Index: scripts/slave/launcher.py
diff --git a/scripts/slave/launcher.py b/scripts/slave/launcher.py
new file mode 100755
index 0000000000000000000000000000000000000000..970c7447b3de46b7f31ad35dd1a397974454f276
--- /dev/null
+++ b/scripts/slave/launcher.py
@@ -0,0 +1,418 @@
+#!/usr/bin/env python
+# Copyright (c) 2013 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+import argparse
+import collections
+import contextlib
+import datetime
+import hashlib
+import itertools
+import logging
+import os
+import platform
+import shutil
+import subprocess
+import sys
+import tempfile
+
+
+# Install Infra environment.
+BUILD_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir,
+                                          os.pardir))
+SCRIPTS = os.path.join(BUILD_ROOT, 'scripts')
+sys.path.insert(0, SCRIPTS)
+import common.env
+common.env.Install()
+
+# Useful to test if we're running on a GCE instance.
+from gerrit_util import GceAuthenticator
+
+PubSubConfig = collections.namedtuple('PubSubConfig', ('project', 'topic'))
+
+# Path to the `annotated_run.py` executable.
+ANNOTATED_RUN = os.path.join(BUILD_ROOT, 'scripts', 'slave', 'annotated_run.py')
+
+# Global logger for launcher.
+LOGGER = logging.getLogger('launcher')
+
+
+# Return codes used by Butler/Annotee to indicate their failure (as opposed to
+# a forwarded return code from the underlying process).
+LOGDOG_ERROR_RETURNCODES = (
+    # Butler runtime error.
+    250,
+    # Annotee runtime error.
+    251,
+)
+
+
+WHITELIST_MASTER_BUILDERS = {
+}
+
+# Sentinel value that, if present in master config, matches all builders
+# underneath that master.
+WHITELIST_ALL = '*'
+
+# Environment will probe this for values.
+# - First, (),
+# - Then, (system,)
+# - Then, (system, platform)
+PLATFORM_CONFIG = {
+  (): {
+    'pubsub': PubSubConfig(
+        project='luci-logdog',
+        topic='chrome-infra-beta',
+    ),
+  },
+  ('Linux',): {
+    'cipd_static_paths': (
+      # XXX: Get this right?
+      '/opt/chrome-infra',
+    ),
+    'credential_paths': (
+      # XXX: Get this right?
+      '/opt/infra/service_accounts',
+    ),
+  },
+  ('Linux', 'x86_64'): {
+    'cipd_packages': {
+      'infra/tools/luci/logdog/butler/linux-amd64': 'testing',
+      'infra/tools/luci/logdog/annotee/linux-amd64': 'testing',
+    },
+  },
+}
+
+
+class LogDogNotBootstrapped(Exception):
+  pass
+
+
+class LogDogBootstrapError(Exception):
+  pass
+
+
+def is_executable(path):
+  return os.path.isfile(path) and os.access(path, os.X_OK)
+
+
+def _assert_executable(path):
+  if not is_executable(path):
+    raise LogDogNotBootstrapped('File is not executable: %s' % (path,))
+  return path
+
+
+def ensure_directory(path):
+  if not os.path.isdir(path):
+    os.makedirs(path)
+
+
+def _run_command(cmd, dry_run=False):
+  LOGGER.info('Executing command: %s', cmd)
+  if dry_run:
+    LOGGER.info('(Dry Run) Not executing command.')
+    return 0, ''
+  proc = subprocess.Popen(cmd, stderr=subprocess.STDOUT)
+  stdout, _ = proc.communicate()
+
+  LOGGER.debug('Process [%s] output:\n%s', cmd, stdout)
+  return proc.returncode, stdout
+
+
+def _check_command(*args, **kwargs):
+  rv, stdout = _run_command(args, **kwargs)
+  if rv != 0:
+    raise ValueError('Process exited with non-zero return code (%d)' % (rv,))
+  return stdout
+
+
+class Environment(object):
+  """Environment is the runtime enviornment of the LogDog bootstrap."""
+
+  _SENTINEL = object()
+
+  def __init__(self, current_platform, **kwargs):
+    self._platform = current_platform
+    self._fields = kwargs
+
+  @classmethod
+  @contextlib.contextmanager
+  def probe(cls, environ, **kw):
+    kw.update({
+        'mastername': environ.get('BUILDBOT_MASTERNAME'),
+        'buildername': environ.get('BUILDBOT_BUILDERNAME'),
+        'slavename': environ.get('BUILDBOT_SLAVENAME'),
+        'buildnumber': environ.get('BUILDBOT_BUILDNUMBER'),
+        'system': platform.system(),
+        'processor': platform.processor(),
+    })
+
+    tdir = None
+    try:
+      tdir = tempfile.mkdtemp()
+      LOGGER.debug('Using temporary directory [%s].', tdir)
+
+      kw['tempdir'] = tdir
+      yield cls((kw['system'], kw['processor']), **kw)
+    finally:
+      if tdir and os.path.isdir(tdir):
+        LOGGER.debug('Cleaning up temporary directory [%s].', tdir)
+        try:
+          shutil.rmtree(tdir)
+        except Exception:
+          LOGGER.exception('Failed to clean up temporary directory [%s].', tdir)
+
+  def __getattr__(self, key):
+    value = getattr(self, 'get')(key)
+    if value is not self._SENTINEL:
+      return value
+    raise KeyError(key)
+
+  def get(self, key, default=None):
+    value = getattr(super(Environment, self), key, self._SENTINEL)
+    if value is not self._SENTINEL:
+      return value
+
+    value = self._fields.get(key, self._SENTINEL)
+    if value is not self._SENTINEL:
+      return value
+
+    for i in xrange(len(self._platform)+1):
+      pcfg = PLATFORM_CONFIG.get(self._platform[:i], {})
+      value = pcfg.get(key, self._SENTINEL)
+      if value is not self._SENTINEL:
+        return value
+    return default
+
+  def __str__(self):
+    return str(self._fields)
+
+  def is_posix(self):
+    """Returns (bool): Whether the current system is a POSIX system."""
+    return self.system in ('Linux', 'Darwin')
+
+  def assert_whitelisted(self):
+    # Key on mastername.
+    bdict = WHITELIST_MASTER_BUILDERS.get(self.mastername)
+    if bdict is not None:
+      # Key on buildername.
+      if WHITELIST_ALL in bdict or self.buildername in bdict:
+        LOGGER.info('Whitelisted master %(mastername)s, builder '
+                    '%(buildername)s.' % self)
+        return
+
+    raise LogDogNotBootstrapped('Master %s, builder %s is not whitelisted.' % (
+                                self.mastername, self.buildername))
+
+  @property
+  def cipd_root(self):
+    return os.path.join(self.tempdir, 'cipd_root')
+
+  def butler_streamserver_uri(self):
+    """Returns (str): The Butler StreamServer URI argument."""
+    return 'unix:%s' % (os.path.join(self.tempdir, 'butler.sock'))
+
+
+class CIPD(object):
+
+  _CIPD_NAME = 'cipd'
+
+  def __init__(self, path, root):
+    self._cipd_path = path
+    self._root = root
+
+  @classmethod
+  def find(cls, env):
+    for p in itertools.chain(
+        iter(os.environ.get('PATH').split(os.pathsep)),
+        env.cipd_static_paths):
+      candidate = os.path.join(p, cls._CIPD_NAME)
+      if os.path.isfile(candidate) and os.access(candidate, os.X_OK):
+        return cls(candidate, env.cipd_root)
+    return None
+
+  def __call__(self, *args):
+    cmd = [self._cipd_path]
+    cmd.extend(args)
+    _check_command(*cmd)
+
+  def path(self, *components):
+    return os.path.join(self._root, *components)
+
+  def ensure(self, **packages):
+    if len(packages) == 0:
+      return
+
+    # Emit package list.
+    package_list = self.path('package_list.txt')
+    lines = [
+      '# Automatically generated CIPD package list (launcher.py)',
+      '# Generated at: %s' % (datetime.datetime.now().isoformat(),),
+      '',
+    ]
+    for pkg, version in sorted(packages.iteritems()):
+      lines.append('%s %s' % (pkg, version))
+
+    ensure_directory(self._root)
+    with open(package_list, 'w+') as fd:
+      fd.write('\n'.join(lines))
+
+    # Ensure against the package list.
+    args = [
+      'ensure',
+      '-root', self._root,
+      '-list', package_list,
+    ]
+    self(*args)
+
+
+def _get_service_account_json(env):
+  """Returns (str/None): If specified, the path to the service account JSON.
+
+  This method probes the local environemnt and returns a (possibly empty) list
+  of arguments to add to the Butler command line for authentication.
+
+  If we're running on a GCE instance, no arguments will be returned, as GCE
+  service account is implicitly authenticated. If we're running on Baremetal,
+  a path to those credentials will be returned.
+
+  Args:
+    env (Environment): The execution environment.
+  Raises:
+    |LogDogBootstrapError| if no credentials could be found.
+  """
+  path = env.get('service_account_json')
+  if path:
+    return path
+
+  if GceAuthenticator.is_gce():
+    LOGGER.info('Running on GCE. No credentials necessary.')
+    return None
+
+  for credential_path in env.get('credential_paths', ()):
+    candidate = os.path.join(credential_path, 'logdog_service_account.json')
+    if os.path.isfile(candidate):
+      return candidate
+
+  raise LogDogBootstrapError('Could not find service account credentials.')
+
+
+def _logdog_bootstrap(env, opts, cmd):
+  butler, annotee = opts.logdog_butler_path, opts.logdog_annotee_path
+  if not (butler and annotee):
+    # Load packages via CIPD.
+    cipd = CIPD.find(env)
+    if env.cipd_packages:
+      if not cipd:
+        raise LogDogBootstrapError('Could not find CIPD binary.')
+      cipd.ensure(**env.cipd_packages)
+    if not butler:
+      butler = cipd.path('logdog_butler')
+    if not annotee:
+      annotee = cipd.path('logdog_annotee')
+
+  # Determine LogDog verbosity.
+  if opts.logdog_verbose == 0:
+    logdog_verbosity = 'warning'
+  elif opts.logdog_verbose == 1:
+    logdog_verbosity = 'info'
+  else:
+    logdog_verbosity = 'debug'
+
+  service_account_args = []
+  service_account_json = _get_service_account_json(env)
+  if service_account_json:
+    service_account_args += ['-service-account-json', service_account_json]
+
+  cmd = [
+      _assert_executable(butler),
+      '-log_level', logdog_verbosity,
+      ] + service_account_args + [
+      '-output', 'gcps,project="%s",topic="%s"' % (env.pubsub.project,
+                                                   env.pubsub.topic),
+      'run',
+      '-streamserver-uri', env.butler_streamserver_uri(),
+      '--',
+      _assert_executable(annotee),
+      '-log_level', logdog_verbosity,
+      '--'
+      ] + cmd
+  rv, _ = _run_command(cmd, dry_run=opts.dry_run)
+  if rv in LOGDOG_ERROR_RETURNCODES:
+    raise LogDogBootstrapError('LogDog Error (%d)' % (rv,))
+  return rv
+
+
+def _parse_args(args):
+  parser = argparse.ArgumentParser()
+  parser.add_argument('-v', '--verbose',
+      action='count', default=0,
+      help='Increase verbosity. This can be specified multiple times.')
+  parser.add_argument('-V', '--logdog-verbose',
+      action='count', default=0,
+      help='Increase LogDog verbosity. This can be specified multiple times.')
+  parser.add_argument('-d', '--dry-run', action='store_true',
+      help='Perform configuration, but refrain from actually executing the '
+           'command.')
+  parser.add_argument('-f', '--force', action='store_true',
+      help='Force LogDog bootstrapping, even if the system is not configured.')
+  parser.add_argument('--logdog-butler-path',
+      help='Path to the LogDog Butler. If empty, one will be probed/downloaded '
+           'from CIPD.')
+  parser.add_argument('--logdog-annotee-path',
+      help='Path to the LogDog Annotee. If empty, one will be '
+           'probed/downloaded from CIPD.')
+  parser.add_argument('--service-account-json',
+      help='Path to the service account JSON. If one is not provided, the '
+           'local system credentials will be used.')
+  parser.add_argument('args', nargs=argparse.REMAINDER)
+
+  opts, extras = parser.parse_known_args(args)
+  if opts.args and opts.args[0] == '--':
+    opts.args.pop(0)
+  if extras:
+    opts.args = extras + opts.args
+  return opts
+
+
+def main(args):
+  opts = _parse_args(args)
+  cmd = [ANNOTATED_RUN] + opts.args
+
+  if opts.verbose == 0:
+    level = logging.WARNING
+  elif opts.verbose == 1:
+    level = logging.INFO
+  else:
+    level = logging.DEBUG
+  logging.getLogger().setLevel(level)
+
+  ran_bootstrap = False
+  rv = 1
+  try:
+    with Environment.probe(
+        os.environ,
+        service_account_json=opts.service_account_json) as env:
+      LOGGER.debug('Loaded environment: %s', env)
+      if not opts.force:
+        env.assert_whitelisted()
+
+      rv = _logdog_bootstrap(env, opts, cmd)
+      ran_bootstrap = True
+  except LogDogNotBootstrapped as e:
+    LOGGER.info('Not bootstrapped: %s', e.message)
+  except LogDogBootstrapError as e:
+    LOGGER.warning('Could not bootstrap LogDog: %s', e.message)
+  except Exception:
+    LOGGER.exception('Exception while bootstrapping LogDog.')
+  finally:
+    if not ran_bootstrap:
+      LOGGER.info('Not using LogDog. Invoking `annotated_run.py` directly.')
+      rv, _ = _run_command(cmd, opts)
+  return rv
+
+
+if __name__ == '__main__':
+  logging.basicConfig()
+  sys.exit(main(sys.argv[1:]))