Agent/examples/auto_put_ad_mini/logging_setup.py

"""Logging setup owned by the auto_put_ad_mini example.

The example is also run as standalone scripts, so this module deliberately
does not depend on ``agent.utils.logging``.  Logs go to the console and,
when configured, to separate Alibaba Cloud SLS Logstores (INFO/WARNING and
ERROR/CRITICAL).  No local file logging is performed.
"""

from __future__ import annotations

import atexit
import logging
import os
import random
import string
import sys
import threading
from datetime import datetime
from typing import Iterable, Optional


_FORMAT = "%Y-%m-%d %H:%M:%S"
_LOCAL_LOG_FORMAT = "%(asctime)s | %(levelname)s | %(name)s | %(message)s"
_SLS_LOG_FORMAT = "%(message)s"
_SLS_HANDLERS: list[logging.Handler] = []
_CAPTURE_INSTALLED = False
_CONFIGURED = False
_EXCEPTION_HOOKS_INSTALLED = False
_ATEXIT_REGISTERED = False
_TRACE_ID: str | None = None


def get_trace_id() -> str | None:
    """Return the trace_id of the current session, or None if not yet initialised."""
    return _TRACE_ID


class _LevelFilter(logging.Filter):
    def __init__(self, minimum: int, maximum: Optional[int] = None):
        super().__init__()
        self.minimum = minimum
        self.maximum = maximum

    def filter(self, record: logging.LogRecord) -> bool:
        return record.levelno >= self.minimum and (
            self.maximum is None or record.levelno <= self.maximum
        )


class _ExcludeCapturedOutput(logging.Filter):
    def filter(self, record: logging.LogRecord) -> bool:
        return not getattr(record, "captured_output", False)


class _TraceIdFilter(logging.Filter):
    """Inject ``trace_id`` into every LogRecord for structured SLS queries.

    The filter reads from the module-level ``_TRACE_ID`` so that it works even
    when ``setup_logging()`` is called after the filter is added.
    """

    def filter(self, record: logging.LogRecord) -> bool:
        record.trace_id = _TRACE_ID or "-"
        return True


class _SafeQueuedLogHandler:
    """Create an SDK handler whose emit failures cannot break the app logger,
    and whose timestamps are forced to second-level precision.

    The SLS protobuf schema includes ``optional fixed32 Time_ns`` — when set
    (even to 0) the console renders ``.000000000``.  We strip it so only
    integer-second timestamps appear.
    """

    @staticmethod
    def create(handler_cls, **kwargs) -> logging.Handler:
        handler = handler_cls(**kwargs)

        # -- safe emit ----------------------------------------------------
        original_emit = handler.emit

        def safe_emit(record):
            try:
                original_emit(record)
            except Exception:
                logging.Handler.handleError(handler, record)

        handler.emit = safe_emit

        # -- strip Time_ns from protobuf before serialization -------------
        # The SDK populates LogItem.time_nano_part from ``time.time()``
        # rather than ``record.created``, injecting a random 9-digit value.
        # Even when forced to 0 the optional protobuf field *is* serialised
        # (fixed32 0 ≠ absent) and the console renders ``.000000000``.
        # Monkey-patching ``LogGroup.SerializeToString`` (called once per
        # batch inside ``LogClient.put_logs``) clears the field right before
        # encoding without duplicating the ~40-line SDK method.
        import aliyun.log.log_logs_pb2 as _sls_pb

        _orig_serialize = _sls_pb.LogGroup.SerializeToString
        original_send = handler.send

        def patched_send(req):
            def _no_nano_serialize(self):
                for log in self.Logs:
                    if log.HasField("Time_ns"):
                        log.ClearField("Time_ns")
                return _orig_serialize(self)

            _sls_pb.LogGroup.SerializeToString = _no_nano_serialize
            try:
                return original_send(req)
            finally:
                _sls_pb.LogGroup.SerializeToString = _orig_serialize

        handler.send = patched_send
        return handler


class _CapturedStream:
    """Mirror a process stream into logging while preserving interactive output."""

    def __init__(self, original, level: int, logger_name: str):
        self._original = original
        self._level = level
        self._logger = logging.getLogger(logger_name)
        self._buffer = ""

    def write(self, value: str) -> int:
        if not value:
            return 0
        self._original.write(value)
        self._original.flush()
        self._buffer += value
        while "\n" in self._buffer:
            line, self._buffer = self._buffer.split("\n", 1)
            line = line.rstrip("\r")
            if line.strip():
                self._logger.log(
                    self._level,
                    "%s",
                    line,
                    extra={"captured_output": True},
                )
        return len(value)

    def flush(self) -> None:
        self._original.flush()
        if self._buffer.strip():
            self._logger.log(
                self._level,
                "%s",
                self._buffer.strip(),
                extra={"captured_output": True},
            )
            self._buffer = ""

    def isatty(self) -> bool:
        return self._original.isatty()

    def fileno(self) -> int:
        return self._original.fileno()

    @property
    def encoding(self):
        return getattr(self._original, "encoding", "utf-8")


def _configured_sls_handlers(log_level: int) -> Iterable[logging.Handler]:
    try:
        from aliyun.log import QueuedLogHandler
    except ImportError:
        logging.getLogger(__name__).warning(
            "[sls] aliyun-log-python-sdk 未安装，跳过 SLS 上报"
        )
        return []

    sls_endpoint = os.getenv("SLS_ENDPOINT", "").strip()
    access_key_id = os.getenv("SLS_ACCESS_KEY_ID", "").strip()
    access_key_secret = os.getenv("SLS_ACCESS_KEY_SECRET", "").strip()
    project = os.getenv("SLS_PROJECT", "auto-put-tecent").strip()
    info_logstore = os.getenv(
        "SLS_INFO_LOGSTORE", os.getenv("SLS_LOGSTORE", "info-log")
    ).strip()
    error_logstore = os.getenv("SLS_ERROR_LOGSTORE", "error-log").strip()
    batch_size = int(os.getenv("SLS_BATCH_SIZE_MAX", "1024"))
    put_wait_ms = int(os.getenv("SLS_PUT_WAIT_MS", "2000"))
    sls_level = getattr(
        logging, os.getenv("SLS_LOG_LEVEL", "INFO").upper(), logging.INFO
    )

    if not all(
        (sls_endpoint, access_key_id, access_key_secret,
         project, info_logstore, error_logstore)
    ):
        logging.getLogger(__name__).warning(
            "[sls] SLS 凭证或 Logstore 未完整配置，跳过上报"
        )
        return []

    # SLS uses its own __time__ field, and level/name/func/file/line/thread are
    # already sent as separate fields via the ``fields`` list. Keep only the
    # message body to avoid redundant data cluttering the SLS console.
    sls_formatter = logging.Formatter(_SLS_LOG_FORMAT)
    common = dict(
        end_point=sls_endpoint,
        access_key_id=access_key_id,
        access_key=access_key_secret,
        project=project,
        fields=[
            "record_name", "level", "func_name", "module", "file_path", "line_no",
            "process_id", "process_name", "thread_id", "thread_name",
        ],
        extract_kv=True,
        batch_size=batch_size,
        put_wait=max(float(put_wait_ms) / 1000, 0.1),
    )
    handlers = []
    for store, level_filter in (
        (info_logstore, _LevelFilter(sls_level, logging.WARNING)),
        (error_logstore, _LevelFilter(logging.ERROR)),
    ):
        handler = _SafeQueuedLogHandler.create(
            QueuedLogHandler, log_store=store, **common
        )
        handler.setLevel(sls_level if store == info_logstore else logging.ERROR)
        handler.addFilter(level_filter)
        handler.setFormatter(sls_formatter)
        handlers.append(handler)
    return handlers


def attach_sls_handler(root_logger: Optional[logging.Logger] = None) -> bool:
    """Attach the example-owned SLS handlers once, returning whether enabled."""
    global _SLS_HANDLERS
    root_logger = root_logger or logging.getLogger()
    if _SLS_HANDLERS and all(handler in root_logger.handlers for handler in _SLS_HANDLERS):
        return True
    # 清理残留的旧 SLS handler（如 setup_logging 被多次调用导致部分 handler 残留）
    for old_handler in list(_SLS_HANDLERS):
        if old_handler in root_logger.handlers:
            root_logger.removeHandler(old_handler)
            try:
                old_handler.close()
            except Exception:
                pass
    _SLS_HANDLERS = []
    try:
        handlers = list(_configured_sls_handlers(root_logger.level or logging.INFO))
        for handler in handlers:
            root_logger.addHandler(handler)
        _SLS_HANDLERS = handlers
        if handlers:
            logging.getLogger(__name__).info(
                "[sls] 已挂载双 Logstore: project=%s info-log/error-log",
                os.getenv("SLS_PROJECT", "auto-put-tecent"),
            )
        return bool(handlers)
    except Exception:
        logging.getLogger(__name__).exception("[sls] handler 挂载失败，降级为本地日志")
        return False


def _close_sls_handlers() -> None:
    for handler in _SLS_HANDLERS:
        try:
            handler.close()
        except Exception:
            pass


def _install_exception_hooks() -> None:
    global _EXCEPTION_HOOKS_INSTALLED
    if _EXCEPTION_HOOKS_INSTALLED:
        return

    original_sys_hook = sys.excepthook

    def log_uncaught_exception(exc_type, exc_value, exc_traceback) -> None:
        if issubclass(exc_type, KeyboardInterrupt):
            original_sys_hook(exc_type, exc_value, exc_traceback)
            return
        logging.getLogger("auto_put_ad_mini.uncaught").critical(
            "event=uncaught_exception thread=main result=failed",
            exc_info=(exc_type, exc_value, exc_traceback),
        )

    def log_thread_exception(args: threading.ExceptHookArgs) -> None:
        logging.getLogger("auto_put_ad_mini.uncaught").critical(
            "event=uncaught_exception thread=%s result=failed",
            args.thread.name if args.thread else "unknown",
            exc_info=(args.exc_type, args.exc_value, args.exc_traceback),
        )

    sys.excepthook = log_uncaught_exception
    threading.excepthook = log_thread_exception
    _EXCEPTION_HOOKS_INSTALLED = True


def setup_logging(
    level: str = "INFO",
    capture_output: bool = True,
) -> None:
    """Configure root logging: console + optional SLS, no local file."""
    global _CAPTURE_INSTALLED, _CONFIGURED, _TRACE_ID

    # Generate a per-session trace_id once per process lifetime.
    if _TRACE_ID is None:
        _TRACE_ID = _generate_trace_id()

    log_level = getattr(logging, str(level).upper(), logging.INFO)
    stdout = sys.__stdout__
    stderr = sys.__stderr__
    trace_filter = _TraceIdFilter()
    console = logging.StreamHandler(stdout)
    console.setLevel(log_level)
    console.addFilter(_ExcludeCapturedOutput())
    console.addFilter(trace_filter)
    formatter = logging.Formatter(_LOCAL_LOG_FORMAT, datefmt=_FORMAT)
    console.setFormatter(formatter)

    root_logger = logging.getLogger()
    # Handler failures (for example an SLS network outage) must never write
    # recursive logging tracebacks to stderr or interrupt the business flow.
    logging.raiseExceptions = False
    if not _CONFIGURED:
        logging.basicConfig(
            level=log_level,
            handlers=[console],
            force=True,
        )
        _CONFIGURED = True
    else:
        root_logger.setLevel(log_level)
    for noisy in ("httpx", "httpcore", "urllib3", "apscheduler", "aliyun.log"):
        logging.getLogger(noisy).setLevel(logging.WARNING)

    attach_sls_handler(root_logger)
    _install_exception_hooks()
    if capture_output and not _CAPTURE_INSTALLED:
        sys.stdout = _CapturedStream(stdout, logging.INFO, "auto_put_ad_mini.stdout")
        sys.stderr = _CapturedStream(stderr, logging.ERROR, "auto_put_ad_mini.stderr")
        _CAPTURE_INSTALLED = True
    global _ATEXIT_REGISTERED
    if not _ATEXIT_REGISTERED:
        atexit.register(_close_sls_handlers)
        _ATEXIT_REGISTERED = True
    logging.getLogger(__name__).info("[logging] 日志初始化完成 trace_id=%s", _TRACE_ID)


def _generate_trace_id() -> str:
    """Generate a compact, sortable, unique trace_id for one process invocation.

    Format: ``YYYYMMDD-HHMMSS-{8 random hex chars}``.
    ~16M combinations per second; sufficient for cron-triggered processes.
    """
    now = datetime.now()
    ts = now.strftime("%Y%m%d-%H%M%S")
    suffix = "".join(random.choices(string.ascii_lowercase + string.digits, k=8))
    return f"{ts}-{suffix}"