Изящное завершение работы приложения uvicorn starlette с помощью веб-сокетов

Question

Учитывая этот образец приложения Starlette с открытым подключением к веб-сокету, как закрыть приложение Starlette? Я бегаю на увикорне. Всякий раз, когда я нажимаю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 уже устанавливает для своего собственного процесса выключения.
Вместо этого вы можете 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!')

Answer 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 метод.