New websockets, handling disconnects

[He Shiming]

Is there an existing issue for this?

• I have searched the existing issues

Is your feature request related to a problem? Please describe.

I'm using version 23.6.0, so it's new websockets. I have followed the guide to build something like this:

async def handler(request, ws):
    while True:
        data = "hello!"
        try:
            await ws.send(data)
            data = await ws.recv()
        except Exception as e:
            print(e)
            break
    clean_up()

bp.add_websocket_route(handler, '/ws')

I discovered that when the client disconnects, no exceptions are thrown in this above code. Making it difficult for me to clean up things. What's worse is that the execution of this handler somehow stops in a disconnect, no subsequent code gets run.

I read the source code and figured out the following mechanism to tackle a disconnect:

def ws_disconnect(my_mess, fut):
    clean_up(my_mess)

async def handler(request, ws):
    ws.connection_lost_waiter.add_done_callback(functools.partial(ws_disconnect, "my_mess"))
    ...

It works, as the WebsocketImplProtocol calls it upon disconnect. But I could help but thinking there should be a better mechanism to handle websocket clean ups. What am I supposed to do to handle disconnects?

Describe the solution you'd like

Keep the handler running upon disconnect, and give it a chance to handle exception or check for "connected" attribute.

Additional context

No response

[Adam Hopkins]

I'd suggest this:

from asyncio import shield

async def cleanup():
    for i in range(10):
        print("Cleaning up")
        await sleep(1)


@app.websocket("/")
async def handler(request: Request, ws):
    try:
        while True:
            ...
    finally:
        print("Connection closed")
        await shield(cleanup())

You might not need the shield, but it may be helpful

[He Shiming]

Thank you Adam! finally works!