perf: bound recent rollout history reads

loopx/rollout_event_log.py

@@ -5,7 +5,7 @@
 from collections import Counter, deque
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import Any, Callable, Iterable, Iterator, Mapping, Sequence
+from typing import Any, BinaryIO, Callable, Iterable, Iterator, Mapping, Sequence
 
 from .file_lock import exclusive_cross_runtime_file_lock
 
@@ -513,43 +513,98 @@ def append_rollout_event_once(
     return payload, True
 
 
+def _parse_rollout_event_line(
+    line: str, *, line_number: int = 0, strict: bool = False,
+) -> dict[str, Any] | None:
+    if not line.strip():
+        return None
+    try:
+        parsed = json.loads(line)
+    except json.JSONDecodeError:
+        if strict:
+            raise ValueError(
+                f"rollout event log line {line_number} must contain valid JSON"
+            ) from None
+        return None
+    if not isinstance(parsed, dict):
+        if strict:
+            raise ValueError(
+                f"rollout event log line {line_number} must contain an object"
+            )
+        return None
+    if parsed.get("schema_version") != ROLLOUT_EVENT_SCHEMA_VERSION:
+        if strict:
+            raise ValueError(
+                f"rollout event log line {line_number} must use "
+                f"{ROLLOUT_EVENT_SCHEMA_VERSION}"
+            )
+        return None
+    return parsed
+
+
 def iter_rollout_events(
     log_path: Path,
     *,
     strict: bool = False,
 ) -> Iterator[dict[str, Any]]:
     for line_number, line in _iter_rollout_event_lines(log_path, strict=strict):
-        if not line.strip():
-            continue
-        try:
-            parsed = json.loads(line)
-        except json.JSONDecodeError:
-            if strict:
-                raise ValueError(
-                    f"rollout event log line {line_number} must contain valid JSON"
-                ) from None
-            continue
-        if not isinstance(parsed, dict):
-            if strict:
-                raise ValueError(
-                    f"rollout event log line {line_number} must contain an object"
-                )
-            continue
-        if parsed.get("schema_version") != ROLLOUT_EVENT_SCHEMA_VERSION:
-            if strict:
-                raise ValueError(
-                    f"rollout event log line {line_number} must use "
-                    f"{ROLLOUT_EVENT_SCHEMA_VERSION}"
-                )
+        event = _parse_rollout_event_line(line, line_number=line_number, strict=strict)
+        if event is not None:
+            yield event
+
+
+def _reverse_rollout_lines(handle: BinaryIO) -> Iterator[bytes]:
+    """Read from the opened EOF, assembling complete UTF-8 rows before decoding."""
+    position = handle.seek(0, 2)
+    fragments: list[bytes] = []
+    while position:
+        size = min(position, 64 * 1024)
+        position -= size
+        handle.seek(position)
+        parts = handle.read(size).split(b"\n")
+        if len(parts) == 1:
+            fragments.append(parts[0])
             continue
-        yield parsed
+        yield b"".join([parts.pop(), *reversed(fragments)])
+        fragments.clear()
+        yield from reversed(parts[1:])
+        fragments.append(parts[0])
+    if fragments:
+        yield b"".join(reversed(fragments))
+
+
+def _recent_rollout_events(log_path: Path, *, limit: int) -> list[dict[str, Any]]:
+    events: deque[dict[str, Any]] = deque(maxlen=limit)
+    try:
+        handle = log_path.open("rb")
+    except OSError:
+        return []
+
+    def parsed_lines(lines: Iterable[bytes]) -> Iterator[dict[str, Any]]:
+        for encoded in lines:
+            try:
+                line = encoded.decode("utf-8")
+            except UnicodeDecodeError:
+                continue
+            event = _parse_rollout_event_line(line)
+            if event is not None:
+                yield event
+
+    with handle:
+        if not handle.seekable():
+            events.extend(parsed_lines(handle))
+        else:
+            for event in parsed_lines(_reverse_rollout_lines(handle)):
+                events.appendleft(event)
+                if len(events) == events.maxlen:
+                    break
+    return list(events)
 
 
 def load_rollout_events(log_path: Path, *, limit: int | None = None) -> list[dict[str, Any]]:
-    events = iter_rollout_events(log_path)
     if limit is None or limit <= 0:
-        return list(events)
-    return list(deque(events, maxlen=limit))
+        return list(iter_rollout_events(log_path))
+    return _recent_rollout_events(log_path, limit=limit)
 
 
 class RolloutEventSnapshot:

tests/test_rollout_event_log_loading.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import io
 import json
 import os
 from pathlib import Path
@@ -60,30 +61,164 @@ def test_limited_rollout_event_load_keeps_only_a_bounded_window(
     tmp_path: Path,
     monkeypatch,
 ) -> None:
-    class TrackedEvent(dict[str, int]):
+    class TrackedEvent(dict[str, object]):
         alive = 0
         peak = 0
 
-        def __init__(self, index: int) -> None:
-            super().__init__(index=index)
+        def __init__(self, value: dict[str, object]) -> None:
+            super().__init__(value)
             type(self).alive += 1
             type(self).peak = max(type(self).peak, type(self).alive)
 
         def __del__(self) -> None:
             type(self).alive -= 1
 
-    monkeypatch.setattr(
-        rollout_event_log,
-        "iter_rollout_events",
-        lambda _path: (TrackedEvent(index) for index in range(100)),
-    )
+    log = tmp_path / "events.jsonl"
+    log.write_text("".join(json.dumps({
+        "schema_version": rollout_event_log.ROLLOUT_EVENT_SCHEMA_VERSION,
+        "index": index,
+    }) + "\n" for index in range(100)), encoding="utf-8")
+    original_loads = json.loads
+
+    def tracked_loads(value):
+        return TrackedEvent(original_loads(value))
+
+    monkeypatch.setattr(rollout_event_log.json, "loads", tracked_loads)
 
-    events = rollout_event_log.load_rollout_events(tmp_path / "unused", limit=3)
+    events = rollout_event_log.load_rollout_events(log, limit=3)
 
     assert [event["index"] for event in events] == [97, 98, 99]
     assert TrackedEvent.peak <= 4
 
 
+def test_recent_rollout_read_does_not_scan_unneeded_history(
+    tmp_path: Path, monkeypatch,
+) -> None:
+    log = tmp_path / "events.jsonl"
+    records = [{"schema_version": rollout_event_log.ROLLOUT_EVENT_SCHEMA_VERSION,
+                "index": index, "summary": "retained history " * 32}
+               for index in range(3000)]
+    log.write_text("".join(json.dumps(row) + "\n" for row in records), encoding="utf-8")
+    original_open = Path.open
+    observed = {"bytes": 0}
+
+    class CountedRead:
+        def __init__(self, handle):
+            self.handle = handle
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *args):
+            return self.handle.__exit__(*args)
+
+        def __getattr__(self, name):
+            return getattr(self.handle, name)
+
+        def __iter__(self):
+            for line in self.handle:
+                observed["bytes"] += len(line)
+                yield line
+
+        def read(self, size=-1):
+            value = self.handle.read(size)
+            observed["bytes"] += len(value)
+            return value
+
+    def counted_open(path, mode="r", *args, **kwargs):
+        handle = original_open(path, mode, *args, **kwargs)
+        return CountedRead(handle) if path == log and mode == "rb" else handle
+
+    monkeypatch.setattr(Path, "open", counted_open)
+
+    assert load_rollout_events(log, limit=3) == records[-3:]
+    assert observed["bytes"] < log.stat().st_size // 4
+
+
+@pytest.mark.parametrize("limit", [1, 2, 7, 100])
+@pytest.mark.parametrize("final_newline", [False, True])
+def test_recent_rollout_read_counts_valid_records_and_preserves_complete_rows(
+    tmp_path: Path, limit: int, final_newline: bool,
+) -> None:
+    log = tmp_path / "events.jsonl"
+    records = [{"schema_version": rollout_event_log.ROLLOUT_EVENT_SCHEMA_VERSION,
+                "index": index, "summary": "雪🙂" * (25000 if index == 4 else 3)}
+               for index in range(8)]
+    invalid = b'\xff\n{"torn":\n[]\nnull\n{"schema_version":"other"}\n\n'
+    contents = b"".join(json.dumps(row, ensure_ascii=False).encode("utf-8")
+                        + b"\n" + invalid for row in records[:-1])
+    contents += json.dumps(records[-1], ensure_ascii=False).encode("utf-8")
+    contents += b"\n" if final_newline else b""
+    log.write_bytes(contents)
+
+    assert load_rollout_events(log, limit=limit) == records[-limit:]
+    assert log.read_bytes() == contents
+
+
+@pytest.mark.parametrize("limit", [None, 0, -1])
+def test_unlimited_rollout_read_retains_its_complete_history(
+    tmp_path: Path, limit: int | None,
+) -> None:
+    log = tmp_path / "events.jsonl"
+    records = [{"schema_version": rollout_event_log.ROLLOUT_EVENT_SCHEMA_VERSION,
+                "index": index} for index in range(5)]
+    log.write_text("\n".join(json.dumps(row) for row in records), encoding="utf-8")
+    assert load_rollout_events(log, limit=limit) == records
+
+
+def test_recent_rollout_read_reopens_the_current_file(tmp_path: Path) -> None:
+    log = tmp_path / "events.jsonl"
+    records = [{"schema_version": rollout_event_log.ROLLOUT_EVENT_SCHEMA_VERSION,
+                "index": index} for index in range(3)]
+    log.write_text(json.dumps(records[0]) + "\n", encoding="utf-8")
+    assert load_rollout_events(log, limit=2) == records[:1]
+    with log.open("a", encoding="utf-8") as handle:
+        handle.write("\n".join(json.dumps(row) for row in records[1:]))
+    assert load_rollout_events(log, limit=2) == records[1:]
+
+
+def test_recent_rollout_read_preserves_nonseekable_input(
+    tmp_path: Path, monkeypatch,
+) -> None:
+    log = tmp_path / "events.jsonl"
+    records = [{"schema_version": rollout_event_log.ROLLOUT_EVENT_SCHEMA_VERSION,
+                "index": index} for index in range(4)]
+    contents = "\n".join(json.dumps(row) for row in records).encode()
+    original_open = Path.open
+
+    class NonseekableRead(io.BytesIO):
+        def seekable(self):
+            return False
+
+        def seek(self, *args):
+            raise io.UnsupportedOperation("not seekable")
+
+    def opened(path, mode="r", *args, **kwargs):
+        return NonseekableRead(contents) if path == log else original_open(path, mode, *args, **kwargs)
+
+    monkeypatch.setattr(Path, "open", opened)
+    assert load_rollout_events(log, limit=2) == records[-2:]
+
+
+@pytest.mark.parametrize("limit", [None, 1])
+def test_rollout_read_keeps_missing_file_behavior(tmp_path: Path, limit: int | None) -> None:
+    assert load_rollout_events(tmp_path / "missing.jsonl", limit=limit) == []
+
+
+@pytest.mark.parametrize("row,message", [
+    ("{", "must contain valid JSON"),
+    ("[]", "must contain an object"),
+    ('{"schema_version":"other"}', "must use loopx_rollout_event_v0"),
+])
+def test_strict_rollout_read_still_reports_the_original_line(
+    tmp_path: Path, row: str, message: str,
+) -> None:
+    log = tmp_path / "events.jsonl"
+    log.write_text("\n" + row + "\n", encoding="utf-8")
+    with pytest.raises(ValueError, match="rollout event log line 2 " + message):
+        list(iter_rollout_events(log, strict=True))
+
+
 @pytest.mark.parametrize("idempotent", [False, True])
 @pytest.mark.parametrize(
     "torn_record",