使用 websockets 优雅关闭 uvicorn starlette 应用程序

Posted 2023-03-08

【中文标题】使用 websockets 优雅关闭 uvicorn starlette 应用程序

【英文标题】：Graceful shutdown of uvicorn starlette app with websockets

【发布时间】：2020-01-27 17:30:47

【问题描述】：

鉴于此示例 Starlette 应用程序具有打开的 websocket 连接，您如何关闭 Starlette 应用程序？我正在运行 uvicorn。每当我按Ctrl+C 时，输出都是Waiting for background tasks to complete.，它会永远挂起。

from starlette.applications import Starlette

app = Starlette()

@app.websocket_route('/ws')
async def ws(websocket):
    await websocket.accept()

    while True:
        # How to interrupt this while loop on the shutdown event?
        await asyncio.sleep(0.1)

    await websocket.close()

我尝试在关机事件上切换一个布尔变量，但该变量永远不会更新。它总是False。

例如。

app.state.is_shutting_down = False


@app.on_event('shutdown')
async def shutdown():
    app.state.is_shutting_down = True


@app.websocket_route('/ws')
async def ws(websocket):
    await websocket.accept()

    while app.state.is_shutting_down is False:

【问题讨论】：

你解决过这个问题吗？我也有类似的问题。

不，我从来没弄明白。对不起。

您是否尝试过等待 con server.should_exit 或 server.force_exit ？

【参考方案1】：

您的变量没有更改的原因是因为“关闭”事件的处理程序是在所有任务都执行后执行的（即无限循环）。 在 asyncio 事件循环上设置信号处理程序可能不起作用，因为我相信只允许一个信号处理程序，uvicorn 已经为其自己的关闭过程设置了该信号处理程序。 相反，您可以 Monkey Patch uvicorn 信号处理程序以检测应用程序关闭并在该新函数中设置您的控制变量。

import asyncio
from starlette.applications import Starlette
from uvicorn.main import Server

original_handler = Server.handle_exit

class AppStatus:
    should_exit = False

    @staticmethod
    def handle_exit(*args, **kwargs):
        AppStatus.should_exit = True
        original_handler(*args, **kwargs)

Server.handle_exit = AppStatus.handle_exit

app = Starlette()

@app.websocket_route('/ws')
async def ws(websocket):
    await websocket.accept()

    while AppStatus.should_exit is False:

        await asyncio.sleep(0.1)

    await websocket.close()
    print('Exited!')

【参考方案2】：

我无法真正重现问题，我明白了

INFO: Started server process [31]
INFO: Waiting for application startup.
INFO: Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)
^CINFO: Shutting down
INFO: Waiting for application shutdown.
INFO: Finished server process [31]

但我在自己的异步应用程序中正常关闭的做法是：

import signal

async def shutdown(signal: signal):
    """
    try to shut down gracefully
    """
    logger.info("Received exit signal %s...", signal.name)
    tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
    [task.cancel() for task in tasks]
    logging.info("Canceling outstanding tasks")
    await asyncio.gather(*tasks)


if __name__ == "__main__":

    loop = asyncio.get_event_loop()
    signals = (signal.SIGHUP, signal.SIGTERM, signal.SIGINT)
    for s in signals:
        loop.add_signal_handler(
            s, lambda s=s: asyncio.create_task(shutdown(s))
        )

您可能必须将await websocket.close() 移动到shutdown 方法中。

【讨论】：

你如何将每个 websocket 的 websocket 实例传递给关闭处理程序？