personal
/
NETPRESSOR
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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
This commit is contained in:
guanjun 2026-06-04 20:13:30 -07:00
parent 20a79b9a23
commit b46d3c3ad9
1 changed files with 101 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

144
stress_test.py
View File

@ -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