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

jettask/utils/rate_limit/config.py

@@ -0,0 +1,145 @@
+#!/usr/bin/env python
+"""限流配置类
+
+定义各种限流策略的配置类，通过不同的类来区分限流模式。
+"""
+
+from abc import ABC, abstractmethod
+from typing import Optional
+
+
+class RateLimitConfig(ABC):
+    """限流配置基类"""
+
+    @abstractmethod
+    def to_dict(self) -> dict:
+        """转换为字典格式，用于存储到 Redis"""
+        pass
+
+    @abstractmethod
+    def get_type(self) -> str:
+        """获取限流类型标识"""
+        pass
+
+    @classmethod
+    @abstractmethod
+    def from_dict(cls, data: dict) -> 'RateLimitConfig':
+        """从字典创建配置实例"""
+        pass
+
+
+class QPSLimit(RateLimitConfig):
+    """QPS 限流配置
+
+    控制每秒最多执行多少个任务（使用滑动窗口算法）
+
+    Args:
+        qps: 每秒允许的最大请求数（所有 workers 总和）
+        window_size: 滑动窗口大小（秒），默认 1.0
+
+    Example:
+        @app.task(rate_limit=QPSLimit(qps=100))
+        async def my_task():
+            pass
+    """
+
+    def __init__(self, qps: int, window_size: float = 1.0):
+        # 转换为正确的类型（处理从 Redis 读取的字符串）
+        qps = int(qps)
+        window_size = float(window_size)
+
+        if qps <= 0:
+            raise ValueError(f"qps must be positive, got {qps}")
+        if window_size <= 0:
+            raise ValueError(f"window_size must be positive, got {window_size}")
+
+        self.qps = qps
+        self.window_size = window_size
+
+    def to_dict(self) -> dict:
+        """转换为字典格式"""
+        return {
+            'type': 'qps',
+            'qps': self.qps,
+            'window_size': self.window_size,
+        }
+
+    def get_type(self) -> str:
+        """获取限流类型标识"""
+        return 'qps'
+
+    @classmethod
+    def from_dict(cls, data: dict) -> 'QPSLimit':
+        """从字典创建配置实例"""
+        return cls(
+            qps=data['qps'],
+            window_size=data.get('window_size', 1.0)
+        )
+
+    def __repr__(self):
+        return f"QPSLimit(qps={self.qps}, window_size={self.window_size})"
+
+
+class ConcurrencyLimit(RateLimitConfig):
+    """并发限流配置
+
+    控制同一时刻最多可以运行多少个任务
+
+    Args:
+        max_concurrency: 最大并发数（所有 workers 总和）
+
+    Example:
+        @app.task(rate_limit=ConcurrencyLimit(max_concurrency=10))
+        async def my_task():
+            pass
+    """
+
+    def __init__(self, max_concurrency: int):
+        # 转换为正确的类型（处理从 Redis 读取的字符串）
+        max_concurrency = int(max_concurrency)
+
+        if max_concurrency <= 0:
+            raise ValueError(f"max_concurrency must be positive, got {max_concurrency}")
+
+        self.max_concurrency = max_concurrency
+
+    def to_dict(self) -> dict:
+        """转换为字典格式"""
+        return {
+            'type': 'concurrency',
+            'max_concurrency': self.max_concurrency,
+        }
+
+    def get_type(self) -> str:
+        """获取限流类型标识"""
+        return 'concurrency'
+
+    @classmethod
+    def from_dict(cls, data: dict) -> 'ConcurrencyLimit':
+        """从字典创建配置实例"""
+        return cls(max_concurrency=data['max_concurrency'])
+
+    def __repr__(self):
+        return f"ConcurrencyLimit(max_concurrency={self.max_concurrency})"
+
+
+def parse_rate_limit_config(data: dict) -> Optional[RateLimitConfig]:
+    """从字典解析限流配置
+
+    Args:
+        data: 限流配置字典
+
+    Returns:
+        RateLimitConfig 实例，如果解析失败则返回 None
+    """
+    if not data or not isinstance(data, dict):
+        return None
+
+    limit_type = data.get('type')
+
+    if limit_type == 'qps':
+        return QPSLimit.from_dict(data)
+    elif limit_type == 'concurrency':
+        return ConcurrencyLimit.from_dict(data)
+    else:
+        return None

jettask/utils/rate_limit/limiter.py

@@ -0,0 +1,41 @@
+"""
+限流器 - 兼容性导入
+
+⚠️ 此文件保持向后兼容，将导入重定向到新的模块化结构。
+原代码已被拆分到以下模块：
+- jettask.rate_limit.qps_limiter: QPSRateLimiter 类
+- jettask.rate_limit.concurrency_limiter: ConcurrencyRateLimiter 类
+- jettask.rate_limit.task_limiter: TaskRateLimiter 类
+- jettask.rate_limit.manager: RateLimiterManager 类
+
+新的模块结构：
+- qps_limiter.py: QPS限流器
+- concurrency_limiter.py: 并发限流器
+- task_limiter.py: 任务级限流器
+- manager.py: 限流器管理器
+
+使用示例：
+    # 旧的导入方式（仍然可用）
+    from jettask.rate_limit.limiter import RateLimiterManager, QPSRateLimiter
+
+    # 新的导入方式（推荐）
+    from jettask.rate_limit import RateLimiterManager, QPSRateLimiter
+"""
+
+# 导入新模块的所有公开接口
+from .config import RateLimitConfig, QPSLimit, ConcurrencyLimit
+from .qps_limiter import QPSRateLimiter
+from .concurrency_limiter import ConcurrencyRateLimiter
+from .task_limiter import TaskRateLimiter
+from .manager import RateLimiterManager
+
+# 保持向后兼容
+__all__ = [
+    'RateLimitConfig',
+    'QPSLimit',
+    'ConcurrencyLimit',
+    'QPSRateLimiter',
+    'ConcurrencyRateLimiter',
+    'TaskRateLimiter',
+    'RateLimiterManager',
+]

jettask/utils/rate_limit/manager.py

@@ -0,0 +1,269 @@
+"""
+限流器管理器
+
+统一管理所有任务的限流器实例
+职责：
+1. 限流器生命周期管理
+2. 配置动态更新
+3. 统一的限流接口
+"""
+
+import asyncio
+import logging
+import time
+from redis.asyncio import Redis
+from typing import Dict, Optional, List, Any
+
+from .task_limiter import TaskRateLimiter
+from .config import RateLimitConfig, parse_rate_limit_config
+
+logger = logging.getLogger('app')
+
+
+class RateLimiterManager:
+    """限流器管理器
+
+    管理多个任务的限流器。
+    """
+
+    def __init__(
+        self,
+        redis_client: Redis,
+        worker_id: str,
+        redis_prefix: str = "jettask",
+        registry_manager = None,
+        worker_state_manager = None
+    ):
+        """初始化限流器管理器
+
+        Args:
+            redis_client: 异步 Redis 客户端
+            worker_id: 当前 worker ID
+            redis_prefix: Redis key 前缀
+            registry_manager: RegistryManager 实例
+            worker_state_manager: WorkerStateManager 实例
+        """
+        self.redis = redis_client
+        self.worker_id = worker_id
+        self.redis_prefix = redis_prefix
+        self.registry_manager = registry_manager
+        self.worker_state_manager = worker_state_manager
+
+        # 任务名 -> 限流器映射
+        self.limiters: Dict[str, TaskRateLimiter] = {}
+
+        logger.debug(f"RateLimiterManager initialized for worker {worker_id}")
+
+    @staticmethod
+    def register_rate_limit_config(redis_client, task_name: str, config: RateLimitConfig, redis_prefix: str = "jettask"):
+        """注册任务的限流配置到 Redis（同步方法）
+
+        Args:
+            redis_client: 同步 Redis 客户端
+            task_name: 任务名称
+            config: RateLimitConfig 对象（ConcurrencyLimit 或 QPSLimit）
+            redis_prefix: Redis key 前缀
+        """
+        try:
+            rate_limit_key = f"{redis_prefix}:RATE_LIMIT:CONFIG:{task_name}"
+            # 将配置对象转换为字典并保存到 Redis
+            config_dict = config.to_dict()
+            redis_client.hset(rate_limit_key, mapping=config_dict)
+
+            # 将任务名添加到索引集合中（用于避免 scan 操作）
+            index_key = f"{redis_prefix}:RATE_LIMIT:INDEX"
+            redis_client.sadd(index_key, task_name)
+
+            logger.debug(f"Registered rate limit config for task '{task_name}': {config}")
+        except Exception as e:
+            logger.error(f"Failed to register rate limit config for task '{task_name}': {e}")
+
+    @staticmethod
+    def unregister_rate_limit_config(redis_client, task_name: str, redis_prefix: str = "jettask"):
+        """从 Redis 中删除任务的限流配置（同步方法）
+
+        Args:
+            redis_client: 同步 Redis 客户端
+            task_name: 任务名称
+            redis_prefix: Redis key 前缀
+
+        Returns:
+            是否成功删除
+        """
+        try:
+            rate_limit_key = f"{redis_prefix}:RATE_LIMIT:CONFIG:{task_name}"
+            deleted = redis_client.delete(rate_limit_key)
+
+            if deleted:
+                # 从索引集合中移除
+                index_key = f"{redis_prefix}:RATE_LIMIT:INDEX"
+                redis_client.srem(index_key, task_name)
+                logger.debug(f"Removed rate limit config for task '{task_name}'")
+                return True
+            return False
+        except Exception as e:
+            logger.error(f"Failed to remove rate limit config for task '{task_name}': {e}")
+            return False
+
+    async def add_limiter(self, task_name: str, config: RateLimitConfig):
+        """添加限流器
+
+        Args:
+            task_name: 任务名称
+            config: 限流配置（QPSLimit 或 ConcurrencyLimit）
+        """
+        if task_name in self.limiters:
+            logger.warning(f"Limiter for {task_name} already exists")
+            return
+
+        limiter = TaskRateLimiter(
+            redis_client=self.redis,
+            task_name=task_name,
+            worker_id=self.worker_id,
+            config=config,
+            redis_prefix=self.redis_prefix,
+            registry_manager=self.registry_manager,
+            worker_state_manager=self.worker_state_manager
+        )
+
+        await limiter.start()
+        self.limiters[task_name] = limiter
+
+        logger.debug(f"Added rate limiter for {task_name}: {config}")
+
+    async def remove_limiter(self, task_name: str, remove_from_redis: bool = False):
+        """移除任务限流器
+
+        Args:
+            task_name: 任务名称
+            remove_from_redis: 是否从 Redis 中删除配置（默认 False）
+        """
+        if task_name not in self.limiters:
+            return
+
+        limiter = self.limiters.pop(task_name)
+        await limiter.stop()
+
+        # 如果需要，从 Redis 中删除配置和索引
+        if remove_from_redis:
+            try:
+                # 删除配置 Hash
+                rate_limit_key = f"{self.redis_prefix}:RATE_LIMIT:CONFIG:{task_name}"
+                await self.redis.delete(rate_limit_key)
+
+                # 从索引集合中移除
+                index_key = f"{self.redis_prefix}:RATE_LIMIT:INDEX"
+                await self.redis.srem(index_key, task_name)
+
+                logger.debug(f"Removed rate limit config and index for {task_name} from Redis")
+            except Exception as e:
+                logger.error(f"Failed to remove rate limit config from Redis: {e}")
+
+        logger.debug(f"Removed rate limiter for {task_name}")
+
+    async def load_config_from_redis(self, task_names: list = None):
+        """从 Redis 加载限流配置
+
+        Args:
+            task_names: 任务名称列表（如果提供，只加载这些任务的配置；否则从索引集合中加载）
+        """
+        try:
+            config_count = 0
+            loaded_limiters = []
+
+            # 如果没有提供 task_names，尝试从索引集合中获取
+            if not task_names:
+                index_key = f"{self.redis_prefix}:RATE_LIMIT:INDEX"
+                task_names_bytes = await self.redis.smembers(index_key)
+                task_names = [
+                    name.decode('utf-8') if isinstance(name, bytes) else name
+                    for name in task_names_bytes
+                ]
+
+                if not task_names:
+                    logger.debug(f"No rate limit configs found in index {index_key}")
+                    return
+
+            # 遍历所有任务名称，加载配置
+            for task_name in task_names:
+                key = f"{self.redis_prefix}:RATE_LIMIT:CONFIG:{task_name}"
+
+                # 检查 key 是否存在
+                exists = await self.redis.exists(key)
+                if not exists:
+                    continue
+
+                # 检查 key 类型
+                key_type = await self.redis.type(key)
+                if key_type != "hash":
+                    logger.debug(f"Skipping non-hash key: {key} (type: {key_type})")
+                    continue
+
+                # 从 Hash 中读取配置
+                config_dict = await self.redis.hgetall(key)
+                if not config_dict:
+                    continue
+
+                # 转换 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()
+                }
+
+                # 解析配置
+                config = parse_rate_limit_config(config_dict)
+                if config and task_name not in self.limiters:
+                    try:
+                        await self.add_limiter(task_name, config)
+                        config_count += 1
+                        loaded_limiters.append(task_name)
+                    except Exception as e:
+                        logger.error(f"Failed to add limiter for {task_name}: {e}")
+
+            logger.debug(f"Loaded {config_count} rate limit configs from Redis")
+            logger.debug(f"Loaded rate limit config from Redis, limiters: {loaded_limiters}")
+        except Exception as e:
+            logger.error(f"Failed to load config from Redis: {e}")
+
+    async def acquire(self, task_name: str, timeout: float = 10.0) -> Optional[str]:
+        """获取指定任务的执行许可
+
+        Returns:
+            成功返回 task_id (或 True), 失败返回 None
+        """
+        limiter = self.limiters.get(task_name)
+        if not limiter:
+            # 没有限流，直接返回 True
+            return True
+
+        return await limiter.acquire(timeout)
+
+    async def release(self, task_name: str, task_id: Optional[str] = None):
+        """释放指定任务的执行许可
+
+        Args:
+            task_name: 任务名称
+            task_id: 任务ID (ConcurrencyLimit 需要)
+        """
+        limiter = self.limiters.get(task_name)
+        if limiter:
+            await limiter.release(task_id)
+
+    async def stop_all(self):
+        """停止所有限流器"""
+        for task_name, limiter in list(self.limiters.items()):
+            await limiter.stop()
+        self.limiters.clear()
+        logger.debug("Stopped all rate limiters")
+
+    def get_all_stats(self) -> Dict[str, dict]:
+        """获取所有限流器的统计信息"""
+        stats = {}
+        for task_name, limiter in self.limiters.items():
+            if hasattr(limiter, 'get_stats'):
+                stats[task_name] = limiter.get_stats()
+        return stats
+
+
+__all__ = ['RateLimiterManager']

jettask/utils/rate_limit/qps_limiter.py

@@ -0,0 +1,154 @@
+"""
+QPS（每秒查询数）限流器
+
+基于滑动窗口算法的QPS限流实现
+"""
+
+import asyncio
+import logging
+import time
+from collections import deque
+from redis.asyncio import Redis
+from typing import Dict, Optional
+
+from jettask.utils.time_sync import get_time_sync
+
+logger = logging.getLogger('app')
+
+
+class QPSRateLimiter:
+    """QPS 限流器 - 本地滑动窗口算法
+
+    使用滑动窗口算法在本地限流,不访问 Redis。
+    记录每个任务的执行时间戳,检查窗口内的执行次数。
+    使用全局时间同步器保证多个 worker 时间一致性。
+
+    特点：
+    - 纯本地计算，零 Redis 访问
+    - 高性能，低延迟
+    - 通过 Redis 协调配额分配
+    """
+
+    def __init__(
+        self,
+        qps: int,
+        window_size: float = 1.0,
+        worker_id: str = "unknown"
+    ):
+        """初始化 QPS 限流器
+
+        Args:
+            qps: 当前 worker 的 QPS 配额
+            window_size: 滑动窗口大小（秒）
+            worker_id: Worker ID，用于日志标识
+        """
+        self.qps_limit = qps
+        self.window_size = window_size
+        self.worker_id = worker_id
+        self.timestamps = deque()  # 记录执行时间戳
+        self._lock = asyncio.Lock()
+        self.time_sync = get_time_sync()
+
+    def _get_time(self) -> float:
+        """获取当前时间"""
+        return self.time_sync.time()
+
+    async def acquire(self, timeout: float = 10.0) -> bool:
+        """获取一个执行许可
+
+        Args:
+            timeout: 等待超时时间（秒）
+
+        Returns:
+            成功获取返回 True，超时返回 False
+        """
+        start_time = self._get_time()
+
+        while True:
+            async with self._lock:
+                now = self._get_time()
+
+                # 清理过期的时间戳（窗口外的）
+                cutoff = now - self.window_size
+                while self.timestamps and self.timestamps[0] < cutoff:
+                    self.timestamps.popleft()
+
+                # 检查 QPS 限制
+                current_count = len(self.timestamps)
+                if current_count < self.qps_limit:
+                    # 未达到限制，允许执行
+                    self.timestamps.append(now)
+                    logger.debug(
+                        f"[QPS] Acquired, count={current_count + 1}/{self.qps_limit}"
+                    )
+                    return True
+
+            # 达到限制，检查超时
+            if timeout is not None and self._get_time() - start_time > timeout:
+                logger.warning(f"[QPS] Acquire timeout after {timeout}s")
+                return False
+
+            # 短暂休眠后重试
+            await asyncio.sleep(0.01)
+
+    async def try_acquire(self) -> bool:
+        """尝试获取执行许可（非阻塞）"""
+        async with self._lock:
+            now = self._get_time()
+
+            # 清理过期的时间戳
+            cutoff = now - self.window_size
+            while self.timestamps and self.timestamps[0] < cutoff:
+                self.timestamps.popleft()
+
+            # 检查是否可以执行
+            current_count = len(self.timestamps)
+            if current_count < self.qps_limit:
+                self.timestamps.append(now)
+                logger.debug(f"[QPS] Try acquired, count={current_count + 1}/{self.qps_limit}")
+                return True
+
+            logger.debug(f"[QPS] Try acquire failed, count={current_count}/{self.qps_limit}")
+            return False
+
+    async def release(self):
+        """释放执行许可（QPS 模式不需要 release）"""
+        pass
+
+    async def update_limit(self, new_limit: int):
+        """动态更新 QPS 限制
+
+        Args:
+            new_limit: 新的 QPS 限制
+        """
+        async with self._lock:
+            if self.qps_limit != new_limit:
+                old_limit = self.qps_limit
+                self.qps_limit = new_limit
+                logger.debug(
+                    f"[WORKER:{self.worker_id}] [QPS] Limit changed: {old_limit} → {new_limit}"
+                )
+
+    async def get_stats(self) -> dict:
+        """获取统计信息"""
+        async with self._lock:
+            now = self._get_time()
+            cutoff = now - self.window_size
+            while self.timestamps and self.timestamps[0] < cutoff:
+                self.timestamps.popleft()
+
+            return {
+                'mode': 'qps',
+                'qps_limit': self.qps_limit,
+                'current_qps': len(self.timestamps),
+                'window_size': self.window_size,
+            }
+
+
+# ============================================================
+# 并发限流器 - 分布式信号量
+# ============================================================
+
+
+
+__all__ = ['QPSRateLimiter']