Async context manager finalization gets canceled in websocket handler

Hello,

I am currently trying to wrap my head around some pitfalls regarding asyncio, handlers and task cancelation in aiohttp.

The solutions I could come up with seem very clunky and unreliable, I hope someone better informed than me could review the approach or point me in a direction of a better solution.

What I am doing

My solution requires for some non-optional finalization code to run when a websocket handler closes. For context certain client data in REDIS is locked using a semaphore at initialization. This lock needs to be cleared again after the client connection is closed.

Essentially something like this

async def test_handler(request: Request):
    app = request.app
    redis_pool = app[APP_REDIS_POOL]
    print("Someone joined.")
    websocket = WebSocketResponse()
    await websocket.prepare(request)

    # run async initialization here e.g fetch user data from redis
    await asyncio.sleep(0.1)

    async for event in websocket:
        pass  # do non-async stuff, process events here

    # run async finalization here e.g write changes to redis
    await asyncio.sleep(0.1)

    await websocket.close()
    # api connection now closed
    print("Someone left.")
    return websocket

Simple JS script that triggers the request:

<script>
        var webSocket;
        let route = 'test_handler'

        try {
            webSocket = new WebSocket(`ws://${window.location.host}/${route}`);
        } catch {
            webSocket = new WebSocket(`wss://${window.location.host}/${route}`);
        }

        webSocket.addEventListener('open', () => {
            console.log("websocket open")
        });
        webSocket.addEventListener('close', () => {
            console.log("websocket closed")
        });
        webSocket.addEventListener('error', (e) => {
            console.log(e)
        });
    </script>

What started the trouble

I noticed that after pressing reload the ‘Someone left’ call is never reached. The websocket spits out a message of type WSMsgType.Closed on the server (code 1001) but after reaching the first await in the finalization code the task gets canceled.
This can be reproduced reliably.

Note:
Adding the following lines in Javascript seem to remove the issue entirely for reloads …

addEventListener("beforeunload", () => {
    webSocket.close()
})

… but this is not reliable as the client might close their browser in some way that prevents this event from being emitted.

My understanding:

In this simple example the code inside of the handling is synchronous. Only after an await control returns to the provider of the message where the handler task gets canceled. This await turns out to be some finalization code.

This is confusing because essentially the Closing message already indicates that the connection needs to be closed, but then the CancelledError occurs. (In the worst case there is not even an error being logged. I was hard stuck on this for 2 days trying to figure out why the last print statement was never reached.)

What I tried

  1. Using an async context (async with ...) seemed like an obvious solution, but I it does not have any privileges in terms of cancelation and the first await inside __aexit__ still gets canceled. So it is only useful to organize code.

  2. Inserting a dummy await to catch the cancelation. This feels like a hack and does not seem reliable at all.

    async for event in websocket:
        pass

    try:
        await asyncio.sleep(1)  # <-- this gets cancelled
    except asyncio.CancelledError as e:
        print(repr(e))

    # run async finalization here e.g write changes to redis
    await asyncio.sleep(0.1)

    await websocket.close()
    # api connection now closed
    print("Someone left.")
    return websocket
  1. I evaluated using aiojobs, because it was recommended in one of the Issues I looked at.

What I am struggling with

My main issue is probably that I expect the finalization code to run without interruptions (“normal” exceptions that might occur there aside). But I am realizing that a CancelledError could occur anywhere anytime inside of a handler.

Additionally when the async for loop exits after a WSMsgType.CLOSED it starts running finalization ‘prematurely’ essentially throwing any await calls under the (CancelationError-) bus.

A construct like a context manager does not help with this and actually makes it harder to wrap finalization inside of a try: ... except CancelledError:.... (maybe it was not meant for this kind of situation?)

The package aiojobs seems neat, but to my understanding this will simply push the CancelationError problem down the call stack. Documentation is very ‘lean’ on this one, so I am honestly confused on how adding more awaitables to the mix is going to help the issue at hand.

The question

  1. Is there a prefered way to tackle async finalization for a client socket going away in aiohttp.
  2. Is there a pythonic way to handle this kind of task cancelation with a guarantee that the task finishes gracefully.

Any help or insight on how to tackle this is appreciated.

Cheers,
I.I.