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

jettask/utils/rate_limit/task_limiter.py

@@ -0,0 +1,384 @@
+"""
+任务级限流器
+
+组合QPS和并发限流，提供统一的任务级限流接口
+"""
+
+import asyncio
+import logging
+import time
+from redis.asyncio import Redis
+from typing import Dict, Optional, List, Tuple, Any
+
+from .qps_limiter import QPSRateLimiter
+from .concurrency_limiter import ConcurrencyRateLimiter
+from .config import RateLimitConfig, QPSLimit, ConcurrencyLimit, parse_rate_limit_config
+
+logger = logging.getLogger('app')
+
+
+class TaskRateLimiter:
+    """任务级别的限流器
+
+    根据配置类型自动选择合适的限流器实现：
+    - QPSLimit: 使用本地滑动窗口 + Redis 协调配额
+    - ConcurrencyLimit: 使用 Redis 分布式信号量
+    """
+
+    def __init__(
+        self,
+        redis_client: Redis,
+        task_name: str,
+        worker_id: str,
+        config: RateLimitConfig,
+        redis_prefix: str = "jettask",
+        sync_interval: float = 5.0,
+        registry_manager = None,
+        worker_state_manager = None
+    ):
+        """初始化任务限流器
+
+        Args:
+            redis_client: 异步 Redis 客户端
+            task_name: 任务名称
+            worker_id: 当前 worker ID
+            config: 限流配置（QPSLimit 或 ConcurrencyLimit）
+            redis_prefix: Redis key 前缀
+            sync_interval: 配额同步间隔（秒），仅用于 QPS 限流
+            registry_manager: RegistryManager 实例
+            worker_state_manager: WorkerStateManager 实例
+        """
+        self.redis = redis_client
+        self.task_name = task_name
+        self.worker_id = worker_id
+        self.redis_prefix = redis_prefix
+        self.global_config = config
+        self.sync_interval = sync_interval
+        self.registry_manager = registry_manager
+        self.worker_state_manager = worker_state_manager
+
+        # Redis key 定义
+        self.rate_limit_key = f"{redis_prefix}:RATE_LIMIT:CONFIG:{task_name}"
+
+        # 根据配置类型创建具体的限流器
+        if isinstance(config, QPSLimit):
+            self.limiter = QPSRateLimiter(
+                qps=config.qps,
+                window_size=config.window_size,
+                worker_id=worker_id
+            )
+            self.mode = 'qps'
+            # QPS 模式需要配额同步
+            self._sync_task: Optional[asyncio.Task] = None
+            self._running = False
+        elif isinstance(config, ConcurrencyLimit):
+            self.limiter = ConcurrencyRateLimiter(
+                redis_client=redis_client,
+                task_name=task_name,
+                worker_id=worker_id,
+                max_concurrency=config.max_concurrency,
+                redis_prefix=redis_prefix
+            )
+            self.mode = 'concurrency'
+            # 并发模式不需要配额同步（分布式信号量自动协调）
+            self._sync_task = None
+            self._running = False
+        else:
+            raise ValueError(f"Unsupported config type: {type(config)}")
+
+        # 统计信息
+        self.stats = {
+            'syncs': 0,
+            'last_sync_time': 0,
+            'current_worker_count': 0,
+        }
+
+    async def start(self):
+        """启动限流器"""
+        if self._running:
+            logger.warning(f"TaskRateLimiter for {self.task_name} already running")
+            return
+
+        logger.debug(
+            f"Starting TaskRateLimiter for task={self.task_name}, "
+            f"config={self.global_config}, mode={self.mode}"
+        )
+
+        # 初始化 Redis 配置
+        await self._initialize_redis()
+
+        # 只有 QPS 模式需要启动配额同步
+        if self.mode == 'qps':
+            self._running = True
+            self._sync_task = asyncio.create_task(self._sync_loop())
+
+            # 注册 worker 状态变更回调
+            if self.worker_state_manager:
+                self.worker_state_manager.register_callback(self._on_worker_state_change)
+                logger.debug(f"Registered worker state change listener for {self.task_name}")
+
+        logger.debug(f"TaskRateLimiter started for {self.task_name}")
+    def __delattr__(self, name):
+        self.stop()
+        
+    async def stop(self):
+        """停止限流器"""
+        
+        print(f'准备退出限流器 {self.task_name} {self.mode}')
+        if not self._running and self.mode == 'qps':
+            return
+
+        logger.debug(f"Stopping TaskRateLimiter for {self.task_name}")
+
+        if self.mode == 'qps':
+            self._running = False
+
+            # 注销状态变更回调
+            if self.worker_state_manager:
+                self.worker_state_manager.unregister_callback(self._on_worker_state_change)
+
+            # 取消后台任务
+            if self._sync_task:
+                self._sync_task.cancel()
+                try:
+                    await self._sync_task
+                except asyncio.CancelledError:
+                    pass
+        elif self.mode == 'concurrency':
+            # 清理并发限流器（释放所有持有的锁）
+            if hasattr(self.limiter, 'stop'):
+                try:
+                    await self.limiter.stop()
+                    logger.debug(f"Concurrency limiter stopped and locks released for {self.task_name}")
+                except Exception as e:
+                    logger.error(f"Error stopping concurrency limiter for {self.task_name}: {e}")
+
+        logger.debug(f"TaskRateLimiter stopped for {self.task_name}")
+
+    async def _initialize_redis(self):
+        """初始化 Redis 配置"""
+        try:
+            exists = await self.redis.exists(self.rate_limit_key)
+            if not exists:
+                config_dict = self.global_config.to_dict()
+                await self.redis.hset(self.rate_limit_key, mapping=config_dict)
+                logger.debug(
+                    f"Initialized rate limit config for {self.task_name}: {self.global_config}"
+                )
+
+            # 将任务名添加到索引集合中（用于避免 scan 操作）
+            index_key = f"{self.redis_prefix}:RATE_LIMIT:INDEX"
+            await self.redis.sadd(index_key, self.task_name)
+            logger.debug(f"Added {self.task_name} to rate limit index")
+
+        except Exception as e:
+            logger.error(f"Failed to initialize Redis: {e}")
+
+    async def _on_worker_state_change(self, state_data: dict):
+        """Worker 状态变更回调（仅用于 QPS 模式）"""
+        if self.mode == 'qps':
+            logger.debug(f"[{self.task_name}] Worker state changed: {state_data}")
+            await self._sync_quota()
+
+    async def _sync_quota(self):
+        """同步配额（仅用于 QPS 模式）"""
+        if self.mode != 'qps':
+            return
+
+        try:
+            # 1. 从 Redis 读取全局配置
+            config_dict = await self.redis.hgetall(self.rate_limit_key)
+            if not config_dict:
+                logger.warning(f"No config found in Redis for {self.task_name}")
+                return
+
+            # 转换 bytes 为 str
+            config_dict = {
+                k.decode() if isinstance(k, bytes) else k:
+                v.decode() if isinstance(v, bytes) else v
+                for k, v in config_dict.items()
+            }
+
+            # 解析配置
+            global_config = parse_rate_limit_config(config_dict)
+            if not isinstance(global_config, QPSLimit):
+                logger.error(f"Invalid QPS config for {self.task_name}: {config_dict}")
+                return
+
+            # 2. 获取当前活跃的 workers
+            worker_ids = [self.worker_id]
+            if self.registry_manager and hasattr(self.registry_manager, 'get_workers_for_task'):
+                try:
+                    workers = await self.registry_manager.get_workers_for_task(self.task_name, only_alive=True)
+                    if workers:
+                        worker_ids = sorted(list(workers))
+                except Exception as e:
+                    logger.debug(f"Failed to get workers from registry: {e}")
+                    # 降级到只使用当前 worker
+
+            worker_count = len(worker_ids)
+
+            # 3. 协商配额分配
+            await self._allocate_quotas(global_config.qps, worker_ids)
+
+            # 4. 读取当前 worker 的配额
+            quota_key = f"quota:{self.worker_id}"
+            quota_str = await self.redis.hget(self.rate_limit_key, quota_key)
+
+            if quota_str:
+                quota_value = int(quota_str)
+            else:
+                # 降级方案：平均分配
+                quota_value = global_config.qps // worker_count if worker_count > 0 else global_config.qps
+
+            # 5. 更新本地限流器的配额
+            await self.limiter.update_limit(quota_value)
+
+            # 6. 更新统计信息
+            self.stats['syncs'] += 1
+            self.stats['last_sync_time'] = time.time()
+            self.stats['current_worker_count'] = worker_count
+
+        except Exception as e:
+            logger.error(f"Quota sync error: {e}")
+
+    async def _sync_loop(self):
+        """后台同步循环（仅用于 QPS 模式）"""
+        while self._running:
+            try:
+                await self._sync_quota()
+                await asyncio.sleep(self.sync_interval)
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"Sync loop error: {e}")
+                await asyncio.sleep(1)
+
+    async def _allocate_quotas(self, total_qps: int, worker_ids: list):
+        """协商配额分配（仅用于 QPS 模式）"""
+        lock_key = f"{self.rate_limit_key}:lock"
+
+        try:
+            async with self.redis.lock(lock_key, timeout=3, blocking_timeout=2):
+                worker_count = len(worker_ids)
+                base_quota = total_qps // worker_count
+                remainder = total_qps % worker_count
+
+                # 分配配额
+                allocations = {}
+                for i, worker_id in enumerate(worker_ids):
+                    quota = base_quota + (1 if i < remainder else 0)
+                    allocations[f"quota:{worker_id}"] = quota
+
+                # 批量写入 Redis
+                if allocations:
+                    await self.redis.hset(self.rate_limit_key, mapping=allocations)
+
+                # 清理不活跃 workers 的配额
+                await self._cleanup_stale_quotas(worker_ids)
+
+                logger.debug(f"[ALLOCATE] Allocated quotas for {worker_count} workers: {allocations}")
+
+        except asyncio.TimeoutError:
+            logger.debug(f"[ALLOCATE] Failed to acquire lock, skipping allocation")
+        except Exception as e:
+            logger.error(f"[ALLOCATE] Error allocating quotas: {e}")
+
+    async def _cleanup_stale_quotas(self, active_worker_ids: list):
+        """清理不活跃 workers 的配额"""
+        try:
+            all_fields = await self.redis.hkeys(self.rate_limit_key)
+            active_quota_keys = {f"quota:{wid}" for wid in active_worker_ids}
+
+            to_delete = []
+            for field in all_fields:
+                if isinstance(field, bytes):
+                    field = field.decode('utf-8')
+
+                if field.startswith("quota:") and field not in active_quota_keys:
+                    to_delete.append(field)
+
+            if to_delete:
+                await self.redis.hdel(self.rate_limit_key, *to_delete)
+                logger.debug(f"[CLEANUP] Removed {len(to_delete)} stale quotas")
+
+        except Exception as e:
+            logger.error(f"[CLEANUP] Error cleaning up quotas: {e}")
+
+    async def acquire(self, timeout: float = 10.0) -> Optional[str]:
+        """获取执行许可
+
+        Returns:
+            成功返回 task_id (ConcurrencyLimit) 或 True (QPSLimit)，失败返回 None
+        """
+        result = await self.limiter.acquire(timeout=timeout)
+        # QPSLimit 返回 bool，ConcurrencyLimit 返回 task_id 或 None
+        if self.mode == 'qps':
+            return True if result else None
+        else:
+            return result  # ConcurrencyLimit 直接返回 task_id 或 None
+
+    async def release(self, task_id: Optional[str] = None):
+        """释放执行许可
+
+        Args:
+            task_id: 任务ID (仅 ConcurrencyLimit 需要)
+        """
+        if self.mode == 'concurrency' and task_id is None:
+            raise ValueError("ConcurrencyLimit requires task_id for release")
+
+        if self.mode == 'concurrency':
+            await self.limiter.release(task_id)
+        else:
+            # QPSLimit 不需要 release（本地计数）
+            pass
+
+    async def try_acquire(self) -> bool:
+        """尝试获取执行许可（非阻塞）"""
+        return await self.limiter.try_acquire()
+
+    def get_stats(self) -> dict:
+        """获取统计信息"""
+        return {
+            **self.stats,
+            'task_name': self.task_name,
+            'worker_id': self.worker_id,
+            'config': str(self.global_config),
+        }
+
+    async def __aenter__(self):
+        """异步上下文管理器入口 - 获取执行许可
+
+        使用示例:
+            async with rate_limiter:
+                # 执行任务
+                await do_something()
+
+        Returns:
+            self: 限流器实例
+        """
+        task_id = await self.acquire()
+        if task_id is None:
+            raise TimeoutError("Failed to acquire rate limit slot")
+        # 保存task_id供__aexit__使用
+        self._current_task_id = task_id
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """异步上下文管理器退出 - 自动释放许可"""
+        if hasattr(self, '_current_task_id'):
+            task_id = self._current_task_id
+            delattr(self, '_current_task_id')
+            if self.mode == 'concurrency':
+                await self.release(task_id)
+            # QPSLimit 不需要显式 release
+        return False  # 不抑制异常
+
+
+# ============================================================
+# 限流器管理器
+# ============================================================
+
+
+
+__all__ = ['TaskRateLimiter']

jettask/utils/serializer.py

@@ -12,6 +12,9 @@ def dumps(obj):
 
 def loads(data):
     """反序列化字节为对象"""
+    if isinstance(data, str):
+        # 兼容性：如果收到字符串，尝试作为 latin-1 解码
+        data = data.encode('latin-1')
     return msgpack.unpackb(data, raw=False, strict_map_key=False)
 
 

jettask/{monitor/stream_backlog_monitor.py → utils/stream_backlog.py}

@@ -37,10 +37,15 @@ class StreamBacklogMonitor:
         self.AsyncSessionLocal = None
         
     async def initialize(self):
-        """初始化连接"""
+        """初始化连接（使用统一的连接池管理）"""
         # 初始化Redis连接
-        self.redis_client = await redis.from_url(self.redis_url, decode_responses=True)
-        
+        from jettask.utils.db_connector import get_async_redis_client
+
+        self.redis_client = get_async_redis_client(
+            redis_url=self.redis_url,
+            decode_responses=True
+        )
+
         # 初始化PostgreSQL连接
         self.engine = create_async_engine(self.pg_url, echo=False)
         self.AsyncSessionLocal = sessionmaker(self.engine, class_=AsyncSession, expire_on_commit=False)
@@ -190,10 +195,13 @@ class StreamBacklogMonitor:
                 for group in groups:
                     group_name = group['name']
                     pending_count = group['pending']  # Redis Stream中的pending数量（已投递未ACK）
-                    
+
                     # 从TASK_OFFSETS获取该组的消费offset
-                    # key格式: f"{queue}:{group_name}"
-                    task_offset_key = f"{stream_name}:{group_name}"
+                    # 从 group_name 中提取 task_name（最后一段）
+                    task_name = group_name.split(':')[-1]
+                    # 构建 field：队列名（含优先级）+ 任务名
+                    # 例如：robust_bench2:8:benchmark_task
+                    task_offset_key = f"{stream_name}:{task_name}"
                     last_acked_offset = int(task_offsets.get(task_offset_key, 0))
                     print(f'{task_offset_key=} {last_acked_offset=}')
                     # 计算各种积压指标

jettask/utils/time_sync.py

@@ -0,0 +1,173 @@
+#!/usr/bin/env python
+"""时间同步模块
+
+提供与 Redis 服务器时间同步的功能，避免多个 worker 因系统时间不一致导致的限流问题。
+在启动时一次性校准与 Redis 服务器的时间差，后续使用本地时间加上时间差来获得标准时间。
+"""
+
+import time
+import logging
+from typing import Optional
+import redis.asyncio as aioredis
+
+logger = logging.getLogger(__name__)
+
+
+class TimeSync:
+    """时间同步类
+
+    在初始化时与 Redis 服务器时间进行一次校准，计算时间差。
+    后续所有时间获取都使用 本地时间 + 时间差 的方式，避免频繁访问 Redis。
+
+    使用方式:
+        # 初始化（异步）
+        time_sync = TimeSync()
+        await time_sync.sync(redis_client)
+
+        # 获取标准时间
+        now = time_sync.time()
+    """
+
+    def __init__(self):
+        """初始化时间同步器"""
+        self._offset = 0.0  # Redis 服务器时间与本地时间的差值（秒）
+        self._synced = False  # 是否已同步
+        self._sync_timestamp = 0.0  # 同步时的本地时间戳
+
+    async def sync(self, redis_client: aioredis.Redis, retry: int = 3) -> bool:
+        """与 Redis 服务器时间同步
+
+        Args:
+            redis_client: 异步 Redis 客户端
+            retry: 重试次数
+
+        Returns:
+            同步是否成功
+        """
+        for attempt in range(retry):
+            try:
+                # 记录请求前的本地时间
+                local_before = time.time()
+
+                # 获取 Redis 服务器时间（微秒级）
+                # TIME 命令返回 [seconds, microseconds]
+                redis_time = await redis_client.time()
+
+                # 记录请求后的本地时间
+                local_after = time.time()
+
+                # 计算网络往返时间的一半（估算单程延迟）
+                network_delay = (local_after - local_before) / 2.0
+
+                # Redis 时间转换为秒（包含微秒）
+                redis_timestamp = float(redis_time[0]) + float(redis_time[1]) / 1_000_000.0
+
+                # 估算 Redis 服务器在请求时刻的时间
+                # 使用请求前时间 + 单程延迟来估算
+                estimated_local_time = local_before + network_delay
+
+                # 计算时间差：Redis 服务器时间 - 本地时间
+                self._offset = redis_timestamp - estimated_local_time
+                self._sync_timestamp = time.time()
+                self._synced = True
+
+                logger.debug(
+                    f"时间同步成功: offset={self._offset:.6f}s, "
+                    f"network_delay={network_delay*1000:.2f}ms, "
+                    f"redis_time={redis_timestamp:.6f}, "
+                    f"local_time={estimated_local_time:.6f}"
+                )
+
+                # 如果时间差超过 1 秒，给出警告
+                if abs(self._offset) > 1.0:
+                    logger.warning(
+                        f"本地时间与 Redis 服务器时间差异较大: {self._offset:.2f}s，"
+                        f"建议检查系统时间配置"
+                    )
+
+                return True
+
+            except Exception as e:
+                logger.error(f"时间同步失败 (attempt {attempt + 1}/{retry}): {e}")
+                if attempt < retry - 1:
+                    await asyncio.sleep(0.1)
+
+        logger.error("时间同步失败，使用本地时间")
+        self._synced = False
+        return False
+
+    def time(self) -> float:
+        """获取标准时间
+
+        Returns:
+            标准时间戳（秒）
+        """
+        if not self._synced:
+            # 未同步时使用本地时间
+            return time.time()
+
+        # 使用本地时间 + 时间差
+        return time.time() + self._offset
+
+    def is_synced(self) -> bool:
+        """检查是否已同步
+
+        Returns:
+            是否已同步
+        """
+        return self._synced
+
+    def get_offset(self) -> float:
+        """获取时间偏移量
+
+        Returns:
+            时间偏移量（秒）
+        """
+        return self._offset
+
+    def get_sync_info(self) -> dict:
+        """获取同步信息
+
+        Returns:
+            同步信息字典
+        """
+        return {
+            'synced': self._synced,
+            'offset': self._offset,
+            'sync_timestamp': self._sync_timestamp,
+            'time_since_sync': time.time() - self._sync_timestamp if self._synced else None,
+        }
+
+
+# 全局时间同步实例（单例模式）
+_global_time_sync: Optional[TimeSync] = None
+
+
+def get_time_sync() -> TimeSync:
+    """获取全局时间同步实例
+
+    Returns:
+        TimeSync 实例
+    """
+    global _global_time_sync
+    if _global_time_sync is None:
+        _global_time_sync = TimeSync()
+    return _global_time_sync
+
+
+async def init_time_sync(redis_client: aioredis.Redis) -> TimeSync:
+    """初始化全局时间同步
+
+    Args:
+        redis_client: 异步 Redis 客户端
+
+    Returns:
+        TimeSync 实例
+    """
+    time_sync = get_time_sync()
+    await time_sync.sync(redis_client)
+    return time_sync
+
+
+# 添加缺失的 asyncio 导入
+import asyncio

jettask/webui/__init__.py

@@ -0,0 +1,27 @@
+"""
+Web UI 模块
+
+提供 Web 界面和 API 接口。
+"""
+
+# 不在 __init__ 中导入 app，避免循环导入
+# 使用时直接: from jettask.webui.app import app
+from .config import *
+
+# 异常类直接从主模块导入（webui/exceptions.py已废弃并删除）
+from jettask.exceptions import (
+    JetTaskException,
+    TaskTimeoutError,
+    TaskExecutionError,
+    TaskNotFoundError,
+    RetryableError
+)
+
+__all__ = [
+    # 'app',  # 移除，避免循环导入
+    'JetTaskException',
+    'TaskTimeoutError',
+    'TaskExecutionError',
+    'TaskNotFoundError',
+    'RetryableError',
+]

jettask/{api/v1 → webui/api}/alerts.py

@@ -7,7 +7,7 @@ from typing import Optional
 import logging
 
 from jettask.schemas import AlertRuleRequest
-from jettask.services.alert_service import AlertService
+from jettask.webui.services.alert_service import AlertService
 
 router = APIRouter(prefix="/alerts", tags=["alerts"])
 logger = logging.getLogger(__name__)

jettask/{api/v1 → webui/api}/analytics.py

@@ -8,8 +8,8 @@ from datetime import datetime
 import logging
 
 from jettask.schemas import TimeRangeQuery
-from jettask.services.analytics_service import AnalyticsService
-from jettask.backend.namespace_data_access import get_namespace_data_access
+from jettask.webui.services.analytics_service import AnalyticsService
+from jettask.persistence.namespace import get_namespace_data_access
 
 router = APIRouter(prefix="/analytics", tags=["analytics"])
 logger = logging.getLogger(__name__)

jettask/{api/v1 → webui/api}/namespaces.py

@@ -8,7 +8,7 @@ import logging
 import traceback
 
 from jettask.schemas import NamespaceCreate, NamespaceUpdate, NamespaceResponse
-from jettask.services.settings_service import SettingsService
+from jettask.webui.services.settings_service import SettingsService
 
 logger = logging.getLogger(__name__)
 

jettask/{api/v1 → webui/api}/overview.py

@@ -8,7 +8,7 @@ import logging
 import traceback
 
 from jettask.schemas import TimeRangeQuery
-from jettask.services.overview_service import OverviewService
+from jettask.webui.services.overview_service import OverviewService
 
 logger = logging.getLogger(__name__)
 

jettask/{api/v1 → webui/api}/queues.py

@@ -14,9 +14,9 @@ from jettask.schemas import (
     TaskActionRequest,
     BacklogTrendRequest
 )
-from jettask.services.queue_service import QueueService
-from jettask.services.task_service import TaskService
-from jettask.services.redis_monitor_service import RedisMonitorService
+from jettask.webui.services.queue_service import QueueService
+from jettask.webui.services.task_service import TaskService
+from jettask.utils.redis_monitor import RedisMonitorService
 
 router = APIRouter(prefix="/queues", tags=["queues"])
 logger = logging.getLogger(__name__)