MCP Client 的并发与异步设计

import asyncio
import httpx
from typing import Dict, Any, Optional, List

class AsyncHttpClient:
    """异步 HTTP 客户端"""
    
    def __init__(self, base_url: str, timeout: float = 30.0, max_connections: int = 100):
        """
        初始化异步 HTTP 客户端
        
        Args:
            base_url: 基础 URL
            timeout: 超时时间，秒
            max_connections: 最大连接数
        """
        self.base_url = base_url
        self.timeout = timeout
        self.max_connections = max_connections
        
        # 创建异步 HTTP 客户端
        self.client = httpx.AsyncClient(
            base_url=base_url,
            timeout=timeout,
            limits=httpx.Limits(max_connections=max_connections),
            http2=True,  # 启用 HTTP/2
        )
    
    async def get(self, path: str, params: Optional[Dict[str, Any]] = None, headers: Optional[Dict[str, str]] = None) -> Dict[str, Any]:
        """异步 GET 请求"""
        try:
            response = await self.client.get(
                path,
                params=params,
                headers=headers,
            )
            response.raise_for_status()
            return response.json()
        except httpx.HTTPError as e:
            print(f"HTTP GET 请求失败: {e}")
            raise
        except Exception as e:
            print(f"GET 请求失败: {e}")
            raise
    
    async def post(self, path: str, json: Optional[Dict[str, Any]] = None, headers: Optional[Dict[str, str]] = None) -> Dict[str, Any]:
        """异步 POST 请求"""
        try:
            response = await self.client.post(
                path,
                json=json,
                headers=headers,
            )
            response.raise_for_status()
            return response.json()
        except httpx.HTTPError as e:
            print(f"HTTP POST 请求失败: {e}")
            raise
        except Exception as e:
            print(f"POST 请求失败: {e}")
            raise
    
    async def close(self):
        """关闭异步 HTTP 客户端"""
        await self.client.aclose()
    
    async def __aenter__(self):
        """进入上下文管理器"""
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """退出上下文管理器"""
        await self.close()

import asyncio
from typing import Dict, Any, Optional, List, Callable, Tuple
from dataclasses import dataclass, field
import time

@dataclass
class Task:
    """任务信息"""
    id: str
    coro: Callable
    priority: int = 0  # 优先级，值越高优先级越高
    delay: float = 0.0  # 延迟执行时间，秒
    created_at: float = field(default_factory=time.time)
    started_at: Optional[float] = None
    completed_at: Optional[float] = None
    result: Optional[Any] = None
    exception: Optional[Exception] = None

class AsyncTaskScheduler:
    """异步任务调度器"""
    
    def __init__(self, max_concurrent_tasks: int = 10):
        """
        初始化异步任务调度器
        
        Args:
            max_concurrent_tasks: 最大并发任务数
        """
        self.max_concurrent_tasks = max_concurrent_tasks
        self.pending_tasks: List[Task] = []  # 待处理任务
        self.running_tasks: Dict[str, asyncio.Task] = {}  # 正在运行的任务
        self.task_results: Dict[str, Any] = {}  # 任务结果
        self.is_running = False
        self.scheduler_task: Optional[asyncio.Task] = None
    
    def add_task(self, task_id: str, coro: Callable, priority: int = 0, delay: float = 0.0) -> str:
        """添加任务"""
        task = Task(
            id=task_id,
            coro=coro,
            priority=priority,
            delay=delay,
        )
        self.pending_tasks.append(task)
        
        # 按优先级和延迟时间排序
        self.pending_tasks.sort(key=lambda x: (-x.priority, x.delay))
        
        return task_id
    
    async def _execute_task(self, task: Task):
        """执行单个任务"""
        task_id = task.id
        try:
            # 等待延迟时间
            if task.delay > 0:
                await asyncio.sleep(task.delay)
            
            task.started_at = time.time()
            
            # 执行任务
            result = await task.coro()
            
            task.completed_at = time.time()
            task.result = result
            self.task_results[task_id] = result
        except Exception as e:
            task.completed_at = time.time()
            task.exception = e
            self.task_results[task_id] = e
        finally:
            # 移除正在运行的任务
            if task_id in self.running_tasks:
                del self.running_tasks[task_id]
    
    async def _scheduler_loop(self):
        """调度器主循环"""
        while self.is_running or self.pending_tasks:
            # 检查是否有足够的并发容量
            if len(self.running_tasks) < self.max_concurrent_tasks and self.pending_tasks:
                # 获取优先级最高的任务
                task = self.pending_tasks.pop(0)
                
                # 创建异步任务
                async_task = asyncio.create_task(self._execute_task(task))
                self.running_tasks[task.id] = async_task
            else:
                # 等待一段时间后再次检查
                await asyncio.sleep(0.1)
    
    async def start(self):
        """启动调度器"""
        if self.is_running:
            return
        
        self.is_running = True
        self.scheduler_task = asyncio.create_task(self._scheduler_loop())
    
    async def stop(self):
        """停止调度器"""
        self.is_running = False
        
        if self.scheduler_task:
            await self.scheduler_task
            self.scheduler_task = None
        
        # 等待所有正在运行的任务完成
        if self.running_tasks:
            await asyncio.gather(*self.running_tasks.values(), return_exceptions=True)
    
    def get_task_result(self, task_id: str) -> Optional[Any]:
        """获取任务结果"""
        return self.task_results.get(task_id)
    
    def get_status(self) -> Dict[str, Any]:
        """获取调度器状态"""
        return {
            "pending_tasks": len(self.pending_tasks),
            "running_tasks": len(self.running_tasks),
            "completed_tasks": len(self.task_results),
            "is_running": self.is_running,
        }

import asyncio
from typing import Dict, Any, Optional

class AsyncConcurrentDict:
    """异步并发安全字典"""
    
    def __init__(self):
        """初始化异步并发安全字典"""
        self._data: Dict[str, Any] = {}
        self._lock = asyncio.Lock()  # 异步锁
    
    async def get(self, key: str, default: Optional[Any] = None) -> Any:
        """获取字典值"""
        async with self._lock:
            return self._data.get(key, default)
    
    async def set(self, key: str, value: Any) -> None:
        """设置字典值"""
        async with self._lock:
            self._data[key] = value
    
    async def delete(self, key: str) -> None:
        """删除字典键"""
        async with self._lock:
            if key in self._data:
                del self._data[key]
    
    async def contains(self, key: str) -> bool:
        """检查字典是否包含键"""
        async with self._lock:
            return key in self._data
    
    async def keys(self) -> List[str]:
        """获取所有键"""
        async with self._lock:
            return list(self._data.keys())
    
    async def values(self) -> List[Any]:
        """获取所有值"""
        async with self._lock:
            return list(self._data.values())
    
    async def items(self) -> List[Tuple[str, Any]]:
        """获取所有键值对"""
        async with self._lock:
            return list(self._data.items())
    
    async def clear(self) -> None:
        """清空字典"""
        async with self._lock:
            self._data.clear()
    
    async def __len__(self) -> int:
        """获取字典长度"""
        async with self._lock:
            return len(self._data)
    
    async def __getitem__(self, key: str) -> Any:
        """获取字典值（[] 操作符）"""
        async with self._lock:
            return self._data[key]
    
    async def __setitem__(self, key: str, value: Any) -> None:
        """设置字典值（[] 操作符）"""
        async with self._lock:
            self._data[key] = value
    
    async def __delitem__(self, key: str) -> None:
        """删除字典键（del 操作符）"""
        async with self._lock:
            del self._data[key]

import asyncio
import time
from typing import Dict, Any, Optional, List
from dataclasses import dataclass, field

@dataclass
class ThreadPoolStats:
    """线程池统计信息"""
    current_threads: int = 0
    idle_threads: int = 0
    active_threads: int = 0
    max_threads: int = 0
    task_queue_size: int = 0
    completed_tasks: int = 0
    rejected_tasks: int = 0
    avg_task_time: float = 0.0
    last_updated: float = field(default_factory=time.time)

class DynamicThreadPool:
    """动态线程池"""
    
    def __init__(self,
                 min_threads: int = 5,
                 max_threads: int = 100,
                 idle_timeout: float = 60.0,
                 scaling_factor: float = 0.5):
        """
        初始化动态线程池
        
        Args:
            min_threads: 最小线程数
            max_threads: 最大线程数
            idle_timeout: 线程空闲超时时间，秒
            scaling_factor: 缩放因子，用于动态调整线程数
        """
        self.min_threads = min_threads
        self.max_threads = max_threads
        self.idle_timeout = idle_timeout
        self.scaling_factor = scaling_factor
        
        # 线程池状态
        self.threads: List[asyncio.Task] = []
        self.idle_threads: asyncio.Queue = asyncio.Queue()
        self.task_queue: asyncio.Queue = asyncio.Queue()
        self.is_running = False
        
        # 统计信息
        self.stats = ThreadPoolStats(
            current_threads=0,
            idle_threads=0,
            active_threads=0,
            max_threads=max_threads,
            task_queue_size=0,
            completed_tasks=0,
            rejected_tasks=0,
            avg_task_time=0.0,
        )
        
        # 动态调整任务
        self.scaling_task: Optional[asyncio.Task] = None
    
    async def _worker(self):
        """工作线程"""
        thread_id = id(asyncio.current_task())
        last_activity = time.time()
        
        while self.is_running:
            try:
                # 尝试从任务队列获取任务，设置超时
                task = await asyncio.wait_for(self.task_queue.get(), timeout=self.idle_timeout)
                
                last_activity = time.time()
                self.stats.active_threads += 1
                self.stats.idle_threads -= 1
                
                # 执行任务
                start_time = time.time()
                await task()
                end_time = time.time()
                
                # 更新统计信息
                self.stats.completed_tasks += 1
                task_time = end_time - start_time
                # 更新平均任务时间
                self.stats.avg_task_time = (
                    self.stats.avg_task_time * (self.stats.completed_tasks - 1) + task_time
                ) / self.stats.completed_tasks
                
                self.stats.active_threads -= 1
                self.stats.idle_threads += 1
                
                # 任务完成
                self.task_queue.task_done()
            except asyncio.TimeoutError:
                # 检查线程是否应该关闭
                if time.time() - last_activity > self.idle_timeout and len(self.threads) > self.min_threads:
                    # 关闭空闲线程
                    break
            except Exception as e:
                print(f"工作线程执行任务失败: {e}")
                # 任务完成（即使失败）
                self.task_queue.task_done()
        
        # 移除线程
        self._remove_thread()
    
    def _add_thread(self):
        """添加线程"""
        if len(self.threads) < self.max_threads:
            worker = asyncio.create_task(self._worker())
            self.threads.append(worker)
            self.idle_threads.put_nowait(worker)
            
            self.stats.current_threads += 1
            self.stats.idle_threads += 1
    
    def _remove_thread(self):
        """移除线程"""
        if self.threads:
            # 从线程列表中移除当前线程
            current_task = asyncio.current_task()
            if current_task in self.threads:
                self.threads.remove(current_task)
                
                self.stats.current_threads -= 1
                # 线程可能处于活跃或空闲状态，需要调整统计信息
                if self.stats.active_threads > 0:
                    self.stats.active_threads -= 1
                else:
                    self.stats.idle_threads -= 1
    
    async def _dynamic_scaling(self):
        """动态调整线程数"""
        while self.is_running:
            # 获取当前状态
            queue_size = self.task_queue.qsize()
            current_threads = len(self.threads)
            idle_threads = self.idle_threads.qsize()
            
            # 更新统计信息
            self.stats.task_queue_size = queue_size
            
            # 动态调整线程数
            if queue_size > current_threads * self.scaling_factor and current_threads < self.max_threads:
                # 需要增加线程
                new_threads = min(
                    int(queue_size * self.scaling_factor) - current_threads,
                    self.max_threads - current_threads
                )
                for _ in range(new_threads):
                    self._add_thread()
            elif idle_threads > self.min_threads and current_threads > self.min_threads:
                # 需要减少线程，但保持最小线程数
                pass  # 线程会在空闲超时后自动关闭
            
            # 每 5 秒调整一次
            await asyncio.sleep(5.0)
    
    async def submit(self, task: callable) -> None:
        """提交任务到线程池"""
        if not self.is_running:
            raise RuntimeError("线程池已关闭")
        
        # 检查任务队列大小
        if self.task_queue.qsize() > self.max_threads * 10:
            # 任务队列已满，拒绝任务
            self.stats.rejected_tasks += 1
            raise RuntimeError("任务队列已满，拒绝任务")
        
        # 提交任务
        await self.task_queue.put(task)
        
        # 如果没有足够的空闲线程，添加线程
        if self.idle_threads.qsize() == 0 and len(self.threads) < self.max_threads:
            self._add_thread()
    
    async def start(self):
        """启动线程池"""
        if self.is_running:
            return
        
        self.is_running = True
        
        # 初始化最小线程数
        for _ in range(self.min_threads):
            self._add_thread()
        
        # 启动动态调整任务
        self.scaling_task = asyncio.create_task(self._dynamic_scaling())
    
    async def stop(self):
        """停止线程池"""
        self.is_running = False
        
        # 等待动态调整任务完成
        if self.scaling_task:
            await self.scaling_task
            self.scaling_task = None
        
        # 等待所有任务完成
        await self.task_queue.join()
        
        # 等待所有线程完成
        if self.threads:
            await asyncio.gather(*self.threads, return_exceptions=True)
            self.threads.clear()
            self.idle_threads = asyncio.Queue()
        
        # 重置统计信息
        self.stats = ThreadPoolStats(
            current_threads=0,
            idle_threads=0,
            active_threads=0,
            max_threads=self.max_threads,
            task_queue_size=0,
            completed_tasks=0,
            rejected_tasks=0,
            avg_task_time=0.0,
        )
    
    def get_stats(self) -> ThreadPoolStats:
        """获取线程池统计信息"""
        self.stats.last_updated = time.time()
        return self.stats
    
    async def __aenter__(self):
        """进入上下文管理器"""
        await self.start()
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """退出上下文管理器"""
        await self.stop()

# 示例：并发与异步使用示例
import asyncio
import time
from async_http_client import AsyncHttpClient
from async_task_scheduler import AsyncTaskScheduler
from async_concurrent_dict import AsyncConcurrentDict
from dynamic_thread_pool import DynamicThreadPool

async def example_async_http_client():
    """异步 HTTP 客户端示例"""
    print("=== 异步 HTTP 客户端示例 ===")
    
    async with AsyncHttpClient(base_url="https://httpbin.org") as client:
        # 并发发送多个请求
        tasks = [
            client.get("/get"),
            client.post("/post", json={"test": "data"}),
            client.get("/delay/2"),  # 延迟 2 秒的请求
        ]
        
        start_time = time.time()
        results = await asyncio.gather(*tasks)
        end_time = time.time()
        
        print(f"发送 3 个请求耗时: {end_time - start_time:.2f} 秒")
        for i, result in enumerate(results):
            print(f"请求 {i+1} 结果: {list(result.keys())[:5]}...")

async def example_async_task_scheduler():
    """异步任务调度器示例"""
    print("\n=== 异步任务调度器示例 ===")
    
    scheduler = AsyncTaskScheduler(max_concurrent_tasks=3)
    
    # 启动调度器
    await scheduler.start()
    
    # 定义任务函数
    async def task_func(task_id, delay):
        print(f"任务 {task_id} 开始执行，延迟 {delay} 秒")
        await asyncio.sleep(delay)
        print(f"任务 {task_id} 执行完成")
        return f"任务 {task_id} 结果"
    
    # 添加任务
    task_ids = []
    for i in range(5):
        delay = i % 3 + 1
        task_id = scheduler.add_task(
            task_id=f"task-{i+1}",
            coro=lambda i=i, delay=delay: task_func(i+1, delay),
            priority=i % 2,  # 交替优先级
            delay=delay * 0.5,  # 延迟执行
        )
        task_ids.append(task_id)
    
    # 等待所有任务完成
    await asyncio.sleep(10)
    
    # 获取任务结果
    for task_id in task_ids:
        result = scheduler.get_task_result(task_id)
        print(f"任务 {task_id} 结果: {result}")
    
    # 停止调度器
    await scheduler.stop()

async def example_async_concurrent_dict():
    """异步并发安全字典示例"""
    print("\n=== 异步并发安全字典示例 ===")
    
    concurrent_dict = AsyncConcurrentDict()
    
    # 定义并发操作任务
    async def set_operation(key, value, delay):
        await asyncio.sleep(delay)
        await concurrent_dict.set(key, value)
        print(f"设置 {key}: {value}")
    
    async def get_operation(key, delay):
        await asyncio.sleep(delay)
        value = await concurrent_dict.get(key, "默认值")
        print(f"获取 {key}: {value}")
    
    # 并发执行操作
    tasks = [
        set_operation("key1", "value1", 0.1),
        get_operation("key1", 0.2),
        set_operation("key2", "value2", 0.15),
        get_operation("key2", 0.1),
        set_operation("key1", "new_value1", 0.25),
        get_operation("key1", 0.3),
    ]
    
    await asyncio.gather(*tasks)
    
    # 获取字典内容
    items = await concurrent_dict.items()
    print(f"最终字典内容: {items}")

async def example_dynamic_thread_pool():
    """动态线程池示例"""
    print("\n=== 动态线程池示例 ===")
    
    async with DynamicThreadPool(
        min_threads=2,
        max_threads=10,
        idle_timeout=10.0,
        scaling_factor=0.5
    ) as thread_pool:
        
        # 定义任务函数
        async def workload(task_id, duration):
            print(f"任务 {task_id} 开始执行，持续 {duration} 秒")
            await asyncio.sleep(duration)
            print(f"任务 {task_id} 执行完成")
        
        # 提交大量任务
        start_time = time.time()
        
        # 提交 20 个任务，每个任务持续 1-3 秒
        for i in range(20):
            duration = (i % 3) + 1
            await thread_pool.submit(lambda i=i, duration=duration: workload(i+1, duration))
        
        # 等待所有任务完成
        await asyncio.sleep(15)
        
        end_time = time.time()
        
        # 获取统计信息
        stats = thread_pool.get_stats()
        print(f"\n线程池统计信息:")
        print(f"  当前线程数: {stats.current_threads}")
        print(f"  空闲线程数: {stats.idle_threads}")
        print(f"  活跃线程数: {stats.active_threads}")
        print(f"  完成任务数: {stats.completed_tasks}")
        print(f"  拒绝任务数: {stats.rejected_tasks}")
        print(f"  平均任务时间: {stats.avg_task_time:.2f} 秒")
        print(f"  总耗时: {end_time - start_time:.2f} 秒")

async def main():
    """主函数"""
    # 依次运行所有示例
    await example_async_http_client()
    await example_async_task_scheduler()
    await example_async_concurrent_dict()
    await example_dynamic_thread_pool()
    
    print("\n所有示例执行完成！")

if __name__ == "__main__":
    asyncio.run(main())