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..04261217440221470910cc1da1887e1e269b80d6 100755
--- a/scripts/slave/annotated_run.py
+++ b/scripts/slave/annotated_run.py
@@ -26,10 +26,43 @@ 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 +70,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,9 +108,71 @@ PLATFORM_CONFIG = {
 # the recipe engine.
 Config = collections.namedtuple('Config', (
     'run_cmd',
+    'logdog_pubsub',
+    'logdog_platform',
 ))
 
 
+class Runtime(object):
+  """Runtime is the runtime context of the recipe execution.
+
+  It is a ContextManager that tracks generated files and cleans them up at
+  exit.
+  """
+
+  def __init__(self, leak=False):
+    self._tempdirs = []
+    self._leak = leak
+
+  def cleanup(self, path):
+    self._tempdirs.append(path)
+
+  def tempdir(self, base=None):
+    """Creates a temporary recipe-local working directory and yields it.
+
+    This creates a temporary directory for this annotation run. Directory
+    cleanup is appended to the supplied Runtime.
+
+    This creates two levels of directory:
+      <base>/.recipe_runtime
+      <base>/.recipe_runtime/tmpFOO
+
+    On termination, the entire "<base>/.recipe_runtime" directory is deleted,
+    removing the subdirectory created by this instance as well as cleaning up
+    any other temporary subdirectories leaked by previous executions.
+
+    Args:
+      rt (Runtime): Process-wide runtime.
+      base (str/None): The directory under which the tempdir should be created.
+          If None, the default temporary directory root will be used.
+    """
+    base = base or tempfile.gettempdir()
+    basedir = ensure_directory(base, '.recipe_runtime')
+    self.cleanup(basedir)
+    tdir = tempfile.mkdtemp(dir=basedir)
+    return tdir
+
+  def __enter__(self):
+    return self
+
+  def __exit__(self, _et, _ev, _tb):
+    self.close()
+
+  def close(self):
+    if self._leak:
+      LOGGER.warning('(--leak) Leaking temporary paths: %s', self._tempdirs)
+    else:
+      for path in reversed(self._tempdirs):
+        try:
+          if os.path.isdir(path):
+            LOGGER.debug('Cleaning up temporary directory [%s].', path)
+            chromium_utils.RemoveDirectory(path)
+        except BaseException:
+          LOGGER.exception('Failed to clean up temporary directory [%s].',
+                           path)
+    del(self._tempdirs[:])
+
+
 def get_config():
   """Returns (Config): The constructed Config object.
 
@@ -83,6 +192,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 +204,29 @@ def ensure_directory(*path):
   return path
 
 
+def _logdog_get_streamserver_uri(rt, typ):
+  """Returns (str): The Butler StreamServer URI.
+
+  Args:
+    rt (Runtime): Process-wide runtime.
+    typ (str): The type of URI to generate. One of: ['unix'].
+  Raises:
+    LogDogBootstrapError: if |typ| is not a known type.
+  """
+  if typ == 'unix':
+    # We have to use a custom temporary directory here. This is due to the path
+    # length limitation on UNIX domain sockets, which is generally 104-108
+    # characters. We can't make that assumption about our standard recipe
+    # temporary directory.
+    sockdir = rt.tempdir()
+    uri = 'unix:%s' % (os.path.join(sockdir, 'butler.sock'),)
+    if len(uri) > 104:
+      raise LogDogBootstrapError('Generated URI exceeds UNIX domain socket '
+                                 'name size: %s' % (uri,))
+    return uri
+  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)
@@ -115,33 +249,255 @@ def _check_command(cmd, **kwargs):
   return stdout
 
 
-@contextlib.contextmanager
-def recipe_tempdir(root=None, leak=False):
-  """Creates a temporary recipe-local working directory and yields it.
+class LogDogNotBootstrapped(Exception):
+  pass
+
+
+class LogDogBootstrapError(Exception):
+  pass
+
+
+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,))
+
 
-  This creates a temporary directory for this annotation run that is
-  automatically cleaned up. It returns the directory.
+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:
-    root (str/None): If not None, the root directory. Otherwise, |os.cwd| will
-        be used.
-    leak (bool): If true, don't clean up the temporary directory on exit.
+    path (str): The CIPD installation root.
+    packages (CipdBinary): The set of CIPD binary packages to install.
   """
-  basedir = ensure_directory((root or os.getcwd()), '.recipe_runtime')
+  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:
-    tdir = tempfile.mkdtemp(dir=basedir)
-    yield tdir
-  finally:
-    if basedir and os.path.isdir(basedir):
-      if not leak:
-        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)
+    _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 _build_logdog_prefix(properties):
+  """Constructs a LogDog stream prefix from the supplied properties.
+
+  The returned prefix is of the form:
+  bb/<mastername>/<buildername>/<buildnumber>
+
+  Any path-incompatible characters will be flattened to underscores.
+  """
+  def normalize(s):
+    parts = []
+    for ch in str(s):
+      if ch.isalnum() or ch in ':_-.':
+        parts.append(ch)
       else:
-        LOGGER.warning('(--leak) Leaking temporary directory [%s].', basedir)
+        parts.append('_')
+    if not parts[0].isalnum():
+      parts.insert(0, 's_')
+    return ''.join(parts)
+
+  components = {}
+  for f in ('mastername', 'buildername', 'buildnumber'):
+    prop = properties.get(f)
+    if not prop:
+      raise LogDogBootstrapError('Missing build property [%s].' % (f,))
+    components[f] = normalize(properties.get(f))
+  return 'bb/%(mastername)s/%(buildername)s/%(buildnumber)s' % components
+
+
+def _logdog_bootstrap(rt, opts, tempdir, config, properties, 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:
+    rt (Runtime): Process-wide runtime.
+    opts (argparse.Namespace): Command-line options.
+    tempdir (str): The path to the session temporary directory.
+    config (Config): Recipe runtime configuration.
+    properties (dict): Build properties.
+    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')
+
+  plat = config.logdog_platform
+  if not plat:
+    raise LogDogNotBootstrapped('LogDog platform is not configured.')
+
+  # Determine LogDog prefix.
+  prefix = _build_logdog_prefix(properties)
+
+  # TODO(dnj): Consider moving this to a permanent directory on the bot so we
+  #            don't CIPD-refresh each time.
+  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(rt, plat.streamserver)
+
+  # 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,
+      '-prefix', prefix,
+      '-output', 'pubsub,project="%(project)s",topic="%(topic)s"' % (
+          config.logdog_pubsub._asdict()),
+  ]
+  cmd += logdog_verbose
+  cmd += service_account_args
+  cmd += [
+      'run',
+      '-stdout', 'tee=stdout',
+      '-stderr', 'tee=stderr',
+      '-streamserver-uri', streamserver_uri,
+      '--',
+  ]
+
+  # Annotee Command.
+  cmd += [
+      annotee,
+      '-butler-stream-server', streamserver_uri,
+      '-annotate', 'tee',
+      '-name-base', 'recipes',
+      '-print-summary',
+      '-tee',
+      '-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 _should_run_logdog(properties):
+  """Returns (bool): True if LogDog should be used for this run.
+
+  Args:
+    properties (dict): The factory properties for this recipe run.
+  """
+  mastername = properties.get('mastername')
+  buildername = properties.get('buildername')
+  if not all((mastername, buildername)):
+    LOGGER.warning('Required mastername/buildername is not set.')
+    return False
+
+  # 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 True
+
+  LOGGER.info('Master %s, builder %s is not whitelisted for LogDog.',
+              mastername, buildername)
+  return False
 
 
 def get_recipe_properties(workdir, build_properties,
@@ -297,6 +653,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 +683,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 +796,51 @@ def write_monitoring_event(config, datadir, build_properties):
     LOGGER.warning("Failed to send monitoring event.", exc_info=True)
 
 
+def _exec_recipe(rt, 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 opts.logdog_force or _should_run_logdog(properties):
+      status = _logdog_bootstrap(rt, opts, tdir, config, properties, 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)
 
@@ -436,7 +853,8 @@ def main(argv):
   clean_old_recipe_engine()
 
   # Enter our runtime environment.
-  with recipe_tempdir(leak=opts.leak) as tdir:
+  with Runtime(leak=opts.leak) as rt:
+    tdir = rt.tempdir(os.getcwd())
     LOGGER.debug('Using temporary directory: [%s].', tdir)
 
     # Load factory properties and configuration.
@@ -450,19 +868,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 +875,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(rt, opts, tdir, config, properties)
 
 
 def shell_main(argv):