ringo-task-queue 0.1.0.dev8__py3-none-macosx_13_0_arm64.whl

ringo_task_queue/_supervisor_runtime.py

@@ -0,0 +1,420 @@
+"""Child-side lifecycle bridge, instantiated only by the supervisor bootstrap."""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import logging
+import os
+import signal
+import threading
+import time
+import weakref
+from typing import TYPE_CHECKING, Any
+
+from ._supervisor_ipc import HEALTH_INTERVAL, Channel
+from ._supervisor_policy import SupervisorPolicy, TaskTimeouts
+from .errors import RingoError
+
+if TYPE_CHECKING:
+    from .worker import Worker
+
+_LOG = logging.getLogger(__name__)
+_HEALTH_INTERVAL = HEALTH_INTERVAL
+_MAX_HEALTH_ITEMS = 128
+_CURRENT: ChildRuntime | None = None
+
+
+def current_runtime() -> ChildRuntime | None:
+    return _CURRENT
+
+
+class ChildRuntime:
+    """An isolated generation's private control and health channel.
+
+    The channel thread does not claim that business/event loops are healthy.
+    Loop ticks are generated *on those loops* and remain stale if they block.
+    """
+
+    def __init__(
+        self, channel: Channel, policy: SupervisorPolicy, task_timeouts: TaskTimeouts
+    ) -> None:
+        self.channel = channel
+        self.policy = policy
+        self.task_timeouts = task_timeouts
+        self._lock = threading.Lock()
+        self._loops: dict[asyncio.AbstractEventLoop, dict[str, Any]] = {}
+        self._workers: weakref.WeakSet[Worker] = weakref.WeakSet()
+        self._worker_snapshots: dict[str, dict[str, Any]] = {}
+        # Strongly retain a single-use Worker's identity until its ordered stop
+        # frame is sent; a pending tombstone must never outlive/reuse that id.
+        self._pending_worker_stops: dict[str, Worker] = {}
+        self._stopping = False
+        self._closed = threading.Event()
+        self._restart_requested = False
+        self._orphan_timer: threading.Timer | None = None
+        self._orphaned = False
+
+    def start(self) -> None:
+        global _CURRENT
+        if _CURRENT is not None:
+            raise RuntimeError("nested process supervision is not supported")
+        _CURRENT = self
+        threading.Thread(
+            target=self._control, name="ringo-supervisor-control", daemon=True
+        ).start()
+        threading.Thread(
+            target=self._report, name="ringo-supervisor-health", daemon=True
+        ).start()
+
+    def _loop_state(self) -> dict[str, Any]:
+        loop = asyncio.get_running_loop()
+        with self._lock:
+            state = self._loops.get(loop)
+            if state is None:
+                state = {
+                    "clients": weakref.WeakSet(),
+                    "stop": asyncio.Event(),
+                    "last_tick": time.monotonic(),
+                    "closed": False,
+                    "pulse": None,
+                    "drain": None,
+                }
+                self._loops[loop] = state
+                state["pulse"] = loop.call_soon(self._pulse, loop, state)
+            stopping = self._stopping
+        if stopping:
+            state["stop"].set()
+            self._schedule_drain(loop, state)
+        return state
+
+    def _pulse(self, loop: asyncio.AbstractEventLoop, state: dict[str, Any]) -> None:
+        # A timer, not a permanent Task: a sync entrypoint may legitimately own
+        # run_until_complete()/loop.close() without SDK-created pending Tasks.
+        with self._lock:
+            if self._closed.is_set():
+                state["closed"] = True
+                return
+            state["last_tick"] = time.monotonic()
+        state["pulse"] = loop.call_later(_HEALTH_INTERVAL, self._pulse, loop, state)
+
+    def register_client(self, client: Any) -> None:
+        state = self._loop_state()
+        with self._lock:
+            if self._stopping:
+                raise RingoError(
+                    "supervised generation is stopping; client start is fenced"
+                )
+            state["clients"].add(client)
+
+    def register_worker(self, worker: Worker) -> None:
+        self._loop_state()
+        with self._lock:
+            if self._stopping:
+                raise RingoError(
+                    "supervised generation is stopping; Worker admission is fenced"
+                )
+            self._workers.add(worker)
+        # Per-Worker configuration, never per-task IPC. Announce before its
+        # Work stream can admit a handler, including explicit unlimited types.
+        timeouts = worker._task_timeouts or TaskTimeouts()
+        scope = worker._task_types
+        budgets = (
+            [timeouts.for_type(kind) for kind in scope]
+            if scope
+            else [timeouts.default, *timeouts.by_type.values()]
+        )
+        bounded = bool(budgets) and all(value is not None for value in budgets)
+        fallback = None
+        if bounded:
+            fallback = (
+                max(value.total_seconds() for value in budgets if value is not None)
+                + timeouts.cancel_grace.total_seconds()
+                + 3 * _HEALTH_INTERVAL
+            )
+        try:
+            self.channel.send(
+                {
+                    "kind": "worker_budget",
+                    "worker_id": str(id(worker)),
+                    "loop_id": str(id(asyncio.get_running_loop())),
+                    "fallback_timeout": fallback,
+                }
+            )
+        except OSError:
+            self._parent_lost()
+            raise RingoError(
+                "supervisor parent control channel is unavailable"
+            ) from None
+
+    def update_worker(self, worker_id: str, snapshot: dict[str, Any]) -> None:
+        with self._lock:
+            # A bounded snapshot is diagnostic only, never task data authority.
+            if (
+                worker_id in self._worker_snapshots
+                or len(self._worker_snapshots) < _MAX_HEALTH_ITEMS
+            ):
+                self._worker_snapshots[worker_id] = dict(snapshot)
+
+    @staticmethod
+    def _worker_busy(worker: Worker) -> bool:
+        watchdog = worker._watchdog
+        return bool(watchdog is not None and watchdog.executions) or any(
+            not task.done() for task in tuple(worker._handler_tasks)
+        )
+
+    def unregister_worker(self, worker: Worker) -> None:
+        # Public worker_id is an audit label, not a process-unique identity.
+        worker_id = str(id(worker))
+        with self._lock:
+            self._worker_snapshots.pop(worker_id, None)
+            still_running = any(
+                candidate is worker and self._worker_busy(candidate)
+                for candidate in self._workers
+            )
+        if still_running:
+            with self._lock:
+                self._pending_worker_stops[worker_id] = worker
+        else:
+            with self._lock:
+                self._workers.discard(worker)
+            try:
+                self.channel.send({"kind": "worker_stopped", "worker_id": worker_id})
+            except OSError:
+                pass
+
+    def request_restart(
+        self, reason: str, detail: dict[str, Any] | None = None
+    ) -> None:
+        with self._lock:
+            if self._stopping or self._restart_requested:
+                return
+            self._restart_requested = True
+        try:
+            self.channel.send(
+                {"kind": "restart", "reason": reason[:128], "detail": detail or {}}
+            )
+        except (OSError, ValueError):
+            self._parent_lost()
+
+    def _report(self) -> None:
+        while not self._closed.wait(_HEALTH_INTERVAL):
+            with self._lock:
+                for loop in tuple(self._loops):
+                    if loop.is_closed():
+                        state = self._loops.pop(loop)
+                        state["closed"] = True
+                        if state.get("pulse") is not None:
+                            state["pulse"].cancel()
+                loops = [
+                    {"id": str(id(loop)), "last_tick": state["last_tick"]}
+                    for loop, state in self._loops.items()
+                    if not state["closed"] and loop.is_running()
+                ][:_MAX_HEALTH_ITEMS]
+                workers = list(self._worker_snapshots.values())
+                business_workers = list(self._workers)
+                stopping = self._stopping
+            executions = []
+            for worker in business_workers:
+                watchdog = worker._watchdog
+                if watchdog is None:
+                    continue
+                try:
+                    external = watchdog.external_snapshot()
+                except (AttributeError, RuntimeError, TypeError, ValueError):
+                    external = {"unknown": True}
+                if external is not None:
+                    executions.append({**external, "sampled_at": time.monotonic()})
+            # Prioritize actual earliest deadlines, bounding the health frame
+            # independently of the number of Workers (not a per-task channel).
+            executions.sort(key=lambda item: item.get("deadline", float("inf")))
+            executions = executions[:_MAX_HEALTH_ITEMS]
+            with self._lock:
+                pending_stops = self._pending_worker_stops.copy()
+            finished = {
+                worker_id
+                for worker_id, retiring in pending_stops.items()
+                if not self._worker_busy(retiring)
+            }
+            with self._lock:
+                for worker_id in finished:
+                    self._pending_worker_stops.pop(worker_id, None)
+                    self._workers.discard(pending_stops[worker_id])
+            try:
+                for worker_id in finished:
+                    self.channel.send(
+                        {"kind": "worker_stopped", "worker_id": worker_id}
+                    )
+                self.channel.send(
+                    {
+                        "kind": "health",
+                        "loops": loops,
+                        "workers": workers,
+                        "executions": executions,
+                        "stopping": stopping,
+                    }
+                )
+            except (OSError, ValueError):
+                self._parent_lost()
+                return
+            # Do not let a WeakSet snapshot retain retired Workers forever or
+            # release a pending identity before its stop has reached the wire.
+            if business_workers:
+                del worker, watchdog
+            del business_workers, pending_stops
+
+    def _control(self) -> None:
+        try:
+            while not self._closed.is_set():
+                message = self.channel.receive()
+                if message is None:
+                    self._parent_lost()
+                    return
+                if message.get("kind") == "stop":
+                    self._request_stop()
+        except (OSError, ValueError):
+            if not self._closed.is_set():
+                self._parent_lost()
+
+    def _request_stop(self) -> None:
+        with self._lock:
+            self._stopping = True
+            states = list(self._loops.items())
+        for loop, state in states:
+            if loop.is_closed():
+                continue
+            try:
+                loop.call_soon_threadsafe(state["stop"].set)
+                loop.call_soon_threadsafe(self._schedule_drain, loop, state)
+            except RuntimeError:
+                pass
+
+    def _schedule_drain(
+        self, loop: asyncio.AbstractEventLoop, state: dict[str, Any]
+    ) -> None:
+        if state["drain"] is None and not loop.is_closed():
+            state["drain"] = loop.create_task(
+                self._drain_loop(state), name="ringo-supervisor-drain"
+            )
+
+    async def _drain_loop(self, state: dict[str, Any]) -> None:
+        with self._lock:
+            clients = list(state["clients"])
+        # Start all Worker close fences in one event-loop turn before waiting.
+        # Client.close otherwise drains its Workers sequentially.
+        workers = {worker for client in clients for worker in list(client._workers)}
+        if workers:
+            await asyncio.gather(
+                *(worker.close() for worker in workers), return_exceptions=True
+            )
+        if clients:
+            await asyncio.gather(
+                *(client.close() for client in clients), return_exceptions=True
+            )
+
+    def _parent_lost(self) -> None:
+        with self._lock:
+            if self._closed.is_set() or self._orphan_timer is not None:
+                return
+            self._orphaned = True
+            timer = threading.Timer(
+                max(0.05, self.policy.shutdown_grace.total_seconds()), self._orphan_exit
+            )
+            timer.daemon = True
+            self._orphan_timer = timer
+        self._request_stop()
+        timer.start()
+
+    @staticmethod
+    def _orphan_exit() -> None:
+        # POSIX children are bootstrapped in a dedicated session. Shared-local
+        # daemons start in another session and are deliberately not signalled.
+        if os.name != "nt":
+            try:
+                os.killpg(os.getpgrp(), signal.SIGKILL)  # type: ignore[attr-defined]
+            except OSError:
+                pass
+        os._exit(1)
+
+    async def run_awaitable(self, awaitable: Any) -> Any:
+        state = self._loop_state()
+
+        async def application():
+            # SystemExit escaping an asyncio Task has special loop-abort
+            # behavior. Normalize it inside the application boundary so an
+            # intentional exit(0) cannot race a spurious crash notification.
+            try:
+                return await awaitable
+            except SystemExit as exc:
+                return (
+                    exc.code
+                    if isinstance(exc.code, int)
+                    else (0 if exc.code is None else 1)
+                )
+
+        main_task = asyncio.create_task(application())
+        stopped = asyncio.create_task(state["stop"].wait())
+        try:
+            done, _ = await asyncio.wait(
+                (main_task, stopped), return_when=asyncio.FIRST_COMPLETED
+            )
+            if stopped in done or self._stopping:
+                # Even if main completed in the same turn, honor the already
+                # started generation-wide drain before asyncio.run shuts down.
+                self._schedule_drain(asyncio.get_running_loop(), state)
+                await state["drain"]
+                if not main_task.done():
+                    main_task.cancel()
+                    await asyncio.gather(main_task, return_exceptions=True)
+                    result = None
+                elif main_task.cancelled():
+                    result = None
+                else:
+                    result = await main_task
+            else:
+                result = await main_task
+            # Bound asyncio.run's shutdown too: leftover application tasks can
+            # swallow cancellation after the entrypoint has already returned.
+            code = (
+                result
+                if isinstance(result, int) and not isinstance(result, bool)
+                else 0
+            )
+            self.channel.send({"kind": "entrypoint_done", "exit_code": code})
+            return result
+        except BaseException:
+            try:
+                self.channel.send({"kind": "entrypoint_failed"})
+            except OSError:
+                pass
+            raise
+        finally:
+            stopped.cancel()
+            await asyncio.gather(stopped, return_exceptions=True)
+
+    def invoke(self, function: Any) -> int:
+        if inspect.iscoroutinefunction(function):
+            # Construct the coroutine once; no recursive launcher or pickling.
+            result = asyncio.run(self.run_awaitable(function()))
+        else:
+            # A synchronous entrypoint may itself call asyncio.run(). SDK
+            # client registration attaches lifecycle hooks to that actual loop.
+            result = function()
+            if inspect.isawaitable(result):
+                result = asyncio.run(self.run_awaitable(result))
+        code = result if isinstance(result, int) and not isinstance(result, bool) else 0
+        self.channel.send({"kind": "entrypoint_done", "exit_code": code})
+        return code
+
+    def close(self) -> None:
+        global _CURRENT
+        self._closed.set()
+        if self._orphan_timer is not None:
+            self._orphan_timer.cancel()
+        if _CURRENT is self:
+            _CURRENT = None
+        if self._orphaned:
+            # A cooperative main can finish *before* the orphan timer fires.
+            # Do not let its exit cancel the only remaining POSIX tree owner.
+            self._orphan_exit()
+        self.channel.close()