如何从callback回调中断api_websocket_wait_update函数执行？

原问题代码 #

def get_contract_listener(access_token: str, callback: Callable[[int | None], Awaitable[None]]) -> None | Watch:
    logger.debug('get contract listener')

    contract = db.collection('Contract').where(filter=FieldFilter('access_token', '==', access_token)).get()

    async def wrap(document: DocumentSnapshot):
        await callback(123)

    return cast(DocumentReference, contract[0].reference).on_snapshot(lambda doc, _, __: run(wrap(doc[0])))


async def api_websocket_wait_update(socket: WebSocket):
    await socket.accept()
    global first
    first = True

    token = socket.cookies.get('access_token')

    async def callback(balance: int | None) -> None:
        assert listener is not None
        global first
        if first:
            first = False
            return

        await socket.send_json({'message': 'updated', 'balance': balance})
        await socket.close(1000)
        listener.unsubscribe()
        # return ??

    listener = get_contract_listener(token, callback)
    await socket.send_json({'message': 'accepted'})

    while True:
        sleep(5)

问题描述 #

上述代码中，api_websocket_wait_update通过get_contract_listener注册回调callback，当callback触发时，仅用return只能中断自身，无法终止外层函数末尾的无限循环sleep(5)，需要实现从callback中终止api_websocket_wait_update的执行。

解决方案 #

核心思路是利用异步事件信号实现协程间通信，同时修正同步阻塞问题：

1. 替换同步sleep为异步sleep #

原代码中的sleep(5)是同步阻塞调用，会卡住整个事件循环，导致回调的异步操作无法及时执行，也无法响应终止信号。需替换为await asyncio.sleep(5)。

2. 使用asyncio.Event传递终止信号 #

创建一个asyncio.Event对象，在callback中设置该事件，外层循环通过检查事件状态决定是否退出。

修改后的完整代码 #

import asyncio
from typing import Callable, Awaitable, cast
# 替换为实际项目的导入路径
from your_module import logger, db, FieldFilter, DocumentSnapshot, Watch, WebSocket

def get_contract_listener(access_token: str, callback: Callable[[int | None], Awaitable[None]]) -> None | Watch:
    logger.debug('get contract listener')

    contract = db.collection('Contract').where(filter=FieldFilter('access_token', '==', access_token)).get()

    async def wrap(document: DocumentSnapshot):
        await callback(123)

    # 用asyncio.create_task替代run，避免创建新事件循环
    return cast(DocumentReference, contract[0].reference).on_snapshot(
        lambda doc, _, __: asyncio.create_task(wrap(doc))
    )


async def api_websocket_wait_update(socket: WebSocket):
    await socket.accept()
    first = True  # 改用局部变量，避免global污染
    token = socket.cookies.get('access_token')
    stop_event = asyncio.Event()  # 终止事件信号
    listener = None

    async def callback(balance: int | None) -> None:
        nonlocal first, listener
        if first:
            first = False
            return

        await socket.send_json({'message': 'updated', 'balance': balance})
        await socket.close(1000)
        if listener:
            listener.unsubscribe()
        stop_event.set()  # 设置终止事件，通知外层循环退出

    listener = get_contract_listener(token, callback)
    await socket.send_json({'message': 'accepted'})

    # 循环检查终止事件，未触发则继续等待
    while not stop_event.is_set():
        await asyncio.sleep(5)

关键说明 #

• asyncio.Event是异步环境下安全的信号传递机制，set()标记事件为触发状态，is_set()用于检查状态。
• 使用nonlocal访问外层函数的局部变量，比global更安全，避免全局变量的副作用。
• 原代码中run(wrap(doc[0]))存在两个问题：一是run会创建新事件循环，异步上下文里应改用asyncio.create_task；二是doc本身就是DocumentSnapshot，无需取doc[0]。

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