142 lines
5.4 KiB
Python
142 lines
5.4 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
Redis 适配器
|
||
基于 redis-py >= 4.x 的异步客户端
|
||
支持：Pub/Sub、Stream（持久化）、Key/Value 消息存储
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
from typing import Optional
|
||
|
||
import redis.asyncio as aioredis
|
||
|
||
from uas.comm.base import CommBaseAdapter, ConnectionConfig, MessageHandler
|
||
from uas.comm.message import Message
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class RedisAdapterComm(CommBaseAdapter):
|
||
"""
|
||
Redis 通信适配器
|
||
URL 格式: redis://[:password@]host[:port][/db][?options]
|
||
示例: redis://:secret@127.0.0.1:6379/0?socket_timeout=5
|
||
"""
|
||
|
||
def __init__(self, config: ConnectionConfig):
|
||
super().__init__(config)
|
||
self._client: Optional[aioredis.Redis] = None
|
||
self._pubsub: Optional[aioredis.client.PubSub] = None
|
||
self._listen_task: Optional[asyncio.Task] = None
|
||
|
||
# ── 生命周期 ──────────────────────────────────────────────────
|
||
|
||
async def connect(self) -> None:
|
||
opts = self.config.options
|
||
db = int(self.config.database) if self.config.database.isdigit() else 0
|
||
|
||
self._client = aioredis.Redis(
|
||
host = self.config.host,
|
||
port = self.config.port,
|
||
password = self.config.password or None,
|
||
db = db,
|
||
socket_timeout = float(opts.get("socket_timeout", 5)),
|
||
decode_responses = True,
|
||
)
|
||
# 连通性测试
|
||
await self._client.ping()
|
||
self._pubsub = self._client.pubsub()
|
||
self._running = True
|
||
self.logger.info(f"✅ Redis 已连接: {self.config.host}:{self.config.port} db={db}")
|
||
|
||
async def disconnect(self) -> None:
|
||
self._running = False
|
||
if self._listen_task:
|
||
self._listen_task.cancel()
|
||
try:
|
||
await self._listen_task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
if self._pubsub:
|
||
await self._pubsub.close()
|
||
if self._client:
|
||
await self._client.aclose()
|
||
self.logger.info("🔌 Redis 连接已关闭")
|
||
|
||
# ── 核心功能 ──────────────────────────────────────────────────
|
||
|
||
async def publish(self, message: Message) -> bool:
|
||
if not self._client:
|
||
raise RuntimeError("Redis 未连接")
|
||
try:
|
||
payload = json.dumps(message.to_dict(), ensure_ascii=False)
|
||
receivers = await self._client.publish(message.topic, payload)
|
||
self.logger.debug(f"📤 发布 topic={message.topic!r} receivers={receivers}")
|
||
return True
|
||
except Exception as e:
|
||
self.logger.error(f"发布失败: {e}")
|
||
return False
|
||
|
||
async def subscribe(self, topic: str, handler: MessageHandler) -> None:
|
||
if not self._pubsub:
|
||
raise RuntimeError("Redis 未连接")
|
||
|
||
self._register_handler(topic, handler)
|
||
|
||
# 支持通配符订阅
|
||
if "*" in topic or "?" in topic:
|
||
await self._pubsub.psubscribe(topic)
|
||
else:
|
||
await self._pubsub.subscribe(topic)
|
||
|
||
# 启动监听循环（幂等）
|
||
if self._listen_task is None or self._listen_task.done():
|
||
self._listen_task = asyncio.create_task(self._listen_loop())
|
||
|
||
self.logger.info(f"📥 订阅 Redis topic: {topic!r}")
|
||
|
||
async def unsubscribe(self, topic: str) -> None:
|
||
if self._pubsub:
|
||
if "*" in topic or "?" in topic:
|
||
await self._pubsub.punsubscribe(topic)
|
||
else:
|
||
await self._pubsub.unsubscribe(topic)
|
||
self._handlers.pop(topic, None)
|
||
self.logger.info(f"🚫 取消订阅: {topic!r}")
|
||
|
||
# ── Stream 持久化发布（可选）──────────────────────────────────
|
||
|
||
async def publish_stream(self, message: Message, maxlen: int = 1000) -> str:
|
||
"""使用 Redis Stream 发布（消息持久化，支持消费组）"""
|
||
stream_id = await self._client.xadd(
|
||
message.topic,
|
||
{"data": json.dumps(message.to_dict(), ensure_ascii=False)},
|
||
maxlen=maxlen,
|
||
)
|
||
self.logger.debug(f"📤 Stream 发布 id={stream_id}")
|
||
return stream_id
|
||
|
||
# ── 内部监听循环 ──────────────────────────────────────────────
|
||
|
||
async def _listen_loop(self):
|
||
self.logger.info("🔄 Redis 监听循环启动")
|
||
try:
|
||
async for raw in self._pubsub.listen():
|
||
if raw["type"] not in ("message", "pmessage"):
|
||
continue
|
||
try:
|
||
data = json.loads(raw["data"])
|
||
message = Message.from_dict(data)
|
||
await self._dispatch(message)
|
||
except Exception as e:
|
||
self.logger.error(f"消息解析异常: {e}", exc_info=True)
|
||
except asyncio.CancelledError:
|
||
self.logger.info("🛑 Redis 监听循环已停止")
|
||
|
||
async def health_check(self) -> bool:
|
||
try:
|
||
return bool(await self._client.ping())
|
||
except Exception:
|
||
return False |