@agentunion/kite 1.3.2 → 1.5.0

package/extensions/services/proxy/agentcp/improved_scheduler.py

@@ -0,0 +1,498 @@
+# -*- coding: utf-8 -*-
+"""
+改进的消息调度器
+支持线程池 + 异步任务池的混合架构
+"""
+
+import asyncio
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from typing import Callable, Awaitable, Dict, Any, Optional
+from agentcp.base.log import log_info, log_error, log_exception, log_warning, log_debug
+
+
+class ImprovedMessageScheduler:
+    """
+    改进的消息调度器
+
+    架构:
+    - 核心线程: core_workers 个 (常驻)
+    - 最大线程: max_workers 个 (高峰扩展)
+    - 每线程并发异步任务: max_tasks_per_worker 个
+
+    优势:
+    - 资源高效: 少量线程处理大量异步任务
+    - 高并发: 总并发 = max_workers × max_tasks_per_worker
+    - 负载均衡: 自动选择负载最低的工作线程
+    """
+
+    def __init__(self,
+                 core_workers: int = 20,
+                 max_workers: int = 50,
+                 max_tasks_per_worker: int = 10):
+        """
+        初始化调度器
+
+        Args:
+            core_workers: 核心工作线程数 (常驻)
+            max_workers: 最大工作线程数 (高峰时扩展)
+            max_tasks_per_worker: 每个工作线程的最大并发异步任务数
+        """
+        self.core_workers = core_workers
+        self.max_workers = max_workers
+        self.max_tasks_per_worker = max_tasks_per_worker
+
+        # 线程池
+        self.thread_pool = ThreadPoolExecutor(
+            max_workers=max_workers,
+            thread_name_prefix="agentcp-worker"
+        )
+
+        # 工作线程状态
+        self.worker_loops: Dict[int, asyncio.AbstractEventLoop] = {}  # worker_id -> event loop
+        self.worker_queues: Dict[int, asyncio.Queue] = {}  # worker_id -> message queue
+        self.worker_tasks_count: Dict[int, int] = {}  # worker_id -> active task count
+        self.worker_lock = threading.Lock()
+
+        # 统计信息
+        self.total_messages = 0
+        self.total_processed = 0
+        self.total_errors = 0
+        self.total_rejected = 0  # ✅ P1修复: 添加拒绝计数
+        self.active_workers = 0
+        self.is_running = True
+
+        # ✅ 修复: 添加统计锁,保护计数器的线程安全
+        self.stats_lock = threading.Lock()
+
+        # ✅ P1修复: 队列监控配置
+        self.queue_warn_threshold = 0.8  # 队列使用率警告阈值
+        self.queue_timeout = 5.0  # ✅ 增加: 队列等待超时从2秒增加到5秒
+        self.max_submit_retries = 3  # ✅ 新增: 提交失败时的最大重试次数
+
+        # 初始化核心工作线程
+        self._init_core_workers()
+
+        log_info(f"[Scheduler] 初始化完成: core_workers={core_workers}, "
+                f"max_workers={max_workers}, max_tasks_per_worker={max_tasks_per_worker}")
+
+    def _init_core_workers(self):
+        """初始化核心工作线程"""
+        for worker_id in range(self.core_workers):
+            self._start_worker(worker_id)
+
+    def _start_worker(self, worker_id: int):
+        """
+        ✅ P1修复: 启动一个工作线程 (使用Event同步)
+
+        Args:
+            worker_id: 工作线程ID
+        """
+        # 创建启动就绪事件
+        ready_event = threading.Event()
+
+        # 提交工作线程任务
+        self.thread_pool.submit(self._worker_main, worker_id, ready_event)
+
+        # ✅ 使用Event等待线程就绪 (最多等待5秒)
+        if not ready_event.wait(timeout=5.0):
+            log_error(f"[Worker-{worker_id}] 启动超时")
+            raise RuntimeError(f"Worker-{worker_id} failed to start within 5 seconds")
+
+        with self.worker_lock:
+            self.worker_tasks_count[worker_id] = 0
+            self.active_workers += 1
+
+        log_info(f"[Worker-{worker_id}] 启动成功")
+
+    def _worker_main(self, worker_id: int, ready_event: threading.Event = None):
+        """
+        ✅ P1修复: 工作线程主函数
+        运行一个持久的事件循环,处理异步任务
+
+        Args:
+            worker_id: 工作线程ID
+            ready_event: 启动就绪事件 (用于同步启动)
+        """
+        # 创建新的事件循环
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+
+        # ✅ 增加: 队列大小从1000增加到5000，提高并发容量
+        queue = asyncio.Queue(maxsize=5000)
+
+        # 注册到全局状态
+        with self.worker_lock:
+            self.worker_loops[worker_id] = loop
+            self.worker_queues[worker_id] = queue
+
+        log_info(f"[Worker-{worker_id}] 事件循环启动, thread={threading.current_thread().name}")
+
+        # ✅ 通知启动完成
+        if ready_event:
+            ready_event.set()
+
+        # 运行事件循环
+        try:
+            loop.run_until_complete(self._worker_loop(worker_id, queue))
+        except Exception as e:
+            log_exception(f"[Worker-{worker_id}] 事件循环异常: {e}")
+        finally:
+            # 清理
+            try:
+                # 取消所有待处理任务
+                pending = asyncio.all_tasks(loop)
+                for task in pending:
+                    task.cancel()
+                # 等待取消完成
+                loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
+                loop.close()
+            except Exception as e:
+                log_error(f"[Worker-{worker_id}] 清理异常: {e}")
+
+            # 从全局状态移除
+            with self.worker_lock:
+                if worker_id in self.worker_loops:
+                    del self.worker_loops[worker_id]
+                if worker_id in self.worker_queues:
+                    del self.worker_queues[worker_id]
+                if worker_id in self.worker_tasks_count:
+                    del self.worker_tasks_count[worker_id]
+                self.active_workers -= 1
+
+            log_info(f"[Worker-{worker_id}] 已停止")
+
+    async def _worker_loop(self, worker_id: int, queue: asyncio.Queue):
+        """✅ P0-3改进: 工作线程事件循环，优化任务拒绝逻辑
+
+        改进要点：
+        1. 移除"放回队列"的逻辑（会导致死锁）
+        2. 任务超限时直接跳过（让 submit_message 重试其他 worker）
+        3. 添加详细的日志
+
+        Args:
+            worker_id: 工作线程ID
+            queue: 消息队列
+        """
+        while self.is_running:
+            try:
+                # 等待新消息 (超时检查,避免卡死)
+                try:
+                    message_handler, data = await asyncio.wait_for(queue.get(), timeout=1.0)
+                except asyncio.TimeoutError:
+                    continue
+
+                # ✅ 检查是否超过任务限制
+                with self.worker_lock:
+                    current_tasks = self.worker_tasks_count.get(worker_id, 0)
+
+                    if current_tasks >= self.max_tasks_per_worker:
+                        # ✅ P0-3改进: 不再尝试放回队列，直接记录拒绝
+                        with self.stats_lock:
+                            self.total_rejected += 1
+
+                        message_id = data.get('message_id', 'unknown') if isinstance(data, dict) else 'unknown'
+                        log_warning(
+                            f"⚠️ [Worker-{worker_id}] 任务超限 "
+                            f"({current_tasks}/{self.max_tasks_per_worker}), "
+                            f"拒绝任务 message_id={message_id[:16] if len(message_id) > 16 else message_id}..."
+                        )
+
+                        # 短暂等待后继续取下一个任务
+                        await asyncio.sleep(0.05)
+                        continue
+
+                    # 增加任务计数
+                    self.worker_tasks_count[worker_id] = current_tasks + 1
+
+                # ✅ 关键:创建异步任务(不等待完成)
+                asyncio.create_task(
+                    self._handle_message_wrapper(worker_id, message_handler, data)
+                )
+
+            except asyncio.CancelledError:
+                log_info(f"[Worker-{worker_id}] 收到取消信号")
+                break
+            except Exception as e:
+                log_exception(f"[Worker-{worker_id}] 事件循环异常: {e}")
+                await asyncio.sleep(0.1)
+
+    async def _handle_message_wrapper(self,
+                                      worker_id: int,
+                                      message_handler: Callable[[Dict], Awaitable[None]],
+                                      data: Dict[str, Any]):
+        """
+        消息处理包装器
+        处理完成后减少任务计数
+
+        Args:
+            worker_id: 工作线程ID
+            message_handler: 异步消息处理函数
+            data: 消息数据
+        """
+        try:
+            # 调用实际的消息处理函数
+            await message_handler(data)
+
+            # ✅ 修复: 使用锁保护统计计数器
+            with self.stats_lock:
+                self.total_processed += 1
+
+        except Exception as e:
+            # ✅ 修复: 使用锁保护统计计数器
+            with self.stats_lock:
+                self.total_errors += 1
+            log_exception(f"[Worker-{worker_id}] 消息处理失败: {e}")
+        finally:
+            # 减少任务计数
+            with self.worker_lock:
+                current = self.worker_tasks_count.get(worker_id, 0)
+                self.worker_tasks_count[worker_id] = max(0, current - 1)
+
+    def submit_message(self,
+                       message_handler: Callable[[Dict], Awaitable[None]],
+                       data: Dict[str, Any],
+                       raise_on_reject: bool = False) -> bool:
+        """✅ P0-3改进: 提交消息到调度器（多候选 worker + 队列监控）
+
+        改进要点：
+        1. 选择多个候选 worker 而不是单个
+        2. 依次尝试候选 worker，失败时自动切换
+        3. 检查队列使用率，跳过接近满的 worker
+        4. 优化重试逻辑和等待时间
+
+        Args:
+            message_handler: 异步消息处理函数 async def handler(data)
+            data: 消息数据
+            raise_on_reject: 如果为True，在任务被拒绝时抛出异常；否则返回False
+
+        Returns:
+            bool: True表示提交成功，False表示被拒绝
+
+        Raises:
+            RuntimeError: 当raise_on_reject=True且任务被拒绝时抛出
+        """
+        # ✅ 统计计数
+        with self.stats_lock:
+            self.total_messages += 1
+
+        # ✅ P0-3改进: 重试机制，每次尝试多个候选 worker
+        last_error = None
+        for retry_attempt in range(self.max_submit_retries):
+            try:
+                # ✅ 获取负载最低的 TOP 3 worker
+                candidate_workers = self._select_workers_by_load(top_n=3)
+
+                if not candidate_workers:
+                    error_msg = "[Scheduler] 没有可用的工作线程"
+                    log_error(error_msg)
+                    with self.stats_lock:
+                        self.total_rejected += 1
+                    if raise_on_reject:
+                        raise RuntimeError(error_msg)
+                    return False
+
+                # ✅ 依次尝试候选 worker
+                submitted = False
+                for worker_id in candidate_workers:
+                    loop = self.worker_loops.get(worker_id)
+                    queue = self.worker_queues.get(worker_id)
+
+                    if not loop or not queue or loop.is_closed():
+                        continue
+
+                    # ✅ 检查队列使用率
+                    queue_size = queue.qsize()
+                    queue_maxsize = queue.maxsize
+                    usage_rate = queue_size / queue_maxsize if queue_maxsize > 0 else 0
+
+                    # ✅ 如果队列使用率超过 90%，跳过这个 worker
+                    if usage_rate >= 0.9:
+                        log_warning(
+                            f"⚠️ [Worker-{worker_id}] 队列接近满 "
+                            f"({queue_size}/{queue_maxsize}, {usage_rate*100:.1f}%), 尝试下一个"
+                        )
+                        continue
+
+                    # ✅ 尝试提交
+                    try:
+                        future = asyncio.run_coroutine_threadsafe(
+                            self._put_with_timeout(queue, message_handler, data),
+                            loop
+                        )
+                        future.result(timeout=self.queue_timeout)
+
+                        # ✅ 提交成功
+                        submitted = True
+                        log_debug(f"✅ [Scheduler] 消息已提交到 Worker-{worker_id}")
+                        return True
+
+                    except Exception as e:
+                        log_debug(f"⚠️ [Worker-{worker_id}] 提交失败: {e}, 尝试下一个")
+                        continue
+
+                if submitted:
+                    return True
+
+                # ✅ 所有候选 worker 都失败
+                last_error = "所有候选 worker 都无法接收任务"
+
+            except Exception as e:
+                last_error = str(e)
+
+            # ✅ 重试前等待（指数退避）
+            if retry_attempt < self.max_submit_retries - 1:
+                wait_time = 0.05 * (2 ** retry_attempt)  # 指数退避: 0.05s, 0.1s, 0.2s
+                log_warning(
+                    f"⚠️ [Scheduler] 提交失败 (第{retry_attempt + 1}次), "
+                    f"{wait_time}s 后重试... reason={last_error}"
+                )
+                time.sleep(wait_time)
+
+        # ✅ 所有重试都失败
+        with self.stats_lock:
+            self.total_rejected += 1
+
+        error_msg = f"[Scheduler] 消息提交最终失败: {last_error}"
+        log_error(error_msg)
+
+        if raise_on_reject:
+            raise RuntimeError(error_msg)
+
+        return False
+
+    async def _put_with_timeout(self, queue: asyncio.Queue, message_handler, data):
+        """
+        ✅ P1修复: 带超时的队列put操作
+
+        Args:
+            queue: 目标队列
+            message_handler: 消息处理器
+            data: 消息数据
+        """
+        try:
+            await asyncio.wait_for(
+                queue.put((message_handler, data)),
+                timeout=self.queue_timeout
+            )
+        except asyncio.TimeoutError:
+            raise Exception(f"队列已满,等待超时 ({self.queue_timeout}s)")
+
+    def _select_worker(self) -> Optional[int]:
+        """
+        选择负载最低的工作线程（保留向后兼容）
+
+        Returns:
+            worker_id 或 None (如果没有可用worker)
+        """
+        with self.worker_lock:
+            if not self.worker_tasks_count:
+                return None
+
+            # 找到任务数最少的worker
+            min_tasks = float('inf')
+            selected_worker = None
+
+            for worker_id, task_count in self.worker_tasks_count.items():
+                if task_count < min_tasks:
+                    min_tasks = task_count
+                    selected_worker = worker_id
+
+                    # 如果找到空闲worker,直接使用
+                    if task_count == 0:
+                        break
+
+            return selected_worker
+
+    def _select_workers_by_load(self, top_n: int = 3) -> list:
+        """✅ P0-3新增: 选择负载最低的 TOP N worker
+
+        Args:
+            top_n: 返回前 N 个负载最低的 worker
+
+        Returns:
+            worker_id 列表，按负载从低到高排序
+        """
+        with self.worker_lock:
+            if not self.worker_tasks_count:
+                return []
+
+            # 按任务数排序（从少到多）
+            sorted_workers = sorted(
+                self.worker_tasks_count.items(),
+                key=lambda x: x[1]  # x[1] 是任务数
+            )
+
+            # 返回前 N 个 worker 的 ID
+            return [worker_id for worker_id, _ in sorted_workers[:top_n]]
+
+    def get_stats(self) -> Dict[str, Any]:
+        """
+        ✅ 修复: 获取统计信息 (线程安全)
+
+        Returns:
+            统计数据字典
+        """
+        # 获取 worker 统计
+        with self.worker_lock:
+            total_active_tasks = sum(self.worker_tasks_count.values())
+            worker_details = dict(self.worker_tasks_count)
+            active_workers = self.active_workers
+
+        # 获取全局统计
+        with self.stats_lock:
+            total_messages = self.total_messages
+            total_processed = self.total_processed
+            total_errors = self.total_errors
+            total_rejected = self.total_rejected  # ✅ P1修复: 添加拒绝统计
+
+        return {
+            'total_messages': total_messages,
+            'total_processed': total_processed,
+            'total_errors': total_errors,
+            'total_rejected': total_rejected,  # ✅ P1修复
+            'active_workers': active_workers,
+            'total_active_tasks': total_active_tasks,
+            'worker_tasks': worker_details,
+            'success_rate': f"{(total_processed / max(1, total_messages)) * 100:.2f}%"
+        }
+
+    def print_stats(self):
+        """✅ P1修复: 打印统计信息 (包含拒绝数)"""
+        stats = self.get_stats()
+        log_info(f"[Scheduler Stats] "
+                f"Messages: {stats['total_messages']}, "
+                f"Processed: {stats['total_processed']}, "
+                f"Errors: {stats['total_errors']}, "
+                f"Rejected: {stats['total_rejected']}, "
+                f"Active Workers: {stats['active_workers']}, "
+                f"Active Tasks: {stats['total_active_tasks']}, "
+                f"Success Rate: {stats['success_rate']}")
+
+    def shutdown(self, wait: bool = True):
+        """
+        关闭调度器
+
+        Args:
+            wait: 是否等待所有任务完成
+        """
+        log_info("[Scheduler] 正在关闭...")
+        self.is_running = False
+
+        if wait:
+            # 等待一段时间让任务完成
+            max_wait = 10  # 最多等待10秒
+            for i in range(max_wait):
+                stats = self.get_stats()
+                if stats['total_active_tasks'] == 0:
+                    break
+                log_info(f"[Scheduler] 等待任务完成... 剩余 {stats['total_active_tasks']} 个")
+                time.sleep(1)
+
+        # 关闭线程池
+        self.thread_pool.shutdown(wait=wait)
+
+        # 打印最终统计
+        self.print_stats()
+        log_info("[Scheduler] 已关闭")