jettask 0.2.18__py3-none-any.whl → 0.2.20__py3-none-any.whl

jettask/persistence/queue_discovery.py

@@ -0,0 +1,215 @@
+"""队列发现模块
+
+负责动态发现新队列，并为新队列创建消费者组。
+使用队列注册表替代scan命令以提高性能。
+"""
+
+import asyncio
+import logging
+import traceback
+from typing import Set
+
+import redis.asyncio as redis
+from redis.asyncio import Redis
+
+logger = logging.getLogger(__name__)
+
+
+class QueueDiscovery:
+    """队列发现器
+
+    职责：
+    - 初始队列发现（启动时执行一次）
+    - 定期发现新队列
+    - 为新队列创建消费者组
+    - 更新ConsumerManager的队列列表
+    """
+
+    def __init__(
+        self,
+        redis_client: Redis,
+        redis_prefix: str,
+        consumer_group: str,
+        consumer_manager=None
+    ):
+        """初始化队列发现器
+
+        Args:
+            redis_client: Redis异步客户端
+            redis_prefix: Redis键前缀
+            consumer_group: 消费者组名称
+            consumer_manager: ConsumerManager实例（可选）
+        """
+        self.redis_client = redis_client
+        self.redis_prefix = redis_prefix
+        self.consumer_group = consumer_group
+        self.consumer_manager = consumer_manager
+
+        # 队列注册表的Redis key
+        self.queue_registry_key = f"{redis_prefix}:QUEUE_REGISTRY"
+
+        # 已知队列集合
+        self._known_queues = set()
+
+        self._running = False
+        self._discovery_task = None
+
+    async def initial_queue_discovery(self) -> Set[str]:
+        """初始队列发现，在启动时执行一次 - 使用队列注册表替代scan
+
+        Returns:
+            发现的队列集合
+        """
+        try:
+            new_queues = set()
+            logger.info(f"Starting initial queue discovery from queue registry: {self.queue_registry_key}")
+
+            # 从队列注册表获取所有队列
+            queue_members = await self.redis_client.smembers(self.queue_registry_key.encode())
+            for queue_name_bytes in queue_members:
+                queue_name = queue_name_bytes.decode('utf-8') if isinstance(queue_name_bytes, bytes) else str(queue_name_bytes)
+                new_queues.add(queue_name)
+                logger.info(f"Found registered queue: {queue_name}")
+
+            # 如果注册表为空，使用 RegistryManager 初始化
+            if not new_queues:
+                logger.warning(f"Queue registry is empty, initializing from RegistryManager...")
+                from jettask.messaging.registry import QueueRegistry
+                registry = QueueRegistry(
+                    redis_client=None,  # 同步客户端，这里不需要
+                    async_redis_client=self.redis_client,
+                    redis_prefix=self.redis_prefix
+                )
+
+                # 初始化注册表（仅在首次运行时需要）
+                await registry.initialize_from_existing_data()
+
+                # 从注册表获取队列
+                new_queues = await registry.get_all_queues()
+                logger.info(f"Got {len(new_queues)} queues from registry manager")
+
+            if new_queues:
+                logger.info(f"Initial queue discovery found {len(new_queues)} queues: {new_queues}")
+                # 合并所有队列：TASK_CHANGES + 动态发现的队列
+                # 转换 bytes 为字符串
+                string_queues = []
+                for q in new_queues:
+                    if isinstance(q, bytes):
+                        string_queues.append(q.decode('utf-8'))
+                    else:
+                        string_queues.append(str(q))
+                all_queues = string_queues + ['TASK_CHANGES']
+
+                # 更新ConsumerManager的配置
+                if self.consumer_manager:
+                    self.consumer_manager.config['queues'] = all_queues
+
+                    # 更新worker的队列信息
+                    await self._update_worker_queues(all_queues)
+
+                self._known_queues = new_queues
+
+            return new_queues
+
+        except Exception as e:
+            logger.error(f"Error in initial queue discovery: {e}")
+            logger.error(traceback.format_exc())
+            return set()
+
+    async def start_discovery(self):
+        """启动定期队列发现"""
+        self._running = True
+        self._discovery_task = asyncio.create_task(self._discover_queues_loop())
+        logger.debug("QueueDiscovery started")
+
+    async def stop_discovery(self):
+        """停止队列发现"""
+        self._running = False
+        if self._discovery_task:
+            self._discovery_task.cancel()
+            try:
+                await self._discovery_task
+            except asyncio.CancelledError:
+                pass
+        logger.debug("QueueDiscovery stopped")
+
+    async def _discover_queues_loop(self):
+        """定期发现新队列 - 使用队列注册表替代scan"""
+        while self._running:
+            try:
+                new_queues = set()
+
+                # 从队列注册表获取所有队列
+                queue_members = await self.redis_client.smembers(self.queue_registry_key)
+                for queue_name_bytes in queue_members:
+                    queue_name = queue_name_bytes.decode('utf-8') if isinstance(queue_name_bytes, bytes) else str(queue_name_bytes)
+                    new_queues.add(queue_name)
+
+                # 优化：添加日志，只在队列数量或内容发生变化时记录
+                if len(new_queues) != len(self._known_queues) or new_queues != self._known_queues:
+                    logger.debug(f"Queue registry contains {len(new_queues)} queues: {sorted(new_queues)}")
+
+                # 为新发现的队列创建消费者组（注意：新队列应该通过生产者自动注册）
+                new_discovered = new_queues - self._known_queues
+                if new_discovered:
+                    for queue in new_discovered:
+                        # 正确构建stream_key，保留优先级部分
+                        stream_key = f"{self.redis_prefix}:QUEUE:{queue}"
+                        try:
+                            await self.redis_client.xgroup_create(
+                                stream_key, self.consumer_group, id='0', mkstream=True
+                            )
+                            logger.info(f"Created consumer group for new queue: {queue} with stream_key: {stream_key}")
+                        except redis.ResponseError:
+                            pass
+
+                # 更新ConsumerManager的队列列表（同步操作）
+                if new_queues != self._known_queues:
+                    logger.info(f"Queue discovery: found {len(new_queues)} queues: {new_queues}")
+                    # 合并所有队列：TASK_CHANGES + 动态发现的队列
+                    all_queues = list(new_queues) + ['TASK_CHANGES']
+
+                    # 更新ConsumerManager的配置
+                    if self.consumer_manager:
+                        self.consumer_manager.config['queues'] = all_queues
+
+                        # 更新worker的队列信息
+                        await self._update_worker_queues(all_queues)
+
+                self._known_queues = new_queues
+                await asyncio.sleep(10)  # 保持较短的检查间隔，确保新队列能及时发现
+
+            except Exception as e:
+                logger.error(f"Error discovering queues: {e}")
+                logger.error(traceback.format_exc())
+                await asyncio.sleep(10)
+
+    async def _update_worker_queues(self, all_queues: list):
+        """更新worker的队列信息到Redis"""
+        try:
+            # ConsumerStrategy 已移除，现在只使用 HEARTBEAT 策略
+
+            # 获取实际的consumer_id（从心跳策略中）
+            if hasattr(self.consumer_manager, '_heartbeat_strategy'):
+                actual_consumer_id = self.consumer_manager._heartbeat_strategy.consumer_id
+            else:
+                # 从config中获取或使用默认值
+                actual_consumer_id = self.consumer_manager.config.get('consumer_id', 'unknown')
+
+            worker_key = f"{self.redis_prefix}:{self.consumer_manager.config.get('worker_prefix', 'PG_CONSUMER')}:{actual_consumer_id}"
+
+            # 使用同步Redis客户端更新
+            self.consumer_manager.redis_client.hset(
+                worker_key,
+                'queues',
+                ','.join(all_queues)
+            )
+            logger.debug(f"Updated worker queues: {all_queues}")
+
+        except Exception as e:
+            logger.error(f"Error updating worker queues: {e}")
+            logger.error(traceback.format_exc())
+
+    def get_known_queues(self) -> Set[str]:
+        """获取已知队列集合"""
+        return self._known_queues.copy()

jettask/persistence/task_persistence.py

@@ -0,0 +1,218 @@
+"""任务持久化模块
+
+负责解析Redis Stream消息，并将任务数据批量插入PostgreSQL数据库。
+"""
+
+import json
+import logging
+import traceback
+from typing import Dict, List, Optional, Any
+from datetime import datetime, timezone
+
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import sessionmaker
+
+logger = logging.getLogger(__name__)
+
+
+class TaskPersistence:
+    """任务持久化处理器
+
+    职责：
+    - 解析Stream消息为任务信息
+    - 批量插入任务到PostgreSQL的tasks表
+    - 处理插入失败的降级策略
+    """
+
+    def __init__(
+        self,
+        async_session_local: sessionmaker,
+        namespace_id: str,
+        namespace_name: str
+    ):
+        """初始化任务持久化处理器
+
+        Args:
+            async_session_local: SQLAlchemy会话工厂
+            namespace_id: 命名空间ID
+            namespace_name: 命名空间名称
+        """
+        self.AsyncSessionLocal = async_session_local
+        self.namespace_id = namespace_id
+        self.namespace_name = namespace_name
+
+    def parse_stream_message(self, task_id: str, data: dict) -> Optional[dict]:
+        """解析Stream消息为任务信息（返回完整的字段）
+
+        Args:
+            task_id: 任务ID（Redis Stream ID）
+            data: 消息数据
+
+        Returns:
+            解析后的任务信息字典，失败返回None
+        """
+        try:
+            from jettask.utils.serializer import loads_str
+
+            if b'data' in data:
+                task_data = loads_str(data[b'data'])
+            else:
+                task_data = {}
+                for k, v in data.items():
+                    key = k.decode('utf-8') if isinstance(k, bytes) else k
+                    if isinstance(v, bytes):
+                        try:
+                            value = loads_str(v)
+                        except:
+                            value = str(v)
+                    else:
+                        value = v
+                    task_data[key] = value
+
+            # 如果配置了命名空间，检查消息是否属于该命名空间
+            # if self.namespace_id:
+            #     msg_namespace_id = task_data.get('__namespace_id')
+            #     # 如果消息没有namespace_id且当前不是默认命名空间，跳过
+            #     if msg_namespace_id != self.namespace_id:
+            #         if not (msg_namespace_id is None and self.namespace_id == 'default'):
+            #             logger.debug(f"Skipping message from different namespace: {msg_namespace_id} != {self.namespace_id}")
+            #             return None
+
+            queue_name = task_data['queue']
+            task_name = task_data.get('name', task_data.get('task', 'unknown'))
+
+            created_at = None
+            if 'trigger_time' in task_data:
+                try:
+                    timestamp = float(task_data['trigger_time'])
+                    created_at = datetime.fromtimestamp(timestamp, tz=timezone.utc)
+                except:
+                    pass
+
+            # 返回完整的字段，包括所有可能为None的字段
+            return {
+                'id': task_id,
+                'queue_name': queue_name,
+                'task_name': task_name,
+                'task_data': json.dumps(task_data),
+                'priority': int(task_data.get('priority', 0)),
+                'retry_count': int(task_data.get('retry', 0)),
+                'max_retry': int(task_data.get('max_retry', 3)),
+                'status': 'pending',
+                'result': None,  # 新任务没有结果
+                'error_message': None,  # 新任务没有错误信息
+                'created_at': created_at,
+                'started_at': None,  # 新任务还未开始
+                'completed_at': None,  # 新任务还未完成
+                'scheduled_task_id': task_data.get('scheduled_task_id'),  # 调度任务ID
+                'metadata': json.dumps(task_data.get('metadata', {})),
+                'worker_id': None,  # 新任务还未分配worker
+                'execution_time': None,  # 新任务还没有执行时间
+                'duration': None,  # 新任务还没有持续时间
+                'namespace_id': self.namespace_id  # 添加命名空间ID
+            }
+
+        except Exception as e:
+            logger.error(f"Error parsing stream message for task {task_id}: {e}")
+            logger.error(traceback.format_exc())
+            return None
+
+    async def insert_tasks(self, tasks: List[Dict[str, Any]]) -> int:
+        """批量插入任务到PostgreSQL（只处理tasks表）
+
+        Args:
+            tasks: 任务信息列表
+
+        Returns:
+            实际插入的记录数
+        """
+        if not tasks:
+            return 0
+
+        logger.info(f"Attempting to insert {len(tasks)} tasks to tasks table")
+
+        try:
+            async with self.AsyncSessionLocal() as session:
+                # 插入tasks表 - 使用批量INSERT忽略冲突
+                # 由于stream_id在实践中是唯一的，我们可以简单地忽略重复
+                tasks_query = text("""
+                    INSERT INTO tasks (stream_id, queue, namespace, scheduled_task_id,
+                                      payload, priority, created_at, source, metadata)
+                    VALUES (:stream_id, :queue, :namespace, :scheduled_task_id,
+                           CAST(:payload AS jsonb), :priority, :created_at, :source, CAST(:metadata AS jsonb))
+                    ON CONFLICT DO NOTHING
+                    RETURNING stream_id;
+                """)
+
+                # 准备tasks表的数据
+                tasks_data = []
+                for task in tasks:
+                    task_data = json.loads(task['task_data'])
+
+                    # 从task_data中获取scheduled_task_id
+                    scheduled_task_id = task_data.get('scheduled_task_id') or task.get('scheduled_task_id')
+
+                    # 根据是否有scheduled_task_id来判断任务来源
+                    if scheduled_task_id:
+                        source = 'scheduler'  # 定时任务
+                    else:
+                        source = 'redis_stream'  # 普通任务
+
+                    tasks_data.append({
+                        'stream_id': task['id'],  # Redis Stream ID作为stream_id
+                        'queue': task['queue_name'],
+                        'namespace': self.namespace_name,
+                        'scheduled_task_id': str(scheduled_task_id) if scheduled_task_id else None,
+                        'payload': task['task_data'],  # 完整的任务数据
+                        'priority': task['priority'],
+                        'created_at': task['created_at'],
+                        'source': source,
+                        'metadata': task.get('metadata', '{}')
+                    })
+
+                # 批量插入 - 使用executemany提高性能
+                logger.debug(f"Executing batch insert with {len(tasks_data)} tasks")
+
+                try:
+                    # 使用executemany批量插入
+                    result = await session.execute(tasks_query, tasks_data)
+
+                    # 获取实际插入的记录数
+                    inserted_count = result.rowcount
+
+                    await session.commit()
+                    logger.debug("Tasks table batch insert transaction completed")
+                    return inserted_count
+
+                except Exception as e:
+                    logger.error(f"Error in batch insert, trying fallback: {e}")
+                    await session.rollback()
+
+                    # 如果批量插入失败，降级为小批量插入（每批10条）
+                    batch_size = 10
+                    total_inserted = 0
+
+                    for i in range(0, len(tasks_data), batch_size):
+                        batch = tasks_data[i:i+batch_size]
+                        try:
+                            result = await session.execute(tasks_query, batch)
+                            batch_inserted = result.rowcount
+                            if batch_inserted > 0:
+                                total_inserted += batch_inserted
+                            await session.commit()
+                        except Exception as batch_error:
+                            logger.error(f"Batch {i//batch_size + 1} failed: {batch_error}")
+                            await session.rollback()
+
+                    if total_inserted > 0:
+                        logger.info(f"Fallback insert completed: {total_inserted} tasks inserted")
+                    else:
+                        logger.info(f"No new tasks inserted in fallback mode")
+
+                    return total_inserted
+
+        except Exception as e:
+            logger.error(f"Error inserting tasks to PostgreSQL: {e}")
+            logger.error(traceback.format_exc())
+            return 0