使用 websockets 正常关闭 uvicorn starlette 应用程序

Question

鉴于此示例 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()

我尝试在关闭事件上切换 bool 变量，但该变量永远不会更新。它总是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:

Answer 1

你的变量没有改变的原因是因为“关闭”事件的处理程序是在所有任务都执行后执行的（即你的无限循环）。

在 asyncio 事件循环上设置信号处理程序可能不起作用，因为我相信只允许 uvicorn 已经为其自己的关闭过程设置的单个信号处理程序。
相反，您可以对 uvicorn 信号处理程序进行 Monkey Patch 以检测应用程序关闭并在该新函数中设置您的控制变量。

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!')