"""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_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()