Python 会话数压测脚本

完整代码

python 复制代码

import asyncio
import time
import argparse
import sys

if sys.platform == 'win32':
    asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())

class AdvancedFeedbackProbe:
    def __init__(self, target, port, start_rate=300, min_rate=1):
        self.target = target
        self.port = port
        self.current_rate = start_rate
        self.start_rate = start_rate
        self.min_rate = min_rate
        self.timeout = 5.0
        
        self.active_writers = {}
        self.conn_counter = 0
        self.results = {"success": 0, "failure": 0, "errors": {}}
        
        self.max_active_observed = 0 
        self.is_running = True
        self.start_time = None
        
        # 反馈控制
        self.recent_failures = 0
        self.failure_threshold = 3

    async def keep_alive_with_retry(self, conn_id):
        while self.is_running:
            writer = None
            try:
                reader, writer = await asyncio.wait_for(
                    asyncio.open_connection(self.target, self.port),
                    timeout=self.timeout
                )
                self.active_writers[conn_id] = writer
                self.results["success"] += 1
                self.recent_failures = 0 # 成功则重置"近期失败"
                
                # 记录最高并发
                current_count = len(self.active_writers)
                if current_count > self.max_active_observed:
                    self.max_active_observed = current_count
                
                while self.is_running:
                    if writer.transport.is_closing(): break
                    await asyncio.sleep(1)
                
            except Exception as e:
                self.recent_failures += 1
                self.results["failure"] += 1
                err_name = type(e).__name__
                self.results["errors"][err_name] = self.results["errors"].get(err_name, 0) + 1
                
                if conn_id in self.active_writers: del self.active_writers[conn_id]
                if writer: 
                    try:
                        writer.close()
                    except: pass
                
                # 触发限速逻辑
                if self.recent_failures >= self.failure_threshold:
                    self.adjust_rate(down=True)
                
                if self.is_running:
                    # 失败后的冷却时间，防止瞬间重试把本地 CPU 跑满
                    await asyncio.sleep(2.0)
                else:
                    break
            finally:
                if conn_id in self.active_writers: del self.active_writers[conn_id]

    def adjust_rate(self, down=True):
        if down:
            # 遇到阻力，油门减半
            new_rate = max(self.min_rate, self.current_rate * 0.5)
            if new_rate != self.current_rate:
                self.current_rate = new_rate
        else:
            # 路况良好，缓慢加速
            if self.current_rate < self.start_rate:
                self.current_rate += 1

    async def monitor(self):
        while self.is_running:
            current_active = len(self.active_writers)
            # 打印包含：实时活跃、最高纪录、当前速率、累计失败、近期连续失败
            print(f"\r[+] 活跃: {current_active} (最高: {self.max_active_observed}) | 速率: {self.current_rate:.1f}/s | 累计失败: {self.results['failure']} | 近期失败: {self.recent_failures}    ", end="", flush=True)
            
            # 如果近期没有新失败，尝试恢复速率
            if self.recent_failures == 0:
                self.adjust_rate(down=False)
            await asyncio.sleep(0.5)

    async def run(self):
        print(f"[*] 目标: {self.target}:{self.port}")
        print(f"[*] 策略: 自适应降速 (连续 {self.failure_threshold} 次失败触发)")
        print("-" * 75)
        self.start_time = time.time()
        asyncio.create_task(self.monitor())

        while self.is_running:
            batch_size = int(self.current_rate)
            for _ in range(batch_size):
                if not self.is_running: break
                self.conn_counter += 1
                asyncio.create_task(self.keep_alive_with_retry(self.conn_counter))
            await asyncio.sleep(1.0) 

    def report(self, final_active):
        self.is_running = False
        duration = time.time() - self.start_time
        print("\n" + "="*50)
        print(f"{'自适应探测最终报告':^46}")
        print("-" * 50)
        print(f" 探测历史最高并发 (峰值): {self.max_active_observed}")
        print(f" 终止时实时活跃会话: {final_active}")
        print(f" 运行总时长: {duration:.1f} 秒")
        print(f" 累计重试/失败总数: {self.results['failure']}")
        print(f" 错误分布: {self.results['errors']}")
        print("-" * 50)
        print(" [提示] 当'累计失败'持续跳动而'活跃'不再增长时，")
        print("        即代表已到达物理连接或防火墙策略上限。")
        print("="*50)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("host")
    parser.add_argument("port", type=int)
    parser.add_argument("--start_rate", type=int, default=100)  # 初始探测速率
    args = parser.parse_args()

    probe = AdvancedFeedbackProbe(args.host, args.port, args.start_rate)
    
    try:
        asyncio.run(probe.run())
    except KeyboardInterrupt:
        snapshot = len(probe.active_writers)
        probe.report(snapshot)

如何使用

此脚本建议压测内网环境服务百度等网站域名具备单IP特定时间会话数保护

arduino 复制代码

python run.py www.baidu.com 443