Add LogDog bootstrapping to `annotated_run.py`.

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 724ed9d02a068e0fae2362f7a642db1f3fc61901..174bfd373dd85ce13cb00ffb24a669dbd1cc353d 100755
--- a/scripts/slave/annotated_run.py
+++ b/scripts/slave/annotated_run.py
@@ -26,10 +26,31 @@ 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)
@@ -37,11 +58,21 @@ LOGGER = logging.getLogger('annotated_run')
 # - 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
@@ -61,6 +92,9 @@ PLATFORM_CONFIG = {
 # the recipe engine.
 Config = collections.namedtuple('Config', (
     'run_cmd',
+    'logdog_pubsub',
+    'logdog_butler_streamserver_gen',
+    'credential_paths',
 ))
 
 
@@ -83,6 +117,10 @@ def get_config():
   # 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', ()),
       )
 
 
@@ -144,6 +182,256 @@ 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 _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.
@@ -291,6 +579,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('-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)
 
 
@@ -477,7 +781,21 @@ def main(argv):
         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