187 lines
7.0 KiB
Python
187 lines
7.0 KiB
Python
from __future__ import annotations
|
||
import gc
|
||
import platform
|
||
import sys
|
||
import time
|
||
import types
|
||
|
||
import cProfile
|
||
|
||
from engine.schema import Calibration, Environment
|
||
|
||
|
||
def collect_environment() -> Environment:
|
||
"""采集运行环境指纹，用于结果可复现性与 UI 展示。"""
|
||
return Environment(
|
||
python=platform.python_version(),
|
||
platform=sys.platform,
|
||
processor=platform.processor() or platform.machine(),
|
||
timerResolution=time.get_clock_info("perf_counter").resolution,
|
||
)
|
||
|
||
|
||
def load_source(script_path: str) -> str:
|
||
"""读取脚本文本，剥离 BOM。
|
||
|
||
utf-8-sig 自动剥离 BOM（Windows Notepad 默认带 BOM 的 UTF-8 文件）；
|
||
用 utf-8 会留下 前缀，compile() 抛 SyntaxError: invalid character。
|
||
引擎内所有读源码的地方都必须走这里，否则会出现"harness 能跑但
|
||
runner 的语法预检先失败"这类不一致（历史 bug）。
|
||
"""
|
||
with open(script_path, "r", encoding="utf-8-sig") as f:
|
||
return f.read()
|
||
|
||
|
||
def _load_code(script_path: str):
|
||
return compile(load_source(script_path), script_path, "exec")
|
||
|
||
|
||
def _build_user_globals(script_path: str) -> dict:
|
||
"""构造用户代码 exec 的 globals。
|
||
|
||
两件事一起处理:
|
||
1) sys.argv 临时清成 [script_path] —— 防用户代码读 --out 注册 atexit 改写
|
||
result.json(关键安全洞 C1)。exec 完后恢复。
|
||
2) 装一个 __main__ module —— 让 pickle / multiprocessing / joblib 这类
|
||
依赖 sys.modules['__main__'] 的库能正常工作。
|
||
之前只把 __name__ 设成 '__main__' 但 sys.modules['__main__'] 仍指向
|
||
engine.runner,直接 break pickle.dumps(本地定义类) ——
|
||
PicklingError: ... not found as __main__.Point。
|
||
"""
|
||
# 装 synthetic __main__ module,跟 exec 的 globals 共享同一份 dict
|
||
main_mod = types.ModuleType("__main__")
|
||
main_mod.__file__ = script_path
|
||
main_mod.__spec__ = None
|
||
g = main_mod.__dict__
|
||
g["__name__"] = "__main__"
|
||
g["__file__"] = script_path
|
||
g["__builtins__"] = __builtins__
|
||
sys.modules["__main__"] = main_mod
|
||
return g
|
||
|
||
|
||
def _scrub_argv_for_user_code(script_path: str) -> list:
|
||
"""把 sys.argv 收成 [script_path] 后备一份原值,exec 完后由调用方还原。
|
||
|
||
关键:用户代码能 sys.argv.index('--out') 拿到结果文件路径然后
|
||
atexit.register(lambda: open(out, 'w').write('{...forged...}')),engine
|
||
写完 result.json 紧接着被覆盖。父进程拿到的就是 attacker-controlled JSON。
|
||
"""
|
||
saved = sys.argv
|
||
sys.argv = [script_path]
|
||
return saved
|
||
|
||
|
||
def _restore_argv(saved):
|
||
sys.argv = saved
|
||
|
||
|
||
# 校准负载：tight-loop,典型的 CPU-bound 工作量。
|
||
# 选 50_000 次是因为典型机器上耗时 ~10-30ms —— 既让 cProfile 开销稳定可测,
|
||
# 又不会让用户感到"启动变慢"。如果换成 1 次循环,ratio 会被调度噪声主导;
|
||
# 换成 1M 次,校准本身就要 ~1s,用户能感知。
|
||
_CALIB_SOURCE = (
|
||
"def _pyrof_calib():\n"
|
||
" s = 0\n"
|
||
" for i in range(50_000):\n"
|
||
" s += i\n"
|
||
" return s\n"
|
||
"_pyrof_calib()\n"
|
||
)
|
||
|
||
# 在 module load 时编译一次 —— 校准负载是常量字符串,之前 calibrate 跑两遍
|
||
# (clean + instrumented) 就 compile 两次,纯浪费。co_filename 用合成路径,
|
||
# 不需要落到磁盘 —— _build_user_globals 只取它当 __file__,不校验文件存在。
|
||
_CALIB_CODE = compile(_CALIB_SOURCE, "<pyrof_calibration>", "exec")
|
||
# 同样用于 _build_user_globals(script_path) —— synthetic __main__.__file__
|
||
# 用同一条合成路径,跟 code.co_filename 保持一致。
|
||
_CALIB_PATH = "<pyrof_calibration>"
|
||
|
||
|
||
def _exec_once_for_calibration() -> None:
|
||
"""校准负载的纯 exec —— 不开 cProfile。
|
||
|
||
GC 状态由 caller 负责 —— calibrate_cprofile_overhead 已经 disable 后
|
||
re-enable,这里再包一层是冗余 no-op,徒增两个 isenabled() 系统调用。
|
||
"""
|
||
g = _build_user_globals(_CALIB_PATH)
|
||
saved_argv = _scrub_argv_for_user_code(_CALIB_PATH)
|
||
try:
|
||
exec(_CALIB_CODE, g)
|
||
finally:
|
||
_restore_argv(saved_argv)
|
||
|
||
|
||
def _exec_once_under_cprofile() -> None:
|
||
"""同一段校准负载,套 cProfile 跑一次 —— 用于测 cProfile 自身开销。
|
||
|
||
GC 状态由 caller 负责（同上理由）。"""
|
||
g = _build_user_globals(_CALIB_PATH)
|
||
saved_argv = _scrub_argv_for_user_code(_CALIB_PATH)
|
||
pr = cProfile.Profile()
|
||
pr.enable()
|
||
try:
|
||
exec(_CALIB_CODE, g)
|
||
finally:
|
||
pr.disable()
|
||
_restore_argv(saved_argv)
|
||
|
||
|
||
def calibrate_cprofile_overhead(workdir: str | None = None) -> Calibration:
|
||
"""跑同一段 tight-loop 各一次（裸 vs 仪器化），返回膨胀系数。
|
||
|
||
为什么需要这步:
|
||
- 用户脚本只 exec 一次 (单跑架构,v4 起),这次 exec 必须开 cProfile 才能拿归因
|
||
- 但 cProfile 是插桩式的,wall-time 会膨胀 1.5~3x;直接报给用户不真实
|
||
- 思路是「cProfile 开销 ≈ 跟用户代码调用次数成正比」,跟代码时间无关。
|
||
所以可以在一段**已知调用次数**的 tight-loop 上测一次 ratio,再把
|
||
instrumented_user / ratio 当作估计的干净耗时。
|
||
- 这一步与用户代码 exec 解耦:校准不会被用户的 import / 全局副作用干扰,
|
||
也不会因为校准改了用户的状态而报错。
|
||
|
||
边界情况:
|
||
- 校准本身耗时 <5ms (机器极快):perf_counter 精度+调度抖动会让 ratio
|
||
跳到 0.5x/5x 这种离谱值,此时回退 ratio=1.0 (即相信 instrumented 时间)。
|
||
- 校准本身耗时 5~50ms:正常,ratio 在 1.5~3 之间。
|
||
|
||
workdir:历史参数,现已废弃 —— 校准脚本不再落盘(改成内存里 compile +
|
||
合成 __file__),workdir 不再被本函数读取。为保持 profiler-service /
|
||
老调用方不破,参数保留并忽略。
|
||
"""
|
||
del workdir # 明确：workdir 已废弃，调用方传过来也不再消费
|
||
|
||
try:
|
||
# 干净跑 —— 不开 cProfile
|
||
gc_was = gc.isenabled(); gc.disable()
|
||
try:
|
||
t0 = time.perf_counter()
|
||
_exec_once_for_calibration()
|
||
clean = time.perf_counter() - t0
|
||
finally:
|
||
if gc_was: gc.enable()
|
||
|
||
# 仪器化跑 —— 开 cProfile
|
||
gc_was = gc.isenabled(); gc.disable()
|
||
try:
|
||
t0 = time.perf_counter()
|
||
_exec_once_under_cprofile()
|
||
inst = time.perf_counter() - t0
|
||
finally:
|
||
if gc_was: gc.enable()
|
||
finally:
|
||
pass # 无文件要清
|
||
|
||
if clean < 0.005 or inst <= 0 or clean <= 0:
|
||
# 校准耗时过短或退化:ratio 不可信,回退 1.0 (即不补偿,UI 显示原始 instrumented 时间)
|
||
return Calibration(
|
||
ratio=1.0,
|
||
workloadName="tight-loop",
|
||
instrumentedWorkloadSec=inst,
|
||
cleanWorkloadSec=clean,
|
||
)
|
||
return Calibration(
|
||
ratio=inst / clean,
|
||
workloadName="tight-loop",
|
||
instrumentedWorkloadSec=inst,
|
||
cleanWorkloadSec=clean,
|
||
) |