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

jettask/persistence/message_consumer.py

@@ -1,259 +0,0 @@
-"""消息消费模块
-
-负责从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/persistence/models.py

@@ -1,49 +0,0 @@
-"""SQLAlchemy models for JetTask WebUI database."""
-from datetime import datetime
-from typing import Optional, Dict, Any, List
-from sqlalchemy import (
-    Column, String, Integer, Float, DateTime, Text, JSON, 
-    ARRAY, UniqueConstraint, Index, func
-)
-from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.dialects.postgresql import JSONB
-
-Base = declarative_base()
-
-
-class Task(Base):
-    """任务表模型"""
-    __tablename__ = 'tasks'
-    
-    id = Column(String(255), primary_key=True)  # Redis Stream的事件ID
-    queue_name = Column(String(255), nullable=False)
-    task_name = Column(String(255), nullable=False)
-    task_data = Column(JSONB)  # 任务的原始数据
-    priority = Column(Integer, default=0)
-    retry_count = Column(Integer, default=0)
-    max_retry = Column(Integer, default=3)
-    status = Column(String(50), default='pending')  # pending, running, success, failed, timeout
-    result = Column(JSONB)  # 执行结果
-    error_message = Column(Text)
-    created_at = Column(DateTime(timezone=True), default=func.current_timestamp())
-    started_at = Column(DateTime(timezone=True))
-    completed_at = Column(DateTime(timezone=True))
-    worker_id = Column(String(255))
-    execution_time = Column(Float)  # 任务执行时间（秒）
-    duration = Column(Float)  # 任务总持续时间（秒）
-    task_metadata = Column('metadata', JSONB)  # 额外的元数据，在数据库中仍叫metadata
-    
-    __table_args__ = (
-        Index('idx_tasks_queue_name', 'queue_name'),
-        Index('idx_tasks_status', 'status'),
-        # 组合索引：优化按队列和状态查询
-        Index('idx_tasks_queue_status', 'queue_name', 'status'),
-        # 时间索引：优化时间范围查询
-        Index('idx_tasks_created_at', 'created_at'),
-        # Worker索引：优化查询特定worker的任务
-        Index('idx_tasks_worker_id', 'worker_id', 
-              postgresql_where=(worker_id.isnot(None))),
-    )
-
-
-# QueueStats 和 Worker 表已废弃，不再使用

jettask/persistence/offline_recovery.py

@@ -1,196 +0,0 @@
-"""离线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())