From b46d3c3ad9c3fdca04b2d9a7f9143de6adf63b58 Mon Sep 17 00:00:00 2001 From: guanjun Date: Thu, 4 Jun 2026 20:13:30 -0700 Subject: [PATCH] Refactor color output, fix request limiting, and improve stats - Extract ANSI color helpers and format_bytes to module level - Make StatsCollector async-safe with locked completion counter - Use linear interpolation for percentile calculation - Fix total_requests limiting in both worker and control loop - Add limit_per_host to TCPConnector --- stress_test.py | 144 ++++++++++++++++++++++++++++++++++--------------- 1 file changed, 101 insertions(+), 43 deletions(-) diff --git a/stress_test.py b/stress_test.py index 6ab255c..8cb3fcb 100755 --- a/stress_test.py +++ b/stress_test.py @@ -56,6 +56,61 @@ except ImportError: sys.exit(1) +# --------------------------- ANSI 颜色支持 --------------------------- + +_NO_COLOR = False + +def _supports_color() -> bool: + """检测终端是否支持 ANSI 颜色""" + if os.environ.get("NO_COLOR"): + return False + if not hasattr(sys.stdout, "isatty"): + return False + return sys.stdout.isatty() + +def c(text: str, code: str) -> str: + """条件着色：仅在支持颜色时添加 ANSI 转义""" + if _NO_COLOR: + return text + return f"\033[{code}m{text}\033[0m" + +def cbold(text: str) -> str: + return c(text, "1") + +def ccyan(text: str) -> str: + return c(text, "1;36") + +def cyellow(text: str) -> str: + return c(text, "1;33") + +def cgreen(text: str) -> str: + return c(text, "1;32") + +def cred(text: str) -> str: + return c(text, "1;31") + +def cmagenta(text: str) -> str: + return c(text, "1;35") + +def cdim(text: str) -> str: + return c(text, "32") # 暗绿/灰色，用于进度条 + +def cdim_yellow(text: str) -> str: + return c(text, "33") + +def cdim_red(text: str) -> str: + return c(text, "31") + + +def format_bytes(n: float) -> str: + """格式化字节数为人类可读字符串""" + for unit in ["B", "KB", "MB", "GB"]: + if abs(n) < 1024: + return f"{n:.1f} {unit}" + n /= 1024 + return f"{n:.1f} TB" + + # --------------------------- 数据结构 --------------------------- @dataclass @@ -99,7 +154,7 @@ class TestConfig: # --------------------------- 统计收集器 --------------------------- class StatsCollector: - """线程安全的统计数据收集器""" + """异步安全的统计数据收集器（仅在单事件循环内使用）""" def __init__(self): self.results: list[RequestResult] = [] @@ -109,16 +164,22 @@ class StatsCollector: self._error_counts: dict[str, int] = defaultdict(int) self._total_bytes: int = 0 self._latencies: list[float] = [] + self._completed: int = 0 # 原子计数器，用于控制总请求数 async def add(self, result: RequestResult): async with self._lock: self.results.append(result) + self._completed += 1 self._status_counts[result.status] += 1 self._total_bytes += result.size self._latencies.append(result.elapsed) if result.error: self._error_counts[result.error] += 1 + async def completed_count(self) -> int: + async with self._lock: + return self._completed + def total_requests(self) -> int: return len(self.results) @@ -128,25 +189,20 @@ class StatsCollector: def error_count(self) -> int: return sum(1 for r in self.results if r.status >= 400 or r.error) - def current_qps(self, window: float = 1.0) -> float: - """计算最近 window 秒的 QPS""" - if not self.results: - return 0.0 - now = time.time() - recent = [r for r in self.results if now - r.timestamp < window] - return len(recent) / window - def avg_latency(self) -> float: return statistics.mean(self._latencies) if self._latencies else 0.0 def percentile(self, p: float) -> float: - """计算百分位延迟""" + """计算百分位延迟（线性插值）""" if not self._latencies: return 0.0 sorted_lats = sorted(self._latencies) - idx = int(len(sorted_lats) * p / 100) - idx = min(idx, len(sorted_lats) - 1) - return sorted_lats[idx] + n = len(sorted_lats) + k = (n - 1) * p / 100 + f = int(k) + c = min(f + 1, n - 1) + d = k - f + return sorted_lats[f] + d * (sorted_lats[c] - sorted_lats[f]) def elapsed_time(self) -> float: return time.time() - self.start_time if self.start_time else 0.0 @@ -169,20 +225,14 @@ class RealtimeDisplay: self.config = config self._frame = 0 self._running = False - self._last_qps_samples: list[float] = [] self.critical_concurrency: int = 0 # find-limit 模式下探测到的临界并发数 - def _format_bytes(self, n: float) -> str: - for unit in ["B", "KB", "MB", "GB"]: - if abs(n) < 1024: - return f"{n:.1f} {unit}" - n /= 1024 - return f"{n:.1f} TB" - - def _format_duration(self, seconds: float) -> str: + @staticmethod + def _format_duration(seconds: float) -> str: return str(timedelta(seconds=int(seconds))) - def _progress_bar(self, current: int, total: int, width: int = 30) -> str: + @staticmethod + def _progress_bar(current: int, total: int, width: int = 30) -> str: if total <= 0: return "∞" filled = int(width * current / total) @@ -203,9 +253,9 @@ class RealtimeDisplay: qps = done / elapsed if elapsed > 0 else 0 lines = [] - lines.append(f"\033[1;36m{'=' * 60}\033[0m") - lines.append(f" {spinner} \033[1mStress Test Running\033[0m") - lines.append(f"\033[1;36m{'-' * 60}\033[0m") + lines.append(ccyan("=" * 60)) + lines.append(f" {spinner} {cbold('Stress Test Running')}") + lines.append(ccyan("-" * 60)) # 进度 if total > 0: @@ -217,33 +267,33 @@ class RealtimeDisplay: lines.append("") # 核心指标 - lines.append(f" \033[1;33mQPS:\033[0m {qps:>10.1f} req/s") - lines.append(f" \033[1;33m并发数:\033[0m {self.config.concurrency:>10}") + lines.append(f" {cyellow('QPS:')} {qps:>10.1f} req/s") + lines.append(f" {cyellow('并发数:')} {self.config.concurrency:>10}") if self.config.find_limit and self.critical_concurrency > 0: remaining = max(0, self.config.duration - elapsed) - lines.append(f" \033[1;35m临界并发:\033[0m {self.critical_concurrency:>10}") - lines.append(f" \033[1;35m剩余时间:\033[0m {self._format_duration(remaining):>10}") - lines.append(f" \033[1;33m总请求:\033[0m {done:>10}") - lines.append(f" \033[1;33m耗时:\033[0m {self._format_duration(elapsed):>10}") + lines.append(f" {cmagenta('临界并发:')} {self.critical_concurrency:>10}") + lines.append(f" {cmagenta('剩余时间:')} {self._format_duration(remaining):>10}") + lines.append(f" {cyellow('总请求:')} {done:>10}") + lines.append(f" {cyellow('耗时:')} {self._format_duration(elapsed):>10}") lines.append("") # 延迟 - lines.append(f" \033[1;32m平均延迟:\033[0m {s.avg_latency() * 1000:>10.1f} ms") - lines.append(f" \033[1;32mP50:\033[0m {s.percentile(50) * 1000:>10.1f} ms") - lines.append(f" \033[1;32mP90:\033[0m {s.percentile(90) * 1000:>10.1f} ms") - lines.append(f" \033[1;32mP99:\033[0m {s.percentile(99) * 1000:>10.1f} ms") + lines.append(f" {cgreen('平均延迟:')} {s.avg_latency() * 1000:>10.1f} ms") + lines.append(f" {cgreen('P50:')} {s.percentile(50) * 1000:>10.1f} ms") + lines.append(f" {cgreen('P90:')} {s.percentile(90) * 1000:>10.1f} ms") + lines.append(f" {cgreen('P99:')} {s.percentile(99) * 1000:>10.1f} ms") lines.append("") # 成功率 success = s.success_count() errors = s.error_count() rate = (success / done * 100) if done > 0 else 0 - color = "\033[1;32m" if rate >= 99 else "\033[1;33m" if rate >= 95 else "\033[1;31m" - lines.append(f" 成功率: {color}{rate:.2f}%\033[0m " + rate_str = cgreen if rate >= 99 else cyellow if rate >= 95 else cred + lines.append(f" 成功率: {rate_str(f'{rate:.2f}%')} " f"(成功: {success} 失败: {errors})") - lines.append(f" 吞吐量: {self._format_bytes(s.throughput())}/s") + lines.append(f" 吞吐量: {format_bytes(s.throughput())}/s") - lines.append(f"\033[1;36m{'=' * 60}\033[0m") + lines.append(ccyan("=" * 60)) return "\n".join(lines) @@ -321,6 +371,11 @@ class StressTester: async def _worker(self, session: aiohttp.ClientSession, url_iter): """单个并发 worker：不断从 url_iter 取 URL 发请求""" while not self._stop_event.is_set(): + # 检查总请求数限制（使用锁内计数器避免超发） + if self.config.total_requests: + completed = await self.stats.completed_count() + if completed >= self.config.total_requests: + break try: url = next(url_iter) except StopIteration: @@ -349,11 +404,12 @@ class StressTester: target = max(1, int(self.config.concurrency * ratio)) return target - async def _control_loop(self, display: RealtimeDisplay, stop_display: asyncio.Event): + async def _control_loop(self, display: RealtimeDisplay | None, stop_display: asyncio.Event): """控制循环：管理任务生命周期和停止条件""" url_iter = iter(self._url_generator()) connector = aiohttp.TCPConnector( limit=self.config.concurrency + 10, + limit_per_host=self.config.concurrency + 10, keepalive_timeout=30 if self.config.keep_alive else 0, enable_cleanup_closed=True, ) @@ -366,8 +422,10 @@ class StressTester: # 主循环 while not self._stop_event.is_set(): # 检查停止条件 - if self.config.total_requests and self.stats.total_requests() >= self.config.total_requests: - break + if self.config.total_requests: + completed = await self.stats.completed_count() + if completed >= self.config.total_requests: + break if self.config.duration and self.stats.elapsed_time() >= self.config.duration: break