Stuck on an issue?

Lightrun Answers was designed to reduce the constant googling that comes with debugging 3rd party libraries. It collects links to all the places you might be looking at while hunting down a tough bug.

And, if you’re still stuck at the end, we’re happy to hop on a call to see how we can help out.

[QUESTION] How to properly shut down websockets waiting forever

See original GitHub issue

Description

What is the proper way to shut down the server when you have a long-running websocket reading data from e.g. Redis or a DB waiting in a while True loop for new messages? For example, given this complete example:

import asyncio
from typing import AsyncGenerator

from fastapi import FastAPI
from starlette.responses import HTMLResponse
from starlette.websockets import WebSocket

app = FastAPI()

html = """
<!DOCTYPE html>
<html>
    <body>
        <script>
            var ws = new WebSocket("ws://localhost:8000/notifications");
            // Processing of messages goes here...
        </script>
    </body>
</html>
"""


@app.get("/")
async def get() -> HTMLResponse:
    return HTMLResponse(html)


async def subscribe(channel: str) -> AsyncGenerator[dict, None]:
    """Fake method which would normally go to Redis/DB and wait."""
    while True:
        await asyncio.sleep(600)
        yield {"hello": "world"}


@app.websocket("/notifications")
async def notifications_handler(websocket: WebSocket) -> None:
    await websocket.accept()

    async for msg in subscribe("*"):
        await websocket.send_json(msg)

If you run it, then go to http://localhost:8000/ in a browser (it’ll load a blank page), then try to <kbd>ctrl</kbd>+<kbd>c</kbd> you get this:

$ uvicorn demo:app
INFO:     Started server process [48444]
INFO:     Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     127.0.0.1:54369 - "GET / HTTP/1.1" 200 OK
INFO:     ('127.0.0.1', 54371) - "WebSocket /notifications" [accepted]
^CINFO:     Shutting down
INFO:     Waiting for background tasks to complete. (CTRL+C to force quit)

It just hangs and never shuts down properly. What’s the right way to handle this? The shutdown server event isn’t fired until after all background tasks have completed so it’s not useful for this either.

Additional context

So far I have come up with the following code, which creates an additional asyncio task that is not linked to uvicorn/starlette and so is ignored in the shutdown path, allowing a shutdown handler to be called that can close the DB server.

import asyncio
from typing import AsyncGenerator

from fastapi import FastAPI
from starlette.responses import HTMLResponse
from starlette.websockets import WebSocket

app = FastAPI()

html = """
<!DOCTYPE html>
<html>
    <body>
        <script>
            var ws = new WebSocket("ws://localhost:8000/notifications");
            // Processing of messages goes here...
        </script>
    </body>
</html>
"""


@app.get("/")
async def get() -> HTMLResponse:
    return HTMLResponse(html)


@app.on_event("shutdown")
async def shutdown() -> None:
    # Close the Redis or DB connection here
    # await redis.close()
    pass


async def subscribe(channel: str) -> AsyncGenerator[dict, None]:
    """Fake method which would normally go to Redis/DB and wait."""
    while True:
        # When using a real Redis or DB this await will return immediately
        # when the server shuts down.
        await asyncio.sleep(600)
        yield {"hello": "world"}


@app.websocket("/notifications")
async def notifications_handler(websocket: WebSocket) -> None:
    await websocket.accept()

    async def _handler() -> None:
        async for msg in subscribe("*"):
            print("Got message!", msg)
            await websocket.send_json(msg)

    # Create the handler as an unmanaged background task so the server won't
    # wait on it forever during shutdown.
    asyncio.create_task(_handler())

    # Use a long-blocking read to keep the socket alive in memory while the
    # above background task works. During shutdown this loop will exit, which
    # in turn will cause the above handler to exit as well.
    while True:
        try:
            print(await websocket.receive())
        except Exception:
            break

This is not ideal since:

It creates two tasks for each incoming request
It won’t work for bidirectional websockets (it works for my use case of push notifications)

So what is the right way to ensure a quick and clean shutdown?

Issue Analytics

State:
Created 4 years ago
Reactions:10
Comments:7 (3 by maintainers)

Top GitHub Comments

4reactions

euri10commented, Nov 14, 2019

Back on a proper pc, I think I’m doing kind of the same of what you’re trying to do but with rabbitmq in https://gitlab.com/euri10/celeryasgimon/tree/master/backend/app/main.py Here you’ll find a snippet of a ws waiting indefinitely to consume the broker’s messages, it handles shutdown gently

3reactions

matiuszkacommented, Dec 15, 2021

Hi,

For Fast API this works for me, when client is disconnected:

@router.websocket("/tst")
async def tst(websocket: WebSocket) -> None:
    async def send_data():
        while websocket.client_state == WebSocketState.CONNECTED:
           await websocket.send_text("hi")
           await asyncio.sleep(1)

    async def watch_status():
        while websocket.client_state == WebSocketState.CONNECTED:
            await websocket.receive()

    await websocket.accept()
    await asyncio.gather(watch_status(), send_data())