Agent/examples/auto_put_ad_mini/logging_setup.py

"""auto_put_ad_mini 生产任务共用的日志配置。

这些任务也会作为独立脚本执行，因此本模块不依赖 ``agent.utils.logging``。
日志始终输出到控制台；配置阿里云 SLS 后，普通日志与错误日志分别写入对应
Logstore，不再额外写本地日志文件。
"""

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
_TRACE_ID_ENV = "AUTO_PUT_TRACE_ID"


def get_trace_id() -> str | None:
    """返回当前进程的 trace_id；日志尚未初始化时返回 None。"""
    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):
    """向每条日志注入 ``trace_id``，供 SLS 按一次任务执行链路检索。

    过滤器从模块级 ``_TRACE_ID`` 读取值，因此先挂载过滤器、后执行
    ``setup_logging()`` 的调用顺序也能正常工作。
    """

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


class _SafeQueuedLogHandler:
    """创建不会因 SLS 发送失败而中断业务的 SDK handler。

    同时移除 SLS protobuf 中的 ``Time_ns`` 可选字段，避免控制台时间戳显示
    无意义的九位纳秒尾数，只保留秒级时间。
    """

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

        # 包装 emit，SLS 网络异常只进入 logging 自身错误处理，不影响业务线程。
        original_emit = handler.emit

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

        handler.emit = safe_emit

        # SDK 使用 ``time.time()`` 填充 LogItem.time_nano_part，而不是使用
        # ``record.created``，会生成随机的九位尾数。即使强制设为 0，该可选字段
        # 仍会被序列化并显示 ``.000000000``。这里在每批日志序列化前临时清除
        # ``Time_ns``，避免复制 SDK 内部约 40 行的发送实现。
        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:
    """保留终端原始输出，同时把 stdout/stderr 镜像到统一日志。"""

    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 自带 __time__ 字段，级别、模块、文件、行号和线程也已通过 fields
    # 单独上报，因此正文只保留 message，避免控制台出现重复字段。
    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:
    """幂等挂载本项目的 SLS handlers，并返回是否已启用 SLS。"""
    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:
    """配置根日志：控制台加可选 SLS，不写本地文件。"""
    global _CAPTURE_INSTALLED, _CONFIGURED, _TRACE_ID

    # 调度器启动的子进程继承任务专属 trace_id；独立运行的脚本自行生成。
    if _TRACE_ID is None:
        _TRACE_ID = _inherited_trace_id() or _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 失败（例如 SLS 网络中断）不能递归向 stderr 打印异常，更不能中断业务。
    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:
    """为一次进程执行生成紧凑、可排序且近似唯一的 trace_id。

    格式为 ``YYYYMMDD-HHMMSS-{8位随机字符}``，足以区分定时任务的并发执行。
    """
    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}"


def _inherited_trace_id() -> str | None:
    """读取并校验调度器传入的 trace_id；不存在或不安全时返回 None。"""

    raw = os.getenv(_TRACE_ID_ENV, "").strip()
    if not raw or len(raw) > 160:
        return None
    allowed_punctuation = frozenset("-_.:")
    if not all(
        character.isalnum() or character in allowed_punctuation
        for character in raw
    ):
        return None
    return raw