annotated_run: extract robust_tempdir so it can be shared with kitchen_run

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 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
 from slave import infra_platform
+from slave import robust_tempdir
 from slave import update_scripts
 
 # 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,
@@ skipping 125 matching lines @@
 # the recipe engine.
 Config = collections.namedtuple('Config', (
     'run_cmd',
     'logdog_host',
     'logdog_pubsub_topic',
     'logdog_max_buffer_age',
     '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.
 
   The Config object is constructed by 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.
@@ skipping 16 matching lines @@
   path = os.path.join(*path)
   if not os.path.isdir(path):
     os.makedirs(path)
   return path
 
 
 def _logdog_get_streamserver_uri(rt, typ):
   """Returns (str): The Butler StreamServer URI.
 
   Args:
-    rt (Runtime): Process-wide runtime.
+    rt (RobustTempdir): context for temporary directories.
     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.
     #
@@ skipping 56 matching lines @@
 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
@@ skipping 77 matching lines @@
 
   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.
+    rt (RobustTempdir): context for temporary directories.
     opts (argparse.Namespace): Command-line options.
     basedir (str): The base (non-temporary) recipe directory.
     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:
@@ skipping 417 matching lines @@
   if opts.verbose == 0:
     level = logging.INFO
   else:
     level = logging.DEBUG
   logging.getLogger().setLevel(level)
 
   clean_old_recipe_engine()
 
   # Enter our runtime environment.
   basedir = os.getcwd()
-  with Runtime(leak=opts.leak) as rt:
+  with robust_tempdir.RobustTempdir(
+      prefix='.recipe_runtime', leak=opts.leak) as rt:
     tdir = rt.tempdir(base=basedir)
     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)
@@ skipping 17 matching lines @@
     # 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))