annotated_run.py: Add LogDog bootstrapping.

scripts/slave/annotated_run.py

Index: scripts/slave/annotated_run.py
diff --git a/scripts/slave/annotated_run.py b/scripts/slave/annotated_run.py
index d5bfd101ce26de0ba26dc5d26f813ebde43d1a31..37fd1c3112d8ef6c03438b60a962da25271f150c 100755
--- a/scripts/slave/annotated_run.py
+++ b/scripts/slave/annotated_run.py
@@ -26,10 +26,44 @@ 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)
@@ -37,11 +71,25 @@ LOGGER = logging.getLogger('annotated_run')
 # - 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
@@ -61,6 +109,8 @@ PLATFORM_CONFIG = {
 # the recipe engine.
 Config = collections.namedtuple('Config', (
     'run_cmd',
+    'logdog_pubsub',
+    'logdog_platform',
 ))
 
 
@@ -83,6 +133,8 @@ def get_config():
   # 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'),
       )
 
 
@@ -93,6 +145,20 @@ def ensure_directory(*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)
@@ -144,6 +210,216 @@ def recipe_tempdir(root=None, leak=False):
         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.
@@ -297,6 +573,22 @@ def get_args(argv):
       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)
 
 
@@ -311,7 +603,7 @@ def update_scripts():
     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:
@@ -424,6 +716,53 @@ def write_monitoring_event(config, datadir, build_properties):
     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)
 
@@ -450,19 +789,6 @@ def main(argv):
     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
@@ -470,21 +796,8 @@ def main(argv):
     # 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):