我已经用 FastAPI 实现了一个简单的 WebSocket 代理（使用这个例子 https://fastapi.tiangolo.com/advanced/websockets/)
应用程序的目标是将其收到的所有消息传递到其活动连接（代理）。
它仅适用于单个实例，因为它将活动的 WebSocket 连接保留在内存中。当存在多个实例时，内存不会共享。
我天真的方法是通过在某些共享存储（Redis）中保持活动连接来解决这个问题。但我坚持腌制它。
这是完整的应用程序：
import pickle
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from collections import defaultdict
import redis
app = FastAPI()
rds = redis.StrictRedis('localhost')
class ConnectionManager:
def __init__(self):
self.active_connections = defaultdict(dict)
async def connect(self, websocket: WebSocket, application: str, client_id: str):
await websocket.accept()
if application not in self.active_connections:
self.active_connections[application] = defaultdict(list)
self.active_connections[application][client_id].append(websocket)
#### this is my attempt to store connections ####
rds.set('connections', pickle.dumps(self.active_connections))
def disconnect(self, websocket: WebSocket, application: str, client_id: str):
self.active_connections[application][client_id].remove(websocket)
async def broadcast(self, message: dict, application: str, client_id: str):
for connection in self.active_connections[application][client_id]:
try:
await connection.send_json(message)
print(f"sent {message}")
except Exception as e:
pass
manager = ConnectionManager()
@app.websocket("/ws/channel/{application}/{client_id}/")
async def websocket_endpoint(websocket: WebSocket, application: str, client_id: str):
await manager.connect(websocket, application, client_id)
while True:
try:
data = await websocket.receive_json()
print(f"received: {data}")
await manager.broadcast(data, application, client_id)
except WebSocketDisconnect:
manager.disconnect(websocket, application, client_id)
except RuntimeError:
break
if __name__ == '__main__':
import uvicorn
uvicorn.run(app, host='0.0.0.0', port=8005)
但是，pickle websocket 连接不成功：
AttributeError: Can't pickle local object 'FastAPI.setup.<locals>.openapi'
在应用程序实例之间存储 WebSocket 连接的正确方法是什么？
UPD每个@AKX 答案的实际解决方案。
服务器的每个实例都订阅了 Redis pubsub 并尝试将收到的消息发送到其所有连接的客户端。
由于一个客户端无法连接到多个实例 - 每条消息只能传递给每个客户端一次
import json
import asyncio
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from collections import defaultdict
import redis
app = FastAPI()
rds = redis.StrictRedis('localhost')
class ConnectionManager:
def __init__(self):
self.active_connections = defaultdict(dict)
async def connect(self, websocket: WebSocket, application: str, client_id: str):
await websocket.accept()
if application not in self.active_connections:
self.active_connections[application] = defaultdict(list)
self.active_connections[application][client_id].append(websocket)
def disconnect(self, websocket: WebSocket, application: str, client_id: str):
self.active_connections[application][client_id].remove(websocket)
async def broadcast(self, message: dict, application: str, client_id: str):
for connection in self.active_connections[application][client_id]:
try:
await connection.send_json(message)
print(f"sent {message}")
except Exception as e:
pass
async def consume(self):
print("started to consume")
sub = rds.pubsub()
sub.subscribe('channel')
while True:
await asyncio.sleep(0.01)
message = sub.get_message(ignore_subscribe_messages=True)
if message is not None and isinstance(message, dict):
msg = json.loads(message.get('data'))
await self.broadcast(msg['message'], msg['application'], msg['client_id'])
manager = ConnectionManager()
@app.on_event("startup")
async def subscribe():
asyncio.create_task(manager.consume())
@app.websocket("/ws/channel/{application}/{client_id}/")
async def websocket_endpoint(websocket: WebSocket, application: str, client_id: str):
await manager.connect(websocket, application, client_id)
while True:
try:
data = await websocket.receive_json()
print(f"received: {data}")
rds.publish(
'channel',
json.dumps({
'application': application,
'client_id': client_id,
'message': data
})
)
except WebSocketDisconnect:
manager.disconnect(websocket, application, client_id)
except RuntimeError:
break
if __name__ == '__main__': # pragma: no cover
import uvicorn
uvicorn.run(app, host='0.0.0.0', port=8005)
