"""Busy-host replies and reading an in-flight job notice are races the daemon E2E wake test cannot reliably force.""" import json import os import select import struct import subprocess import sys import tempfile import time import unittest from pathlib import Path ROOT = Path(__file__).resolve().parents[2] KERNEL = ROOT / "priv" / "python" / "albedo_kernel.py" sys.path.insert(0, str(ROOT / "test")) import scratch # noqa: E402 # Temporary workspaces and homes go under this run's scratch directory, which # is removed at exit. scratch.claim("harness") MAX_FRAME = 8 * 1024 * 1024 RETRY_WINDOW = 3.2 # NOTICE_RETRY plus slack; must not mistake a retry for silence class Owner: """A minimal owner that answers host calls and buffers every other frame.""" def __init__(self, modules, env=None): self.workspace = tempfile.mkdtemp(prefix="albedo-wake-") environment = dict(os.environ) environment.update(env or {}) self.process = subprocess.Popen( [sys.executable, "-u", str(KERNEL), json.dumps(modules)], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=self.workspace, env=environment, bufsize=0, ) self.buffered = [] def close(self): self.send({"type": "shutdown"}) self.process.wait(timeout=10) for stream in (self.process.stdin, self.process.stdout, self.process.stderr): assert stream is not None stream.close() self.buffered.clear() def send(self, message): data = json.dumps(message).encode() assert self.process.stdin is not None self.process.stdin.write(struct.pack(">I", len(data)) + data) self.process.stdin.flush() def recv(self, timeout=0.05): """One frame, or None when the kernel stays quiet for the window. stdout is unbuffered: a frame read ahead into a buffer is invisible to select(), and would stall until the next one arrives.""" assert self.process.stdout is not None ready, _, _ = select.select([self.process.stdout], [], [], timeout) if not ready: return None size = struct.unpack(">I", self.read_exactly(4))[0] assert size <= MAX_FRAME, "control frame exceeds the ceiling" return json.loads(self.read_exactly(size)) def read_exactly(self, size): assert self.process.stdout is not None data = b"" while len(data) < size: chunk = self.process.stdout.read(size - len(data)) assert chunk, "kernel closed its control stream" data += chunk return data def wait_for(self, predicate, timeout=10.0): deadline = time.monotonic() + timeout while time.monotonic() < deadline: for index, frame in enumerate(self.buffered): if predicate(frame): return self.buffered.pop(index) frame = self.recv(timeout=max(0.05, deadline - time.monotonic())) if frame is not None: self.buffered.append(frame) raise AssertionError( f"no matching frame; buffered: {[f.get('type') for f in self.buffered]}" ) def expect_no_call(self, method, window=RETRY_WINDOW): deadline = time.monotonic() + window while time.monotonic() < deadline: try: frame = self.wait_for( lambda f: f.get("type") == "call" and f.get("method") == method, timeout=max(0.05, deadline - time.monotonic()), ) except AssertionError: return raise AssertionError(f"unexpected {method} call: {frame}") def invoke(self, call_id, **frame): self.send({"type": "invoke", "id": call_id, **frame}) while True: reply = self.wait_for( lambda f: f.get("id") == call_id and f["type"] == "invoked" ) if reply.get("ok") is True or "error" in reply: return reply class WakeProtocolTest(unittest.TestCase): def setUp(self): self.owner = Owner(["run"]) self.addCleanup(self.owner.close) self.assertEqual( self.owner.wait_for(lambda f: f.get("type") == "ready")["type"], "ready" ) def start_job(self, *argv): reply = self.owner.invoke("j1", name="run", args=list(argv)) self.assertTrue(reply["ok"]) return reply["handle"] def wait_job_done(self): """Wait on the raw completion frame; awaiting the job would retire the wake.""" start = self.owner.wait_for(lambda f: f.get("type") == "job_start") done = self.owner.wait_for( lambda f: f.get("type") == "job" and f.get("id") == start["id"] ) return done def test_a_busy_session_is_retried_and_a_refusal_gives_up(self): self.start_job("echo", "retry-echo") self.wait_job_done() # Busy is retried, not dropped; a permanent refusal ends the attempts. frame = self.owner.wait_for(lambda f: f.get("type") == "call") self.owner.send( { "type": "reply", "id": frame["id"], "value": {"ok": False, "code": "busy", "message": "session is busy"}, } ) retry = self.owner.wait_for( lambda f: f.get("type") == "call" and f.get("method") == "jobs.completed", timeout=RETRY_WINDOW, ) self.owner.send( { "type": "reply", "id": retry["id"], "value": {"ok": False, "code": "unavailable", "message": "no session"}, } ) self.owner.expect_no_call("jobs.completed") def test_jobs_owed_when_a_busy_session_frees_share_one_notice(self): self.start_job("echo", "first-echo") self.wait_job_done() frame = self.owner.wait_for(lambda f: f.get("type") == "call") self.assertEqual(len(frame["args"]["jobs"]), 1) self.start_job("echo", "second-echo") self.wait_job_done() self.owner.send( { "type": "reply", "id": frame["id"], "value": {"ok": False, "code": "busy", "message": "session is busy"}, } ) retry = self.owner.wait_for( lambda f: f.get("type") == "call" and f.get("method") == "jobs.completed", timeout=RETRY_WINDOW, ) self.assertEqual(len(retry["args"]["jobs"]), 2) self.assertIn("2 background jobs finished", retry["args"]["text"]) self.assertIn("second-echo", retry["args"]["display"]) self.owner.send( {"type": "reply", "id": retry["id"], "value": {"ok": True, "value": None}} ) self.owner.expect_no_call("jobs.completed") def test_a_read_result_retires_a_pending_notice(self): handle = self.start_job("echo", "read-echo") self.wait_job_done() # The notice is in flight (busy) when the result is read; the retry must stop. frame = self.owner.wait_for(lambda f: f.get("type") == "call") self.owner.send( { "type": "reply", "id": frame["id"], "value": {"ok": False, "code": "busy", "message": "session is busy"}, } ) tail = self.owner.invoke("t1", target={"handle": handle}, name="tail") self.assertIn("read-echo", tail["value"]) self.owner.expect_no_call("jobs.completed") if __name__ == "__main__": unittest.main()