如何实现强制退出异步上下文管理器并抛出异常？

要实现连接断开时强制退出上下文并抛出异常，核心是在上下文管理器中跟踪当前异步任务，并在连接断开回调中通过任务对象主动抛出异常。以下是具体实现步骤和代码示例：

1. 定义自定义异常类 #

首先创建专属的异常类型，让消费者可以明确捕获连接丢失的情况：

class ConnectionLostException(Exception):
    """当服务连接意外断开时抛出的异常"""
    pass

2. 改造异步上下文管理器 #

修改SomeServiceContextManager，添加任务跟踪逻辑，在连接断开时主动向当前任务抛出异常：

import asyncio

class SomeServiceContextManager:
    def __init__(self, service):
        self.service = service
        self._task = None  # 保存当前异步任务引用

    async def __aenter__(self):
        # 获取当前正在执行的异步任务
        self._task = asyncio.current_task()
        await self.service.connect(self.connection_state_callback)
        return self.service

    async def __aexit__(self, exc_type, exc, tb):
        # 确保无论是否发生异常，都断开连接
        self.service.disconnect()
        # 返回False表示不抑制异常，让异常继续向上传播
        return False

    def connection_state_callback(self, state):
        if state == "connection lost" and self._task is not None:
            # 向当前异步任务抛出连接丢失异常
            self._task.throw(ConnectionLostException("连接意外断开"))

关键细节说明： #

• asyncio.current_task()：在__aenter__中获取当前执行的异步任务，保存为实例属性，这样回调就能定位到需要抛出异常的上下文。
• task.throw()：主动向任务抛出异常，会立即中断当前任务中正在执行的代码（比如消费者的循环或消息处理逻辑），触发上下文管理器的__aexit__方法完成资源清理，同时将异常传递给外层的try-except块。

3. 消费者代码的异常处理 #

消费者只需正常捕获ConnectionLostException即可处理连接断开的情况，嵌套上下文场景下，异常会自动触发内层上下文退出，外层上下文也会在函数执行完成时正常关闭：

简单场景示例 #

class Consumer:
    async def send_stuff(self):
        try:
            async with SomeServiceContextManager(self.service) as connected_service:
                while True:
                    await asyncio.sleep(1)
                    connected_service.send("hello")
        except ConnectionLostException as e:
            print(f"Oh no my connection was lost!! {e}")

嵌套上下文场景示例 #

class Consumer:
    async def translate_stuff_stuff(self):
        async with SomeOtherServiceContextManager(self.otherservice) as connected_other_service:
            try:
                async with SomeServiceContextManager(self.service) as connected_service:
                    # 假设messages()是异步迭代器，使用async for符合异步场景规范
                    async for message in connected_other_service.messages():
                        connected_service.send("message received: " + message.text)
            except ConnectionLostException as e:
                print(f"Oh no my connection was lost - I'll also drop out of the other service connection!! {e}")
        # 外层的SomeOtherServiceContextManager会在函数退出时自动调用__aexit__断开连接

4. 线程安全注意事项 #

如果服务的connection_state_callback是在非事件循环线程中触发的（比如某些WebSocket库的回调在独立线程），直接调用task.throw()会导致线程安全问题，此时需要用asyncio.call_soon_threadsafe来触发异常：

def connection_state_callback(self, state):
    if state == "connection lost" and self._task is not None:
        # 线程安全地向事件循环中的任务抛出异常
        asyncio.call_soon_threadsafe(
            self._task.throw,
            ConnectionLostException("连接意外断开")
        )

核心原理 #

结构化并发中，异步上下文管理器的__aexit__会在上下文退出时自动执行，无论正常退出还是异常退出。通过向当前任务主动抛出异常，我们可以强制中断上下文内的代码执行，触发__aexit__完成资源清理，同时将异常传递给上层调用者处理，完美契合结构化并发的资源管理原则。

内容的提问来源于stack exchange，提问作者Sam