src/tools/swarming_client/utils/tools.py

   1 # Copyright 2013 The Swarming Authors. All rights reserved.
   2 # Use of this source code is governed under the Apache License, Version 2.0 that
   3 # can be found in the LICENSE file.
   4
   5 """Various utility functions and classes not specific to any single area."""
   6
   7 import cStringIO
   8 import json
   9 import logging
  10 import logging.handlers
  11 import optparse
  12 import os
  13 import re
  14 import sys
  15 import threading
  16 import time
  17
  18 import utils
  19 from utils import zip_package
  20
  21
  22 # Path to (possibly extracted from zip) cacert.pem bundle file.
  23 # See get_cacerts_bundle().
  24 _ca_certs = None
  25 _ca_certs_lock = threading.Lock()
  26
  27
  28 class OptionParserWithLogging(optparse.OptionParser):
  29   """Adds --verbose option."""
  30
  31   # Set to True to enable --log-file options.
  32   enable_log_file = True
  33
  34   def __init__(self, verbose=0, log_file=None, **kwargs):
  35     kwargs.setdefault('description', sys.modules['__main__'].__doc__)
  36     optparse.OptionParser.__init__(self, **kwargs)
  37     self.group_logging = optparse.OptionGroup(self, 'Logging')
  38     self.group_logging.add_option(
  39         '-v', '--verbose',
  40         action='count',
  41         default=verbose,
  42         help='Use multiple times to increase verbosity')
  43     if self.enable_log_file:
  44       self.group_logging.add_option(
  45           '-l', '--log-file',
  46           default=log_file,
  47           help='The name of the file to store rotating log details')
  48       self.group_logging.add_option(
  49           '--no-log', action='store_const', const='', dest='log_file',
  50           help='Disable log file')
  51
  52   def parse_args(self, *args, **kwargs):
  53     # Make sure this group is always the last one.
  54     self.add_option_group(self.group_logging)
  55
  56     options, args = optparse.OptionParser.parse_args(self, *args, **kwargs)
  57     levels = [logging.ERROR, logging.INFO, logging.DEBUG]
  58     level = levels[min(len(levels) - 1, options.verbose)]
  59
  60     logging_console = logging.StreamHandler()
  61     logging_console.setFormatter(logging.Formatter(
  62         '%(levelname)5s %(module)15s(%(lineno)3d): %(message)s'))
  63     logging_console.setLevel(level)
  64     logging.getLogger().setLevel(level)
  65     logging.getLogger().addHandler(logging_console)
  66
  67     if self.enable_log_file and options.log_file:
  68       # This is necessary otherwise attached handler will miss the messages.
  69       logging.getLogger().setLevel(logging.DEBUG)
  70
  71       logging_rotating_file = logging.handlers.RotatingFileHandler(
  72           options.log_file,
  73           maxBytes=10 * 1024 * 1024,
  74           backupCount=5,
  75           encoding='utf-8')
  76       # log files are always at DEBUG level.
  77       logging_rotating_file.setLevel(logging.DEBUG)
  78       logging_rotating_file.setFormatter(logging.Formatter(
  79           '%(asctime)s %(levelname)-8s %(module)15s(%(lineno)3d): %(message)s'))
  80       logging.getLogger().addHandler(logging_rotating_file)
  81
  82     return options, args
  83
  84
  85 class Profiler(object):
  86   """Context manager that records time spend inside its body."""
  87   def __init__(self, name):
  88     self.name = name
  89     self.start_time = None
  90
  91   def __enter__(self):
  92     self.start_time = time.time()
  93     return self
  94
  95   def __exit__(self, _exc_type, _exec_value, _traceback):
  96     time_taken = time.time() - self.start_time
  97     logging.info('Profiling: Section %s took %3.3f seconds',
  98                  self.name, time_taken)
  99
 100
 101 class Unbuffered(object):
 102   """Disable buffering on a file object."""
 103   def __init__(self, stream):
 104     self.stream = stream
 105
 106   def write(self, data):
 107     self.stream.write(data)
 108     if '\n' in data:
 109       self.stream.flush()
 110
 111   def __getattr__(self, attr):
 112     return getattr(self.stream, attr)
 113
 114
 115 def disable_buffering():
 116   """Makes this process and child processes stdout unbuffered."""
 117   if not os.environ.get('PYTHONUNBUFFERED'):
 118     # Since sys.stdout is a C++ object, it's impossible to do
 119     # sys.stdout.write = lambda...
 120     sys.stdout = Unbuffered(sys.stdout)
 121     os.environ['PYTHONUNBUFFERED'] = 'x'
 122
 123
 124 def fix_python_path(cmd):
 125   """Returns the fixed command line to call the right python executable."""
 126   out = cmd[:]
 127   if out[0] == 'python':
 128     out[0] = sys.executable
 129   elif out[0].endswith('.py'):
 130     out.insert(0, sys.executable)
 131   return out
 132
 133
 134 def read_json(filepath):
 135   with open(filepath, 'r') as f:
 136     return json.load(f)
 137
 138
 139 def write_json(filepath_or_handle, data, dense):
 140   """Writes data into filepath or file handle encoded as json.
 141
 142   If dense is True, the json is packed. Otherwise, it is human readable.
 143   """
 144   if dense:
 145     kwargs = {'sort_keys': True, 'separators': (',',':')}
 146   else:
 147     kwargs = {'sort_keys': True, 'indent': 2}
 148
 149   if hasattr(filepath_or_handle, 'write'):
 150     json.dump(data, filepath_or_handle, **kwargs)
 151   else:
 152     with open(filepath_or_handle, 'wb') as f:
 153       json.dump(data, f, **kwargs)
 154
 155
 156 def format_json(data, dense):
 157   """Returns a string with json encoded data.
 158
 159   If dense is True, the json is packed. Otherwise, it is human readable.
 160   """
 161   buf = cStringIO.StringIO()
 162   write_json(buf, data, dense)
 163   return buf.getvalue()
 164
 165
 166 def gen_blacklist(regexes):
 167   """Returns a lambda to be used as a blacklist."""
 168   compiled = [re.compile(i) for i in regexes]
 169   return lambda f: any(j.match(f) for j in compiled)
 170
 171
 172 def get_bool_env_var(name):
 173   """Return True if integer environment variable |name| value is non zero.
 174
 175   If environment variable is missing or is set to '0', returns False.
 176   """
 177   return bool(int(os.environ.get(name, '0')))
 178
 179
 180 def is_headless():
 181   """True if running in non-interactive mode on some bot machine.
 182
 183   Examines os.environ for presence of SWARMING_HEADLESS var.
 184   """
 185   headless_env_keys = (
 186     # This is Chromium specific. Set when running under buildbot slave.
 187     'CHROME_HEADLESS',
 188     # Set when running under swarm bot.
 189     'SWARMING_HEADLESS',
 190   )
 191   return any(get_bool_env_var(key) for key in headless_env_keys)
 192
 193
 194 def get_cacerts_bundle():
 195   """Returns path to a file with CA root certificates bundle.
 196
 197   Python's ssl module needs a real file on disk, so if code is running from
 198   a zip archive, we need to extract the file first.
 199   """
 200   global _ca_certs
 201   with _ca_certs_lock:
 202     if _ca_certs is not None and os.path.exists(_ca_certs):
 203       return _ca_certs
 204     # Some rogue process clears /tmp and causes cacert.pem to disappear. Extract
 205     # to current directory instead. We use our own bundled copy of cacert.pem.
 206     _ca_certs = zip_package.extract_resource(utils, 'cacert.pem', temp_dir='.')
 207     return _ca_certs