okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

dawnkernel: a job_stop frame lets /jobs stop remote background jobs57a82be

main
branchesmainall refs →
12.8 KiB313 linesraw
1"""Exercise owner frames the daemon cannot generate, including malformed messages.
2
3The real kernel must clean up children on invalid input while preserving pending
4call races, live references and typed values without an SSH server.
5"""
6
7import json
8import select
9import struct
10import subprocess
11import sys
12import tempfile
13import time
14import unittest
15from pathlib import Path
16
17ROOT = Path(__file__).resolve().parents[2]
18KERNEL = ROOT / "priv" / "python" / "albedo_kernel.py"
19
20sys.path.insert(0, str(ROOT / "test"))
21import scratch  # noqa: E402
22
23# Temporary workspaces and homes go under this run's scratch directory, which
24# is removed at exit.
25scratch.claim("harness")
26MAX_FRAME = 8 * 1024 * 1024
27
28
29class OwnerChannel:
30    """A minimal owner: speak the kernel's control frames over stdio.
31
32    Frames the owner is not looking for stay buffered, so looking for one kind
33    never discards another (mirrors and job events interleave with replies).
34    """
35
36    def __init__(self, modules):
37        self.workspace = tempfile.mkdtemp(prefix="albedo-owner-")
38        self.process = subprocess.Popen(
39            [sys.executable, "-u", str(KERNEL), json.dumps(modules)],
40            stdin=subprocess.PIPE,
41            stdout=subprocess.PIPE,
42            stderr=subprocess.PIPE,
43            bufsize=0,
44            cwd=self.workspace,
45        )
46        self.buffered = []
47
48    def close(self):
49        if self.process.poll() is None:
50            self.send({"type": "shutdown"})
51        self.process.wait(timeout=10)
52        assert self.process.stdin is not None
53        assert self.process.stdout is not None
54        assert self.process.stderr is not None
55        self.process.stdin.close()
56        self.process.stdout.close()
57        self.process.stderr.close()
58        self.buffered.clear()
59
60    def send(self, message):
61        data = json.dumps(message).encode()
62        assert self.process.stdin is not None
63        self.process.stdin.write(struct.pack(">I", len(data)) + data)
64        self.process.stdin.flush()
65
66    def recv(self):
67        assert self.process.stdout is not None
68        if not select.select([self.process.stdout], [], [], 10)[0]:
69            raise AssertionError("kernel did not send a frame within 10 seconds")
70        header = self.read_exact(4)
71        size = struct.unpack(">I", header)[0]
72        assert size <= MAX_FRAME, "control frame exceeds the ceiling"
73        return json.loads(self.read_exact(size))
74
75    def read_exact(self, size):
76        assert self.process.stdout is not None
77        data = bytearray()
78        while len(data) < size:
79            chunk = self.process.stdout.read(size - len(data))
80            if not chunk:
81                raise AssertionError("kernel closed the channel during a frame")
82            data.extend(chunk)
83        return bytes(data)
84
85    def wait_for(self, predicate, timeout=10.0):
86        deadline = time.monotonic() + timeout
87        while time.monotonic() < deadline:
88            for index, frame in enumerate(self.buffered):
89                if predicate(frame):
90                    return self.buffered.pop(index)
91            self.buffered.append(self.recv())
92        raise AssertionError(
93            f"no matching frame; buffered: {[f.get('type') for f in self.buffered]}"
94        )
95
96    def recv_reply(self, call_id):
97        return self.wait_for(
98            lambda f: (
99                f.get("id") == call_id and f["type"] in ("invoked", "introspected")
100            )
101        )
102
103    def invoke(self, call_id, **frame):
104        self.send({"type": "invoke", "id": call_id, **frame})
105        return self.recv_reply(call_id)
106
107
108def python(source):
109    """run() arguments for a few steps of Python in a child interpreter."""
110    return [sys.executable, "-c", source]
111
112
113class KernelInvokeTest(unittest.TestCase):
114    def setUp(self):
115        self.channel = OwnerChannel(["run", "files"])
116        self.addCleanup(self.channel.close)
117        ready = self.channel.recv()
118        self.assertEqual(ready["type"], "ready")
119
120    def test_malformed_frames_shutdown_and_clean_up_a_child(self):
121        frames = [
122            [],
123            None,
124            {},
125            {"type": "unknown"},
126            {"type": "execute", "id": "x"},
127            {"type": "execute", "id": "x", "code": 7},
128            {"type": "execute", "id": "x", "code": "", "durable": 1},
129            {"type": "execute", "id": "x", "code": "", "max_edge": True},
130            {"type": "invoke", "id": []},
131            {"type": "invoke", "id": "x", "args": {}},
132            {"type": "invoke", "id": "x", "kwargs": []},
133            {"type": "invoke", "id": "x", "await": 1},
134            {"type": "invoke", "id": "x", "target": {"handle": "a", "pending": "b"}},
135            {"type": "invoke", "id": "x", "target": {"handle": 1}},
136            {"type": "reply", "id": "x", "value": []},
137            {"type": "reply", "id": "x", "value": {"ok": 1, "value": None}},
138            {"type": "reply", "id": "x", "value": {"ok": True}},
139            {"type": "reply", "id": "x", "value": {"ok": False, "code": "bad"}},
140            {"type": "interrupt", "id": "x", "reason": "other"},
141            {"type": "restore", "id": "x", "path": None},
142            {"type": "release", "handle": []},
143            {"type": "job_stop", "id": []},
144        ]
145        payloads = [json.dumps(frame).encode() for frame in frames] + [b"{bad json"]
146        wire_frames = [struct.pack(">I", len(data)) + data for data in payloads]
147        wire_frames.append(struct.pack(">I", MAX_FRAME + 1))
148        for wire in wire_frames:
149            with self.subTest(wire=wire[:160]):
150                channel = OwnerChannel(["run"])
151                try:
152                    self.assertEqual(channel.recv()["type"], "ready")
153                    channel.invoke(
154                        "child", name="run", args=python("import time; time.sleep(60)")
155                    )
156                    started = channel.wait_for(
157                        lambda frame: frame["type"] == "job_start"
158                    )
159                    assert channel.process.stdin is not None
160                    channel.process.stdin.write(wire)
161                    channel.process.stdin.flush()
162                    channel.process.wait(timeout=5)
163                    # Cleanup reaps the child before the kernel kills its own group.
164                    self.assertFalse(Path(f"/proc/{started['pgid']}").exists())
165                finally:
166                    channel.close()
167
168    def test_extra_fields_and_missing_methods_keep_the_channel_open(self):
169        reply = self.channel.invoke(
170            "extra", name="files.read", args=["missing"], extra=True
171        )
172        self.assertFalse(reply["ok"])
173        self.assertFalse(self.channel.invoke("missing")["ok"])
174        self.channel.send({"type": "introspect", "id": "still-open", "extra": []})
175        self.assertIn("run", self.channel.recv_reply("still-open")["names"])
176
177    def test_pending_targets_resolve_calls_raced_ahead_of_their_reply(self):
178        self.channel.send(
179            {
180                "type": "invoke",
181                "id": "r1",
182                "name": "run",
183                "args": python(
184                    "import time; print('raced', flush=True); time.sleep(0.3); print('done')"
185                ),
186            }
187        )
188        # tail() is called before the run reply can exist; the pending target waits
189        raced = self.channel.invoke("r2", target={"pending": "r1"}, name="tail")
190        self.assertTrue(raced["ok"])
191        self.assertEqual(raced["value"], "")
192        reply = self.channel.recv_reply("r1")
193        self.assertTrue(reply["ok"])
194        done = self.channel.invoke(
195            "r3", target={"handle": reply["handle"]}, **{"await": True}
196        )
197        self.assertEqual(done["state"]["tail"], "raced\ndone\n")
198
199    def test_interrupt_cancels_an_await_without_losing_the_reference(self):
200        reply = self.channel.invoke(
201            "job",
202            name="run",
203            args=python("import time; time.sleep(1); print('finished')"),
204        )
205        handle = reply["handle"]
206        self.channel.send(
207            {
208                "type": "invoke",
209                "id": "wait",
210                "target": {"handle": handle},
211                "await": True,
212            }
213        )
214        # A mirror arrives after the await task has had a loop turn to start.
215        self.channel.wait_for(lambda frame: frame["type"] == "mirror")
216        self.channel.send({"type": "interrupt", "id": "wait"})
217        cancelled = self.channel.recv_reply("wait")
218        self.assertFalse(cancelled["ok"])
219        self.assertTrue(cancelled["cancelled"])
220        finished = self.channel.invoke(
221            "finish", target={"handle": handle}, **{"await": True}
222        )
223        self.assertTrue(finished["ok"])
224        self.assertEqual(finished["state"]["tail"], "finished\n")
225
226    def test_interrupt_for_a_queued_cell_lands_when_it_starts(self):
227        self.channel.send(
228            {"type": "execute", "id": "busy", "code": "import time; time.sleep(0.5)"}
229        )
230        self.channel.send({"type": "execute", "id": "queued", "code": "ran = True"})
231        self.channel.send({"type": "interrupt", "id": "queued", "reason": "deadline"})
232        busy = self.channel.wait_for(lambda f: f["type"] == "done")
233        self.assertEqual((busy["id"], busy["status"]), ("busy", "ok"))
234        queued = self.channel.wait_for(lambda f: f["type"] == "done")
235        self.assertEqual((queued["id"], queued["status"]), ("queued", "interrupted"))
236        self.assertIn("deadline exceeded", queued["output"])
237        self.channel.send({"type": "execute", "id": "check", "code": "'ran' in dir()"})
238        check = self.channel.wait_for(lambda f: f["type"] == "done")
239        self.assertEqual(check["value"], "False")
240
241    def test_a_failing_pending_call_propagates_its_error(self):
242        self.channel.send({"type": "invoke", "id": "f1", "name": "nope"})
243        failure = self.channel.invoke("f2", target={"pending": "f1"}, name="tail")
244        self.assertFalse(failure["ok"])
245        self.assertEqual(failure["error"]["ename"], "LookupError")
246
247    def test_typed_composites_carry_a_reference_and_their_value(self):
248        Path(self.channel.workspace, "note.txt").write_text(
249            "needle one\nplain\nneedle two\n"
250        )
251        reply = self.channel.invoke("i1", name="files.find", args=["needle", "."])
252        self.assertTrue(reply["ok"])
253        self.assertIn("handle", reply)
254        value = reply["value"]
255        self.assertEqual(value["__list__"], "albedo_plugins.files.Rows")
256        self.assertEqual(value["items"][0]["__class__"], "albedo_plugins.files.Match")
257        self.assertEqual(value["items"][0]["fields"]["line"], 1)
258        self.assertEqual(len(value["items"]), 2)
259
260    def test_reference_arguments_resolve_back_to_live_objects(self):
261        reply = self.channel.invoke("i1", name="run", args=["echo", "ref-roundtrip"])
262        handle = reply["handle"]
263        bad = self.channel.invoke(
264            "i2", name="files.read", args=[{"__ref__": "not-a-handle"}]
265        )
266        self.assertFalse(bad["ok"])
267        self.assertEqual(bad["error"]["ename"], "LookupError")
268        # A live reference resolves to the object itself; the failure that follows
269        # is the call's own, proving the marker never leaked through as a dict.
270        answered = self.channel.invoke(
271            "i3", name="files.read", args=[{"__ref__": handle}]
272        )
273        self.assertFalse(answered["ok"])
274        self.assertNotEqual(answered["error"]["ename"], "LookupError")
275
276    def test_release_forgets_the_reference(self):
277        reply = self.channel.invoke("i1", name="run", args=["true"])
278        handle = reply["handle"]
279        self.channel.send({"type": "release", "handle": handle})
280        gone = self.channel.invoke("i2", target={"handle": handle}, name="poll")
281        self.assertFalse(gone["ok"])
282
283    def test_job_stop_unknown_job_answers_not_found(self):
284        self.channel.send({"type": "job_stop", "id": "unknown-job"})
285        reply = self.channel.wait_for(lambda f: f["type"] == "job_stop")
286        self.assertEqual(
287            reply,
288            {
289                "type": "job_stop",
290                "id": "unknown-job",
291                "ok": False,
292                "error": "job not found",
293            },
294        )
295
296    def test_job_stop_running_job_stops_and_confirms(self):
297        self.channel.invoke(
298            "child", name="run", args=python("import time; time.sleep(60)")
299        )
300        started = self.channel.wait_for(lambda f: f["type"] == "job_start")
301        job_id = started["id"]
302        self.channel.send({"type": "job_stop", "id": job_id})
303        reply = self.channel.wait_for(lambda f: f["type"] == "job_stop")
304        self.assertEqual(reply, {"type": "job_stop", "id": job_id, "ok": True})
305        # The job also announces its exit
306        ended = self.channel.wait_for(
307            lambda f: f["type"] == "job" and f["id"] == job_id
308        )
309        self.assertTrue(ended["cleanup"]["gone"])
310
311
312if __name__ == "__main__":
313    unittest.main()