Move goma_teardown, goma_setup and determine_goma_jobs to goma_utils.py

scripts/slave/goma_utils.py

Index: scripts/slave/goma_utils.py
diff --git a/scripts/slave/goma_utils.py b/scripts/slave/goma_utils.py
index 1485300d2419c98feea00853759c1d97ec89cc95..19f28fc7dd9aca7e5ce3f7cf3fe98ca7e309f1b3 100644
--- a/scripts/slave/goma_utils.py
+++ b/scripts/slave/goma_utils.py
@@ -11,10 +11,12 @@ import getpass
 import glob
 import gzip
 import json
+import multiprocessing
 import os
 import re
 import shutil
 import socket
+import subprocess
 import sys
 import tempfile
 import time
@@ -35,6 +37,21 @@ PLATFORM_RUN_CMD = {
 TIMESTAMP_PATTERN = re.compile('(\d{4}/\d{2}/\d{2} \d{2}:\d{2}:\d{2})')
 TIMESTAMP_FORMAT = '%Y/%m/%d %H:%M:%S'
 
+# Define a bunch of directory paths (same as bot_update.py)
+CURRENT_DIR = os.path.abspath(os.getcwd())
+BUILDER_DIR = os.path.dirname(CURRENT_DIR)
+SLAVE_DIR = os.path.dirname(BUILDER_DIR)
+# GOMA_CACHE_DIR used for caching long-term data.
+
+DEFAULT_GOMA_CACHE_DIR = os.path.join(SLAVE_DIR, 'goma_cache')
+# Path of the scripts/slave/ checkout on the slave, found by looking at the
+# current compile.py script's path's dirname().
+SLAVE_SCRIPTS_DIR = os.path.dirname(os.path.abspath(__file__))
+# Path of the build/ checkout on the slave, found relative to the
+# scripts/slave/ directory.
+BUILD_DIR = os.path.dirname(os.path.dirname(SLAVE_SCRIPTS_DIR))
+
+
 
 def GetShortHostname():
   """Get this machine's short hostname in lower case."""
@@ -355,3 +372,359 @@ def SendGomaTsMon(json_file, exit_status):
 
   except Exception as ex:
     print('error while sending ts mon json_file=%s: %s' % (json_file, ex))
+
+
+def UpdateWindowsEnvironment(envfile_dir, env):
+  """Update windows environment in environment.{x86,x64}.
+
+  Args:
+    envfile_dir: a directory name environment.{x86,x64} are stored.
+    env: an instance of EchoDict that represents environment.
+  """
+  # envvars_to_save come from _ExtractImportantEnvironment in
+  # https://chromium.googlesource.com/external/gyp/+/\
+  # master/pylib/gyp/msvs_emuation.py
+  # You must update this when the original code is updated.
+  envvars_to_save = (
+      'goma_.*', # TODO(scottmg): This is ugly, but needed for goma.
+      'include',
+      'lib',
+      'libpath',
+      'path',
+      'pathext',
+      'systemroot',
+      'temp',
+      'tmp',
+  )
+  env_to_store = {}
+  for envvar in envvars_to_save:
+    compiled = re.compile(envvar, re.IGNORECASE)
+    for key in env.overrides:
+      if compiled.match(key):
+        if envvar == 'path':
+          env_to_store[key] = (os.path.dirname(sys.executable) +
+                               os.pathsep + env[key])
+        else:
+          env_to_store[key] = env[key]
+
+  if not env_to_store:
+    return
+
+  nul = '\0'
+  for arch in ['x86', 'x64']:
+    path = os.path.join(envfile_dir, 'environment.%s' % arch)
+    print '%s will be updated with %s.' % (path, env_to_store)
+    env_in_file = {}
+    with open(path) as f:
+      for entry in f.read().split(nul):
+        if not entry:
+          continue
+        key, value = entry.split('=', 1)
+        env_in_file[key] = value
+    env_in_file.update(env_to_store)
+    with open(path, 'wb') as f:
+      f.write(nul.join(['%s=%s' % (k, v) for k, v in env_in_file.iteritems()]))
+      f.write(nul * 2)
+
+
+def goma_setup(options, env):
+  """Sets up goma if necessary.
+
+  If using the Goma compiler, first call goma_ctl  to ensure the proxy is
+  available, and returns (True, instance of cloudtail subprocess).
+  If it failed to start up compiler_proxy, modify options.compiler and
+  options.goma_dir and returns (False, None)
+
+  Args:
+    options (Option): options for goma
+    env (dict): env for compiler_proxy
+
+  Returns:
+    (bool, subprocess.Popen):
+      first element represents whether compiler_proxy starts successfully
+      second element is subprocess instance of cloudtail
+
+  """
+  if options.compiler not in ('goma', 'goma-clang'):
+    # Unset goma_dir to make sure we'll not use goma.
+    options.goma_dir = None
+    return False, None
+
+  hostname = GetShortHostname()
+  # HACK(shinyak, yyanagisawa, goma): Windows NO_NACL_GOMA (crbug.com/390764)
+  # Building NaCl untrusted code using goma brings large performance
+  # improvement but it sometimes cause build failure by race condition.
+  # Let me enable goma build on goma canary buildslaves to confirm the issue
+  # has been fixed by a workaround.
+  # vm*-m4 are trybots. build*-m1 and vm*-m1 are all goma canary bots.
+  if hostname in ['build28-m1', 'build58-m1', 'vm191-m1', 'vm480-m1',
+                  'vm820-m1', 'vm821-m1', 'vm848-m1']:
+    env['NO_NACL_GOMA'] = 'false'
+
+  if options.goma_fail_fast:
+    # startup fails when initial ping failed.
+    env['GOMA_FAIL_FAST'] = 'true'
+  else:
+    # If a network error continues 30 minutes, compiler_proxy make the compile
+    # failed.  When people use goma, they expect using goma is faster than
+    # compile locally. If goma cannot guarantee that, let it make compile
+    # as error.
+    env['GOMA_ALLOWED_NETWORK_ERROR_DURATION'] = '1800'
+
+  # HACK(yyanagisawa): reduce GOMA_BURST_MAX_PROCS crbug.com/592306
+  # Recently, I sometimes see buildbot slave time out, one possibility I come
+  # up with is burst mode use up resource.
+  # Let me temporary set small values to GOMA_BURST_MAX_PROCS to confirm
+  # the possibility is true or false.
+  max_subprocs = '3'
+  max_heavy_subprocs = '1'
+  number_of_processors = 0
+  try:
+    number_of_processors = multiprocessing.cpu_count()
+  except NotImplementedError:
+    print 'cpu_count() is not implemented, using default value.'
+    number_of_processors = 1
+  if number_of_processors > 3:
+    max_subprocs = str(number_of_processors - 1)
+    max_heavy_subprocs = str(number_of_processors / 2)
+  env['GOMA_BURST_MAX_SUBPROCS'] = max_subprocs
+  env['GOMA_BURST_MAX_SUBPROCS_LOW'] = max_subprocs
+  env['GOMA_BURST_MAX_SUBPROCS_HEAVY'] = max_heavy_subprocs
+
+  # Caches CRLs in GOMA_CACHE_DIR.
+  # Since downloading CRLs is usually slow, caching them may improves
+  # compiler_proxy start time.
+  if not os.path.exists(options.goma_cache_dir):
+    os.mkdir(options.goma_cache_dir, 0700)
+  env['GOMA_CACHE_DIR'] = options.goma_cache_dir
+
+  # Enable DepsCache. DepsCache caches the list of files to send goma server.
+  # This will greatly improve build speed when cache is warmed.
+  # The cache file is stored in the target output directory.
+  env['GOMA_DEPS_CACHE_DIR'] = (
+      options.goma_deps_cache_dir or options.target_output_dir)
+
+  if not env.get('GOMA_HERMETIC'):
+    env['GOMA_HERMETIC'] = options.goma_hermetic
+  if options.goma_enable_remote_link:
+    env['GOMA_ENABLE_REMOTE_LINK'] = options.goma_enable_remote_link
+  if options.goma_store_local_run_output:
+    env['GOMA_STORE_LOCAL_RUN_OUTPUT'] = options.goma_store_local_run_output
+  if options.goma_enable_compiler_info_cache:
+    # Will be stored in GOMA_CACHE_DIR.
+    env['GOMA_COMPILER_INFO_CACHE_FILE'] = 'goma-compiler-info.cache'
+
+  if options.build_data_dir:
+    env['GOMA_DUMP_STATS_FILE'] = os.path.join(options.build_data_dir,
+                                               'goma_stats_proto')
+
+  # goma is requested.
+  goma_key = os.path.join(options.goma_dir, 'goma.key')
+  if os.path.exists(goma_key):
+    env['GOMA_API_KEY_FILE'] = goma_key
+  if options.goma_service_account_json_file:
+    env['GOMA_SERVICE_ACCOUNT_JSON_FILE'] = \
+        options.goma_service_account_json_file
+  if chromium_utils.IsWindows():
+    env['GOMA_RPC_EXTRA_PARAMS'] = '?win'
+
+  # HACK(yyanagisawa): update environment files on |env| update.
+  # For compiling on Windows, environment in environment files are used.
+  # It means even if enviroment such as GOMA_DISABLED is updated in
+  # compile.py, the update will be ignored.
+  # We need to update environment files to reflect the update.
+  if chromium_utils.IsWindows() and NeedEnvFileUpdateOnWin(env):
+    print 'Updating environment.{x86,x64} files.'
+    UpdateWindowsEnvironment(options.target_output_dir, env)
+
+
+  goma_start_command = ['restart']
+  goma_ctl_cmd = [sys.executable,
+                  os.path.join(options.goma_dir, 'goma_ctl.py')]
+  result = chromium_utils.RunCommand(goma_ctl_cmd + goma_start_command, env=env)
+  if not result:
+    # goma started sucessfully.
+    # Making cloudtail to upload the latest log.
+    # TODO(yyanagisawa): install cloudtail from CIPD.
+    cloudtail_path = '/opt/infra-tools/cloudtail'
+    if chromium_utils.IsWindows():
+      cloudtail_path = 'C:\\infra-tools\\cloudtail'
+    try:
+      cloudtail_proc = subprocess.Popen(
+          [cloudtail_path, 'tail', '--log-id', 'goma_compiler_proxy', '--path',
+           GetLatestGomaCompilerProxyInfo()])
+    except Exception as e:
+      print 'failed to invoke cloudtail: %s' % e
+      return True, None
+    return True, cloudtail_proc
+
+  if options.goma_jsonstatus:
+    chromium_utils.RunCommand(
+        goma_ctl_cmd + ['jsonstatus', options.goma_jsonstatus], env=env)
+    SendGomaTsMon(options.goma_jsonstatus, -1)
+
+  # Try to stop compiler_proxy so that it flushes logs and stores
+  # GomaStats.
+  if options.build_data_dir:
+    env['GOMACTL_CRASH_REPORT_ID_FILE'] = os.path.join(options.build_data_dir,
+                                                       'crash_report_id_file')
+  chromium_utils.RunCommand(goma_ctl_cmd + ['stop'], env=env)
+
+  override_gsutil = None
+  if options.gsutil_py_path:
+    override_gsutil = [sys.executable, options.gsutil_py_path]
+
+  # Upload compiler_proxy.INFO to investigate the reason of compiler_proxy
+  # start-up failure.
+  UploadGomaCompilerProxyInfo(override_gsutil=override_gsutil)
+  # Upload GomaStats to make it monitored.
+  if env.get('GOMA_DUMP_STATS_FILE'):
+    SendGomaStats(env['GOMA_DUMP_STATS_FILE'],
+                  env.get('GOMACTL_CRASH_REPORT_ID_FILE'),
+                  options.build_data_dir)
+
+  if options.goma_disable_local_fallback:
+    print 'error: failed to start goma; fallback has been disabled'
+    raise Exception('failed to start goma')
+
+  print 'warning: failed to start goma. falling back to non-goma'
+  # Drop goma from options.compiler
+  options.compiler = options.compiler.replace('goma-', '')
+  if options.compiler == 'goma':
+    options.compiler = None
+  # Reset options.goma_dir.
+  options.goma_dir = None
+  env['GOMA_DISABLED'] = '1'
+  return False, None
+
+
+def goma_teardown(options, env, exit_status, cloudtail_proc):
+  """Tears down goma if necessary. """
+  if (options.compiler in ('goma', 'goma-clang') and
+      options.goma_dir):
+    override_gsutil = None
+    if options.gsutil_py_path:
+      override_gsutil = [sys.executable, options.gsutil_py_path]
+
+    # If goma compiler_proxy crashes during the build, there could be crash
+    # dump.
+    if options.build_data_dir:
+      env['GOMACTL_CRASH_REPORT_ID_FILE'] = os.path.join(options.build_data_dir,
+                                                         'crash_report_id_file')
+    goma_ctl_cmd = [sys.executable,
+                    os.path.join(options.goma_dir, 'goma_ctl.py')]
+    if options.goma_jsonstatus:
+      chromium_utils.RunCommand(
+          goma_ctl_cmd + ['jsonstatus', options.goma_jsonstatus], env=env)
+      SendGomaTsMon(options.goma_jsonstatus, exit_status)
+    # Always stop the proxy for now to allow in-place update.
+    chromium_utils.RunCommand(goma_ctl_cmd + ['stop'], env=env)
+    UploadGomaCompilerProxyInfo(override_gsutil=override_gsutil)
+    if env.get('GOMA_DUMP_STATS_FILE'):
+      SendGomaStats(env['GOMA_DUMP_STATS_FILE'],
+                    env.get('GOMACTL_CRASH_REPORT_ID_FILE'),
+                    options.build_data_dir)
+  if cloudtail_proc:
+    cloudtail_proc.terminate()
+    cloudtail_proc.wait()
+
+
+def determine_goma_jobs():
+  # We would like to speed up build on Windows a bit, since it is slowest.
+  number_of_processors = 0
+  try:
+    number_of_processors = multiprocessing.cpu_count()
+  except NotImplementedError:
+    print 'cpu_count() is not implemented, using default value 50.'
+    return 50
+
+  assert number_of_processors > 0
+
+  # When goma is used, 10 * number_of_processors is basically good in
+  # various situations according to our measurement. Build speed won't
+  # be improved if -j is larger than that.
+  #
+  # Since Mac had process number limitation before, we had to set
+  # the upper limit to 50. Now that the process number limitation is 2000,
+  # so we would be able to use 10 * number_of_processors.
+  # For the safety, we'd like to set the upper limit to 200.
+  #
+  # Note that currently most try-bot build slaves have 8 processors.
+  if chromium_utils.IsMac() or chromium_utils.IsWindows():
+    return min(10 * number_of_processors, 200)
+
+  # For Linux, we also would like to use 10 * cpu. However, not sure
+  # backend resource is enough, so let me set Linux and Linux x64 builder
+  # only for now.
+  hostname = GetShortHostname()
+  if hostname in (
+      ['build14-m1', 'build48-m1'] +
+      # Also increasing cpus for v8/blink trybots.
+      ['build%d-m4' % x for x in xrange(45, 48)] +
+      # Also increasing cpus for LTO buildbots.
+      ['slave%d-c1' % x for x in [20, 33] + range(78, 108)]):
+    return min(10 * number_of_processors, 200)
+
+  return 50
+
+
+def NeedEnvFileUpdateOnWin(env):
+  """Returns true if environment file need to be updated."""
+  # Following GOMA_* are applied to compiler_proxy not gomacc,
+  # you do not need to update environment files.
+  ignore_envs = (
+      'GOMA_API_KEY_FILE',
+      'GOMA_DEPS_CACHE_DIR',
+      'GOMA_HERMETIC',
+      'GOMA_RPC_EXTRA_PARAMS',
+      'GOMA_ALLOWED_NETWORK_ERROR_DURATION'
+  )
+  for key in env.overrides:
+    if key not in ignore_envs:
+      return True
+  return False
+
+
+def construct_optionparser(option_parser):
+  option_parser.add_option('--goma-cache-dir',
+                           default=DEFAULT_GOMA_CACHE_DIR,
+                           help='specify goma cache directory')
+  option_parser.add_option('--goma-deps-cache-dir',
+                           help='specify goma deps cache directory')
+  option_parser.add_option('--goma-hermetic', default='error',
+                           help='Set goma hermetic mode')
+  option_parser.add_option('--goma-enable-remote-link', default=None,
+                           help='Enable goma remote link.')
+  option_parser.add_option('--goma-enable-compiler-info-cache',
+                           action='store_true',
+                           help='Enable goma CompilerInfo cache')
+  option_parser.add_option('--goma-store-local-run-output', default=None,
+                           help='Store local run output to goma servers.')
+  option_parser.add_option('--goma-fail-fast', action='store_true')
+  option_parser.add_option('--goma-disable-local-fallback', action='store_true')
+  option_parser.add_option('--goma-jsonstatus',
+                           help='Specify a file to dump goma_ctl jsonstatus.')
+  option_parser.add_option('--goma-service-account-json-file',
+                           help='Specify a file containing goma service account'
+                           ' credentials')
+  option_parser.add_option('--build-data-dir', action='store',
+                           help='specify a build data directory.')
+
+  if not option_parser.has_option('--goma-dir'):
+    option_parser.add_option('--goma-dir',
+                             default=os.path.join(BUILD_DIR, 'goma'),
+                             help='specify goma directory')
+
+  if not option_parser.has_option('--goma-jobs'):
+    option_parser.add_option('--goma-jobs', default=None,
+                             help='The number of jobs for ninja -j.')
+
+  if not option_parser.has_option('--gsutil-py-path'):
+    option_parser.add_option('--gsutil-py-path',
+                             help='Specify path to gsutil.py script.')
+
+  if not option_parser.has_option('--compiler'):
+    option_parser.add_option('--compiler', default=None,
+                             help='specify alternative compiler (e.g. clang)')
+
+  return option_parser