在Python多线程上下文中监控函数调用

import time
import threading
from dataclasses import dataclass
from collections import UserList # 用于后续的LocalList

@dataclass
class MonitorRecord:
    function: str
    time: float

class MonitorContext:
    """
    监控上下文管理器。
    当进入上下文时，注册自身到全局处理器；退出时，注销。
    """
    def __init__(self):
        self._records: list[MonitorRecord] = []

    def add_record(self, record: MonitorRecord) -> None:
        """向当前上下文添加一条监控记录。"""
        self._records.append(record)

    def __enter__(self) -> 'MonitorContext':
        """进入上下文时，将当前上下文实例注册到全局处理器。"""
        handlers.register(self)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """退出上下文时，从全局处理器中删除当前上下文实例。"""
        handlers.delete(self)
        return

def monitor_decorator(f):
    """
    函数监控装饰器。
    记录被装饰函数的执行时间，并将记录添加到所有当前活跃的监控上下文中。
    """
    def _(*args, **kwargs):
        start = time.time()
        result = f(*args, **kwargs) # 执行原始函数
        handlers.add_record(
            MonitorRecord(
                function=f.__name__,
                time=time.time() - start,
            )
        )
        return result
    return _

# 初始的 MonitorHandlers (存在多线程问题)
class MonitorHandlersInitial:
    def __init__(self):
        self._handlers: list[MonitorContext] = []

    def register(self, handler: MonitorContext) -> None:
        self._handlers.append(handler)

    def delete(self, handler: MonitorContext) -> None:
        self._handlers.remove(handler)

    def add_record(self, record: MonitorRecord) -> None:
        for h in self._handlers:
            h.add_record(record)

handlers = MonitorHandlersInitial() # 全局实例

# 示例：单线程环境
# @monitor_decorator
# def run_single_thread():
#     time.sleep(0.1)

# with MonitorContext() as m1:
#     run_single_thread()
#     with MonitorContext() as m2:
#         run_single_thread()
#         run_single_thread()

# print(f"m1 records: {len(m1._records)}") # 预期输出 3
# print(f"m2 records: {len(m2._records)}") # 预期输出 2

# @monitor_decorator
# def run_multithread():
#     time.sleep(0.1)

# def nested_multithread_context():
#     with MonitorContext() as m:
#         run_multithread()
#     print(f"Thread {threading.get_ident()} local context records: {len(m._records)}")

# with MonitorContext() as m_main:
#     threads = [threading.Thread(target=nested_multithread_context) for _ in range(3)]
#     [t.start() for t in threads]
#     [t.join() for t in threads]

# print(f"Main context records: {len(m_main._records)}")

class LocalList(threading.local, UserList):
    """
    一个线程本地的列表。
    每个线程都会有自己的独立列表实例。
    UserList 提供列表的基本行为。
    """
    pass

class MonitorHandlers:
    """
    线程安全的监控处理器。
    使用 threading.local 为每个线程维护独立的上下文列表。
    使用 threading.Lock 保护主线程的上下文列表和共享操作。
    """
    def __init__(self):
        self._lock = threading.Lock() # 用于保护 _mainhandlers 的并发访问
        with self._lock:
            # _mainhandlers 存储主线程注册的上下文，由锁保护
            self._mainhandlers: list[MonitorContext] = []
        # _handlers 存储非主线程注册的上下文，每个线程拥有独立的实例
        self._handlers: list[MonitorContext] = LocalList()

    def register(self, handler: MonitorContext) -> None:
        """
        注册一个监控上下文。
        主线程的上下文添加到 _mainhandlers (加锁)。
        非主线程的上下文添加到 _handlers (线程本地)。
        """
        if threading.main_thread().ident == threading.get_ident():
            # 当前是主线程
            with self._lock:
                self._mainhandlers.append(handler)
        else:
            # 当前是非主线程
            self._handlers.append(handler)

    def delete(self, handler: MonitorContext) -> None:
        """
        删除一个监控上下文。
        与 register 逻辑类似，区分主线程和非主线程。
        """
        if threading.main_thread().ident == threading.get_ident():
            with self._lock:
                if handler in self._mainhandlers: # 检查是否存在以避免ValueError
                    self._mainhandlers.remove(handler)
        else:
            if handler in self._handlers: # 检查是否存在以避免ValueError
                self._handlers.remove(handler)

    def add_record(self, record: MonitorRecord) -> None:
        """
        向所有活跃的监控上下文添加记录。
        遍历当前线程的 _handlers，并遍历主线程的 _mainhandlers (加锁)。
        """
        # 向当前线程的上下文添加记录
        for h in self._handlers:
            h.add_record(record)

        # 向主线程的上下文添加记录 (需要加锁保护访问 _mainhandlers)
        with self._lock:
            for h in self._mainhandlers:
                h.add_record(record)

# 替换全局的 handlers 实例为新的线程安全版本
handlers = MonitorHandlers()

@monitor_decorator
def run_multithread_fixed():
    """一个简单的模拟耗时函数。"""
    time.sleep(0.1)

def nested_multithread_context_fixed():
    """
    在子线程中创建上下文并执行监控函数。
    注意：这里的 `m` 是每个子线程自己独立的 `MonitorContext` 实例。
    """
    with MonitorContext() as m:
        run_multithread_fixed()
    print(f"Thread {threading.get_ident()} local context records: {len(m._records)}")

if __name__ == "__main__":
    print("--- 单线程示例 (验证基础功能) ---")
    @monitor_decorator
    def run_single_thread_fixed():
        time.sleep(0.05)

    with MonitorContext() as m_s1:
        run_single_thread_fixed()
        with MonitorContext() as m_s2:
            run_single_thread_fixed()
            run_single_thread_fixed()

    print(f"Single-thread m_s1 records: {len(m_s1._records)}")
    print(f"Single-thread m_s2 records: {len(m_s2._records)}")
    print("\n--- 多线程示例 (验证线程安全) ---")

    with MonitorContext() as m_main_fixed:
        # 在主线程中创建多个子线程
        threads = [threading.Thread(target=nested_multithread_context_fixed) for i in range(3)]
        [t.start() for t in threads]
        [t.join() for t in threads]

    print(f"Main context records: {len(m_main_fixed._records)}")

--- 单线程示例 (验证基础功能) ---
Single-thread m_s1 records: 3
Single-thread m_s2 records: 2

--- 多线程示例 (验证线程安全) ---
Thread XXX local context records: 1
Thread YYY local context records: 1
Thread ZZZ local context records: 1
Main context records: 3