okami.mom/albedo

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

nekomimi.petkernel: run the cell runtime from cached bytecode and precompile the packaged python9f86700

main
branchesmainall refs →
11.6 KiB311 linesraw
1"""OS process-group identity and shutdown races need controlled child processes rather than flaky daemon timing."""
2
3from __future__ import annotations
4
5from collections.abc import Sequence
6from typing import cast
7import asyncio
8import os
9import sys
10import time
11import unittest
12
13ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
14sys.path.insert(0, os.path.join(ROOT, "priv", "python"))
15
16import albedo_api  # noqa: E402
17import albedo_cells  # noqa: E402
18import albedo_proc  # noqa: E402
19from albedo_plugins import run as plugin  # noqa: E402
20
21LOOP = asyncio.new_event_loop()
22asyncio.set_event_loop(LOOP)
23albedo_proc.reap_stopped_safely(LOOP)
24EVENTS: list[dict[str, object]] = []
25
26
27def install() -> None:
28    plugin.setup(
29        albedo_api.PythonApi(
30            version=1,
31            loop=LOOP,
32            host=lambda method, args: (_ for _ in ()).throw(
33                RuntimeError("no host in tests")
34            ),
35            HostError=RuntimeError,
36            capture=albedo_cells.background_capture,
37            preview=albedo_cells.PREVIEW,
38            send=EVENTS.append,
39            on_shutdown=lambda close: None,
40            background_handle=lambda _: None,
41        )
42    )
43
44
45def run(coro):
46    return LOOP.run_until_complete(coro)
47
48
49def start(command: str, timeout: float = 300):
50    """Start a shell line as a job without waiting for it; supervision is the
51    same for any program, and a shell makes descendants easy to arrange."""
52    return plugin.start(["/bin/sh", "-c", command], timeout)
53
54
55def quick(term: float = 0.01, kill: float = 0.02):
56    """Patch the ladder windows so a refusing group is tested without real waiting."""
57    real = albedo_proc.terminate
58    quick_term, quick_kill = term, kill
59
60    async def patched(
61        groups: Sequence[albedo_proc.Group], term: float = 0.0, kill: float = 0.0
62    ) -> list[albedo_proc.Termination]:
63        return await real(groups, quick_term, quick_kill)
64
65    setattr(albedo_proc, "terminate", patched)
66    return real
67
68
69def wait(job):
70    return run(job)
71
72
73def present(pgid: int) -> bool:
74    try:
75        os.killpg(pgid, 0)
76        return True
77    except ProcessLookupError:
78        return False
79    except PermissionError:
80        return True
81
82
83def pid_present(pid: int) -> bool:
84    try:
85        os.kill(pid, 0)
86        return True
87    except ProcessLookupError:
88        return False
89    except PermissionError:
90        return True
91
92
93def ready(job, marker: str = "ready") -> None:
94    """Wait for the job to print `marker`. A shell that has not yet run its
95    first line cannot be said to refuse a signal or to have started a child, and
96    how soon it runs one depends on the machine's load, not on the code."""
97    deadline = time.monotonic() + 30
98    while marker not in job.tail():
99        assert time.monotonic() < deadline, f"job never printed {marker!r}"
100        run(asyncio.sleep(0.01))
101
102
103def pid_from(output: str) -> int:
104    return int(output.strip().splitlines()[0])
105
106
107class SupervisionTest(unittest.TestCase):
108    def setUp(self):
109        EVENTS.clear()
110        install()
111
112    def test_deadline_ends_the_whole_group(self):
113        # The deadline must outlast the shell's first lines on a loaded machine,
114        # so wait for the child's pid before relying on the deadline to end it.
115        job = start("sleep 30 & echo $!; wait", timeout=2)
116        ready(job, "\n")
117        wait(job)
118        self.assertTrue(job.timed_out)
119        self.assertTrue(job.termination.gone)
120        self.assertFalse(present(job.group.pgid))
121        self.assertFalse(pid_present(pid_from(job.tail())))
122        self.assertIn("deadline exceeded after 2s", job.tail())
123        self.assertIn("terminated", job.tail())
124
125    def test_command_exit_ends_descendants_that_hold_no_output(self):
126        started = time.monotonic()
127        job = start("sleep 30 >/dev/null 2>&1 & echo $!; exit 0", timeout=30)
128        wait(job)
129        grandchild = pid_from(job.tail())
130        self.assertLess(time.monotonic() - started, 5)  # no wait for the deadline
131        self.assertFalse(job.timed_out)
132        self.assertTrue(job.termination.gone)
133        self.assertFalse(pid_present(grandchild))
134
135    def test_command_exit_leaves_what_a_process_that_left_the_group_parents(self):
136        # Like ssh with ControlPersist: a master forks its proxy, daemonizes out
137        # of the job (setsid), and outlives it; the proxy stays in the group.
138        # A plain background child is an orphan once the command ends.
139        script = (
140            "import os, sys, time\n"
141            "if os.fork() == 0:\n"
142            "    proxy = os.fork()\n"
143            "    if proxy == 0:\n"
144            "        os.execvp('sleep', ['sleep', '31'])\n"
145            "    os.setsid()\n"
146            "    print('proxy', proxy, 'master', os.getpid(), flush=True)\n"
147            "    os.execvp('sleep', ['sleep', '32'])\n"
148            "orphan = os.fork()\n"
149            "if orphan == 0:\n"
150            "    os.execvp('sleep', ['sleep', '33'])\n"
151            "print('orphan', orphan, flush=True)\n"
152            "time.sleep(0.5)\n"
153        )
154        job = plugin.start([sys.executable, "-c", script], 30)
155        wait(job)
156        pids = dict(zip(*[iter(job.tail().split())] * 2))
157        proxy, master, orphan = (
158            int(pids[name]) for name in ("proxy", "master", "orphan")
159        )
160        self.addCleanup(os.kill, master, 9)
161        self.addCleanup(os.kill, proxy, 9)
162        ending, group = job.termination, job.group
163        assert ending is not None and group is not None
164        self.assertTrue(ending.gone, ending)
165        self.assertIn("kept 1 that a process which left the group", ending.note)
166        self.assertTrue(pid_present(proxy))
167        self.assertTrue(pid_present(master))
168        self.assertFalse(pid_present(orphan))
169        # A stop or a deadline still ends everything in the group.
170        self.assertTrue(run(albedo_proc.end(group)).gone)
171        self.assertEqual(albedo_proc.live_members(group.pgid), [])
172
173    def test_stop_escalates_and_is_retryable(self):
174        job = start("trap '' TERM; echo ready; while :; do sleep 0.2; done", timeout=30)
175        self.addCleanup(lambda: run(job.stop()))
176        ready(job)  # the trap is installed before the first signal is sent
177        self.assertIsNone(job.poll())  # still running
178        ending = run(job.stop())
179        self.assertTrue(ending.gone)
180        self.assertEqual(ending.signals, ("SIGTERM", "SIGKILL"))
181        self.assertIs(run(job.stop()), ending)  # a reached verdict is not re-run
182        run(job)
183        self.assertIn(job.id, plugin.retained)
184
185    def test_surviving_group_stays_owned_and_is_surfaced(self):
186        """A group that outlives KILL stays owned and is reported, not forgotten."""
187        real = (
188            albedo_proc.alive,
189            albedo_proc.current,
190            plugin.SHUTDOWN_TERM,
191            plugin.SHUTDOWN_KILL,
192            quick(),
193        )
194
195        def still_alive(group: albedo_proc.Group) -> bool:
196            return True
197
198        setattr(albedo_proc, "alive", still_alive)
199        setattr(albedo_proc, "current", still_alive)
200        plugin.SHUTDOWN_TERM, plugin.SHUTDOWN_KILL = 0.01, 0.02
201        try:
202            job = start("true")
203            wait(job)
204            self.assertFalse(job.termination.gone)
205            self.assertIn(job.id, plugin.active)  # surviving work stays owned
206            self.assertNotIn(job.id, plugin.retained)
207            self.assertIn("cleanup failed", job.tail())
208            run(plugin.close())
209        finally:
210            (
211                albedo_proc.alive,
212                albedo_proc.current,
213                plugin.SHUTDOWN_TERM,
214                plugin.SHUTDOWN_KILL,
215                albedo_proc.terminate,
216            ) = real
217        cleanup = [event for event in EVENTS if event.get("type") == "cleanup"]
218        self.assertEqual(len(cleanup), 1)
219        self.assertTrue(
220            any(
221                "SURVIVED" in failure
222                for failure in cast(list[str], cleanup[0]["failures"])
223            )
224        )
225
226    def test_ladder_reports_gone_and_refuses_reuse(self):
227        self.assertFalse(albedo_proc.alive(albedo_proc.Group(999_999)))
228        ending = run(albedo_proc.terminate([albedo_proc.Group(999_999)]))
229        self.assertEqual(ending[0].signals, ())
230        self.assertTrue(ending[0].gone)
231        leader = run(
232            asyncio.create_subprocess_shell("sleep 30", start_new_session=True)
233        )
234        try:
235            # Linux and Darwin both give a process an identity beyond its pid.
236            self.assertIsNotNone(albedo_proc.leader_token(leader.pid))
237            group = albedo_proc.Group(leader.pid, "not-the-leader-we-created")
238            self.assertFalse(albedo_proc.alive(group))
239            self.assertTrue(present(leader.pid))  # and it was left alone
240        finally:
241            run(albedo_proc.terminate([albedo_proc.Group(leader.pid)]))
242            run(leader.wait())
243
244    def test_stop_before_spawn_and_immediate_cancellation_keep_ownership(self):
245        async def check():
246            job = start("sleep 30")
247            ending = await job.stop()
248            await job
249            self.assertTrue(ending.gone)
250            if job.group is not None:
251                self.assertFalse(present(job.group.pgid))
252            self.assertNotIn(job.id, plugin.active)
253            cancelled = start("sleep 30")
254            cancelled.task.cancel()
255            await asyncio.sleep(0)
256            await cancelled.stop()
257            if cancelled.group is not None:
258                self.assertFalse(present(cancelled.group.pgid))
259            self.assertNotIn(cancelled.id, plugin.active)
260
261        run(check())
262
263    def test_retrying_failed_cleanup_releases_quota(self):
264        from unittest.mock import patch
265
266        async def refused(groups, *args, **kwargs):
267            return [
268                albedo_proc.Termination(group.pgid, (), False, ("fixture refusal",))
269                for group in groups
270            ]
271
272        job = start("true")
273        with patch.object(albedo_proc, "terminate", side_effect=refused):
274            wait(job)
275        self.assertIn(job.id, plugin.active)
276        self.assertTrue(run(job.stop()).gone)
277        self.assertNotIn(job.id, plugin.active)
278        self.assertTrue(cast(dict[str, object], EVENTS[-1]["cleanup"])["gone"])
279
280    def test_helper_watchdog_terminates_native_code_even_if_alarm_was_ignored(self):
281        import signal
282        import subprocess
283
284        helper = os.path.join(ROOT, "priv", "python", "albedo_signal.py")
285        code = (
286            f"import sys; sys.path.insert(0, {os.path.dirname(helper)!r})\n"
287            "import ctypes, runpy, signal, albedo_proc\n"
288            "signal.signal(signal.SIGALRM, signal.SIG_IGN)\n"
289            # the helper's alarm, fired sooner: it still has to arm it and
290            # restore the default action, or the sleep outlasts the timeout
291            "signal.alarm = lambda seconds: signal.setitimer(signal.ITIMER_REAL, 0.2)\n"
292            "async def stuck(*args, **kwargs):\n    ctypes.PyDLL(None).sleep(30)\n"
293            "albedo_proc.terminate = stuck\n"
294            "sys.argv = ['helper', '{\"targets\": []}']\n"
295            f"runpy.run_path({helper!r}, run_name='__main__')\n"
296        )
297        started = time.monotonic()
298        result = subprocess.run(
299            [sys.executable, "-c", code], capture_output=True, timeout=5
300        )
301        self.assertEqual(result.returncode, -signal.SIGALRM, result.stderr)
302        self.assertLess(time.monotonic() - started, 5)
303
304    def test_start_token_reads_a_kernel_stat_line(self):
305        line = b"1234 (some cmd) S 1 1234 1234 0 -1 4194560 1 0 0 0 5 3 0 0 20 0 3 0 987654 1 2"
306        self.assertEqual(albedo_proc.start_token(line), "987654")
307        self.assertIsNone(albedo_proc.start_token(b"nonsense"))
308
309
310if __name__ == "__main__":
311    unittest.main(verbosity=2)