fastapi实现websocket在线聊天

最近要实现一个在线聊天功能,基于fastapi的websocket实现了这个功能。下面介绍一下遇到的技术问题

1.问题难点

在线上环境部署时,一般是多进程的方式进行部署启动fastapi服务,而每个启动的进程都有自己的独立存储空间。导致存储的连接对象分布在不同的进程中,当进行通信时,可能无法找到已连接的连接对象。

2.解决方案

使用使用redis的订阅发布机制,使所有的进程都能进行消息订阅。这样能保证每个进程收到消息后都会进行相关的信息处理了。

3.方案设计

  • 每个进程启动的时候都进行一个消息的订阅。
  • 通过http请求,进行消息发布。
  • 每个进程收到发布的消息后,进行判断是否由自己进行处理。

4.代码实现

①在服务启动时,进行消息订阅,并一直监听消息通道。当有消息发布时,进行消息处理。

python 复制代码
# 初始化app
app = FastAPI(title="Ws Chat", description="测试", version="1.0.0")
app.openapi_version = "3.0.0"


app.include_router(chat.app, prefix='/api/chat', tags=['Chat'])

@app.on_event('startup')
async def on_startup():
    print(f"订阅初始化:{os.getpid()}")
    # 执行消息订阅机制https://aioredis.readthedocs.io/en/latest/examples/
    loop = asyncio.get_event_loop()
    loop.create_task(register_pubsub())


async def reader(channel):
    # 进行消息的消费
    async for msg in channel.listen():  # 监听通道
        # print(msg)
        msg_data = msg.get("data")
        if msg_data and isinstance(msg_data, str):
            msg_data_dict = json.loads(msg_data)
            print(f"chat:{msg_data_dict}")
            sender = msg_data_dict.get("sender")
            # 进行消息处理
            await chat.cm.handle_websocket_message(msg_data_dict, sender)


async def register_pubsub():
    pool = aioredis.from_url(
        "redis://{}".format(host), db=db, password=password, port=port, encoding="utf-8", decode_responses=True
    )
    psub = pool.pubsub()

    async with psub as p:
        # 消息订阅
        await p.subscribe("chat")
        await reader(p)
        await p.unsubscribe("chat")

②websocket处理类

python 复制代码
from fastapi import WebSocket, WebSocketDisconnect


class ConnectionManager:
    def __init__(self):
        # 保存当前所有的链接的websocket对象
        self.websocket_connections = {}

    async def connect(self, websocket: WebSocket, client_id):
        # 添加连接并发送欢迎消息
        await websocket.accept()
        self.websocket_connections[client_id] = websocket
        await websocket.send_json({"type": "system",
                                   "msg": "Welcome to the chat app!",
                                   "sender": "system",
                                   "recipient": client_id})

        try:
            # 处理消息
            while True:
                # 获取信息
                message = await websocket.receive_json()
                # 处理发送信息
                await self.handle_websocket_message(message, client_id)

        except WebSocketDisconnect:
            # 连接断开时移除连接
            del self.websocket_connections[client_id]

    async def handle_websocket_message(self, message: dict, client_id):

        # 处理私聊消息
        if message.get("type") == "private_message":
            recipient = message.get("recipient")
            msg = message.get("msg")
            recipient_conn = self.websocket_connections.get(recipient)

            if recipient_conn:
                # 在线
                await recipient_conn.send_json({"type": "private_message",
                                                "sender": client_id,
                                                "msg": msg,
                                                "recipient": recipient})

    async def broadcast(self, message: dict):
        # 循环变量给所有在线激活的链接发送消息-全局广播
        for connection in self.websocket_connections:
            await connection.send_text(message)

    async def close(self, websocket: WebSocket, client_id):
        # 断开客户端的链接
        await websocket.close()
        del self.websocket_connections[client_id]

    async def disconnect(self, user_id):
        websocket: WebSocket = self.websocket_connections[user_id]
        await websocket.close()
        del self.websocket_connections[user_id]

③websocket连接

python 复制代码
from app.chat_manager.server import ConnectionManager

cm = ConnectionManager()

@app.websocket("/connect_chat")
async def connect_chat(websocket: WebSocket, user_code: str):
    try:

        await cm.connect(websocket, user_code)

    except WebSocketDisconnect:
        # 连接断开时移除连接
        del cm.websocket_connections[user_code]

④http请求进行消息发布

python 复制代码
@app.post("/create_chat", summary="发起聊天")
async def create_chat(param: DiagnosisChatSch, r=Depends(get_redis)):
    """

    """

    ws_param = {"type": "private_message",
                "msg": param.msg,
                "sender": param.sender,
                "recipient": param.recipient}

    # 进行消息发布
    await r.publish('diagnosis_chat', json.dumps(ws_param))

    return {'code': 200, 'msg': '成功', 'data': ''}

5.源码

github源码地址:https://github.com/zhangyukuo/fastapi_ws_chat

6.参考文章

https://www.cnblogs.com/a00ium/p/16931133.html

相关推荐
悟空胆好小41 分钟前
分音塔科技(BABEL Technology) 的公司背景、股权构成、产品类型及技术能力的全方位解读
网络·人工智能·科技·嵌入式硬件
DoraBigHead1 小时前
《电磁波的浪漫,铜线上的灵魂》——计算机网络·物理层全解版
网络协议
爱吃土豆的马铃薯ㅤㅤㅤㅤㅤㅤㅤㅤㅤ1 小时前
如何使用Java WebSocket API实现客户端和服务器端的通信?
java·开发语言·websocket
cooldream20092 小时前
Python 包管理新时代:深入了解 `uv` 的使用与实践
python·uv·包管理器
之歆2 小时前
Python-魔术方法-创建、初始化与销毁-hash-bool-可视化-运算符重载-容器和大小-可调用对象-上下文管理-反射-描述器-二分-学习笔记
笔记·python·学习
胖达不服输2 小时前
「日拱一码」025 机器学习——评价指标
人工智能·python·机器学习·评价指标
ssswywywht2 小时前
OSPF实验
网络
FCM662 小时前
HCIA第一次实验报告:静态路由综合实验
网络·tcp/ip·信息与通信
apihz2 小时前
VM虚拟机全版本网盘+免费本地网络穿透端口映射实时同步动态家庭IP教程
android·服务器·开发语言·网络·数据库·网络协议·tcp/ip
dog2503 小时前
TCP 传输时 sk_buff 的 clone 和 unclone
网络·网络协议·tcp/ip