Bound import memory with staged streaming

1 files changed, 38 insertions, 9 deletions

diff --git a/scripts/benchmark_p1_cp.py b/scripts/benchmark_p1_cp.py
index f456fbb..c6a9d0a 100644
--- a/scripts/benchmark_p1_cp.py
+++ b/scripts/benchmark_p1_cp.py
@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import argparse
+import ctypes
 import json
 import shutil
 import subprocess
@@ -12,6 +13,21 @@ from pathlib import Path
 FIXTURE_PREFIX = "p1-cp-benchmark-"
 
 
+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),
+    ]
+
+
 def make_ulid(timestamp_ms: int, random_bits: int) -> str:
     alphabet = "0123456789ABCDEFGHJKMNPQRSTVWXYZ"
     bytes_data = bytearray(16)
@@ -59,18 +75,28 @@ def write_fixture(path: Path, count: int) -> None:
             file.write(json.dumps(message, separators=(",", ":")) + "\n")
 
 
-def run_import(executable: Path, data_dir: Path, input_path: Path) -> float:
+def run_import(executable: Path, data_dir: Path, input_path: Path) -> tuple[float, int]:
     shutil.rmtree(data_dir, ignore_errors=True)
     start = time.perf_counter()
-    result = subprocess.run(
+    process = subprocess.Popen(
         [str(executable), "--data-dir", str(data_dir), "--import", str(input_path)],
-        capture_output=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
         text=True,
-        check=True,
     )
+    stdout, stderr = process.communicate()
     elapsed = time.perf_counter() - start
-    print(result.stdout.strip())
-    return elapsed
+    if process.returncode != 0:
+        raise RuntimeError(stderr)
+    counters = ProcessMemoryCounters()
+    counters.cb = ctypes.sizeof(counters)
+    peak = 0
+    if ctypes.windll.psapi.GetProcessMemoryInfo(  # type: ignore[attr-defined]
+        int(process._handle), ctypes.byref(counters), counters.cb  # type: ignore[attr-defined]
+    ):
+        peak = int(counters.PeakWorkingSetSize)
+    print(stdout.strip())
+    return elapsed, peak
 
 
 def run_export(executable: Path, data_dir: Path, output_path: Path) -> float:
@@ -105,15 +131,16 @@ def main() -> None:
             print(f"fixture size: {fixture_size:,} bytes")
 
             import_dir = work_dir / f"import-{count}"
-            import_ms = run_import(executable, import_dir, fixture)
+            import_ms, import_peak = run_import(executable, import_dir, fixture)
             print(f"import: {import_ms:.3f}s ({count / import_ms:,.0f} msg/s)")
+            print(f"import peak working set: {import_peak:,} bytes")
 
             export_path = work_dir / f"exported-{count}.jsonl"
             export_ms = run_export(executable, import_dir, export_path)
             print(f"export: {export_ms:.3f}s ({count / export_ms:,.0f} msg/s)")
 
             roundtrip_dir = work_dir / f"roundtrip-{count}"
-            roundtrip_ms = run_import(executable, roundtrip_dir, export_path)
+            roundtrip_ms, roundtrip_peak = run_import(executable, roundtrip_dir, export_path)
             print(f"round-trip import: {roundtrip_ms:.3f}s ({count / roundtrip_ms:,.0f} msg/s)")
 
             export_size = export_path.stat().st_size
@@ -124,8 +151,10 @@ def main() -> None:
                         "count": count,
                         "fixture_bytes": fixture_size,
                         "import_ms": round(import_ms * 1000, 3),
+                        "import_peak_working_set_bytes": import_peak,
                         "export_ms": round(export_ms * 1000, 3),
                         "roundtrip_ms": round(roundtrip_ms * 1000, 3),
+                        "roundtrip_peak_working_set_bytes": roundtrip_peak,
                         "export_bytes": export_size,
                     }
                 )
@@ -135,4 +164,4 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()