Skip to content

Instantly share code, notes, and snippets.

@wshayes
Created May 27, 2019 20:19
Show Gist options
  • Save wshayes/c22a07e9815d980a9a1d0bd1ab56a690 to your computer and use it in GitHub Desktop.
Save wshayes/c22a07e9815d980a9a1d0bd1ab56a690 to your computer and use it in GitHub Desktop.
[Websocket demo for fastapi] example of broadcast using websockets for fastapi #fastapi #websockets
# From https://github.com/tiangolo/fastapi/issues/258
from typing import List
from fastapi import FastAPI
from starlette.responses import HTMLResponse
from starlette.websockets import WebSocket, WebSocketDisconnect
app = FastAPI()
html = """
<!DOCTYPE html>
<html>
<head>
<title>Chat</title>
</head>
<body>
<h1>WebSocket Chat</h1>
<form action="" onsubmit="sendMessage(event)">
<input type="text" id="messageText" autocomplete="off"/>
<button>Send</button>
</form>
<ul id='messages'>
</ul>
<script>
var ws = new WebSocket("ws://localhost:8000/ws");
ws.onmessage = function(event) {
var messages = document.getElementById('messages')
var message = document.createElement('li')
var content = document.createTextNode(event.data)
message.appendChild(content)
messages.appendChild(message)
};
function sendMessage(event) {
var input = document.getElementById("messageText")
ws.send(input.value)
input.value = ''
event.preventDefault()
}
</script>
</body>
</html>
"""
@app.get("/")
async def get():
return HTMLResponse(html)
class Notifier:
def __init__(self):
self.connections: List[WebSocket] = []
self.generator = self.get_notification_generator()
async def get_notification_generator(self):
while True:
message = yield
await self._notify(message)
async def push(self, msg: str):
await self.generator.asend(msg)
async def connect(self, websocket: WebSocket):
await websocket.accept()
self.connections.append(websocket)
def remove(self, websocket: WebSocket):
self.connections.remove(websocket)
async def _notify(self, message: str):
living_connections = []
while len(self.connections) > 0:
# Looping like this is necessary in case a disconnection is handled
# during await websocket.send_text(message)
websocket = self.connections.pop()
await websocket.send_text(message)
living_connections.append(websocket)
self.connections = living_connections
notifier = Notifier()
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
await notifier.connect(websocket)
try:
while True:
data = await websocket.receive_text()
await websocket.send_text(f"Message text was: {data}")
except WebSocketDisconnect:
notifier.remove(websocket)
@app.get("/push/{message}")
async def push_to_connected_websockets(message: str):
await notifier.push(f"! Push notification: {message} !")
@app.on_event("startup")
async def startup():
# Prime the push notification generator
await notifier.generator.asend(None)
If you want to test it out, copy the above to a file main.py and start the server:
uvicorn main:app --reload
Then, open a few tabs at http://localhost:8000/ and send some chat messages (this should work the same as the base tutorial app). Then open http://localhost:8000/push/hello%20world and you should receive a push notification in each of your open tabs showing the message hello world.
@LowinLi
Copy link

LowinLi commented Feb 9, 2023

but how to avoid block main thread?

@wshayes
Copy link
Author

wshayes commented Feb 9, 2023

As long as everything you wrote in your code is async, it will hand off control to the next task and run fine with a single thread. If you have blocking file, database or other long-running tasks as part of your code that are not through async libraries, they will block your async tasks.

@CrasCris
Copy link

@wshayes Did you not face momery leak or stack issues using fastapi websocket ?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment