大規模履歴向けread_messages pagination対応

1 files changed, 180 insertions, 0 deletions

diff --git a/scripts/python/p2_benchmark.py b/scripts/python/p2_benchmark.py
new file mode 100644
index 0000000..792673c
--- /dev/null
+++ b/scripts/python/p2_benchmark.py
@@ -0,0 +1,180 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import ctypes
+import json
+import shutil
+import subprocess
+import tempfile
+import time
+from pathlib import Path
+from typing import Any
+
+
+class ProcessMemoryCounters(ctypes.Structure):
+    _fields_ = [
+        ("cb", ctypes.c_ulong),
+        ("PageFaultCount", ctypes.c_ulong),
+        ("PeakWorkingSetSize", ctypes.c_size_t),
+        ("WorkingSetSize", ctypes.c_size_t),
+        ("QuotaPeakPagedPoolUsage", ctypes.c_size_t),
+        ("QuotaPagedPoolUsage", ctypes.c_size_t),
+        ("QuotaPeakNonPagedPoolUsage", ctypes.c_size_t),
+        ("QuotaNonPagedPoolUsage", ctypes.c_size_t),
+        ("PagefileUsage", ctypes.c_size_t),
+        ("PeakPagefileUsage", ctypes.c_size_t),
+    ]
+
+
+class Server:
+    def __init__(self, exe: Path, data_dir: Path) -> None:
+        self.process = subprocess.Popen(
+            [str(exe), "--data-dir", str(data_dir)],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            encoding="utf-8",
+            bufsize=1,
+        )
+        self.request_id = 1
+        self.max_response_bytes = 0
+        self.request(
+            "initialize",
+            {
+                "protocolVersion": "2024-11-05",
+                "capabilities": {},
+                "clientInfo": {"name": "p2-benchmark", "version": "1"},
+            },
+        )
+        self.notify("notifications/initialized", {})
+
+    def request(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
+        assert self.process.stdin and self.process.stdout
+        request_id = self.request_id
+        self.request_id += 1
+        self.process.stdin.write(
+            json.dumps(
+                {"jsonrpc": "2.0", "id": request_id, "method": method, "params": params},
+                separators=(",", ":"),
+            )
+            + "\n"
+        )
+        self.process.stdin.flush()
+        line = self.process.stdout.readline()
+        self.max_response_bytes = max(self.max_response_bytes, len(line.encode("utf-8")))
+        response = json.loads(line)
+        if "error" in response:
+            raise RuntimeError(response)
+        return response["result"]
+
+    def call(self, name: str, arguments: dict[str, Any]) -> dict[str, Any]:
+        result = self.request("tools/call", {"name": name, "arguments": arguments})
+        payload = json.loads(result["content"][0]["text"])
+        if result.get("isError"):
+            raise RuntimeError(payload)
+        return payload
+
+    def notify(self, method: str, params: dict[str, Any]) -> None:
+        assert self.process.stdin
+        self.process.stdin.write(
+            json.dumps({"jsonrpc": "2.0", "method": method, "params": params}) + "\n"
+        )
+        self.process.stdin.flush()
+
+    def peak_working_set(self) -> int:
+        counters = ProcessMemoryCounters()
+        counters.cb = ctypes.sizeof(counters)
+        handle = ctypes.windll.kernel32.OpenProcess(0x1000 | 0x0400, False, self.process.pid)
+        try:
+            if handle and ctypes.windll.psapi.GetProcessMemoryInfo(
+                handle, ctypes.byref(counters), counters.cb
+            ):
+                return int(counters.PeakWorkingSetSize)
+            return 0
+        finally:
+            if handle:
+                ctypes.windll.kernel32.CloseHandle(handle)
+
+    def close(self) -> None:
+        if self.process.stdin:
+            self.process.stdin.close()
+        code = self.process.wait(timeout=30)
+        if code != 0:
+            raise RuntimeError(self.process.stderr.read() if self.process.stderr else code)
+
+
+def read_all(server: Server, session_id: str) -> tuple[int, int]:
+    cursor = None
+    count = pages = 0
+    while True:
+        arguments: dict[str, Any] = {"session_id": session_id, "limit": 1000}
+        if cursor:
+            arguments["cursor"] = cursor
+        payload = server.call("read_messages", arguments)
+        count += len(payload["history"])
+        pages += 1
+        cursor = payload.get("nextCursor")
+        if not cursor:
+            return count, pages
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--server", type=Path, required=True)
+    parser.add_argument("--messages", type=int, choices=(10000, 50000), required=True)
+    parser.add_argument("--sessions", type=int, choices=(1, 100), required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    args = parser.parse_args()
+    data_dir = Path(tempfile.mkdtemp(prefix="tiny-convmem-p2-"))
+    try:
+        server = Server(args.server.resolve(), data_dir)
+        write_started = time.perf_counter()
+        for index in range(args.messages):
+            server.call(
+                "write_message",
+                {
+                    "session_id": f"session-{index % args.sessions}",
+                    "role": "user",
+                    "content": f"history item {index}",
+                },
+            )
+        write_ms = (time.perf_counter() - write_started) * 1000
+        server.close()
+
+        restart_started = time.perf_counter()
+        server = Server(args.server.resolve(), data_dir)
+        restart_ms = (time.perf_counter() - restart_started) * 1000
+        read_started = time.perf_counter()
+        total = pages = 0
+        for session in range(args.sessions):
+            session_count, session_pages = read_all(server, f"session-{session}")
+            total += session_count
+            pages += session_pages
+        read_ms = (time.perf_counter() - read_started) * 1000
+        rebuild_started = time.perf_counter()
+        server.call("rebuild_index", {})
+        rebuild_ms = (time.perf_counter() - rebuild_started) * 1000
+        result = {
+            "messages": args.messages,
+            "sessions": args.sessions,
+            "messages_read": total,
+            "pages": pages,
+            "write_ms": write_ms,
+            "restart_ms": restart_ms,
+            "full_read_ms": read_ms,
+            "rebuild_ms": rebuild_ms,
+            "peak_working_set_bytes": server.peak_working_set(),
+            "max_response_bytes": server.max_response_bytes,
+        }
+        server.close()
+        args.output.write_text(json.dumps(result, indent=2), encoding="utf-8")
+        print(json.dumps(result, indent=2))
+        return 0 if total == args.messages and server.max_response_bytes <= 8 * 1024 * 1024 else 1
+    finally:
+        shutil.rmtree(data_dir, ignore_errors=True)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())