153 lines
4.9 KiB
Python
153 lines
4.9 KiB
Python
"""
|
||
RabbitMQ 可靠消息生产者模块
|
||
"""
|
||
|
||
import asyncio
|
||
import aio_pika
|
||
import json
|
||
import logging
|
||
from datetime import datetime
|
||
from typing import Dict, Any, Optional
|
||
|
||
from .config import config
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class ReliableProducer:
|
||
"""可靠消息生产者"""
|
||
|
||
def __init__(self,
|
||
exchange_name: Optional[str] = None,
|
||
queue_name: Optional[str] = None):
|
||
"""
|
||
初始化生产者
|
||
|
||
Args:
|
||
exchange_name: 交换器名称，默认使用配置中的值
|
||
queue_name: 队列名称，默认使用配置中的值
|
||
"""
|
||
self.exchange_name = exchange_name or config.exchange_name
|
||
self.queue_name = queue_name or config.queue_name
|
||
self.connection = None
|
||
self.channel = None
|
||
self.exchange = None
|
||
self.queue = None
|
||
|
||
async def connect(self):
|
||
"""建立连接并设置确认机制"""
|
||
try:
|
||
# 使用 robust 连接，支持自动重连
|
||
connection_config = config.get_connection_config()
|
||
self.connection = await aio_pika.connect_robust(connection_config['uri'])
|
||
self.channel = await self.connection.channel()
|
||
|
||
# 开启发布确认机制 - 确保消息成功发送到队列
|
||
await self.channel.set_qos(prefetch_count=connection_config['prefetch_count'])
|
||
|
||
# 声明持久化交换器
|
||
self.exchange = await self.channel.declare_exchange(
|
||
self.exchange_name,
|
||
aio_pika.ExchangeType.DIRECT,
|
||
durable=True # 交换器持久化
|
||
)
|
||
|
||
# 声明持久化队列
|
||
self.queue = await self.channel.declare_queue(
|
||
self.queue_name,
|
||
durable=True, # 队列持久化
|
||
auto_delete=False, # 队列不自动删除
|
||
)
|
||
|
||
# 绑定队列到交换器
|
||
await self.queue.bind(self.exchange, routing_key="reliable")
|
||
|
||
logger.info(f"[生产者] 已连接，队列: {self.queue_name}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"[生产者] 连接失败: {e}")
|
||
raise
|
||
|
||
def _generate_message_id(self, message_data: Dict[str, Any]) -> str:
|
||
"""
|
||
为消息生成消息ID
|
||
对于 duplicate_test 类型的消息，生成固定的ID用于测试幂等性
|
||
|
||
Args:
|
||
message_data: 消息数据字典
|
||
|
||
Returns:
|
||
str: 消息ID
|
||
"""
|
||
message_type = message_data.get('type', '')
|
||
content = message_data.get('content', '')
|
||
|
||
# 对于 duplicate_test 类型的消息，基于内容生成固定ID
|
||
if message_type == 'duplicate_test':
|
||
# 使用内容生成固定的消息ID
|
||
import hashlib
|
||
content_hash = hashlib.md5(content.encode('utf-8')).hexdigest()
|
||
return f"duplicate_{content_hash[:8]}"
|
||
else:
|
||
# 其他消息使用时间戳生成唯一ID
|
||
return f"msg_{asyncio.get_running_loop().time()}"
|
||
|
||
async def publish_reliable_message(self, message_data: Dict[str, Any]) -> bool:
|
||
"""
|
||
发送可靠消息
|
||
|
||
Args:
|
||
message_data: 消息数据字典
|
||
|
||
Returns:
|
||
bool: 发送是否成功
|
||
"""
|
||
try:
|
||
# 生成消息ID
|
||
message_id = self._generate_message_id(message_data)
|
||
|
||
# 添加消息元数据
|
||
message_data.update({
|
||
'timestamp': datetime.now().isoformat(),
|
||
'message_id': message_id
|
||
})
|
||
|
||
# 创建持久化消息
|
||
message = aio_pika.Message(
|
||
body=json.dumps(message_data, ensure_ascii=False).encode('utf-8'),
|
||
delivery_mode=aio_pika.DeliveryMode.PERSISTENT, # 消息持久化
|
||
message_id=message_id,
|
||
timestamp=datetime.now()
|
||
)
|
||
|
||
# 发送消息并等待确认
|
||
await self.exchange.publish(
|
||
message,
|
||
routing_key="reliable"
|
||
)
|
||
|
||
logger.info(f"[生产者] 消息已发送: {message_id} (类型: {message_data.get('type', 'N/A')}, 内容: {message_data.get('content', 'N/A')})")
|
||
return True
|
||
|
||
except Exception as e:
|
||
logger.error(f"[生产者] 发送消息失败: {e}")
|
||
return False
|
||
|
||
async def close(self):
|
||
"""关闭连接"""
|
||
try:
|
||
if self.connection:
|
||
await self.connection.close()
|
||
logger.info("[生产者] 连接已关闭")
|
||
except Exception as e:
|
||
logger.error(f"[生产者] 关闭连接时出错: {e}")
|
||
|
||
async def __aenter__(self):
|
||
"""异步上下文管理器入口"""
|
||
await self.connect()
|
||
return self
|
||
|
||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||
"""异步上下文管理器出口"""
|
||
await self.close()
|