mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-10-29 03:18:14 +08:00
124 lines
3.9 KiB
Python
124 lines
3.9 KiB
Python
"""测量 TaskRegistry 跨线程 pending submission 的关停终态与提交吞吐。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import asyncio
|
||
import inspect
|
||
import json
|
||
import statistics
|
||
import sys
|
||
import time
|
||
from pathlib import Path
|
||
|
||
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
||
if str(PROJECT_ROOT) not in sys.path:
|
||
sys.path.insert(0, str(PROJECT_ROOT))
|
||
|
||
from app.runtime.tasks import TaskRegistry
|
||
|
||
|
||
async def _value(value: int) -> int:
|
||
"""返回输入值，提供最小可完成协程。"""
|
||
return value
|
||
|
||
|
||
def pending_shutdown_probe() -> dict[str, object]:
|
||
"""在目标 loop 尚未分发 callback 时关闭 Registry，并报告提交终态。"""
|
||
registry = TaskRegistry()
|
||
target_loop = asyncio.new_event_loop()
|
||
coroutine = _value(1)
|
||
completion = registry.submit_threadsafe(
|
||
coroutine,
|
||
loop=target_loop,
|
||
owner="probe.pending",
|
||
)
|
||
shutdown_result = asyncio.run(registry.shutdown(timeout_seconds=0.001))
|
||
result = {
|
||
"shutdown_result": shutdown_result,
|
||
"completion_done": completion.done(),
|
||
"completion_cancelled": completion.cancelled(),
|
||
"coroutine_state": inspect.getcoroutinestate(coroutine),
|
||
}
|
||
completion.cancel()
|
||
coroutine.close()
|
||
target_loop.close()
|
||
return result
|
||
|
||
|
||
async def throughput_sample(iterations: int) -> dict[str, float]:
|
||
"""从工作线程提交一组最小协程，并测量提交和完整完成时间。"""
|
||
registry = TaskRegistry()
|
||
loop = asyncio.get_running_loop()
|
||
started = time.perf_counter()
|
||
|
||
def submit_all():
|
||
"""在同一宿主线程连续提交，保持各轮工作负载一致。"""
|
||
return [
|
||
registry.submit_threadsafe(
|
||
_value(index),
|
||
loop=loop,
|
||
owner="probe.throughput",
|
||
)
|
||
for index in range(iterations)
|
||
]
|
||
|
||
completions = await asyncio.to_thread(submit_all)
|
||
submitted = time.perf_counter()
|
||
values = await asyncio.gather(
|
||
*(asyncio.wrap_future(completion) for completion in completions)
|
||
)
|
||
finished = time.perf_counter()
|
||
assert sum(values) == iterations * (iterations - 1) // 2
|
||
assert await registry.shutdown(timeout_seconds=1.0) is True
|
||
return {
|
||
"submit_ms": (submitted - started) * 1000,
|
||
"total_ms": (finished - started) * 1000,
|
||
}
|
||
|
||
|
||
async def run_samples(iterations: int, samples: int) -> list[dict[str, float]]:
|
||
"""顺序执行多轮样本，避免并行样本互相争抢事件循环。"""
|
||
return [await throughput_sample(iterations) for _ in range(samples)]
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
"""解析探针负载规模。"""
|
||
parser = argparse.ArgumentParser(description=__doc__)
|
||
parser.add_argument("--iterations", type=int, default=2000)
|
||
parser.add_argument("--samples", type=int, default=7)
|
||
return parser.parse_args()
|
||
|
||
|
||
def main() -> None:
|
||
"""运行终态探针和吞吐样本并输出 JSON。"""
|
||
args = parse_args()
|
||
if args.iterations < 1 or args.samples < 1:
|
||
raise SystemExit("iterations 和 samples 必须大于 0")
|
||
samples = asyncio.run(run_samples(args.iterations, args.samples))
|
||
print(
|
||
json.dumps(
|
||
{
|
||
"pending": pending_shutdown_probe(),
|
||
"throughput": {
|
||
"iterations": args.iterations,
|
||
"samples": args.samples,
|
||
"submit_ms": [sample["submit_ms"] for sample in samples],
|
||
"total_ms": [sample["total_ms"] for sample in samples],
|
||
"submit_median_ms": statistics.median(
|
||
sample["submit_ms"] for sample in samples
|
||
),
|
||
"total_median_ms": statistics.median(
|
||
sample["total_ms"] for sample in samples
|
||
),
|
||
},
|
||
},
|
||
ensure_ascii=False,
|
||
indent=2,
|
||
)
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|