Add LogDog bootstrapping to `annotated_run.py`.

scripts/slave/annotated_run.py

 #!/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 json
 import logging
 import os
 import platform
 import shutil
 import socket
 import subprocess
 import sys
 import tempfile
 
 
 # Install Infra build environment.
 BUILD_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(
                              os.path.abspath(__file__))))
 sys.path.insert(0, os.path.join(BUILD_ROOT, 'scripts'))
 
 from common import annotator
 from common import chromium_utils
 from common import env
 from common import master_cfg_utils
+from slave import gce
 
 # Logging instance.
 LOGGER = logging.getLogger('annotated_run')
 
+# 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 of {master}=>[{builder}|WHITELIST_ALL] whitelisting specific masters
+# and builders for experimental LogDog/Annotee export.
+LOGDOG_WHITELIST_MASTER_BUILDERS = {
+}
+
+# Sentinel value that, if present in master config, matches all builders
+# underneath that master.
+WHITELIST_ALL = '*'
+
+# Configuration for a Pub/Sub topic.
+PubSubConfig = collections.namedtuple('PubSubConfig', ('project', 'topic'))
 
 # RecipeRuntime will probe this for values.
 # - First, (system, platform)
 # - Then, (system,)
 # - Finally, (),
 PLATFORM_CONFIG = {
   # All systems.
-  (): {},
+  (): {
+    'logdog_pubsub': PubSubConfig(
+        project='luci-logdog',
+        topic='chrome-infra-beta',
+    ),
+  },
 
   # Linux
   ('Linux',): {
     'run_cmd': ['/opt/infra-python/run.py'],
+    'credential_paths': (
+      # XXX: Get this right?
+      '/opt/infra/service_accounts',
+    ),
+    'logdog_butler_streamserver_gen': lambda d: os.path.join(d, 'butler.sock'),
   },
 
   # Mac OSX
   ('Darwin',): {
     'run_cmd': ['/opt/infra-python/run.py'],
   },
 
   # Windows
   ('Windows',): {
     'run_cmd': ['C:\\infra-python\\ENV\\Scripts\\python.exe',
                 'C:\\infra-python\\run.py'],
   },
 }
 
 
 # Config is the runtime configuration used by `annotated_run.py` to bootstrap
 # the recipe engine.
 Config = collections.namedtuple('Config', (
     'run_cmd',
+    'logdog_pubsub',
+    'logdog_butler_streamserver_gen',
+    'credential_paths',
 ))
 
 
 def get_config():
   """Returns (Config): The constructed Config object.
 
   The Config object is constructed from:
   - Cascading the PLATFORM_CONFIG fields together based on current
         OS/Architecture.
 
   Raises:
     KeyError: if a required configuration key/parameter is not available.
   """
   # Cascade the platform configuration.
   p = (platform.system(), platform.processor())
   platform_config = {}
   for i in xrange(len(p)+1):
     platform_config.update(PLATFORM_CONFIG.get(p[:i], {}))
 
   # Construct runtime configuration.
   return Config(
       run_cmd=platform_config.get('run_cmd'),
+      logdog_pubsub=platform_config.get('logdog_pubsub'),
+      logdog_butler_streamserver_gen=platform_config.get(
+          'logdog_butler_streamserver_gen'),
+      credential_paths=platform_config.get('credential_paths', ()),
       )
 
 
 def ensure_directory(*path):
   path = os.path.join(*path)
   if not os.path.isdir(path):
     os.makedirs(path)
   return path
 
 
@@ skipping 41 matching lines @@
         LOGGER.debug('Cleaning up temporary directory [%s].', basedir)
         try:
           chromium_utils.RemoveDirectory(basedir)
         except Exception:
           LOGGER.exception('Failed to clean up temporary directory [%s].',
                            basedir)
       else:
         LOGGER.warning('(--leak) Leaking temporary directory [%s].', basedir)
 
 
+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 ensure_directory(*path):
+  path = os.path.join(*path)
+  if not os.path.isdir(path):
+    os.makedirs(path)
+  return path
+
+
+def _run_command(cmd, **kwargs):
+  dry_run = kwargs.pop('dry_run', False)
+
+  LOGGER.debug('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] returned [%d] with output:\n%s',
+               cmd, proc.returncode, 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 RecipeRuntime(object):
+  """RecipeRuntime is the platform-specific runtime enviornment.
+
+  The runtime is loaded with a set of read-only attributes that are a
+  combination of plaetform and runtime values used in the setup and execution of
+  the recipe engine.
+  """
+
+  _SENTINEL = object()
+
+  def __init__(self, **kwargs):
+    self._attrs = kwargs
+
+  @classmethod
+  @contextlib.contextmanager
+  def enter(cls, leak, **kw):
+    """Enters the annotated_run environment.
+
+    This creates a temporary directory for this annotation run that is
+    automatically cleaned up. It returns a RecipeRuntime object containing a
+    combination of the supplied keyword arguments and the platform-specific
+    configuration.
+
+    Args:
+      leak (bool): If true, don't clean up the temporary directory on exit.
+      kw (dict): Key/value pairs to add as attributes to the RecipeRuntime.
+    """
+    # Build our platform attributes.
+    p = (platform.system(), platform.processor())
+    attrs = {}
+    for i in xrange(len(p)+1):
+      attrs.update(PLATFORM_CONFIG.get(p[:i], {}))
+    attrs.update(kw)
+
+    basedir = ensure_directory(os.getcwd(), '.recipe_runtime')
+    try:
+      tdir = tempfile.mkdtemp(dir=basedir)
+      LOGGER.debug('Using temporary directory [%s].', tdir)
+
+      attrs['workdir'] = tdir
+      yield cls(**attrs)
+    finally:
+      if basedir and os.path.isdir(basedir):
+        if not leak:
+          LOGGER.debug('Cleaning up temporary directory [%s].', basedir)
+          try:
+            # TODO(pgervais): use infra_libs.rmtree instead.
+            shutil.rmtree(basedir)
+          except Exception:
+            LOGGER.exception('Failed to clean up temporary directory [%s].',
+                             basedir)
+        else:
+          LOGGER.warning('(--leak) Leaking temporary directory [%s].', basedir)
+
+  def __getattr__(self, key):
+    # Class methods/variables.
+    value = getattr(super(RecipeRuntime, self), key, self._SENTINEL)
+    if value is not self._SENTINEL:
+      return value
+
+    value = getattr(self, 'get')(key, self._SENTINEL)
+    if value is not self._SENTINEL:
+      return value
+    raise KeyError(key)
+
+  def get(self, key, default=None):
+    value = self._attrs.get(key, self._SENTINEL)
+    if value is not self._SENTINEL:
+      return value
+    return default
+
+  def __str__(self):
+    return str(self._attrs)
+
+
+def _get_service_account_json(opts, credential_paths):
+  """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:
+    rt (RecipeRuntime): The runtime environment.
+  Raises:
+    |LogDogBootstrapError| if no credentials could be found.
+  """
+  path = opts.get('service_account_json')
+  if path:
+    return path
+
+  if gce.Authenticator.is_gce():
+    LOGGER.info('Running on GCE. No credentials necessary.')
+    return None
+
+  for credential_path in 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. '
+                             'Tried: %s' % (credential_paths,))
+
+
+def _logdog_bootstrap(tempdir, config, opts, cmd):
+  """Executes the recipe engine, bootstrapping it through LogDog/Annotee.
+
+  This method executes the recipe engine, bootstrapping it through
+  LogDog/Annotee so its output and annotations are streamed to LogDog. The
+  bootstrap is configured to tee the annotations through STDOUT/STDERR so they
+  will still be sent to BuildBot.
+
+  The overall setup here is:
+  [annotated_run.py] => [logdog_butler] => [logdog_annotee] => [recipes.py]
+
+  Args:
+    config (Config): Recipe runtime configuration.
+    opts (argparse.Namespace): Command-line options.
+    cmd (list): The recipe runner command list to bootstrap.
+
+  Returns (int): The return code of the recipe runner process.
+
+  Raises:
+    LogDogNotBootstrapped: if the recipe engine was not executed because the
+        LogDog bootstrap requirements are not available.
+    LogDogBootstrapError: if there was an error bootstrapping the recipe runner
+        through LogDog.
+  """
+  bootstrap_dir = ensure_directory(tempdir, 'logdog_bootstrap')
+  butler, annotee = opts.logdog_butler_path, opts.logdog_annotee_path
+
+  if not is_executable(annotee):
+    raise LogDogNotBootstrapped('Annotee is not executable: %s' % (annotee,))
+  if not is_executable(butler):
+    raise LogDogNotBootstrapped('Butler is not executable: %s' % (butler,))
+
+  # Determine LogDog verbosity.
+  logdog_verbose = []
+  if opts.logdog_verbose == 0:
+    pass
+  elif opts.logdog_verbose == 1:
+    logdog_verbose.extend('-log_level=info')
+  else:
+    logdog_verbose.extend('-log_level=debug')
+
+  service_account_args = []
+  service_account_json = _get_service_account_json(
+      opts, config.credential_paths)
+  if service_account_json:
+    service_account_args += ['-service-account-json', service_account_json]
+
+  streamserver_uri_gen = config.logdog_butler_streamserver_gen
+  if not streamserver_uri_gen:
+    raise LogDogBootstrapError('No streamserver URI generator.')
+  streamserver_uri = streamserver_uri_gen(tempdir)
+
+  # Dump Annotee command to JSON.
+  cmd_json = os.path.join(bootstrap_dir, 'annotee_cmd.json')
+  with open(cmd_json, 'w') as fd:
+    json.dump(cmd, fd)
+
+  cmd = [
+      # Butler Command.
+      butler,
+      ] + logdog_verbose + service_account_args + [
+      '-output', 'gcps,project="%s",topic="%s"' % (config.logdog_pubsub.project,
+                                                   config.logdog_pubsub.topic),
+      'run',
+      '-streamserver-uri', streamserver_uri,
+      '--',
+
+      # Annotee Command.
+      annotee,
+      ] + logdog_verbose + [
+      '-json-args-path', cmd_json,
+      ]
+  rv, _ = _run_command(cmd, dry_run=opts.dry_run)
+  if rv in LOGDOG_ERROR_RETURNCODES:
+    raise LogDogBootstrapError('LogDog Error (%d)' % (rv,))
+  return rv
+
+
+def _assert_logdog_whitelisted(mastername, buildername):
+  """Asserts that the runtime environment is whitelisted for LogDog bootstrap.
+
+  Args:
+    mastername (str): The master name string.
+    buildername (str): The builder name.
+  Raises:
+    LogDogNotBootstrapped: if the runtime is not whitelisted.
+  """
+  if not all((mastername, buildername)):
+    raise LogDogNotBootstrapped('Required mastername/buildername is not set.')
+
+  # Key on mastername.
+  bdict = LOGDOG_WHITELIST_MASTER_BUILDERS.get(mastername)
+  if bdict is not None:
+    # Key on buildername.
+    if WHITELIST_ALL in bdict or buildername in bdict:
+      LOGGER.info('Whitelisted master %s, builder %s.',
+                  mastername, buildername)
+      return
+  raise LogDogNotBootstrapped('Master %s, builder %s is not whitelisted.' % (
+                              mastername, buildername))
+
+
 def get_recipe_properties(workdir, build_properties,
                           use_factory_properties_from_disk):
   """Constructs the recipe's properties from buildbot's properties.
 
   This retrieves the current factory properties from the master_config
   in the slave's checkout (no factory properties are handed to us from the
   master), and merges in the build properties.
 
   Using the values from the checkout allows us to do things like change
   the recipe and other factory properties for a builder without needing
@@ skipping 127 matching lines @@
       help='factory properties in b64 gz JSON format')
   parser.add_argument('--keep-stdin', action='store_true', default=False,
       help='don\'t close stdin when running recipe steps')
   parser.add_argument('--master-overrides-slave', action='store_true',
       help='use the property values given on the command line from the master, '
            'not the ones looked up on the slave')
   parser.add_argument('--use-factory-properties-from-disk',
       action='store_true', default=False,
       help='use factory properties loaded from disk on the slave')
 
+  group = parser.add_argument_group('LogDog Bootstrap')
+  group.add_argument('-V', '--logdog-verbose',
+      action='count', default=0,
+      help='Increase LogDog verbosity. This can be specified multiple times.')
+  group.add_argument('-f', '--logdog-force', action='store_true',
+      help='Force LogDog bootstrapping, even if the system is not configured.')
+  group.add_argument('--logdog-butler-path',
+      help='Path to the LogDog Butler. If empty, one will be probed/downloaded '
+           'from CIPD.')
+  group.add_argument('--logdog-annotee-path',
+      help='Path to the LogDog Annotee. If empty, one will be '
+           'probed/downloaded from CIPD.')
+  group.add_argument('--logdog-service-account-json',
+      help='Path to the service account JSON. If one is not provided, the '
+           'local system credentials will be used.')
+
   return parser.parse_args(argv)
 
 
 def update_scripts():
   if os.environ.get('RUN_SLAVE_UPDATED_SCRIPTS'):
     os.environ.pop('RUN_SLAVE_UPDATED_SCRIPTS')
     return False
 
   stream = annotator.StructuredAnnotationStream()
 
@@ skipping 166 matching lines @@
     with open(props_file, 'w') as fh:
       json.dump(properties, fh)
     cmd = [
         sys.executable, '-u', recipe_runner,
         'run',
         '--workdir=%s' % os.getcwd(),
         '--properties-file=%s' % props_file,
         properties['recipe'],
     ]
 
-    status, _ = _run_command(cmd, dry_run=opts.dry_run)
+    status = None
+    try:
+      if not opts.logdog_force:
+        _assert_logdog_whitelisted(config.mastername, config.buildername)
+      status = _logdog_bootstrap(tdir, config, opts, cmd)
+    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 status is None:
+        LOGGER.info('Not using LogDog. Invoking `annotated_run.py` directly.')
+        status, _ = _run_command(cmd, dry_run=opts.dry_run)
 
   return status
 
 
 def shell_main(argv):
   if update_scripts():
     # Re-execute with the updated annotated_run.py.
     rv, _ = _run_command([sys.executable] + argv)
     return rv
   else:
     return main(argv[1:])
 
 
 if __name__ == '__main__':
   logging.basicConfig(level=logging.INFO)
   sys.exit(shell_main(sys.argv))