Merge 5476d6fbdd80aeba9367b4c8b5bcaa5ac2d903a6 into 9e145d9ceb6268b6957fef6aa1889cb99e9f35be

This commit is contained in:
NIKHIL PENTAPALLI 2026-08-29 17:04:24 -07:00 committed by GitHub
commit c64183521d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 51 additions and 7 deletions

View File

@ -11,6 +11,7 @@ owns the wake records + the due/complete logic; the scheduler tick consumes ``du
from __future__ import annotations
import json
import os
import threading
import uuid
from dataclasses import asdict, dataclass, field
@ -50,20 +51,31 @@ class WakeStore:
self._lock = threading.Lock()
self._wakes: dict[str, Wake] = {}
if self.path and self.path.is_file():
for raw in json.loads(self.path.read_text(encoding="utf-8")).get(
"wakes", []
):
w = Wake(**raw)
# A corrupt/partial file must not crash startup — better to lose pending
# wakes than wedge the whole server on every launch.
try:
raw_wakes = json.loads(self.path.read_text(encoding="utf-8")).get(
"wakes", []
)
except (OSError, json.JSONDecodeError):
raw_wakes = []
for raw in raw_wakes:
try:
w = Wake(**raw)
except TypeError:
continue
self._wakes[w.id] = w
def _save(self) -> None:
if not self.path:
return
self.path.parent.mkdir(parents=True, exist_ok=True)
self.path.write_text(
tmp = self.path.with_name(self.path.name + ".tmp")
tmp.write_text(
json.dumps({"wakes": [asdict(w) for w in self._wakes.values()]}, indent=2),
encoding="utf-8",
)
os.replace(tmp, self.path)
def add_timer(self, session_id: str, fire_at: datetime, *, note: str = "") -> Wake:
w = Wake(
@ -99,8 +111,10 @@ class WakeStore:
def due(self, now: Optional[datetime] = None) -> list[Wake]:
"""Timer wakes whose fire time has passed, plus completion/event wakes marked due."""
now = now or _now()
with self._lock:
wakes = list(self._wakes.values())
out = []
for w in self._wakes.values():
for w in wakes:
if w.state != STATE_PENDING and w.state != STATE_DUE:
continue
if (
@ -144,9 +158,11 @@ class WakeStore:
self._save()
def pending(self, session_id: Optional[str] = None) -> list[Wake]:
with self._lock:
wakes = list(self._wakes.values())
return [
w
for w in self._wakes.values()
for w in wakes
if w.state != STATE_FIRED
and (session_id is None or w.session_id == session_id)
]

View File

@ -2,6 +2,7 @@
from __future__ import annotations
import json
from datetime import datetime, timedelta, timezone
from coworker.selfwake import WakeStore, selfwake_tools
@ -64,3 +65,30 @@ def test_selfwake_tools(tmp_path):
pend = store.pending("s1")
assert len(pend) == 4
assert {w.kind for w in pend} == {"timer", "completion", "event"}
def test_corrupt_state_file_does_not_crash_load(tmp_path):
path = tmp_path / "wakes.json"
path.write_text('{"wakes": [truncated', encoding="utf-8")
store = WakeStore(path) # must not raise
assert store.pending() == []
w = store.add_completion("s1", "job-1")
assert any(x.id == w.id for x in WakeStore(path).pending("s1"))
def test_unknown_wake_fields_are_skipped_not_fatal(tmp_path):
path = tmp_path / "wakes.json"
good = WakeStore(path).add_completion("s1", "job-1")
data = json.loads(path.read_text(encoding="utf-8"))
data["wakes"].append({"id": "x", "future_field": True})
path.write_text(json.dumps(data), encoding="utf-8")
reloaded = WakeStore(path)
assert [w.id for w in reloaded.pending("s1")] == [good.id]
def test_save_is_atomic_no_tmp_left_behind(tmp_path):
path = tmp_path / "wakes.json"
store = WakeStore(path)
store.add_timer("s1", _now() + timedelta(seconds=60))
assert not path.with_name(path.name + ".tmp").exists()
assert len(json.loads(path.read_text(encoding="utf-8"))["wakes"]) == 1