chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,186 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Stress test for EditorStateCache.GetSnapshot() to reproduce GC allocation spikes.
|
||||
|
||||
This script rapidly polls the editor state to simulate an MCP client that frequently
|
||||
checks Unity's readiness state. Run this while profiling in Unity to see if it
|
||||
causes the GC spikes reported in GitHub issue #577.
|
||||
|
||||
Usage:
|
||||
python tools/stress_editor_state.py --duration 30 --interval 0.05
|
||||
|
||||
While this runs, open Unity Profiler and look for:
|
||||
- EditorStateCache.OnUpdate
|
||||
- EditorStateCache.GetSnapshot
|
||||
- GC.Alloc spikes
|
||||
"""
|
||||
import asyncio
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import struct
|
||||
import time
|
||||
from pathlib import Path
|
||||
import sys
|
||||
|
||||
|
||||
TIMEOUT = 5.0
|
||||
|
||||
|
||||
def find_status_files() -> list[Path]:
|
||||
home = Path.home()
|
||||
status_dir = Path(os.environ.get("UNITY_MCP_STATUS_DIR", home / ".unity-mcp"))
|
||||
if not status_dir.exists():
|
||||
return []
|
||||
return sorted(status_dir.glob("unity-mcp-status-*.json"), key=lambda p: p.stat().st_mtime, reverse=True)
|
||||
|
||||
|
||||
def discover_port(project_path: str | None) -> int:
|
||||
default_port = 6400
|
||||
files = find_status_files()
|
||||
for f in files:
|
||||
try:
|
||||
data = json.loads(f.read_text())
|
||||
port = int(data.get("unity_port", 0) or 0)
|
||||
if 0 < port < 65536:
|
||||
return port
|
||||
except Exception:
|
||||
pass
|
||||
return default_port
|
||||
|
||||
|
||||
async def read_exact(reader: asyncio.StreamReader, n: int) -> bytes:
|
||||
buf = b""
|
||||
while len(buf) < n:
|
||||
chunk = await reader.read(n - len(buf))
|
||||
if not chunk:
|
||||
raise ConnectionError("Connection closed while reading")
|
||||
buf += chunk
|
||||
return buf
|
||||
|
||||
|
||||
async def read_frame(reader: asyncio.StreamReader) -> bytes:
|
||||
header = await read_exact(reader, 8)
|
||||
(length,) = struct.unpack(">Q", header)
|
||||
if length <= 0 or length > (64 * 1024 * 1024):
|
||||
raise ValueError(f"Invalid frame length: {length}")
|
||||
return await read_exact(reader, length)
|
||||
|
||||
|
||||
async def write_frame(writer: asyncio.StreamWriter, payload: bytes) -> None:
|
||||
header = struct.pack(">Q", len(payload))
|
||||
writer.write(header)
|
||||
writer.write(payload)
|
||||
await asyncio.wait_for(writer.drain(), timeout=TIMEOUT)
|
||||
|
||||
|
||||
async def do_handshake(reader: asyncio.StreamReader) -> None:
|
||||
line = await reader.readline()
|
||||
if not line or b"WELCOME UNITY-MCP" not in line:
|
||||
raise ConnectionError(f"Unexpected handshake from server: {line!r}")
|
||||
|
||||
|
||||
def make_get_editor_state_frame() -> bytes:
|
||||
payload = {"type": "get_editor_state", "params": {}}
|
||||
return json.dumps(payload).encode("utf-8")
|
||||
|
||||
|
||||
async def stress_loop(host: str, port: int, duration: float, interval: float, verbose: bool):
|
||||
stop_time = time.time() + duration
|
||||
stats = {"requests": 0, "errors": 0, "reconnects": 0}
|
||||
|
||||
print(f"Starting editor state stress test...")
|
||||
print(f" Target: {host}:{port}")
|
||||
print(f" Duration: {duration}s")
|
||||
print(f" Interval: {interval}s ({1/interval:.1f} requests/sec)")
|
||||
print(f" Press Ctrl+C to stop early")
|
||||
print()
|
||||
|
||||
writer = None
|
||||
reader = None
|
||||
|
||||
try:
|
||||
while time.time() < stop_time:
|
||||
try:
|
||||
# Connect if needed
|
||||
if writer is None:
|
||||
reader, writer = await asyncio.wait_for(
|
||||
asyncio.open_connection(host, port), timeout=TIMEOUT
|
||||
)
|
||||
await asyncio.wait_for(do_handshake(reader), timeout=TIMEOUT)
|
||||
if verbose:
|
||||
print(f"[{time.time():.2f}] Connected")
|
||||
|
||||
# Send get_editor_state request
|
||||
await write_frame(writer, make_get_editor_state_frame())
|
||||
response = await asyncio.wait_for(read_frame(reader), timeout=TIMEOUT)
|
||||
stats["requests"] += 1
|
||||
|
||||
if verbose and stats["requests"] % 20 == 0:
|
||||
try:
|
||||
data = json.loads(response.decode("utf-8", errors="ignore"))
|
||||
seq = data.get("data", {}).get("sequence", "?")
|
||||
print(f"[{time.time():.2f}] Request #{stats['requests']}, sequence={seq}")
|
||||
except Exception:
|
||||
print(f"[{time.time():.2f}] Request #{stats['requests']}")
|
||||
|
||||
await asyncio.sleep(interval)
|
||||
|
||||
except (ConnectionError, OSError, asyncio.TimeoutError) as e:
|
||||
stats["errors"] += 1
|
||||
stats["reconnects"] += 1
|
||||
if verbose:
|
||||
print(f"[{time.time():.2f}] Connection error: {e}, reconnecting...")
|
||||
if writer:
|
||||
try:
|
||||
writer.close()
|
||||
await writer.wait_closed()
|
||||
except Exception:
|
||||
pass
|
||||
writer = None
|
||||
reader = None
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\nStopped by user")
|
||||
finally:
|
||||
if writer:
|
||||
try:
|
||||
writer.close()
|
||||
await writer.wait_closed()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
elapsed = duration - max(0, stop_time - time.time())
|
||||
print()
|
||||
print("=" * 50)
|
||||
print("Results:")
|
||||
print(f" Total requests: {stats['requests']}")
|
||||
print(f" Errors: {stats['errors']}")
|
||||
print(f" Reconnects: {stats['reconnects']}")
|
||||
print(f" Elapsed: {elapsed:.1f}s")
|
||||
print(f" Rate: {stats['requests']/elapsed:.1f} requests/sec")
|
||||
print("=" * 50)
|
||||
|
||||
|
||||
async def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Stress test EditorStateCache.GetSnapshot() to reproduce GC spikes"
|
||||
)
|
||||
parser.add_argument("--host", default="127.0.0.1", help="Unity bridge host")
|
||||
parser.add_argument("--port", type=int, default=0, help="Unity bridge port (0=auto-discover)")
|
||||
parser.add_argument("--duration", type=float, default=30.0, help="Test duration in seconds")
|
||||
parser.add_argument("--interval", type=float, default=0.05, help="Interval between requests (0.05 = 20/sec)")
|
||||
parser.add_argument("-v", "--verbose", action="store_true", help="Verbose output")
|
||||
args = parser.parse_args()
|
||||
|
||||
port = args.port if args.port > 0 else discover_port(None)
|
||||
|
||||
await stress_loop(args.host, port, args.duration, args.interval, args.verbose)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
asyncio.run(main())
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
Reference in New Issue
Block a user