fastapi
fastapi copied to clipboard
tiangolo
WebSocket disconnected state is not propagated to the application code (proper closures, ping timeouts)
Code first, explanation below.
import asyncio
import uvicorn
from fastapi import FastAPI, WebSocket
from starlette.websockets import WebSocketState, WebSocketDisconnect
app = FastAPI(debug=True)
@app.websocket('/ws')
async def registration(websocket: WebSocket):
await websocket.accept()
try:
while True:
if websocket.client_state != WebSocketState.CONNECTED:
print('Connection broken') # This line is never reached even if the connection is gone
break
some_condition = False # will be set to True when I need to send something to the client
if some_condition:
await websocket.send_text('some text')
await asyncio.sleep(1)
except WebSocketDisconnect as e:
print(f'Connection closed {e.code}') # This line is never reached unless I send anything to the client
if __name__ == '__main__':
uvicorn.run("main:app", port=5000)
I used the code from the WS web chat example from here: https://fastapi.tiangolo.com/advanced/websockets/
The issue is that unless you call websocket.receive_text() there is no way to get if the connection is already dead or not, even if this fact is already established. An idle WS connection can get closed via three ways that I can think of 1) proper connection procedure (client sends close frame, server confirms) 2) TCP connection dies and ping fails to be sent 3) Ping-pong response timed out. All of them are handled by the websockets library that is used inside FastAPI/Starlette/uvicorn. Since the ping-pong is enabled by default, the actual state of connection is known.
But I don't see any way to check the underlying connection state, websocket.client_state is alwaysWebSocketState.CONNECTED even after it is closed.
The only way to "probe" the state of the connection and trigger WebSocketDisconnect if it is closed is to periodically run this code:
try:
await asyncio.wait_for(
websocket.receive_text(), 0.0001
)
except asyncio.TimeoutError:
pass
It works in a sense that it triggers the check of the actual state and raises an exception if it is gone, but it is a wasteful call is the connection is actually still okay (needless receive call). Meanwhile if using websockets library raw there is websocket.closed which is actually updated once the connection is gone with a reason as well.
Is there a way to get the actual connection state without this polling receive_text?
the same happens in case of wsproto is used.
this piece of code
manager.disconnect(websocket)
await websocket.close()
print(websocket.client_state)
results in
disconnecting client `ec25918c-e735-4f39-b390-0b8959324fd0`
WebSocketState.CONNECTED
WebSocketState.CONNECTED seems to be wrongly reported.
I had a similar problem never getting any disconnect event. Maybe this will help someone:
I'm using the lower-ish level method websocket.receive() instead of websocket.receive_text() etc. and I never got the WebSocketDisconnect error but a RuntimeError instead. The comment in the OP about receive_text made me check the code and indeed only the higher level version are checking connection state by raising the proper WebSocketDisconnect error.
What works for me though isreplacing while True with:
while websocket.client_state == WebSocketState.CONNECTED:
(this is from starlette.websockets import WebSocketState btw)
I admit I'm not a big fan of changing the classic WebSocket style of onopen -> onmessage -> onclose to a loop style syntax with while True ... it just feels wrong somehow ... but thats just me I guess :-p
What should be fixed though is to throw the proper error when using websocket.receive().
@fquirin for me websocket.client_state == WebSocketState.CONNECTED check never worked, because client_state stays connected even if the client is long gone.
@asyschikov how does your client leave? Some error event? Because the receive method should usually see the message_type == "websocket.disconnect" event and if the internal states don't fit it should throw the disconnect or runtime exception ... I guess.
I think it can happen that no events are triggered when you kill your client by force (e.g. kill the whole browser) but this happens for every WebSocket implementation.
@fquirin I just tested the code that I pasted in the original message with the client gracefully closing the connection (btw I am not using a browser, this is python to python communication).
Because the receive method should usually see the message_type == "websocket.disconnect"
Indeed if I call receive I get that result and websocket.client_state changes as well, but I don't want to call receive because I don't need it. For my use case I am not expecting client to send anything first, I use websockets for the server to stream messages to the client and I want to know when the connection is gone (for any reason) without trying to send something to the client. As I explained in the ticket, I can "poll" the connection via a short lived receive to force the WebSocketDisconnect exception to be raised but this seems unnecessary because the fact that the connection is closed is known and should be exposed via some field/property.
@asyschikov I'm in essentially the same situation - can see the ping/pong messages, but can't detect if the client disconnects without closing the session via the starlette websocket. Did you find a workaround? Is there a way to attach the lower level python websockets class to the method to detect the connection state?
I would also love to see a solution for this problem. I have the same situation as @asyschikov: A Websocket without any need to receive something from the clients. Shouldn't the ping/pong messages be able to set the state the right way?
@dm-intropic the only workaround I found (and tested in production and it works) is to occasionally "poll" the connection this way:
try:
await asyncio.wait_for(
websocket.receive_text(), 0.0001
)
except asyncio.TimeoutError:
pass
Somewhere under the hood websocket.receive_text actually checks if the connection is closed or not and it it is closed it will raise WebSocketDisconnect that you can catch and act on it.
My solution is
async def _alive_task(websocket: WebSocket):
try:
await websocket.receive_text()
except (WebSocketDisconnect, ConnectionClosedError):
pass
async def _send_data(websocket: WebSocket):
try:
while True:
data = await wait_data()
await websocket.send_json(data)
except (WebSocketDisconnect, ConnectionClosedError):
pass
@router.websocket_route("/api/ws")
async def handle_something(websocket: WebSocket):
await websocket.accept()
loop = asyncio.get_running_loop()
alive_task = loop.create_task(
_alive_task(websocket),
name=f"WS alive check: {websocket.client}",
)
send_task: asyncio.Task = loop.create_task(
_send_data(websocket),
name=f"WS data sending: {websocket.client}",
)
alive_task.add_done_callback(send_task.cancel)
send_task.add_done_callback(alive_task.cancel)
await asyncio.wait({alive_task, send_task})
Same problem for me. Sometimes a client loses connection, but the websocket is unaware and will not close by itself, resulting in exceptions on send data.
This is normal. If you want to have an early crash, set the
--ws-ping-intervalon Uvicorn.
Interesting, didn't know that option. I'm using the UvicornWorker through Gunicorn though, and it doesn't seem to have that option?
You need to subclass the UvicornWorker and pass the CONFIG_SOMETHING (see how the worker is implemented) yourself, then you use that one with gunicorn --worker-class your.path.CustomUvicornWorker.
(I'm on my phone)
EDIT: Like this: https://www.uvicorn.org/deployment/#gunicorn
Why this isn't documented properly? And why if send fails it doesn't change connection state?
@Kludex the ticket is about the fact that there is no way to cause early crash. The ping pong is enabled and it fails at some point but nothing in the code indicates that it happened. There is no way to detect that connection is dead using whatever public methods FastAPI offer.
I've tried --ws-ping-interval 1. BUt websocket.client_state.name is not changing on client disconnect.