-
-
Save wshayes/c22a07e9815d980a9a1d0bd1ab56a690 to your computer and use it in GitHub Desktop.
# 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. |
It may be blocking - try running uvicorn with multiple workers and see if that works better. (Added later as this comment wasn't clear) it may be blocking because you added code that wasn't async (e.g. blocking IO - file access, database access, etc that wasn't managed by an async library). The code above should not block as long as you are using an async runner (e.g. uvicorn) as all of the main tasks are async.
but how to avoid block main thread?
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.
@wshayes Did you not face momery leak or stack issues using fastapi websocket ?
This code not emitting sometimes on Ubuntu Server 20.04. What I missed?