5

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]
Tom Carrick
  • 6,349
  • 13
  • 54
  • 78
drhagen
  • 8,331
  • 8
  • 53
  • 82

1 Answers1

2

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" }')
Gabriel Cappelli
  • 3,632
  • 1
  • 17
  • 31