logs.py

# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Set up for logging."""
from enum import Enum

import logging
import os
import sys
import time
import traceback

import google.cloud.logging
from google.cloud.logging_v2.handlers.handlers import CloudLoggingHandler
from google.cloud import error_reporting

# Disable this check since we have a bunch of non-constant globals in this file.
# pylint: disable=invalid-name

from common import utils

_default_logger = None
_log_client = None
_error_reporting_client = None
_default_extras = {}

LOG_LENGTH_LIMIT = 250 * 1000

NUM_ATTEMPTS = 4
RETRY_DELAY = 1
BACKOFF = 2


def _initialize_cloud_clients():
    """Initialize clients for Google Cloud Logging and Error reporting."""
    assert not utils.is_local()
    global _log_client
    if _log_client:
        return
    _log_client = google.cloud.logging.Client()
    logging_handler = CloudLoggingHandler(_log_client)
    logging.getLogger().addHandler(logging_handler)
    global _error_reporting_client
    _error_reporting_client = error_reporting.Client()


def initialize(name='fuzzbench', default_extras=None, log_level=logging.INFO):
    """Initializes stackdriver logging if running on Google Cloud."""
    logging.getLogger().setLevel(log_level)
    logging.getLogger().addFilter(LengthFilter())

    # Don't log so much with SQLalchemy to avoid stressing the logging library.
    # See crbug.com/1044343.
    logging.getLogger('sqlalchemy').setLevel(logging.ERROR)

    if utils.is_local():
        return
    _initialize_cloud_clients()
    global _default_logger
    _default_logger = _log_client.logger(name)

    default_extras = {} if default_extras is None else default_extras

    _set_instance_name(default_extras)
    _set_experiment(default_extras)

    # pylint: disable=global-variable-not-assigned
    global _default_extras
    _default_extras.update(default_extras)


def _set_instance_name(extras: dict):
    """Set instance_name in |extras| if it is provided by the environment and
    not already set."""
    if 'instance_name' in extras:
        return

    instance_name = os.getenv('INSTANCE_NAME')
    if instance_name is None:
        return

    extras['instance_name'] = instance_name


def _set_experiment(extras: dict):
    """Set experiment in |extras| if it is provided by the environment and
    not already set."""
    if 'experiment' in extras:
        return

    experiment = os.getenv('EXPERIMENT')
    if experiment is None:
        return

    extras['experiment'] = experiment


class Logger:
    """Wrapper around logging.Logger that allows it to be used like we use the
    root logger for stackdriver."""

    _LOGGER_NAME = 'fuzzbench'

    def __init__(self, default_extras=None, log_level=logging.INFO):
        if not utils.is_local():
            _initialize_cloud_clients()
            self.logger = _log_client.logger(self._LOGGER_NAME)
        else:
            self.logger = logging.getLogger(self._LOGGER_NAME)

        logging.getLogger(self._LOGGER_NAME).setLevel(log_level)
        logging.getLogger(self._LOGGER_NAME).addFilter(LengthFilter())
        self.default_extras = default_extras if default_extras else {}

    def error(self, *args, **kwargs):
        """Wrapper that uses _log_function_wrapper to call error."""
        self._log_function_wrapper(error, *args, **kwargs)

    def warning(self, *args, **kwargs):
        """Wrapper that uses _log_function_wrapper to call warning."""
        self._log_function_wrapper(warning, *args, **kwargs)

    def info(self, *args, **kwargs):
        """Wrapper that uses _log_function_wrapper to call info."""
        self._log_function_wrapper(info, *args, **kwargs)

    def debug(self, *args, **kwargs):
        """Wrapper that uses _log_function_wrapper to call debug."""
        self._log_function_wrapper(debug, *args, **kwargs)

    def _log_function_wrapper(self, log_function, message, *args, extras=None):
        """Wrapper around log functions that passes extras and the logger this
        object wraps (self.logger)."""
        extras = {} if extras is None else extras
        extras = extras.copy()
        extras.update(self.default_extras)
        log_function(message, *args, extras=extras, logger=self.logger)


class LogSeverity(Enum):
    """Enum for different levels of log severity."""
    ERROR = logging.ERROR
    WARNING = logging.WARNING
    INFO = logging.INFO
    DEBUG = logging.DEBUG


def log(logger, severity, message, *args, extras=None):
    """Log a message with severity |severity|. If using stack driver logging
    then |extras| is also logged (in addition to default extras)."""
    # Custom retry logic to avoid circular dependency as retry from
    # retry.py uses log.
    for num_try in range(1, NUM_ATTEMPTS + 1):
        try:
            message = str(message)
            if args:
                message = message % args

            if utils.is_local():
                if extras:
                    message += ' Extras: ' + str(extras)
                logging.log(severity, message)
                return

            if logger is None:
                logger = _default_logger
            assert logger

            struct_message = {
                'message': message,
            }
            all_extras = _default_extras.copy()
            extras = extras or {}
            all_extras.update(extras)
            struct_message.update(all_extras)
            severity = LogSeverity(severity).name
            logger.log_struct(struct_message, severity=severity)
            break
        except Exception:  # pylint: disable=broad-except
            # We really dont want do to do anything here except sleep here,
            # since we cant log it out as log itself is already failing
            time.sleep(utils.get_retry_delay(num_try, RETRY_DELAY, BACKOFF))


def error(message, *args, extras=None, logger=None):
    """Logs |message| to stackdriver logging and error reporting (including
    exception if there was one."""

    def _report_error_with_retries(message):
        if utils.is_local():
            return

        # Custom retry logic to avoid circular dependency
        # as retry from retry.py uses log
        for num_try in range(1, NUM_ATTEMPTS + 1):
            try:
                _error_reporting_client.report(message)
                break
            except Exception:  # pylint: disable=broad-except
                # We really dont want do to do anything here except sleep here,
                # since we cant log it out as log itself is already failing
                time.sleep(utils.get_retry_delay(num_try, RETRY_DELAY, BACKOFF))

    if not any(sys.exc_info()):
        _report_error_with_retries(message % args)
        log(logger, logging.ERROR, message, *args, extras=extras)
        return
    # I can't figure out how to include both the message and the exception
    # other than this having the exception message preceed the log message
    # (without using private APIs).
    _report_error_with_retries(traceback.format_exc() + '\nMessage: ' +
                               message % args)
    extras = {} if extras is None else extras
    extras['traceback'] = traceback.format_exc()
    log(logger, logging.ERROR, message, *args, extras=extras)


def warning(message, *args, extras=None, logger=None):
    """Log a message with severity 'WARNING'."""
    log(logger, logging.WARNING, message, *args, extras=extras)


def info(message, *args, extras=None, logger=None):
    """Log a message with severity 'INFO'."""
    log(logger, logging.INFO, message, *args, extras=extras)


def debug(message, *args, extras=None, logger=None):
    """Log a message with severity 'DEBUG'."""
    log(logger, logging.DEBUG, message, *args, extras=extras)


class LengthFilter(logging.Filter):
    """Filter for truncating log messages that are too long for stackdriver."""

    def filter(self, record):
        if len(record.msg) > LOG_LENGTH_LIMIT:
            record.msg = ('TRUNCATED: ' + record.msg)[:LOG_LENGTH_LIMIT]
        return True