annotated_run.py: Add LogDog bootstrapping.

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,
+)
+
+# Sentinel value that, if present in master config, matches all builders
+# underneath that master.
+WHITELIST_ALL = '*'
+
+# Whitelist of {master}=>[{builder}|WHITELIST_ALL] whitelisting specific masters
+# and builders for experimental LogDog/Annotee export.
+LOGDOG_WHITELIST_MASTER_BUILDERS = {
+}
+
+# Configuration for a Pub/Sub topic.
+PubSubConfig = collections.namedtuple('PubSubConfig', ('project', 'topic'))
+
+# LogDogPlatform is the set of platform-specific LogDog bootstrapping
+# configuration parameters.
+#
+# See _logdog_get_streamserver_uri for "streamserver" parameter details.
+LogDogPlatform = collections.namedtuple('LogDogPlatform', (
+    'butler', 'annotee', 'credential_path', 'streamserver',
+    ))
+
+# A CIPD binary description, including the package name, version, and relative
+# path of the binary within the package.
+CipdBinary = collections.namedtuple('CipdBinary',
+    ('package', 'version', 'relpath'))
 
 # RecipeRuntime will probe this for values.
 # - First, (system, platform)
 # - Then, (system,)
 # - Finally, (),
 PLATFORM_CONFIG = {
   # All systems.
-  (): {},
+  (): {
+    'logdog_pubsub': PubSubConfig(
+        project='luci-logdog',
+        topic='logs',
+    ),
+  },
 
   # Linux
   ('Linux',): {
     'run_cmd': ['/opt/infra-python/run.py'],
+    'logdog_platform': LogDogPlatform(
+        butler=CipdBinary('infra/tools/luci/logdog/butler/linux-amd64',
+                          'latest', 'logdog_butler'),
+        annotee=CipdBinary('infra/tools/luci/logdog/annotee/linux-amd64',
+                          'latest', 'logdog_annotee'),
+        credential_path=(
+            '/creds/service_accounts/service-account-luci-logdog-pubsub.json'),
+        streamserver='unix',
+    ),
   },
 
   # 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_platform',
 ))
 
 
 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_platform=platform_config.get('logdog_platform'),
       )
 
 
 def ensure_directory(*path):
   path = os.path.join(*path)
   if not os.path.isdir(path):
     os.makedirs(path)
   return path
 
 
+def _logdog_get_streamserver_uri(typ, d):
+  """Returns (str): The Butler StreamServer URI.
+
+  Args:
+    typ (str): The type of URI to generate. One of: ['unix'].
+    d (str): The working directory path.
+  Raises:
+    LogDogBootstrapError: if |typ| is not a known type.
+  """
+  if typ == 'unix':
+    return 'unix:%s' % (os.path.join(d, 'butler.sock'),)
+  raise LogDogBootstrapError('No streamserver URI generator.')
+
+
 def _run_command(cmd, **kwargs):
   if kwargs.pop('dry_run', False):
     LOGGER.info('(Dry Run) Would have executed command: %s', cmd)
     return 0, ''
 
   LOGGER.debug('Executing command: %s', cmd)
   kwargs.setdefault('stderr', subprocess.STDOUT)
   proc = subprocess.Popen(cmd, **kwargs)
   stdout, _ = proc.communicate()
 
@@ skipping 31 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 _get_service_account_json(opts, credential_path):
+  """Returns (str/None): If specified, the path to the service account JSON.
+
+  This method probes the local environment 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.logdog_service_account_json
+  if path:
+    return path
+
+  if gce.Authenticator.is_gce():
+    LOGGER.info('Running on GCE. No credentials necessary.')
+    return None
+
+  if os.path.isfile(credential_path):
+    return credential_path
+
+  raise LogDogBootstrapError('Could not find service account credentials. '
+                             'Tried: %s' % (credential_path,))
+
+
+def _logdog_install_cipd(path, *packages):
+  """Returns (list): The paths to the binaries in each of the packages.
+
+  This method bootstraps CIPD in "path", installing the packages specified
+  by "packages" and returning the paths to their binaries.
+
+  Args:
+    path (str): The CIPD installation root.
+    packages (CipdBinary): The set of CIPD binary packages to install.
+  """
+  verbosity = 0
+  level = logging.getLogger().level
+  if level <= logging.INFO:
+    verbosity += 1
+  if level <= logging.DEBUG:
+    verbosity += 1
+
+  packages_path = os.path.join(path, 'packages.json')
+  pmap = {}
+  cmd = [
+      sys.executable,
+      os.path.join(env.Build, 'scripts', 'slave', 'cipd.py'),
+      '--dest-directory', path,
+      '--json-output', packages_path,
+  ] + (['--verbose'] * verbosity)
+  for p in packages:
+    cmd += ['-P', '%s@%s' % (p.package, p.version)]
+    pmap[p.package] = os.path.join(path, p.relpath)
+
+  try:
+    _check_command(cmd)
+  except subprocess.CalledProcessError:
+    LOGGER.exception('Failed to install LogDog CIPD packages.')
+    raise LogDogBootstrapError()
+
+  # Resolve installed packages.
+  return tuple(pmap[p.package] for p in packages)
+
+
+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')

[Vadim Sh., 2016/01/09 01:25:13]

it will run bootstrap from scratch each time? :( They aren't very large files,
but when >2000 bots fetch them all the time, using single inbound pipe, it adds
up. Why not cache them locally between runs? 

[dnj (Google), 2016/01/09 02:50:29]

No specific reason other than this is nice and hermetic. In future-land this
will probably be isolated (as opposed to machine-native) so less pressure. This
sort of emulates isolate setup circumstances it in the meantime.

+
+  plat = config.logdog_platform
+  if not plat:
+    raise LogDogNotBootstrapped('LogDog platform is not configured.')
+
+  cipd_path = os.path.join(bootstrap_dir, 'cipd')
+  butler, annotee = _logdog_install_cipd(cipd_path, plat.butler, plat.annotee)
+  if opts.logdog_butler_path:
+    butler = opts.logdog_butler_path
+  if opts.logdog_annotee_path:
+    annotee = opts.logdog_annotee_path
+
+  if not config.logdog_pubsub:
+    raise LogDogNotBootstrapped('No Pub/Sub configured.')
+  if not config.logdog_pubsub.project:
+    raise LogDogNotBootstrapped('No Pub/Sub project configured.')
+  if not config.logdog_pubsub.topic:
+    raise LogDogNotBootstrapped('No Pub/Sub topic configured.')
+
+  # Determine LogDog verbosity.
+  logdog_verbose = []
+  if opts.logdog_verbose == 0:
+    pass
+  elif opts.logdog_verbose == 1:
+    logdog_verbose.append('-log_level=info')
+  else:
+    logdog_verbose.append('-log_level=debug')
+
+  service_account_args = []
+  service_account_json = _get_service_account_json(opts, plat.credential_path)
+  if service_account_json:
+    service_account_args += ['-service-account-json', service_account_json]
+
+  # Generate our Butler stream server URI.
+  streamserver_uri = _logdog_get_streamserver_uri(plat.streamserver, tempdir)
+
+  # Dump the bootstrapped Annotee command to JSON for Annotee to load.
+  #
+  # Annotee can run accept bootstrap parameters through either JSON or
+  # command-line, but using JSON effectively steps around any sort of command-
+  # line length limits such as those experienced on Windows.
+  cmd_json = os.path.join(bootstrap_dir, 'annotee_cmd.json')
+  with open(cmd_json, 'w') as fd:
+    json.dump(cmd, fd)
+
+  # Butler Command.
+  cmd = [
+      butler,
+      '-output', 'pubsub,project="%(project)s",topic="%(topic)s"' % (
+          config.logdog_pubsub._asdict()),
+  ]
+  cmd += logdog_verbose
+  cmd += service_account_args
+  cmd += [
+      'run',
+      '-streamserver-uri', streamserver_uri,
+      '--',
+  ]
+
+  # Annotee Command.
+  cmd += [
+      annotee,
+      '-butler-stream-server', streamserver_uri,
+      '-json-args-path', cmd_json,
+  ]
+  cmd += logdog_verbose
+
+  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 133 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('--logdog-verbose',
+      action='count', default=0,
+      help='Increase LogDog verbosity. This can be specified multiple times.')
+  group.add_argument('--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()
 
   with stream.step('update_scripts') as s:
     gclient_name = 'gclient'
     if sys.platform.startswith('win'):
       gclient_name += '.bat'
-    gclient_path = os.path.join(env.Build, '..', 'depot_tools',
+    gclient_path = os.path.join(env.Build, os.pardir, 'depot_tools',
                                 gclient_name)
     gclient_cmd = [gclient_path, 'sync', '--force', '--verbose', '--jobs=2']
     try:
       fd, output_json = tempfile.mkstemp()
       os.close(fd)
       gclient_cmd += ['--output-json', output_json]
     except Exception:
       # Super paranoia try block.
       output_json = None
     cmd_dict = {
@@ skipping 92 matching lines @@
      ]
     # Add this conditionally so that we get an error in
     # send_monitoring_event log files in case it isn't present.
     if hostname:
       cmd += ['--build-event-hostname', hostname]
     _check_command(cmd)
   except Exception:
     LOGGER.warning("Failed to send monitoring event.", exc_info=True)
 
 
+def _exec_recipe(opts, tdir, config, properties):
+  # Find out if the recipe we intend to run is in build_internal's recipes. If
+  # so, use recipes.py from there, otherwise use the one from build.
+  recipe_file = properties['recipe'].replace('/', os.path.sep) + '.py'
+
+  # Use the standard recipe runner unless the recipes are explicitly in the
+  # "build_limited" repository.
+  recipe_runner = os.path.join(env.Build,
+                               'scripts', 'slave', 'recipes.py')
+  if env.BuildInternal:
+    build_limited = os.path.join(env.BuildInternal, 'scripts', 'slave')
+    if os.path.exists(os.path.join(build_limited, 'recipes', recipe_file)):
+      recipe_runner = os.path.join(build_limited, 'recipes.py')
+
+  # Dump properties to JSON and build recipe command.
+  props_file = os.path.join(tdir, 'recipe_properties.json')
+  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 = None
+  try:
+    if not opts.logdog_force:
+      _assert_logdog_whitelisted(properties.get('mastername'),

[Vadim Sh., 2016/01/09 01:25:13]

I was confused by this... Do not use exception as a normal flow control.
Exceptions should be exceptional. e.g. change it to is_logdog_whitelisted.

[dnj (Google), 2016/01/09 02:50:28]

Done.

+                                 properties.get('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 as e:
+    LOGGER.exception('Exception while bootstrapping LogDog.')
+  finally:
+    if status is None:
+      LOGGER.info('Not using LogDog. Invoking `recipes.py` directly.')
+      status, _ = _run_command(cmd, dry_run=opts.dry_run)
+
+  return status
+
+
 def main(argv):
   opts = get_args(argv)
 
   if opts.verbose == 0:
     level = logging.INFO
   else:
     level = logging.DEBUG
   logging.getLogger().setLevel(level)
 
   clean_old_recipe_engine()
 
   # Enter our runtime environment.
   with recipe_tempdir(leak=opts.leak) as tdir:
     LOGGER.debug('Using temporary directory: [%s].', tdir)
 
     # Load factory properties and configuration.
     # TODO(crbug.com/551165): remove flag "factory_properties".
     use_factory_properties_from_disk = (opts.use_factory_properties_from_disk or
                                         bool(opts.factory_properties))
     properties = get_recipe_properties(
         tdir, opts.build_properties, use_factory_properties_from_disk)
     LOGGER.debug('Loaded properties: %s', properties)
 
     config = get_config()
     LOGGER.debug('Loaded runtime configuration: %s', config)
 
-    # Find out if the recipe we intend to run is in build_internal's recipes. If
-    # so, use recipes.py from there, otherwise use the one from build.
-    recipe_file = properties['recipe'].replace('/', os.path.sep) + '.py'
-
-    # Use the standard recipe runner unless the recipes are explicitly in the
-    # "build_limited" repository.
-    recipe_runner = os.path.join(env.Build,
-                                 'scripts', 'slave', 'recipes.py')
-    if env.BuildInternal:
-      build_limited = os.path.join(env.BuildInternal, 'scripts', 'slave')
-      if os.path.exists(os.path.join(build_limited, 'recipes', recipe_file)):
-        recipe_runner = os.path.join(build_limited, 'recipes.py')
-
     # Setup monitoring directory and send a monitoring event.
     build_data_dir = ensure_directory(tdir, 'build_data')
     properties['build_data_dir'] = build_data_dir
 
     # Write our annotated_run.py monitoring event.
     write_monitoring_event(config, build_data_dir, properties)
 
-    # Dump properties to JSON and build recipe command.
-    props_file = os.path.join(tdir, 'recipe_properties.json')
-    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)
-
-  return status
+    # Execute our recipe.
+    return _exec_recipe(opts, tdir, config, properties)
 
 
 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))