Shutdown immediately closes websocket connections without grace

[tonyay163]

I'm using uvicorn with fastapi. On SIGTERM, I can see that existing websocket connections are closed with 1012 which comes from uvicorn. This seems to differ from all the other request types where they are allowed to drain. Is there a way I can allow a graceful timeout for these connections instead of killing them abruptly?

Dependencies:

python = "3.10"
fastapi = "0.115.6"
requests = "2.32.3"
uvicorn = "0.34.0"
websockets = "14.1"

Code:

import asyncio

import requests
import websockets
from fastapi import FastAPI, WebSocket

app = FastAPI()


@app.get("/")
async def get():
    await asyncio.sleep(1)
    return "hello"


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    while True:
        data = await websocket.receive_text()
        print(data)

async def main():
    async with websockets.connect("ws://localhost:8080/ws") as websocket:
        while True:
            await websocket.send("hello")
            print(requests.get("http://localhost:8080").text)

if __name__ == '__main__':
    asyncio.get_event_loop().run_until_complete(main())

Repro:

1. Run the server (in a separate tab): uvicorn main:app --host 0.0.0.0 --port 8080
2. Run the client: python main.py
3. Ctrl + C to kill the server

Logs:

INFO:     ('127.0.0.1', 60184) - "WebSocket /ws" [accepted]
INFO:     connection open
hello
INFO:     127.0.0.1:60186 - "GET / HTTP/1.1" 200 OK
hello
INFO:     127.0.0.1:60188 - "GET / HTTP/1.1" 200 OK
hello
INFO:     127.0.0.1:60190 - "GET / HTTP/1.1" 200 OK
hello
INFO:     127.0.0.1:60192 - "GET / HTTP/1.1" 200 OK
hello
INFO:     127.0.0.1:60195 - "GET / HTTP/1.1" 200 OK
hello
^CINFO:     Shutting down
ERROR:    Exception in ASGI application
Traceback (most recent call last):
...
  File "/Users/tony/workspace/websocket-repro/main.py", line 20, in websocket_endpoint
    data = await websocket.receive_text()
  File "/Users/tony/Library/Caches/pypoetry/virtualenvs/websocket-repro-PYTa93q3-py3.10/lib/python3.10/site-packages/starlette/websockets.py", line 119, in receive_text
    self._raise_on_disconnect(message)
  File "/Users/tony/Library/Caches/pypoetry/virtualenvs/websocket-repro-PYTa93q3-py3.10/lib/python3.10/site-packages/starlette/websockets.py", line 113, in _raise_on_disconnect
    raise WebSocketDisconnect(message["code"], message.get("reason"))
starlette.websockets.WebSocketDisconnect: (1012, None)
INFO:     connection closed
INFO:     Waiting for connections to close. (CTRL+C to force quit)
INFO:     127.0.0.1:60197 - "GET / HTTP/1.1" 200 OK
INFO:     Waiting for application shutdown.
INFO:     Application shutdown complete.
INFO:     Finished server process [88662]

As you can see in the logs, the HTTP request finishes gracefully before the server shutdown but it kills the websocket connection immediately.

[Kludex]

MRE?

[Kludex]

Include installed packages please.

[tonyay163]

updated