Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

What happens to existing awaits when WebSocket.close is called

If I open a Starlette/FastAPI WebSocket, what happens to any coroutines currently waiting to receive data from the client if I close the WebSocket from outside the coroutine? Does the call to receive raise an exception or does the coroutine sit in memory forever because it is never going to receive anything?

from fastapi import FastAPI
from fastapi.websockets import WebSocket

app = FastAPI()

open_sockets = {}

@app.websocket('/connection/')
async def connection(websocket: WebSocket):
    await websocket.accept()

    id = await websocket.receive_json()['id']
    open_sockets[id] = websocket

    while True:
        data = await websocket.receive_json()

@app.post('/kill/{id}/')
async def kill(id=str):
    # What does this do to the above `await websocket.receive_json()`?
    await open_sockets[id].close()
    del open_sockets[id]
like image 480
drhagen Avatar asked Oct 26 '22 22:10

drhagen


1 Answers

It raises a starlette.websockets.WebSocketDisconnect with code 1006

  File ".\websocket_close_test.py", line 27, in connection
    data = await websocket.receive_json()
  File "C:\Apps\Python38\lib\site-packages\starlette\websockets.py", line 98, in receive_json
    self._raise_on_disconnect(message)
  File "C:\Apps\Python38\lib\site-packages\starlette\websockets.py", line 80, in _raise_on_disconnect
    raise WebSocketDisconnect(message["code"])

I did note that the .close() call blocked for a little while.

Here's the code I used to test, the second ws.send was never received.

import time

import requests
import websocket

ws = websocket.WebSocket()
ws.connect("ws://localhost:8000/connection/")
print('Sending id 0')
ws.send('{ "id": "0" }')
time.sleep(2)
print('Closing id 0')
requests.post('http://localhost:8000/kill/0/')
print('id 0 is closed')
time.sleep(2)
print('Trying to send data on closed connection')
ws.send('{ "id": "10" }')
like image 142
Gabriel Cappelli Avatar answered Nov 15 '22 07:11

Gabriel Cappelli