初始化提交

2026-04-07 14:11:45 +08:00
parent 46342134b6
commit c412e4de5f
69 changed files with 19321 additions and 1 deletions
@@ -0,0 +1,119 @@
+# ///
+# 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()