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

jettask/messaging/scanner.py

@@ -122,6 +122,36 @@ class DelayedMessageScanner:
 
         logger.info(f"DelayedMessageScanner started with {len(self._scan_tasks)} scan tasks")
 
+    async def add_queues(self, queues: List[str]):
+        """
+        动态添加队列到扫描器（用于通配符队列发现）
+
+        Args:
+            queues: 要添加的队列名列表
+        """
+        if not self._running:
+            logger.warning("Cannot add queues: DelayedMessageScanner is not running")
+            return
+
+        # 记录已经存在的队列任务
+        existing_queues = {task.get_name().replace('_scan_queue_loop_', '')
+                          for task in self._scan_tasks if not task.done()}
+
+        # 只为新队列创建扫描任务
+        new_queues = [q for q in queues if q not in existing_queues]
+
+        if not new_queues:
+            logger.debug(f"No new queues to add (already scanning: {existing_queues})")
+            return
+
+        logger.info(f"Adding {len(new_queues)} new queues to DelayedMessageScanner: {new_queues}")
+
+        for queue in new_queues:
+            task = asyncio.create_task(self._scan_queue_loop(queue), name=f'_scan_queue_loop_{queue}')
+            self._scan_tasks.append(task)
+
+        logger.info(f"DelayedMessageScanner now has {len(self._scan_tasks)} scan tasks")
+
     async def stop(self):
         """停止扫描器"""
         if not self._running:

jettask/persistence/__init__.py

@@ -1,118 +1,43 @@
-"""PostgreSQL Consumer模块
+"""PostgreSQL Consumer 模块 - V4 重构版本
 
-从Redis队列消费任务并持久化到PostgreSQL数据库。
+基于 Jettask 通配符队列的 PostgreSQL 消费者实现。
+
+核心特性：
+- 使用通配符队列 (*) 自动发现和监听所有队列
+- 独立的 TASK_CHANGES 队列处理状态更新
+- 批量 INSERT 和 UPDATE 提升性能
+- 支持多命名空间管理
 
 模块结构：
-- consumer.py: 主消费者类（协调器）
-- backlog_monitor.py: Stream积压监控
-- task_updater.py: 任务状态更新
-- offline_recovery.py: 离线Worker恢复
-- task_persistence.py: 任务数据持久化
-- queue_discovery.py: 队列发现
-- message_consumer.py: 消息消费
-- maintenance.py: 数据库维护
+- consumer.py: PostgreSQL 消费者（通配符队列实现）
+- manager.py: 统一的消费者管理器
+- buffer.py: 批量缓冲区
+- persistence.py: 任务持久化逻辑
 
 使用示例：
-    from jettask.services.pg_consumer import PostgreSQLConsumer, run_pg_consumer, main
+    from jettask.persistence import PostgreSQLConsumer, UnifiedConsumerManager
 
-    # 使用Consumer类
-    consumer = PostgreSQLConsumer(pg_config, redis_config)
+    # 使用 Consumer 类（单命名空间）
+    consumer = PostgreSQLConsumer(
+        pg_config={'url': 'postgresql://...'},
+        redis_config={'url': 'redis://...'},
+        prefix='jettask',
+        namespace_name='default'
+    )
     await consumer.start()
 
-    # 或直接运行
-    await run_pg_consumer(pg_config, redis_config)
-
-    # 或使用main函数（从环境变量读取配置）
-    main()
+    # 使用 UnifiedConsumerManager（多命名空间）
+    manager = UnifiedConsumerManager(
+        task_center_url='http://localhost:8001',
+        check_interval=30
+    )
+    await manager.run()
 """
 
-import asyncio
-import logging
-import os
-
-from jettask.webui.config import PostgreSQLConfig, RedisConfig
-# ConsumerStrategy 已移除，现在只使用 HEARTBEAT 策略
-
 from .consumer import PostgreSQLConsumer
+from .manager import UnifiedConsumerManager
 
-logger = logging.getLogger(__name__)
-
-# 导出主要类和函数
 __all__ = [
     'PostgreSQLConsumer',
-    'run_pg_consumer',
-    'main'
+    'UnifiedConsumerManager',
 ]
-
-
-async def run_pg_consumer(
-    pg_config: PostgreSQLConfig,
-    redis_config: RedisConfig,
-    # consumer_strategy 参数已移除，现在只使用 HEARTBEAT 策略
-):
-    """运行PostgreSQL消费者
-
-    Args:
-        pg_config: PostgreSQL配置
-        redis_config: Redis配置
-        consumer_strategy: 消费者策略
-    """
-    # 从环境变量读取监控配置
-    enable_backlog_monitor = os.getenv('JETTASK_ENABLE_BACKLOG_MONITOR', 'true').lower() == 'true'
-    backlog_monitor_interval = int(os.getenv('JETTASK_BACKLOG_MONITOR_INTERVAL', '60'))
-
-    logger.info(f"Backlog monitor config: enabled={enable_backlog_monitor}, interval={backlog_monitor_interval}s")
-
-    consumer = PostgreSQLConsumer(
-        pg_config,
-        redis_config,
-        consumer_strategy=consumer_strategy,
-        enable_backlog_monitor=enable_backlog_monitor,
-        backlog_monitor_interval=backlog_monitor_interval
-    )
-
-    try:
-        await consumer.start()
-        while True:
-            await asyncio.sleep(1)
-
-    except KeyboardInterrupt:
-        logger.debug("Received interrupt signal")
-    finally:
-        await consumer.stop()
-
-
-def main():
-    """主入口函数（从环境变量读取配置）"""
-    from dotenv import load_dotenv
-
-    load_dotenv()
-
-    logging.basicConfig(
-        level=logging.INFO,
-        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-    )
-
-    pg_config = PostgreSQLConfig(
-        host=os.getenv('JETTASK_PG_HOST', 'localhost'),
-        port=int(os.getenv('JETTASK_PG_PORT', '5432')),
-        database=os.getenv('JETTASK_PG_DB', 'jettask'),
-        user=os.getenv('JETTASK_PG_USER', 'jettask'),
-        password=os.getenv('JETTASK_PG_PASSWORD', '123456'),
-    )
-
-    redis_config = RedisConfig(
-        host=os.getenv('REDIS_HOST', 'localhost'),
-        port=int(os.getenv('REDIS_PORT', '6379')),
-        db=int(os.getenv('REDIS_DB', '0')),
-        password=os.getenv('REDIS_PASSWORD'),
-    )
-
-    # 使用 HEARTBEAT 策略（唯一支持的策略）
-    logger.debug("Using consumer strategy: HEARTBEAT")
-
-    asyncio.run(run_pg_consumer(pg_config, redis_config))
-
-
-if __name__ == '__main__':
-    main()

jettask/persistence/buffer.py

@@ -0,0 +1,170 @@
+"""批量缓冲区管理器
+
+负责收集任务数据和ACK信息，批量写入数据库并ACK。
+支持 INSERT 和 UPDATE 两种操作类型。
+"""
+
+import time
+import asyncio
+import logging
+from typing import List, Dict, Any
+
+logger = logging.getLogger(__name__)
+
+
+class BatchBuffer:
+    """批量缓冲区管理器
+
+    负责：
+    1. 收集任务数据和ACK信息
+    2. 判断是否应该刷新（批量大小或超时）
+    3. 批量写入数据库并ACK
+    """
+
+    def __init__(
+        self,
+        max_size: int = 1000,
+        max_delay: float = 5.0,
+        operation_type: str = 'insert'  # 'insert' 或 'update'
+    ):
+        """初始化缓冲区
+
+        Args:
+            max_size: 缓冲区最大容量（条数）
+            max_delay: 最大延迟时间（秒）
+            operation_type: 操作类型，'insert' 或 'update'
+        """
+        self.max_size = max_size
+        self.max_delay = max_delay
+        self.operation_type = operation_type
+
+        # 任务数据缓冲区
+        self.records: List[Dict[str, Any]] = []
+        self.contexts: List[Any] = []  # 保存 TaskContext 用于 ACK
+
+        # 刷新控制
+        self.last_flush_time = time.time()
+        self.flush_lock = asyncio.Lock()
+
+        # 统计信息
+        self.total_flushed = 0
+        self.flush_count = 0
+
+    def add(self, record: dict, context: Any = None):
+        """添加到缓冲区
+
+        Args:
+            record: 任务数据或更新数据
+            context: TaskContext（用于 ACK）
+        """
+        self.records.append(record)
+        if context:
+            self.contexts.append(context)
+
+    def should_flush(self) -> bool:
+        """判断是否应该刷新
+
+        Returns:
+            是否需要刷新
+        """
+        if not self.records:
+            return False
+
+        # 缓冲区满了
+        if len(self.records) >= self.max_size:
+            logger.info(
+                f"[{self.operation_type.upper()}] 缓冲区已满 "
+                f"({len(self.records)}/{self.max_size})，触发刷新"
+            )
+            return True
+
+        # 超时了
+        elapsed = time.time() - self.last_flush_time
+        if elapsed >= self.max_delay:
+            logger.info(
+                f"[{self.operation_type.upper()}] 缓冲区超时 "
+                f"({elapsed:.1f}s >= {self.max_delay}s)，触发刷新"
+            )
+            return True
+
+        return False
+
+    async def flush(self, db_manager):
+        """刷新缓冲区到数据库
+
+        1. 批量写入数据库
+        2. 批量ACK（如果有context）
+        3. 清空缓冲区
+
+        Args:
+            db_manager: 数据库管理器，需要有 batch_insert_tasks 或 batch_update_tasks 方法
+        """
+        async with self.flush_lock:
+            if not self.records:
+                return 0
+
+            count = len(self.records)
+            start_time = time.time()
+
+            try:
+                logger.info(f"[{self.operation_type.upper()}] 开始批量刷新 {count} 条记录...")
+
+                # 1. 批量写入数据库
+                if self.operation_type == 'insert':
+                    await db_manager.batch_insert_tasks(self.records)
+                    logger.info(f"  ✓ 批量插入 {count} 条任务记录")
+                else:  # update
+                    print(f'{self.records=}')
+                    await db_manager.batch_update_tasks(self.records)
+                    logger.info(f"  ✓ 批量更新 {count} 条任务状态")
+
+                # 2. 批量ACK（如果使用 Jettask 的 context）
+                if self.contexts:
+                    for ctx in self.contexts:
+                        if hasattr(ctx, 'ack'):
+                            await ctx.ack()
+                    logger.info(f"  ✓ 批量确认 {len(self.contexts)} 条消息")
+
+                # 3. 清空缓冲区
+                self.records.clear()
+                self.contexts.clear()
+                self.last_flush_time = time.time()
+
+                # 4. 统计
+                self.total_flushed += count
+                self.flush_count += 1
+                elapsed = time.time() - start_time
+
+                logger.info(
+                    f"[{self.operation_type.upper()}] ✓ 批量刷新完成! "
+                    f"本次: {count}条, "
+                    f"耗时: {elapsed:.3f}s, "
+                    f"总计: {self.total_flushed}条 ({self.flush_count}次刷新)"
+                )
+
+                return count
+
+            except Exception as e:
+                logger.error(
+                    f"[{self.operation_type.upper()}] ✗ 批量刷新失败: {e}",
+                    exc_info=True
+                )
+                # 失败时清空缓冲区，避免无限重试
+                self.records.clear()
+                self.contexts.clear()
+                raise
+
+    def get_stats(self) -> dict:
+        """获取统计信息
+
+        Returns:
+            统计信息字典
+        """
+        return {
+            'operation_type': self.operation_type,
+            'current_size': len(self.records),
+            'max_size': self.max_size,
+            'total_flushed': self.total_flushed,
+            'flush_count': self.flush_count,
+            'avg_per_flush': self.total_flushed // self.flush_count if self.flush_count > 0 else 0
+        }