75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
import aio_pika
|
||
from config import RABBITMQ_URI
|
||
|
||
|
||
async def setup_multi_queue_balance(
|
||
exchange_name="demo.direct.multi.queue",
|
||
queue_prefix="task.queue.",
|
||
route_prefix="route.",
|
||
queue_count=3 # 3个队列
|
||
):
|
||
"""创建多个队列，每个队列绑定不同的路由键"""
|
||
connection = await aio_pika.connect_robust(RABBITMQ_URI)
|
||
channel = await connection.channel()
|
||
|
||
# 声明Direct交换器
|
||
await channel.declare_exchange(
|
||
exchange_name,
|
||
aio_pika.ExchangeType.DIRECT,
|
||
durable=True
|
||
)
|
||
|
||
# 创建N个队列，每个队列绑定独立的路由键
|
||
for i in range(queue_count):
|
||
queue_name = f"{queue_prefix}{i + 1}"
|
||
route_key = f"{route_prefix}{i + 1}"
|
||
|
||
# 声明队列
|
||
queue = await channel.declare_queue(
|
||
queue_name,
|
||
durable=True,
|
||
auto_delete=False
|
||
)
|
||
|
||
# 绑定路由键（每个队列对应唯一路由键）
|
||
await queue.bind(exchange_name, routing_key=route_key)
|
||
print(f"Queue {queue_name} bound to routing key: {route_key}")
|
||
|
||
await connection.close()
|
||
|
||
|
||
class BalancedProducer:
|
||
def __init__(self, exchange_name="demo.direct.multi.queue", queue_count=3):
|
||
self.exchange_name = exchange_name
|
||
self.queue_count = queue_count
|
||
self.current_index = 0 # 轮询索引
|
||
|
||
async def connect(self):
|
||
"""建立连接（可复用连接提升性能）"""
|
||
self.connection = await aio_pika.connect_robust(RABBITMQ_URI)
|
||
self.channel = await self.connection.channel()
|
||
self.exchange = await self.channel.declare_exchange(
|
||
self.exchange_name,
|
||
aio_pika.ExchangeType.DIRECT,
|
||
durable=True
|
||
)
|
||
|
||
async def publish(self, message: str):
|
||
"""轮询选择路由键，发送消息到对应的队列"""
|
||
# 轮询算法：每次发送后切换到下一个路由键
|
||
self.current_index = (self.current_index + 1) % self.queue_count
|
||
route_key = f"route.{self.current_index + 1}" # 对应 route_1/2/3
|
||
|
||
message_obj = aio_pika.Message(
|
||
body=message.encode("utf-8"),
|
||
delivery_mode=aio_pika.DeliveryMode.PERSISTENT
|
||
)
|
||
|
||
await self.exchange.publish(message_obj, routing_key=route_key)
|
||
print(f"Message sent: {message} (routed to {route_key})")
|
||
|
||
async def close(self):
|
||
"""关闭连接"""
|
||
await self.connection.close()
|
||
|