Python多线程环境中上下文内函数调用监控的线程安全实现

import time
from dataclasses import dataclass

# ... (MonitorRecord, MonitorContext, MonitorHandlers, monitor_decorator 定义略，与下文完整代码一致)

# 假设上述类和装饰器已定义
# handlers = MonitorHandlers() # 全局实例

@monitor_decorator
def run_task():
    time.sleep(0.1)

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

import threading
import time
from dataclasses import dataclass
from collections import UserList

@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 LocalList(threading.local, UserList):
    """
    一个线程局部存储的列表，每个线程都有其独立的列表实例。
    继承UserList以便拥有标准列表的所有方法。
    """
    def __init__(self, initlist=None):
        super().__init__() # 初始化UserList部分
        if initlist is not None:
            self.data = list(initlist)
        else:
            self.data = []

class MonitorHandlers:
    """
    负责管理监控上下文的处理器。
    使用threading.local实现线程隔离，并使用锁确保主线程处理器的线程安全。
    """
    def __init__(self):
        self._lock = threading.Lock()
        with self._lock:
            # _mainhandlers 存储主线程的上下文，所有线程的记录都会汇总到这里
            self._mainhandlers: list[MonitorContext] = []
        # _handlers 存储非主线程的上下文，每个线程有其独立的列表
        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:
        """
        向所有活跃的监控上下文添加记录。
        记录会添加到当前线程的局部上下文，以及主线程的全局上下文。
        """
        # 1. 添加到当前线程的局部上下文（如果有）
        for h in self._handlers:
            h.add_record(record)
        # 2. 添加到主线程的全局上下文 (需要锁保护，因为_mainhandlers是共享资源)
        with self._lock:
            for h in self._mainhandlers:
                h.add_record(record)

# 创建全局的 MonitorHandlers 实例
handlers = MonitorHandlers()

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 _

# 假设上述 MonitorRecord, MonitorContext, LocalList, MonitorHandlers, monitor_decorator 已定义

@monitor_decorator
def run_monitored_task():
    """一个模拟耗时操作的函数。"""
    time.sleep(0.05) # 缩短睡眠时间以便快速测试

def thread_nested_context():
    """在子线程中创建监控上下文并执行被监控函数。"""
    with MonitorContext() as m_thread:
        run_monitored_task()
    # 打印当前子线程内部上下文的记录数。
    # 注意：m_thread._records 只包含该特定上下文实例接收到的记录，
    # 而不是所有线程的记录。
    print(f"线程 {threading.get_ident()} 内部上下文记录数: {len(m_thread._records)}")

print("--- 单线程测试 ---")
with MonitorContext() as m1_single:
    run_monitored_task()
    with MonitorContext() as m2_single:
        run_monitored_task()
        run_monitored_task()

print(f"单线程外层上下文 m1_single 记录数: {len(m1_single._records)}")
print(f"单线程内层上下文 m2_single 记录数: {len(m2_single._records)}")

print("\n--- 多线程测试 ---")
with MonitorContext() as m1_multi: # 这是主线程创建的上下文
    threads = []
    # 创建并启动5个子线程
    for i in range(5):
        t = threading.Thread(target=thread_nested_context)
        threads.append(t)
        t.start()