Python上下文中的函数调用监控与多线程兼容性实现

import time
import threading
from dataclasses import dataclass
from collections import UserList # UserList for LocalList inheritance

@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

class MonitorHandlers:
    def __init__(self):
        # _handlers 列表存储所有注册的 MonitorContext
        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 = MonitorHandlers()

def monitor_decorator(f):
    def wrapper(*args, **kwargs):
        start = time.time()
        result = f(*args, **kwargs) # 执行被装饰的函数
        handlers.add_record(
            MonitorRecord(
                function=f.__name__,
                time=time.time() - start,
            )
        )
        return result
    return wrapper

# 单线程示例
@monitor_decorator
def run_task():
    time.sleep(0.1) # 模拟任务执行

if __name__ == '__main__':
    print("--- 单线程测试 ---")
    with MonitorContext() as m1:
        run_task()
        with MonitorContext() as m2:
            run_task()
            run_task()

    print(f"m1 记录数: {len(m1._records)}") # 预期 3
    print(f"m2 记录数: {len(m2._records)}") # 预期 2

    # 单线程输出:
    # m1 记录数: 3
    # m2 记录数: 2

import time
import threading
from dataclasses import dataclass
from collections import UserList # UserList for LocalList inheritance

@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

# 继承 threading.local 和 UserList，实现线程局部的列表
class LocalList(threading.local, UserList):
    """
    一个线程局部列表，每个线程都会有自己的独立列表实例。
    UserList 提供列表的完整接口。
    """
    def __init__(self, initlist=None):
        super().__init__() # 调用 UserList 的 __init__
        if initlist is not None:
            self.data = list(initlist) # 初始化内部数据

class MonitorHandlers:
    def __init__(self):
        # 用于保护 _mainhandlers 列表的锁
        self._lock = threading.Lock()
        # _mainhandlers 存储主线程的上下文，需要锁保护
        with self._lock:
            self._mainhandlers: list[MonitorContext] = []
        # _handlers 存储非主线程的上下文，使用 LocalList 实现线程局部
        self._handlers: list[MonitorContext] = LocalList()

    def register(self, handler: MonitorContext) -> None:
        # 根据当前线程是否为主线程，注册到不同的列表中
        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:
        # 根据当前线程是否为主线程，从不同的列表中删除
        if threading.main_thread().ident == threading.get_ident():
            with self._lock: # 主线程列表需要加锁
                self._mainhandlers.remove(handler)
        else:
            self._handlers.remove(handler) # 非主线程从线程局部列表删除

    def add_record(self, record: MonitorRecord) -> None:
        # 将记录分发给当前线程的所有上下文
        for h in self._handlers:
            h.add_record(record)
        # 同时将记录分发给主线程的所有上下文（如果存在）
        with self._lock: # 访问主线程列表需要加锁
            for h in self._mainhandlers:
                h.add_record(record)

# 全局的监控处理器实例
handlers = MonitorHandlers()

def monitor_decorator(f):
    def wrapper(*args, **kwargs):
        start = time.time()
        result = f(*args, **kwargs)
        handlers.add_record(
            MonitorRecord(
                function=f.__name__,
                time=time.time() - start,
            )
        )
        return result
    return wrapper

# 多线程示例
@monitor_decorator
def run_threaded_task():
    time.sleep(0.1) # 模拟任务执行

def nested_thread_context():
    """
    在子线程中创建监控上下文并执行任务。
    """
    with MonitorContext() as m_thread:
        run_threaded_task()
    print(f"子线程 {threading.get_ident()} 内部上下文记录数: {len(m_thread._records)}")

if __name__ == '__main__':
    print("\n--- 多线程测试 ---")
    threads = []
    # 在主线程中创建监控上下文
    with MonitorContext() as m_main:
        for i in range(5): # 创建 5 个子线程
            t = threading.Thread(target=nested_thread_context, name=f"Thread-{i}")
            threads.append(t)
            t.start()

        for t in threads:
            t.join() # 等待所有子线程完成

    print(f"主线程上下文记录数: {len(m_main._records)}")
    # 预期输出：
    # 子线程 ... 内部上下文记录数: 1 (每个子线程的上下文只记录自己的任务)
    # 主线程上下文记录数: 5 (主线程的上下文记录了所有子线程的任务)