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

jettask/persistence/message_consumer.py

@@ -0,0 +1,259 @@
+"""消息消费模块
+
+负责从Redis Stream队列中消费消息，并持久化到PostgreSQL。
+"""
+
+import asyncio
+import logging
+import traceback
+from typing import List, Dict
+from collections import defaultdict
+
+import redis.asyncio as redis
+from redis.asyncio import Redis
+
+from .task_persistence import TaskPersistence
+
+logger = logging.getLogger(__name__)
+
+
+class MessageConsumer:
+    """消息消费器
+
+    职责：
+    - 从Redis Stream队列消费消息
+    - 解析消息并持久化到数据库
+    - 管理多个队列的消费任务
+    - 处理错误重试和ACK
+    """
+
+    def __init__(
+        self,
+        redis_client: Redis,
+        redis_prefix: str,
+        consumer_group: str,
+        consumer_id: str,
+        task_persistence: TaskPersistence,
+        queue_discovery: 'QueueDiscovery'
+    ):
+        """初始化消息消费器
+
+        Args:
+            redis_client: Redis异步客户端
+            redis_prefix: Redis键前缀
+            consumer_group: 消费者组名称
+            consumer_id: 消费者ID
+            task_persistence: 任务持久化处理器
+            queue_discovery: 队列发现器
+        """
+        self.redis_client = redis_client
+        self.redis_prefix = redis_prefix
+        self.consumer_group = consumer_group
+        self.consumer_id = consumer_id
+        self.task_persistence = task_persistence
+        self.queue_discovery = queue_discovery
+
+        # 错误计数器
+        self._consecutive_errors = defaultdict(int)
+
+        # 已处理任务ID缓存（用于优化查询）
+        self._processed_task_ids = set()
+        self._processed_ids_lock = asyncio.Lock()
+        self._processed_ids_max_size = 100000
+        self._processed_ids_cleanup_interval = 300
+
+        self._running = False
+        self._consume_task = None
+        self._queue_tasks = {}
+
+    async def start(self):
+        """启动消费器"""
+        self._running = True
+        self._consume_task = asyncio.create_task(self._consume_queues())
+        logger.debug("MessageConsumer started")
+
+    async def stop(self):
+        """停止消费器"""
+        self._running = False
+
+        if self._consume_task:
+            self._consume_task.cancel()
+            try:
+                await self._consume_task
+            except asyncio.CancelledError:
+                pass
+
+        # 取消所有队列任务
+        for task in self._queue_tasks.values():
+            task.cancel()
+
+        if self._queue_tasks:
+            await asyncio.gather(*self._queue_tasks.values(), return_exceptions=True)
+
+        logger.debug("MessageConsumer stopped")
+
+    async def _consume_queues(self):
+        """启动所有队列的消费任务"""
+        while self._running:
+            try:
+                # 获取已知队列
+                known_queues = self.queue_discovery.get_known_queues()
+
+                # 为每个队列启动消费任务
+                for queue in known_queues:
+                    if queue not in self._queue_tasks or self._queue_tasks[queue].done():
+                        self._queue_tasks[queue] = asyncio.create_task(self._consume_queue(queue))
+                        logger.debug(f"Started consumer task for queue: {queue}")
+
+                # 移除不存在的队列任务
+                for queue in list(self._queue_tasks.keys()):
+                    if queue not in known_queues:
+                        self._queue_tasks[queue].cancel()
+                        del self._queue_tasks[queue]
+                        logger.debug(f"Stopped consumer task for removed queue: {queue}")
+
+                await asyncio.sleep(10)
+
+            except Exception as e:
+                logger.error(f"Error in consume_queues manager: {e}")
+                logger.error(traceback.format_exc())
+                await asyncio.sleep(5)
+
+    async def _consume_queue(self, queue_name: str):
+        """消费单个队列的任务（包括优先级队列）"""
+        # 判断是否是优先级队列
+        is_priority_queue = ':' in queue_name and queue_name.rsplit(':', 1)[-1].isdigit()
+
+        if is_priority_queue:
+            # 优先级队列格式：base_queue:priority (如 robust_bench2:2)
+            base_queue = queue_name.rsplit(':', 1)[0]
+            priority = queue_name.rsplit(':', 1)[1]
+            stream_key = f"{self.redis_prefix}:QUEUE:{base_queue}:{priority}"
+        else:
+            # 普通队列
+            stream_key = f"{self.redis_prefix}:QUEUE:{queue_name}"
+
+        logger.debug(f"Consuming queue: {queue_name}, stream_key: {stream_key}, is_priority: {is_priority_queue}")
+
+        check_backlog = True
+        lastid = "0-0"
+
+        # pg_consumer 应该使用统一的 consumer_id，而不是为每个队列创建新的
+        # 因为 pg_consumer 的职责是消费所有队列的消息并写入数据库
+        # 它不是真正的任务执行者，所以不需要为每个队列创建独立的 consumer
+        consumer_name = self.consumer_id
+
+        # ConsumerManager会自动处理离线worker的pending消息恢复
+        # 不需要手动恢复
+
+        while self._running and queue_name in self.queue_discovery.get_known_queues():
+            try:
+                myid = lastid if check_backlog else ">"
+
+                messages = await self.redis_client.xreadgroup(
+                    self.consumer_group,
+                    consumer_name,  # 使用ConsumerManager管理的consumer_name
+                    {stream_key: myid},
+                    count=10000,
+                    block=1000 if not check_backlog else 0
+                )
+
+                if not messages or (messages and len(messages[0][1]) == 0):
+                    check_backlog = False
+                    continue
+
+                if messages:
+                    await self._process_messages(messages)
+                    self._consecutive_errors[queue_name] = 0
+
+                    if messages[0] and messages[0][1]:
+                        lastid = messages[0][1][-1][0].decode('utf-8') if isinstance(messages[0][1][-1][0], bytes) else messages[0][1][-1][0]
+                        check_backlog = len(messages[0][1]) >= 2000
+
+            except redis.ResponseError as e:
+                if "NOGROUP" in str(e):
+                    try:
+                        await self.redis_client.xgroup_create(
+                            stream_key, self.consumer_group, id='0', mkstream=True
+                        )
+                        logger.debug(f"Recreated consumer group for queue: {queue_name}")
+                        check_backlog = True
+                        lastid = "0-0"
+                    except:
+                        pass
+                else:
+                    logger.error(f"Redis error for queue {queue_name}: {e}")
+                    logger.error(traceback.format_exc())
+                    self._consecutive_errors[queue_name] += 1
+
+                if self._consecutive_errors[queue_name] > 10:
+                    logger.debug(f"Too many errors for queue {queue_name}, will retry later")
+                    await asyncio.sleep(30)
+                    self._consecutive_errors[queue_name] = 0
+
+            except Exception as e:
+                logger.error(f"Error consuming queue {queue_name}: {e}", exc_info=True)
+                self._consecutive_errors[queue_name] += 1
+                await asyncio.sleep(1)
+
+    async def _process_messages(self, messages: List):
+        """处理消息并保存到PostgreSQL"""
+        tasks_to_insert = []
+        ack_batch = []
+
+        for stream_key, stream_messages in messages:
+            if not stream_messages:
+                continue
+
+            stream_key_str = stream_key.decode('utf-8') if isinstance(stream_key, bytes) else stream_key
+            msg_ids_to_ack = []
+
+            for msg_id, data in stream_messages:
+                try:
+                    if not msg_id or not data:
+                        continue
+
+                    msg_id_str = msg_id.decode('utf-8') if isinstance(msg_id, bytes) else str(msg_id)
+
+                    # 使用TaskPersistence解析消息
+                    task_info = self.task_persistence.parse_stream_message(msg_id_str, data)
+                    if task_info:
+                        tasks_to_insert.append(task_info)
+                        msg_ids_to_ack.append(msg_id)
+
+                except Exception as e:
+                    logger.error(f"Error processing message {msg_id}: {e}")
+                    logger.error(traceback.format_exc())
+
+            if msg_ids_to_ack:
+                ack_batch.append((stream_key, msg_ids_to_ack))
+
+        if tasks_to_insert:
+            # 使用TaskPersistence插入任务
+            inserted_count = await self.task_persistence.insert_tasks(tasks_to_insert)
+
+            # 将成功插入的任务ID添加到内存集合中
+            async with self._processed_ids_lock:
+                for task in tasks_to_insert:
+                    self._processed_task_ids.add(task['id'])
+
+                # 如果集合过大，清理最早的一半
+                if len(self._processed_task_ids) > self._processed_ids_max_size:
+                    # 只保留最新的一半ID
+                    ids_list = list(self._processed_task_ids)
+                    keep_count = self._processed_ids_max_size // 2
+                    self._processed_task_ids = set(ids_list[-keep_count:])
+                    logger.debug(f"Cleaned processed IDs cache, kept {keep_count} most recent IDs")
+
+            # ACK所有消息（即使部分插入失败，也要ACK，避免重复处理）
+            if ack_batch:
+                pipeline = self.redis_client.pipeline()
+                for stream_key, msg_ids in ack_batch:
+                    pipeline.xack(stream_key, self.consumer_group, *msg_ids)
+
+                try:
+                    await pipeline.execute()
+                    total_acked = sum(len(msg_ids) for _, msg_ids in ack_batch)
+                    logger.debug(f"Successfully ACKed {total_acked} messages")
+                except Exception as e:
+                    logger.error(f"Error executing batch ACK: {e}")

jettask/{backend/namespace_data_access.py → persistence/namespace.py}

@@ -11,152 +11,113 @@ from datetime import datetime, timedelta, timezone
 from typing import Dict, List, Optional, Tuple, Any
 import redis.asyncio as redis
 from sqlalchemy import text, bindparam
-from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.ext.asyncio import AsyncSession
 import aiohttp
 
+# 导入统一的数据库连接工具
+from ..utils.db_connector import (
+    get_dual_mode_async_redis_client,
+    get_pg_engine_and_factory
+)
+
 logger = logging.getLogger(__name__)
 
 
 class NamespaceConnection:
     """单个命名空间的数据库连接"""
-    
+
     def __init__(self, namespace_name: str, redis_config: dict, pg_config: dict):
         self.namespace_name = namespace_name
         self.redis_config = redis_config
         self.pg_config = pg_config
         self.redis_prefix = namespace_name  # 使用命名空间名作为Redis前缀
-        
-        # 数据库连接对象
+
+        # 使用全局单例连接池
+        self._text_redis_client: Optional[redis.Redis] = None
+        self._binary_redis_client: Optional[redis.Redis] = None
+        self._initialized = False
+
+        # PostgreSQL 相关
         self.async_engine = None
         self.AsyncSessionLocal = None
-        self._redis_pool = None
-        self._binary_redis_pool = None
-        self._initialized = False
         
     async def initialize(self):
         """初始化数据库连接"""
         if self._initialized:
             return
-            
+
         try:
-            # 初始化PostgreSQL连接
+            # 初始化 PostgreSQL 连接（使用全局单例）
             if self.pg_config:
-                dsn = self._build_pg_dsn()
-                if dsn.startswith('postgresql://'):
-                    dsn = dsn.replace('postgresql://', 'postgresql+psycopg://', 1)
-                
-                print(f'{dsn=}')
-                self.async_engine = create_async_engine(
-                    dsn,
+                self.async_engine, self.AsyncSessionLocal = get_pg_engine_and_factory(
+                    config=self.pg_config,
                     pool_size=10,
                     max_overflow=5,
                     pool_pre_ping=True,
                     echo=False
                 )
-                
-                self.AsyncSessionLocal = sessionmaker(
-                    bind=self.async_engine,
-                    class_=AsyncSession,
-                    expire_on_commit=False
-                )
-            
-            # 初始化Redis连接池
+
+            # 初始化 Redis 连接（使用全局单例，双模式）
             if self.redis_config:
-                # 支持两种格式：url格式或分离的host/port格式
-                redis_url = self.redis_config.get('url')
-                if redis_url:
-                    # 从URL创建连接池
-                    self._redis_pool = redis.ConnectionPool.from_url(
-                        redis_url,
-                        decode_responses=True,
-                        encoding='utf-8'
-                    )
-                    
-                    self._binary_redis_pool = redis.ConnectionPool.from_url(
-                        redis_url,
-                        decode_responses=False
-                    )
-                else:
-                    # 从分离的配置创建连接池
-                    self._redis_pool = redis.ConnectionPool(
-                        host=self.redis_config.get('host', 'localhost'),
-                        port=self.redis_config.get('port', 6379),
-                        db=self.redis_config.get('db', 0),
-                        password=self.redis_config.get('password'),
-                        decode_responses=True,
-                        encoding='utf-8'
-                    )
-                    
-                    self._binary_redis_pool = redis.ConnectionPool(
-                        host=self.redis_config.get('host', 'localhost'),
-                        port=self.redis_config.get('port', 6379),
-                        db=self.redis_config.get('db', 0),
-                        password=self.redis_config.get('password'),
-                        decode_responses=False
-                    )
-            
+                self._text_redis_client, self._binary_redis_client = get_dual_mode_async_redis_client(
+                    redis_url=self.redis_config.get('url') if isinstance(self.redis_config, dict) else self.redis_config,
+                    max_connections=50
+                )
+
             self._initialized = True
             logger.info(f"命名空间 {self.namespace_name} 数据库连接初始化成功")
-            
+
         except Exception as e:
             logger.error(f"初始化命名空间 {self.namespace_name} 数据库连接失败: {e}")
             traceback.print_exc()
             raise
-    
-    def _build_pg_dsn(self) -> str:
-        """构建PostgreSQL DSN"""
-        config = self.pg_config
-        # 支持两种格式：url格式或分离的配置
-        if 'url' in config:
-            return config['url']
-        else:
-            return f"postgresql://{config['user']}:{config['password']}@{config['host']}:{config['port']}/{config['database']}"
-    
+
     async def get_redis_client(self, decode: bool = True) -> redis.Redis:
-        """获取Redis客户端"""
+        """获取 Redis 客户端（使用全局单例）"""
         try:
             if not self._initialized:
                 await self.initialize()
-            
-            pool = self._redis_pool if decode else self._binary_redis_pool
-            if not pool:
-                raise ValueError(f"命名空间 {self.namespace_name} 没有配置Redis")
-            
-            return redis.Redis(connection_pool=pool)
+
+            # 根据 decode 参数选择文本或二进制客户端
+            client = self._text_redis_client if decode else self._binary_redis_client
+            if not client:
+                raise ValueError(f"命名空间 {self.namespace_name} 没有配置 Redis")
+
+            return client
         except Exception as e:
             # 连接异常时重置初始化标志，允许重新初始化
-            logger.error(f"获取Redis客户端失败: {e}")
+            logger.error(f"获取 Redis 客户端失败: {e}")
             traceback.print_exc()
             self._initialized = False
             raise
-    
+
     async def get_pg_session(self) -> AsyncSession:
-        """获取PostgreSQL会话"""
+        """获取 PostgreSQL 会话（使用全局单例）"""
         try:
             if not self._initialized:
                 await self.initialize()
-            
+
             if not self.AsyncSessionLocal:
-                raise ValueError(f"命名空间 {self.namespace_name} 没有配置PostgreSQL")
-            
+                raise ValueError(f"命名空间 {self.namespace_name} 没有配置 PostgreSQL")
+
             return self.AsyncSessionLocal()
         except Exception as e:
             # 连接异常时重置初始化标志，允许重新初始化
-            logger.error(f"获取PostgreSQL会话失败: {e}")
+            logger.error(f"获取 PostgreSQL 会话失败: {e}")
             traceback.print_exc()
             self._initialized = False
             raise
-    
+
     async def close(self):
-        """关闭数据库连接"""
-        if self._redis_pool:
-            await self._redis_pool.aclose()
-        if self._binary_redis_pool:
-            await self._binary_redis_pool.aclose()
-        if self.async_engine:
-            await self.async_engine.dispose()
-        
+        """关闭数据库连接（由于使用全局单例，这里只重置状态）"""
+        # 注意：连接池由全局单例管理，这里只清理引用
+        self._text_redis_client = None
+        self._binary_redis_client = None
+
+        # PostgreSQL engine 也是全局单例，只清理引用
+        self.async_engine = None
+        self.AsyncSessionLocal = None
+
         self._initialized = False
         logger.info(f"命名空间 {self.namespace_name} 数据库连接已关闭")
 
@@ -331,12 +292,19 @@ class NamespaceJetTaskDataAccess:
         redis_client = await conn.get_redis_client()
         
         try:
-            # 获取所有队列
-            queue_pattern = f"{conn.redis_prefix}:QUEUE:*"
-            print(f'{queue_pattern=}')
-            queue_keys = []
-            async for key in redis_client.scan_iter(match=queue_pattern):
-                queue_keys.append(key)
+            # 使用 RegistryManager 获取所有队列，避免 SCAN
+            from jettask.messaging.registry import QueueRegistry
+            registry = QueueRegistry(
+                redis_client=None,
+                async_redis_client=redis_client,
+                redis_prefix=conn.redis_prefix
+            )
+            
+            # 获取所有队列名称
+            queue_names = await registry.get_all_queues()
+            
+            # 构建完整的队列键
+            queue_keys = [f"{conn.redis_prefix}:QUEUE:{queue_name}" for queue_name in queue_names]
             
             stats = []
             for queue_key in queue_keys:

jettask/persistence/offline_recovery.py

@@ -0,0 +1,196 @@
+"""离线Worker恢复模块
+
+负责恢复离线PG_CONSUMER的消息，包括TASK_CHANGES流的离线消息。
+"""
+
+import asyncio
+import logging
+import msgpack
+import traceback
+from typing import Optional
+
+from redis.asyncio import Redis
+from jettask.worker.recovery import OfflineWorkerRecovery
+
+logger = logging.getLogger(__name__)
+
+
+class OfflineRecoveryHandler:
+    """离线Worker恢复处理器
+
+    职责：
+    - 启动离线worker恢复服务
+    - 恢复TASK_CHANGES stream的离线消息
+    - 处理恢复的消息并更新任务状态
+    """
+
+    def __init__(
+        self,
+        redis_client: Redis,
+        redis_prefix: str,
+        consumer_id: str,
+        task_updater: 'TaskUpdater'  # 类型提示使用字符串避免循环导入
+    ):
+        """初始化离线恢复处理器
+
+        Args:
+            redis_client: Redis异步客户端
+            redis_prefix: Redis键前缀
+            consumer_id: 消费者ID
+            task_updater: 任务更新器实例（用于处理恢复的消息）
+        """
+        self.redis_client = redis_client
+        self.redis_prefix = redis_prefix
+        self.consumer_id = consumer_id
+        self.task_updater = task_updater
+
+        # 创建 WorkerState 实例（用于查询 Worker 状态）
+        from jettask.worker.manager import WorkerState
+        self.worker_state = WorkerState(
+            redis_client=None,  # persistence 模块使用异步客户端
+            async_redis_client=redis_client,
+            redis_prefix=redis_prefix
+        )
+
+        # 创建离线worker恢复器（用于恢复TASK_CHANGES stream的离线消息）
+        # 注意：这里不传入consumer_manager，因为需要在start时初始化
+        self.offline_recovery = None
+
+        self._running = False
+        self._recovery_task = None
+
+    def set_consumer_manager(self, consumer_manager):
+        """设置ConsumerManager（延迟初始化）
+
+        Args:
+            consumer_manager: ConsumerManager实例
+        """
+        self.offline_recovery = OfflineWorkerRecovery(
+            async_redis_client=self.redis_client,
+            redis_prefix=self.redis_prefix,
+            worker_prefix='PG_CONSUMER',  # 使用PG_CONSUMER前缀
+            consumer_manager=consumer_manager,
+            worker_state=self.worker_state  # 传入在 __init__ 中创建的 WorkerState
+        )
+
+    async def start(self):
+        """启动离线恢复服务"""
+        if not self.offline_recovery:
+            logger.warning("OfflineRecovery not initialized, please call set_consumer_manager first")
+            return
+
+        self._running = True
+        self._recovery_task = asyncio.create_task(self._recovery_loop())
+        logger.debug("OfflineRecoveryHandler started")
+
+    async def stop(self):
+        """停止离线恢复服务"""
+        self._running = False
+
+        if self.offline_recovery:
+            self.offline_recovery.stop()  # stop() 不是异步方法
+
+        if self._recovery_task:
+            self._recovery_task.cancel()
+            try:
+                await self._recovery_task
+            except asyncio.CancelledError:
+                pass
+
+        logger.debug("OfflineRecoveryHandler stopped")
+
+    async def _recovery_loop(self):
+        """离线恢复循环"""
+        while self._running:
+            try:
+                total_recovered = 0
+
+                # 恢复TASK_CHANGES stream的消息
+                recovered = await self._recover_task_changes_offline_messages()
+                if recovered > 0:
+                    logger.debug(f"Recovered {recovered} TASK_CHANGES messages")
+                    total_recovered += recovered
+
+                if total_recovered > 0:
+                    logger.debug(f"Total recovered {total_recovered} messages in this cycle")
+
+                # 每1秒扫描一次
+                await asyncio.sleep(1)
+
+            except Exception as e:
+                logger.error(f"Error in offline recovery service: {e}")
+                await asyncio.sleep(10)
+
+    async def _recover_task_changes_offline_messages(self) -> int:
+        """恢复TASK_CHANGES stream的离线消息"""
+        # 使用 OfflineWorkerRecovery 的标准接口
+        try:
+            # 为TASK_CHANGES定义自定义的队列格式化器
+            def task_changes_formatter(queue):
+                # 对于TASK_CHANGES，直接返回stream key（不加QUEUE:前缀）
+                if queue == 'TASK_CHANGES':
+                    return f"{self.redis_prefix}:TASK_CHANGES"
+                else:
+                    return f"{self.redis_prefix}:QUEUE:{queue}"
+
+            # 创建专门用于TASK_CHANGES的恢复器
+            task_changes_recovery = OfflineWorkerRecovery(
+                async_redis_client=self.redis_client,
+                redis_prefix=self.redis_prefix,
+                worker_prefix='PG_CONSUMER',
+                queue_formatter=task_changes_formatter,
+                worker_state=self.worker_state  # 传入在 __init__ 中创建的 WorkerState
+            )
+
+            # 调用标准的恢复方法
+            # TASK_CHANGES作为队列名传入，会被正确处理
+            recovered = await task_changes_recovery.recover_offline_workers(
+                queue='TASK_CHANGES',  # 这个队列名会用于查找离线worker
+                current_consumer_name=self.consumer_id,
+                process_message_callback=self._process_recovered_task_change_v2
+            )
+
+            return recovered
+
+        except Exception as e:
+            logger.error(f"Error in recover_task_changes_offline_messages: {e}")
+            return 0
+
+    async def _process_recovered_task_change_v2(self, msg_id, msg_data, queue, consumer_id):
+        """处理恢复的TASK_CHANGES消息（符合OfflineWorkerRecovery的回调接口）"""
+        try:
+            logger.debug(f'处理恢复的TASK_CHANGES消息（符合OfflineWorkerRecovery的回调接口） {msg_data=}')
+            # 解析消息 - 现在使用task_id而不是event_id
+            if b'task_id' in msg_data:
+                # 使用msgpack解压task_id
+                compressed_task_id = msg_data[b'task_id']
+                task_key = msgpack.unpackb(compressed_task_id)
+                task_key = task_key.decode('utf-8') if isinstance(task_key, bytes) else str(task_key)
+
+                # 从完整的task_key格式提取stream_id
+                # 格式: namespace:TASK:stream_id:queue_name
+                stream_id = None
+                if ':TASK:' in task_key:
+                    parts = task_key.split(':TASK:')
+                    if len(parts) == 2:
+                        # 再从右边部分提取stream_id
+                        right_parts = parts[1].split(':')
+                        if right_parts:
+                            stream_id = right_parts[0]  # 提取stream_id
+
+                if stream_id:
+                    logger.debug(f"Processing recovered TASK_CHANGES message: {stream_id} from offline worker {consumer_id}")
+                    # 更新任务状态 - 传入(stream_id, task_key)元组
+                    # 使用task_updater的内部方法
+                    await self.task_updater._update_tasks_by_event([(stream_id, task_key)])
+                else:
+                    logger.warning(f"Cannot extract stream_id from task_key: {task_key}")
+
+                # ACK消息
+                change_stream_key = f"{self.redis_prefix}:TASK_CHANGES"
+                consumer_group = f"{self.redis_prefix}_changes_consumer"
+                await self.redis_client.xack(change_stream_key, consumer_group, msg_id)
+
+        except Exception as e:
+            logger.error(f"Error processing recovered task change {msg_id}: {e}")
+            logger.error(traceback.format_exc())