#!/usr/bin/env python3 """ 并发选区压测脚本 功能： - 批量生成登录链接（拿到 codeNo） - 并发对这些 codeNo 执行“选区”操作 - 统计是否出现多个 codeNo 选到同一设备（设备串号） 使用前请在“配置区”按实际接口改动： - BASE_URL、GENERATE_PATH、SELECT_PATH - 请求方法/参数/头部及字段名（CODE_NO_FIELD、DEVICE_FIELD、STATUS_FIELD） 运行示例： python scripts/concurrency_region_select.py --base http://localhost:18080 --gen-path /api/link/generate --sel-path /api/region/select --links 100 --concurrency 100 --region Q 依赖：pip install httpx anyio """ from __future__ import annotations import argparse import asyncio import json import random import sys from dataclasses import dataclass from typing import Any, Dict, List, Optional, Tuple import anyio import httpx # ===================== 配置区（可通过命令行覆盖） ===================== # 响应字段名 CODE_NO_FIELD = "codeNo" # 生成链接接口返回的 codeNo 字段 DEVICE_FIELD = "device" # 选区接口返回的 设备字段 STATUS_FIELD = "status" # 选区接口返回的 状态字段（可选） # 生成链接：请求方法与构造（默认 POST，无体） DEFAULT_GEN_METHOD = "POST" # 可选：GET/POST DEFAULT_GEN_BODY: Optional[Dict[str, Any]] = None # 如需传 body，可在此设置 DEFAULT_GEN_HEADERS: Dict[str, str] = {} # 如需鉴权：{"Authorization": "Bearer ..."} # 选区：请求方法与构造（默认 POST，body 含 codeNo 与 region） DEFAULT_SEL_METHOD = "POST" DEFAULT_SEL_HEADERS: Dict[str, str] = {} # ===================== 实现 ===================== @dataclass class Config: base_url: str generate_path: str select_path: str links: int concurrency: int region: str gen_method: str = DEFAULT_GEN_METHOD sel_method: str = DEFAULT_SEL_METHOD gen_body: Optional[str] = None # JSON 字符串，命令行传入后解析 gen_headers: Optional[str] = None # JSON 字符串 sel_headers: Optional[str] = None # JSON 字符串 timeout: float = 10.0 jitter_ms: int = 10 # 并发抖动，避免完全同一时刻 def parsed_gen_body(self) -> Optional[Dict[str, Any]]: return json.loads(self.gen_body) if self.gen_body else DEFAULT_GEN_BODY def parsed_gen_headers(self) -> Dict[str, str]: if self.gen_headers: return json.loads(self.gen_headers) return DEFAULT_GEN_HEADERS def parsed_sel_headers(self) -> Dict[str, str]: if self.sel_headers: return json.loads(self.sel_headers) return DEFAULT_SEL_HEADERS async def generate_links(cfg: Config, client: httpx.AsyncClient) -> List[str]: """并发生成链接，返回 codeNo 列表。""" url = cfg.base_url.rstrip("/") + cfg.generate_path sem = anyio.Semaphore(cfg.concurrency) code_nos: List[Optional[str]] = [None] * cfg.links async def one(idx: int): async with sem: # 抖动 await anyio.sleep(random.random() * cfg.jitter_ms / 1000.0) try: if cfg.gen_method.upper() == "GET": resp = await client.get(url, headers=cfg.parsed_gen_headers(), timeout=cfg.timeout) else: resp = await client.post(url, headers=cfg.parsed_gen_headers(), json=cfg.parsed_gen_body(), timeout=cfg.timeout) resp.raise_for_status() data = resp.json() code_no = data.get(CODE_NO_FIELD) if not code_no: # 尝试从 URL 中解析（如果返回的是完整 link） # 例如: ...?codeNo=XXXX link = data if isinstance(data, str) else data.get("link") if isinstance(data, dict) else None if isinstance(link, str) and "codeNo=" in link: code_no = link.split("codeNo=")[-1].split("&")[0] if not code_no: raise ValueError(f"生成链接响应缺少 {CODE_NO_FIELD} 字段: {data}") code_nos[idx] = code_no except Exception as e: print(f"[ERR] 生成链接失败 idx={idx}: {e}") async with anyio.create_task_group() as tg: for i in range(cfg.links): tg.start_soon(one, i) ok = [c for c in code_nos if c] if len(ok) != cfg.links: print(f"[WARN] 生成成功 {len(ok)}/{cfg.links}") return ok async def select_region_for_codes(cfg: Config, client: httpx.AsyncClient, code_nos: List[str]) -> List[Tuple[str, Optional[str], Optional[str], int]]: """并发对 codeNos 执行选区，返回 (codeNo, device, status, http_status) 列表。""" url = cfg.base_url.rstrip("/") + cfg.select_path sem = anyio.Semaphore(cfg.concurrency) results: List[Optional[Tuple[str, Optional[str], Optional[str], int]]] = [None] * len(code_nos) def build_select_body(code_no: str) -> Dict[str, Any]: # 按需修改为你的服务契约 return {"codeNo": code_no, "region": cfg.region} async def one(idx: int, code_no: str): async with sem: await anyio.sleep(random.random() * cfg.jitter_ms / 1000.0) try: if cfg.sel_method.upper() == "GET": # GET 方式：拼接 query params = {"codeNo": code_no, "region": cfg.region} resp = await client.get(url, headers=cfg.parsed_sel_headers(), params=params, timeout=cfg.timeout) else: body = build_select_body(code_no) resp = await client.post(url, headers=cfg.parsed_sel_headers(), json=body, timeout=cfg.timeout) status_code = resp.status_code device = None status = None try: data = resp.json() if isinstance(data, dict): device = data.get(DEVICE_FIELD) status = data.get(STATUS_FIELD) except Exception: pass results[idx] = (code_no, device, status, status_code) except Exception as e: print(f"[ERR] 选区失败 codeNo={code_no}: {e}") results[idx] = (code_no, None, None, 0) async with anyio.create_task_group() as tg: for i, c in enumerate(code_nos): tg.start_soon(one, i, c) return [r for r in results if r] def analyze_collisions(pairs: List[Tuple[str, Optional[str], Optional[str], int]]) -> None: # 汇总设备 -> codeNos device_map: Dict[str, List[str]] = {} success = 0 http_ok = 0 for code_no, device, status, http_status in pairs: if http_status == 200: http_ok += 1 if device: success += 1 device_map.setdefault(device, []).append(code_no) total = len(pairs) print("\n===== 结果统计 =====") print(f"请求总数: {total}") print(f"HTTP 200: {http_ok}") print(f"解析出设备: {success}") collisions = {d: cs for d, cs in device_map.items() if len(cs) > 1} if collisions: print("\n[ALERT] 检测到设备串号（同一设备被多个 codeNo 绑定）：") for device, codes in sorted(collisions.items(), key=lambda x: len(x[1]), reverse=True): sample = ", ".join(codes[:10]) + (" ..." if len(codes) > 10 else "") print(f"- device={device}: {len(codes)} codeNos -> {sample}") else: print("\n[OK] 未发现同一设备被多个 codeNo 同时绑定的情况。") async def main_async(cfg: Config) -> int: async with httpx.AsyncClient(http2=False) as client: print(f"生成链接 {cfg.links} 个 ...") code_nos = await generate_links(cfg, client) if not code_nos: print("未成功生成 codeNo，退出。") return 2 print(f"开始并发选区（并发={cfg.concurrency}, region={cfg.region}）...") pairs = await select_region_for_codes(cfg, client, code_nos) analyze_collisions(pairs) return 0 def parse_args(argv: List[str]) -> Config: p = argparse.ArgumentParser(description="并发选区压测脚本") p.add_argument("--base", dest="base_url", required=True, help="服务基础地址，如 http://localhost:8080") p.add_argument("--gen-path", dest="generate_path", required=True, help="生成链接接口路径，如 /api/link/generate") p.add_argument("--sel-path", dest="select_path", required=True, help="选区接口路径，如 /api/region/select") p.add_argument("--links", dest="links", type=int, default=100, help="生成链接数量") p.add_argument("--concurrency", dest="concurrency", type=int, default=100, help="并发度（同时进行的请求数）") p.add_argument("--region", dest="region", default="CN", help="选区标识，如 CN/US/EU") p.add_argument("--gen-method", dest="gen_method", default=DEFAULT_GEN_METHOD, choices=["GET", "POST"], help="生成链接的 HTTP 方法") p.add_argument("--sel-method", dest="sel_method", default=DEFAULT_SEL_METHOD, choices=["GET", "POST"], help="选区的 HTTP 方法") p.add_argument("--gen-body", dest="gen_body", help="生成链接的 JSON 请求体(字符串)") p.add_argument("--gen-headers", dest="gen_headers", help="生成链接的请求头(JSON字符串)") p.add_argument("--sel-headers", dest="sel_headers", help="选区的请求头(JSON字符串)") p.add_argument("--timeout", dest="timeout", type=float, default=10.0, help="单请求超时时间(秒)") p.add_argument("--jitter-ms", dest="jitter_ms", type=int, default=10, help="请求抖动(毫秒)") args = p.parse_args(argv) return Config( base_url=args.base_url, generate_path=args.generate_path, select_path=args.select_path, links=args.links, concurrency=args.concurrency, region=args.region, gen_method=args.gen_method, sel_method=args.sel_method, gen_body=args.gen_body, gen_headers=args.gen_headers, sel_headers=args.sel_headers, timeout=args.timeout, jitter_ms=args.jitter_ms, ) def main() -> int: cfg = parse_args(sys.argv[1:]) try: return asyncio.run(main_async(cfg)) except KeyboardInterrupt: print("用户中断。") return 130 if __name__ == "__main__": raise SystemExit(main())