utils/logging_utils.py

# Copyright 2015 The LUCI Authors. All rights reserved.
# Use of this source code is governed by the Apache v2.0 license that can be
# found in the LICENSE file.

"""Utility relating to logging."""

import argparse
import codecs
import ctypes
import logging
import logging.handlers
import optparse
import os
import sys
import tempfile
import time


# This works around file locking issue on Windows specifically in the case of
# long lived child processes.
#
# Python opens files with inheritable handle and without file sharing by
# default. This causes the RotatingFileHandler file handle to be duplicated in
# the subprocesses even if the log file is not used in it. Because of this
# handle in the child process, when the RotatingFileHandler tries to os.rename()
# the file in the parent process, it fails with:
#     WindowsError: [Error 32] The process cannot access the file because
#     it is being used by another process
if sys.platform == 'win32':
  import ctypes
  import msvcrt  # pylint: disable=F0401
  import _subprocess  # pylint: disable=F0401

  FILE_ATTRIBUTE_NORMAL = 0x00000080
  FILE_SHARE_READ = 1
  FILE_SHARE_WRITE = 2
  FILE_SHARE_DELETE = 4
  GENERIC_READ = 0x80000000
  GENERIC_WRITE = 0x40000000
  OPEN_ALWAYS = 4

  # TODO(maruel): Make it work in cygwin too if necessary. This would have to
  # use ctypes.cdll.kernel32 instead of _subprocess and msvcrt.


  def shared_open(path):
    """Opens a file with full sharing mode and without inheritance.

    The file is open for both read and write.

    See https://bugs.python.org/issue15244 for inspiration.
    """
    path = unicode(path)
    handle = ctypes.windll.kernel32.CreateFileW(
        path,
        GENERIC_READ|GENERIC_WRITE,
        FILE_SHARE_READ|FILE_SHARE_WRITE|FILE_SHARE_DELETE,
        None,
        OPEN_ALWAYS,
        FILE_ATTRIBUTE_NORMAL,
        None)
    ctr_handle = msvcrt.open_osfhandle(handle, os.O_BINARY | os.O_NOINHERIT)
    return os.fdopen(ctr_handle, 'r+b')


  class NoInheritRotatingFileHandler(logging.handlers.RotatingFileHandler):
    def _open(self):
      """Opens the log file without handle inheritance but with file sharing.

      Ignores self.mode.
      """
      f = shared_open(self.baseFilename)
      if self.encoding:
        # Do the equivalent of
        # codecs.open(self.baseFilename, self.mode, self.encoding)
        info = codecs.lookup(self.encoding)
        f = codecs.StreamReaderWriter(
            f, info.streamreader, info.streamwriter, 'replace')
        f.encoding = self.encoding
      return f


else:  # Not Windows.


  NoInheritRotatingFileHandler = logging.handlers.RotatingFileHandler


# Levels used for logging.
LEVELS = [logging.ERROR, logging.INFO, logging.DEBUG]


class CaptureLogs(object):
  """Captures all the logs in a context."""
  def __init__(self, prefix, root=None):
    handle, self._path = tempfile.mkstemp(prefix=prefix, suffix='.log')
    os.close(handle)
    self._handler = logging.FileHandler(self._path, 'w')
    self._handler.setLevel(logging.DEBUG)
    formatter = UTCFormatter(
        '%(process)d %(asctime)s: %(levelname)-5s %(message)s')
    self._handler.setFormatter(formatter)
    self._root = root or logging.getLogger()
    self._root.addHandler(self._handler)
    assert self._root.isEnabledFor(logging.DEBUG)

  def read(self):
    """Returns the current content of the logs.

    This also closes the log capture so future logs will not be captured.
    """
    self._disconnect()
    assert self._path
    try:
      with open(self._path, 'rb') as f:
        return f.read()
    except IOError as e:
      return 'Failed to read %s: %s' % (self._path, e)

  def close(self):
    """Closes and delete the log."""
    self._disconnect()
    if self._path:
      try:
        os.remove(self._path)
      except OSError as e:
        logging.error('Failed to delete log file %s: %s', self._path, e)
      self._path = None

  def __enter__(self):
    return self

  def __exit__(self, _exc_type, _exc_value, _traceback):
    self.close()

  def _disconnect(self):
    if self._handler:
      self._root.removeHandler(self._handler)
      self._handler.close()
      self._handler = None


class UTCFormatter(logging.Formatter):
  converter = time.gmtime

  def formatTime(self, record, datefmt=None):
    """Change is ',' to '.'."""
    ct = self.converter(record.created)
    if datefmt:
      return time.strftime(datefmt, ct)
    else:
      t = time.strftime("%Y-%m-%d %H:%M:%S", ct)
      return "%s.%03d" % (t, record.msecs)


class Filter(object):
  """Adds fields used by the infra-specific formatter.

  Fields added:
  - 'severity': one-letter indicator of log level (first letter of levelname).
  """

  def filter(self, record):
    record.severity = record.levelname[0]
    return True


def find_stderr(root=None):
  """Returns the logging.handler streaming to stderr, if any."""
  for log in (root or logging.getLogger()).handlers:
    if getattr(log, 'stream', None) is sys.stderr:
      return log


def prepare_logging(filename, root=None):
  """Prepare logging for scripts.

  Makes it log in UTC all the time. Prepare a rotating file based log.
  """
  assert not find_stderr(root)
  formatter = UTCFormatter('%(process)d %(asctime)s %(severity)s: %(message)s')

  # It is a requirement that the root logger is set to DEBUG, so the messages
  # are not lost. It defaults to WARNING otherwise.
  logger = root or logging.getLogger()
  if not logger:
    # Better print insanity than crash.
    print >> sys.stderr, 'OMG NO ROOT'
    return
  logger.setLevel(logging.DEBUG)

  stderr = logging.StreamHandler()
  stderr.setFormatter(formatter)
  stderr.addFilter(Filter())
  # Default to ERROR.
  stderr.setLevel(logging.ERROR)
  logger.addHandler(stderr)

  # Setup up logging to a constant file so we can debug issues where
  # the results aren't properly sent to the result URL.
  if filename:
    try:
      rotating_file = NoInheritRotatingFileHandler(
          filename, maxBytes=10 * 1024 * 1024, backupCount=5,
          encoding='utf-8')
      rotating_file.setLevel(logging.DEBUG)
      rotating_file.setFormatter(formatter)
      rotating_file.addFilter(Filter())
      logger.addHandler(rotating_file)
    except Exception:
      # May happen on cygwin. Do not crash.
      logging.exception('Failed to open %s', filename)


def set_console_level(level, root=None):
  """Reset the console (stderr) logging level."""
  handler = find_stderr(root)
  if not handler:
    # Better print insanity than crash.
    print >> sys.stderr, 'OMG NO STDERR'
    return
  handler.setLevel(level)


class OptionParserWithLogging(optparse.OptionParser):
  """Adds --verbose option."""

  # Set to True to enable --log-file options.
  enable_log_file = True

  # Set in unit tests.
  logger_root = None

  def __init__(self, verbose=0, log_file=None, **kwargs):
    kwargs.setdefault('description', sys.modules['__main__'].__doc__)
    optparse.OptionParser.__init__(self, **kwargs)
    self.group_logging = optparse.OptionGroup(self, 'Logging')
    self.group_logging.add_option(
        '-v', '--verbose',
        action='count',
        default=verbose,
        help='Use multiple times to increase verbosity')
    if self.enable_log_file:
      self.group_logging.add_option(
          '-l', '--log-file',
          default=log_file,
          help='The name of the file to store rotating log details')
      self.group_logging.add_option(
          '--no-log', action='store_const', const='', dest='log_file',
          help='Disable log file')

  def parse_args(self, *args, **kwargs):
    # Make sure this group is always the last one.
    self.add_option_group(self.group_logging)

    options, args = optparse.OptionParser.parse_args(self, *args, **kwargs)
    prepare_logging(self.enable_log_file and options.log_file, self.logger_root)
    set_console_level(
        LEVELS[min(len(LEVELS) - 1, options.verbose)], self.logger_root)
    return options, args


class ArgumentParserWithLogging(argparse.ArgumentParser):
  """Adds --verbose option."""

  # Set to True to enable --log-file options.
  enable_log_file = True

  def __init__(self, verbose=0, log_file=None, **kwargs):
    kwargs.setdefault('description', sys.modules['__main__'].__doc__)
    kwargs.setdefault('conflict_handler', 'resolve')
    self.__verbose = verbose
    self.__log_file = log_file
    super(ArgumentParserWithLogging, self).__init__(**kwargs)

  def _add_logging_group(self):
    group = self.add_argument_group('Logging')
    group.add_argument(
        '-v', '--verbose',
        action='count',
        default=self.__verbose,
        help='Use multiple times to increase verbosity')
    if self.enable_log_file:
      group.add_argument(
          '-l', '--log-file',
          default=self.__log_file,
          help='The name of the file to store rotating log details')
      group.add_argument(
          '--no-log', action='store_const', const='', dest='log_file',
          help='Disable log file')

  def parse_args(self, *args, **kwargs):
    # Make sure this group is always the last one.
    self._add_logging_group()

    args = super(ArgumentParserWithLogging, self).parse_args(*args, **kwargs)
    prepare_logging(self.enable_log_file and args.log_file, self.logger_root)
    set_console_level(
        LEVELS[min(len(LEVELS) - 1, args.verbose)], self.logger_root)
    return args