init

2025-11-28 21:20:40 +08:00
commit f940b95b67
73 changed files with 15721 additions and 0 deletions
--- a/utils/queue.py
+++ b/utils/queue.py
@@ -0,0 +1,141 @@
+import json
+import time
+from datetime import datetime
+from typing import Dict, Optional
+
+# 内存队列（当Redis不可用时使用）
+_memory_queue = []
+_memory_processing = {}
+_memory_results = {}
+
+def get_redis_client():
+    """获取Redis客户端"""
+    try:
+        from app import redis_client
+        return redis_client
+    except:
+        return None
+
+class ParseQueue:
+    """解析队列管理器"""
+
+    QUEUE_KEY = "parse_queue"
+    PROCESSING_KEY = "parse_processing"
+    RESULT_KEY_PREFIX = "parse_result:"
+
+    @staticmethod
+    def add_task(task_id: str, video_url: str, user_id: Optional[int] = None, ip_address: str = ""):
+        """添加任务到队列"""
+        task = {
+            'task_id': task_id,
+            'video_url': video_url,
+            'user_id': user_id,
+            'ip_address': ip_address,
+            'created_at': datetime.utcnow().isoformat(),
+            'status': 'queued'
+        }
+
+        redis_client = get_redis_client()
+        if redis_client:
+            redis_client.rpush(ParseQueue.QUEUE_KEY, json.dumps(task))
+        else:
+            # 使用内存队列
+            _memory_queue.append(task)
+        return task_id
+
+    @staticmethod
+    def get_task() -> Optional[Dict]:
+        """从队列获取任务"""
+        redis_client = get_redis_client()
+        if redis_client:
+            task_json = redis_client.lpop(ParseQueue.QUEUE_KEY)
+            if task_json:
+                task = json.loads(task_json)
+                redis_client.hset(ParseQueue.PROCESSING_KEY, task['task_id'], json.dumps(task))
+                return task
+        else:
+            # 使用内存队列
+            if _memory_queue:
+                task = _memory_queue.pop(0)
+                _memory_processing[task['task_id']] = task
+                return task
+        return None
+
+    @staticmethod
+    def complete_task(task_id: str, result: Dict):
+        """完成任务"""
+        redis_client = get_redis_client()
+        if redis_client:
+            redis_client.hdel(ParseQueue.PROCESSING_KEY, task_id)
+            redis_client.setex(
+                f"{ParseQueue.RESULT_KEY_PREFIX}{task_id}",
+                3600,
+                json.dumps(result)
+            )
+        else:
+            # 使用内存
+            _memory_processing.pop(task_id, None)
+            _memory_results[task_id] = result
+
+    @staticmethod
+    def get_result(task_id: str) -> Optional[Dict]:
+        """获取任务结果"""
+        redis_client = get_redis_client()
+        if redis_client:
+            result_json = redis_client.get(f"{ParseQueue.RESULT_KEY_PREFIX}{task_id}")
+            if result_json:
+                return json.loads(result_json)
+        else:
+            # 使用内存
+            return _memory_results.get(task_id)
+        return None
+
+    @staticmethod
+    def get_queue_length() -> int:
+        """获取队列长度"""
+        redis_client = get_redis_client()
+        if redis_client:
+            return redis_client.llen(ParseQueue.QUEUE_KEY)
+        else:
+            return len(_memory_queue)
+
+    @staticmethod
+    def get_processing_count() -> int:
+        """获取正在处理的任务数"""
+        redis_client = get_redis_client()
+        if redis_client:
+            return redis_client.hlen(ParseQueue.PROCESSING_KEY)
+        else:
+            return len(_memory_processing)
+
+    @staticmethod
+    def get_queue_status() -> Dict:
+        """获取队列状态"""
+        return {
+            'queued': ParseQueue.get_queue_length(),
+            'processing': ParseQueue.get_processing_count()
+        }
+
+
+class ConcurrencyController:
+    """并发控制器"""
+
+    @staticmethod
+    def can_process() -> bool:
+        """检查是否可以处理新任务"""
+        from models import SiteConfig
+        config = SiteConfig.query.filter_by(config_key='max_concurrent').first()
+        max_concurrent = int(config.config_value) if config else 3
+
+        processing_count = ParseQueue.get_processing_count()
+        return processing_count < max_concurrent
+
+    @staticmethod
+    def wait_for_slot(timeout: int = 60) -> bool:
+        """等待可用槽位"""
+        start_time = time.time()
+        while time.time() - start_time < timeout:
+            if ConcurrencyController.can_process():
+                return True
+            time.sleep(0.5)
+        return False