如何结合FastAPI WebSocket端点与RabbitMQ Stream（basic_consume）？

核心问题点 #

1. 同步客户端与异步环境冲突：你使用的pika是同步RabbitMQ客户端，channel.start_consuming()会阻塞FastAPI的异步事件循环，导致WebSocket协程无法正常处理后续消息和异步回调，这是切换到Stream队列后失效的核心原因。
2. 错误的停止消费逻辑：回调函数里的channel.stop_consuming()会让消费在处理完第一条消息后就终止，完全违背流式传输的需求。
3. 未确认消息：设置了auto_ack=False但未调用basic_ack，会导致消息堆积在Stream队列中，无法被正确清理。
4. Stream队列适配问题：普通的basic_consume对Stream队列的支持有限，需要使用适配Stream的消费方式，比如指定消费偏移。

修复步骤 #

1. 替换为异步客户端aio-pika #

pika的同步模型和FastAPI的异步环境不兼容，改用aio-pika（异步RabbitMQ客户端）才能完美适配。先安装依赖：

pip install aio-pika

2. 重构WebSocket消费逻辑 #

以下是修改后的代码，适配Stream队列和异步环境：

import aio_pika
import asyncio
import json
import cv2
import base64
from fastapi import WebSocket, Depends, APIRouter
from your_module import get_current_user_ws, PipelineVisualizer, STREAMING_EXCHANGE

router = APIRouter()

@router.websocket("/ws/messages")
async def websocket_messages(
    websocket: WebSocket,
    device_name: str,
    pipeline_name: str,
    username: str = Depends(get_current_user_ws)
):
    await websocket.accept()
    pipeline = await websocket.receive_text()
    queue_name = pipeline_name

    # 建立异步RabbitMQ连接
    connection = await aio_pika.connect_robust("amqp://user:password@localhost/")  # 替换为你的RabbitMQ地址
    async with connection:
        channel = await connection.channel()
        await channel.set_qos(prefetch_count=1)

        # 声明Stream队列
        queue = await channel.declare_queue(
            queue_name,
            durable=True,
            arguments={
                "x-queue-type": "stream",
                "x-max-length-bytes": 20_000_000_000
            }
        )
        # 绑定交换机
        await queue.bind(STREAMING_EXCHANGE, routing_key=f"{device_name}/{pipeline_name}")

        pipeline_visualizer = PipelineVisualizer(json.loads(pipeline))

        async def process_message(message: aio_pika.IncomingMessage):
            async with message.process():  # 自动确认消息，替代手动basic_ack
                payload_str = message.body.decode("utf-8")
                record_data = json.loads(payload_str)
                data, data_type = pipeline_visualizer.render_frame(record_data)

                output = {"messages": record_data}

                if data_type == "image":
                    _, image = cv2.imencode('.jpg', data)
                    image = image.tobytes()
                    image = base64.b64encode(image).decode('ascii')
                    output.update({
                        "visualizationData": image,
                        "type": "image"
                    })
                elif data_type == "pointcloud":
                    with open("./out.pcd", "rb") as fp:
                        data_bytes = fp.read()
                    pointcloud = base64.b64encode(data_bytes).decode('ascii')
                    output.update({
                        "visualizationData": pointcloud,
                        "type": "pointcloud"
                    })

                await websocket.send_text(json.dumps(output))
                await asyncio.sleep(1/60)

        # 启动异步消费，指定Stream队列的消费偏移（这里用latest，即从最新消息开始）
        await queue.consume(process_message, arguments={"x-stream-offset": "latest"})

        # 保持协程运行，直到WebSocket断开
        await asyncio.Future()
    finally:
        await websocket.close()

关键修改说明 #

• 异步连接与消费：用aio-pika的异步API替代pika的同步调用，避免阻塞事件循环。
• 自动消息确认：message.process()上下文管理器会自动完成消息确认，无需手动调用basic_ack。
• Stream队列消费偏移：通过arguments={"x-stream-offset": "latest"}指定消费起始位置，可选值还有"earliest"或具体的偏移值，符合Stream队列的特性。
• 移除错误的停止消费逻辑：去掉了stop_consuming，让消费持续运行直到WebSocket断开。
• 协程保持运行：await asyncio.Future()让WebSocket协程一直处于运行状态，直到断开连接触发异常进入finally块。

额外注意事项 #

• 确保RabbitMQ服务器开启了Stream插件：执行rabbitmq-plugins enable rabbitmq_stream。
• 若需要精确控制消费位置，可以保存消费的偏移量，下次连接时指定x-stream-offset为该值，实现断点续传。
• 对于点云文件的读取，建议避免每次都读本地文件，最好直接使用render_frame返回的data处理，减少IO开销。

内容的提问来源于stack exchange，提问作者Luis Filipe Araujo