Enhancement: Graceful cancellation of handler on client disconnect

[jorgeecardona]

Summary

I am unsure if this already possible, but it seems that one can't detect a client disconnection in order to prevent costly db operations. The example is taken from a similar question reported for fastapi.

Basic Example

import asyncio
from litestar import Litestar, get
from litestar.openapi import OpenAPIConfig

@get("/")
async def hello_world() -> dict[str, str]:
    """Keeping the tradition alive with hello world."""
    try:
        print("sleep")
        await asyncio.sleep(10)
        print("done")
    except asyncio.CancelledError:
        print("cancelled")
        raise
    return {"hello": "world"}


app = Litestar(
    route_handlers=[hello_world],
    openapi_config=OpenAPIConfig(title="My API", version="1.0.0")
)

Drawbacks and Impact

No response

Unresolved questions

No response

---

Note

While we are open for sponsoring on GitHub Sponsors and
OpenCollective, we also utilize Polar.sh to engage in pledge-based sponsorship.

Check out all issues funded or available for funding on our Polar.sh Litestar dashboard

• If you would like to see an issue prioritized, make a pledge towards it!
• We receive the pledge once the issue is completed & verified
• This, along with engagement in the community, helps us know which features are a priority to our users.

[peterschutt]

Can you link the fastapi issue please?

[provinzkraut]

I generally like this idea. However, I’m not sure how it would interact with background tasks. They might be cancelled / never executed as well, which could lead to some unexpected behaviour.

Maybe we could explicitly check for those and ensure they are properly awaited in any case?

[peterschutt]

Hmm, good thinking, not sure. We could use shield?

We say in the docs that they are called "after the response finishes sending the data" - is it too naive to think that if a request is cancelled by client disconnect that we haven't finished sending the data?

[provinzkraut]

is it too naive to think that if a request is cancelled by client disconnect that we haven't finished sending the data?

Good point.

One other thing to consider is that the behaviour for fully synchronous route handlers might be unexpected. Since we don't call those asynchronously (and even if we did, it wouldn't change anything), the handler function will always run through once called, since there's no way to "cancel" a synchronous function. So only the after request pipeline (hooks, background tasks) would be cancelled in those cases.

We could document this, I'm just concerned that diverging behaviour like this might lead to confusion.

[provinzkraut]

In any case this would be considered a breaking change though, as it differs from the current behaviour.

[peterschutt]

Yeh this will certainly add complexity, and it would be good to keep that away from our base case.

Maybe it could be made an opt-in feature as users that are looking for this will more readily understand the complexity of the case

However, I’m not sure how it would interact with background tasks.

Another thing we could do for this is to leave it up to the user. If they are doing stuff in their background tasks that shouldn't be cancelled, they can use asyncio.shield and catch CancelledError. Actually, this would be my preference as it would mean adding less complexity to the core and we could wait to see if that would be good enough once there was some actual use.

[jorgeecardona]

Probably an easy workaround can be built using a Stream response since they protect for disconnect see

litestar/litestar/response/streaming.py

Line 83 in 0674fba

async def _listen_for_disconnect(self, cancel_scope: CancelScope, receive: Receive) -> None:

[dorinclisu]

As of Litestar 2.12 at least (didn't investigate for which older version this was not the case), the handlers are cancelled on disconnect.

See the tests for proof and clarity: https://github.com/dorinclisu/runner-with-api/blob/main/tests/litestar/test_cancellation.py

EDIT: no they aren't, but it can be easily achieved, see bellow.

[provinzkraut]

This is only true because you are using the test client, which doesn't establish an actual HTTP connection, but rather calls the ASGI app directly. You are cancelling the asyncio task, not an HTTP connection.

[dorinclisu]

Thanks for the heads-up, indeed testing this automatically is quite elusive. In any case, have a look on a working implementation I just copied over from fastapi:
https://github.com/dorinclisu/runner-with-api/blob/main/runner_with_api/litestar/cancellation.py

from runner_with_api.litestar.cancellation import cancel_on_disconnect

@get('/test')
async def test(request: Request):
    async with cancel_on_disconnect(request):
        # this context is cancellable
        await sleep(1)  

Actually tested with $ curl http://localhost:8000/test --max-time 0.1

$ uvicorn tests.litestar.test_cancellation:app
INFO:     Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)
INFO - 2024-11-08 20:20:57,190 - runner_with_api.litestar - cancellation - 127.0.0.1:37532 - "GET /test" 499 DISCONNECTED
INFO - 2024-11-08 20:20:57,192 - root - test_cancellation - Cancelled Handler