当前位置:   article > 正文

fastapi实现websocket在线聊天_fastapi websocket

fastapi websocket

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

1.问题难点

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

2.解决方案

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

3.方案设计

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

4.代码实现

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

# 初始化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")
  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28
  • 29
  • 30
  • 31
  • 32
  • 33
  • 34
  • 35
  • 36
  • 37
  • 38
  • 39

②websocket处理类

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]

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28
  • 29
  • 30
  • 31
  • 32
  • 33
  • 34
  • 35
  • 36
  • 37
  • 38
  • 39
  • 40
  • 41
  • 42
  • 43
  • 44
  • 45
  • 46
  • 47
  • 48
  • 49
  • 50
  • 51
  • 52
  • 53
  • 54
  • 55
  • 56
  • 57
  • 58
  • 59

③websocket连接

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]

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14

④http请求进行消息发布

@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': ''}
  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15

5.源码

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

6.参考文章

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

声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/秋刀鱼在做梦/article/detail/819492
推荐阅读
相关标签
  

闽ICP备14008679号