120 lines
3.6 KiB
Python
120 lines
3.6 KiB
Python
# ///
|
|
# queue_service.py
|
|
# 描述：消息队列服务，支持Redis后端
|
|
# 作者：AI Generated
|
|
# 创建日期：2026-04-05
|
|
# ///
|
|
|
|
import asyncio
|
|
import json
|
|
from typing import Dict, Any, Optional, Callable
|
|
from datetime import datetime
|
|
from enum import Enum
|
|
|
|
from config import get_settings
|
|
|
|
|
|
class QueueName(Enum):
|
|
MESSAGE_SEND = "wxauto:queue:message_send"
|
|
MESSAGE_BROADCAST = "wxauto:queue:message_broadcast"
|
|
|
|
|
|
class QueueService:
|
|
def __init__(self):
|
|
self._redis_client = None
|
|
self._redis_available = False
|
|
self._local_queue = asyncio.Queue()
|
|
self._init_redis()
|
|
self._worker_task = None
|
|
self._running = False
|
|
|
|
def _init_redis(self):
|
|
try:
|
|
import redis
|
|
settings = get_settings()
|
|
redis_config = settings.redis
|
|
self._redis_client = redis.from_url(
|
|
redis_config.url,
|
|
db=redis_config.db,
|
|
password=redis_config.password,
|
|
decode_responses=True
|
|
)
|
|
self._redis_client.ping()
|
|
self._redis_available = True
|
|
except Exception:
|
|
self._redis_available = False
|
|
self._redis_client = None
|
|
|
|
def enqueue(self, queue_name: QueueName, data: Dict[str, Any]) -> bool:
|
|
job_data = {
|
|
"data": data,
|
|
"enqueued_at": datetime.now().isoformat(),
|
|
"status": "pending"
|
|
}
|
|
|
|
if self._redis_available:
|
|
try:
|
|
self._redis_client.rpush(queue_name.value, json.dumps(job_data, ensure_ascii=False))
|
|
return True
|
|
except Exception:
|
|
return False
|
|
else:
|
|
try:
|
|
asyncio.create_task(self._local_enqueue(queue_name, job_data))
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
async def _local_enqueue(self, queue_name: QueueName, job_data: Dict[str, Any]):
|
|
await self._local_queue.put((queue_name, job_data))
|
|
|
|
async def dequeue(self, queue_name: QueueName, timeout: int = 1) -> Optional[Dict[str, Any]]:
|
|
if self._redis_available:
|
|
try:
|
|
result = self._redis_client.blpop(queue_name.value, timeout=timeout)
|
|
if result:
|
|
_, job_data = result
|
|
return json.loads(job_data)
|
|
except Exception:
|
|
pass
|
|
else:
|
|
try:
|
|
queue_item = await asyncio.wait_for(self._local_queue.get(), timeout=timeout)
|
|
if queue_item and queue_item[0] == queue_name:
|
|
return queue_item[1]
|
|
except asyncio.TimeoutError:
|
|
pass
|
|
return None
|
|
|
|
def get_queue_length(self, queue_name: QueueName) -> int:
|
|
if self._redis_available:
|
|
try:
|
|
return self._redis_client.llen(queue_name.value)
|
|
except Exception:
|
|
pass
|
|
return 0
|
|
|
|
def start_worker(self, processor: Callable):
|
|
if self._running:
|
|
return
|
|
self._running = True
|
|
self._worker_task = asyncio.create_task(self._worker_loop(processor))
|
|
|
|
async def _worker_loop(self, processor: Callable):
|
|
while self._running:
|
|
for queue_name in [QueueName.MESSAGE_SEND]:
|
|
job = await self.dequeue(queue_name, timeout=1)
|
|
if job:
|
|
try:
|
|
await processor(job)
|
|
except Exception as e:
|
|
pass
|
|
|
|
def stop_worker(self):
|
|
self._running = False
|
|
if self._worker_task:
|
|
self._worker_task.cancel()
|
|
|
|
|
|
queue_service = QueueService()
|