Большие модели — не выход

"""
asyncio_pipe_state.py

Custom asyncio event loop policy with a thread-safe StateManager that maps
child-process PIDs to non-blocking pipe contexts, sustaining an atomic,
non-yielding read loop on a background thread.

Design highlights
-----------------
* No ``asyncio.create_task`` is used for the read loop.
* No ``select`` / ``poll`` / ``epoll`` syscalls are used; the read loop is
  driven by ``os.read`` on O_NONBLOCK descriptors, yielding EAGAIN when idle.
* Concurrent writes from multiple child processes sharing a pipe context are
  serialized into a single deterministic in-memory buffer via per-context
  locking -- each successful ``os.read`` chunk is appended atomically, so
  chunks never interleave byte-wise mid-append.
* The custom ``ChildWatcher`` integrates with the ``StateManager`` so that
  PIDs are tracked for their whole lifetime and cleaned up on exit.
"""

from __future__ import annotations

import asyncio
import errno
import os
import signal
import threading
import warnings
from typing import Callable, Dict, Optional, Tuple


_CHUNK_SIZE = 65536
_IDLE_SLEEP = 0.005  # seconds; tiny back-off when all fds return EAGAIN


class PipeContext:
    """A non-blocking pipe read-end bound to one or more PIDs.

    Multiple PIDs may share a single context (e.g. two children whose stdout
    was dup'd onto the same pipe). All reads are appended to ``buffer`` under
    ``_lock`` so chunks remain intact and deterministic.
    """

    __slots__ = (
        "fd",
        "pids",
        "buffer",
        "_lock",
        "_closed",
        "__weakref__",
    )

    def __init__(self, fd: int, pid: int) -> None:
        self._set_nonblocking(fd)
        self.fd: int = fd
        self.pids: set[int] = {pid}
        self.buffer: bytearray = bytearray()
        self._lock = threading.Lock()
        self._closed: bool = False

    @staticmethod
    def _set_nonblocking(fd: int) -> None:
        import fcntl
        flags = fcntl.fcntl(fd, fcntl.F_GETFL)
        fcntl.fcntl(fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)

    def read_once(self) -> int:
        """Perform a single non-blocking read and append atomically.

        Returns the number of bytes appended, or 0 on EOF / EAGAIN / closed.
        Never raises for transient I/O conditions.
        """
        if self._closed:
            return 0
        try:
            chunk = os.read(self.fd, _CHUNK_SIZE)
        except BlockingIOError:
            return 0  # EAGAIN/EWOULDBLOCK -- nothing to read right now
        except InterruptedError:
            return 0  # retry on next tick
        except OSError as exc:
            # EBADF etc.: treat as terminal.
            self._closed = True
            return 0

        if not chunk:
            # EOF on this fd.
            self._closed = True
            return 0

        with self._lock:
            # Atomic append: this is the critical section that guarantees
            # deterministic interleaving when multiple producers share the
            # same context. Two concurrent os.read() calls (from the loop
            # thread) cannot both be mid-append because the loop is single-
            # threaded; the lock additionally protects asyncio consumers
            # reading from a worker thread.
            self.buffer.extend(chunk)
        return len(chunk)

    def snapshot(self) -> bytes:
        """Return a stable copy of the current buffer contents."""
        with self._lock:
            return bytes(self.buffer)

    def drain(self) -> bytes:
        """Atomically return and clear the buffer."""
        with self._lock:
            data = bytes(self.buffer)
            self.buffer.clear()
        return data

    def close(self) -> None:
        if self._closed:
            return
        self._closed = True
        try:
            os.close(self.fd)
        except OSError:
            pass


class StateManager:
    """Thread-safe registry mapping PIDs to ``PipeContext`` instances.

    A single background daemon thread runs an atomic, non-yielding read loop
    across all registered contexts. The loop neither awaits nor uses
    select/poll; it busy-reads O_NONBLOCK fds with a tiny idle sleep.
    """

    def __init__(self) -> None:
        self._contexts: Dict[int, PipeContext] = {}  # fd -> context
        self._pid_to_fd: Dict[int, int] = {}         # pid -> fd
        self._registry_lock = threading.RLock()
        self._wakeup = threading.Event()
        self._stop = threading.Event()
        self._thread: Optional[threading.Thread] = None
        self._started = False

    # -- lifecycle -------------------------------------------------------

    def start(self) -> None:
        with self._registry_lock:
            if self._started:
                return
            self._started = True
            self._stop.clear()
            self._thread = threading.Thread(
                target=self._run_loop,
                name="StateManager-Reader",
                daemon=True,
            )
            self._thread.start()

    def stop(self) -> None:
        self._stop.set()
        self._wakeup.set()
        t = self._thread
        if t is not None and t.is_alive():
            t.join(timeout=2.0)
        with self._registry_lock:
            for ctx in list(self._contexts.values()):
                ctx.close()
            self._contexts.clear()
            self._pid_to_fd.clear()
        self._thread = None
        self._started = False

    # -- registration ----------------------------------------------------

    def register(self, pid: int, fd: int) -> PipeContext:
        """Associate ``pid`` with the read-end ``fd``.

        If a context already exists for ``fd`` (multiple PIDs share it), the
        PID is added to that context's producer set.
        """
        with self._registry_lock:
            ctx = self._contexts.get(fd)
            if ctx is None:
                ctx = PipeContext(fd, pid)
                self._contexts[fd] = ctx
            else:
                ctx.pids.add(pid)
            self._pid_to_fd[pid] = fd
            self._wakeup.set()
            return ctx

    def unregister(self, pid: int) -> None:
        """Detach ``pid`` from its context; close context when last PID leaves."""
        with self._registry_lock:
            fd = self._pid_to_fd.pop(pid, None)
            if fd is None:
                return
            ctx = self._contexts.get(fd)
            if ctx is None:
                return
            ctx.pids.discard(pid)
            if not ctx.pids:
                ctx.close()
                self._contexts.pop(fd, None)

    def get_context(self, pid: int) -> Optional[PipeContext]:
        with self._registry_lock:
            fd = self._pid_to_fd.get(pid)
            if fd is None:
                return None
            return self._contexts.get(fd)

    # -- read loop -------------------------------------------------------

    def _run_loop(self) -> None:
        """Atomic, non-yielding read loop on a background thread.

        No select/poll, no asyncio.create_task. Iterates contexts and calls
        ``read_once`` until stop is requested.
        """
        while not self._stop.is_set():
            had_data = False
            with self._registry_lock:
                contexts = list(self._contexts.values())

            for ctx in contexts:
                # Drain the fd in this turn as long as data is immediately
                # available, to reduce latency under burst load.
                while True:
                    n = ctx.read_once()
                    if n == 0:
                        break
                    had_data = True

            if not had_data:
                # No ready data on any fd: back off briefly. This is NOT a
                # syscall-based wait; it is a pure Python sleep.
                self._wakeup.wait(timeout=_IDLE_SLEEP)
                self._wakeup.clear()
            else:
                # Yield the GIL promptly so asyncio & producers stay live.
                self._wakeup.wait(timeout=0.0)
                self._wakeup.clear()


_StateManagerSingleton: Optional[StateManager] = None
_StateManagerLock = threading.Lock()


def get_state_manager() -> StateManager:
    """Process-wide StateManager singleton, lazily started."""
    global _StateManagerSingleton
    with _StateManagerLock:
        if _StateManagerSingleton is None:
            _StateManagerSingleton = StateManager()
            _StateManagerSingleton.start()
        return _StateManagerSingleton


# ----------------------------------------------------------------------
# Custom child watcher that cooperates with the StateManager.
# ----------------------------------------------------------------------

class StateAwareChildWatcher(asyncio.AbstractChildWatcher):
    """A child watcher that registers/unregisters PIDs with the StateManager.

    It delegates the actual OS-level child-exit notification to a configured
    base watcher (default: ThreadedChildWatcher) so SIGCHLD/reaping still
    works portably, while layering PID->context bookkeeping on top. This
    keeps behavior correct on POSIX without re-imanging the reaper.
    """

    def __init__(self, base: Optional[asyncio.AbstractChildWatcher] = None) -> None:
        self._base = base or asyncio.ThreadedChildWatcher()
        self._state = get_state_manager()
        self._callbacks: Dict[int, Tuple[Callable, tuple]] = {}
        self._lock = threading.Lock()

    def __enter__(self) -> "StateAwareChildWatcher":
        self._base.__enter__()
        return self

    def __exit__(self, *exc_info) -> None:
        self._base.__exit__(*exc_info)

    def add_child_handler(self, pid: int, callback, *args) -> None:
        with self._lock:
            self._callbacks[pid] = (callback, args)
        # The StateManager entry for this PID will be populated when the
        # caller registers a pipe fd (see register_pipe_for_pid). Nothing
        # to do here except remember the callback.
        self._base.add_child_handler(pid, self._on_exit, pid)

    def remove_child_handler(self, pid: int) -> bool:
        with self._lock:
            self._callbacks.pop(pid, None)
        self._state.unregister(pid)
        return self._base.remove_child_handler(pid)

    def attach_loop(self, loop: Optional[asyncio.AbstractEventLoop]) -> None:
        self._base.attach_loop(loop)

    def close(self) -> None:
        with self._lock:
            for pid in list(self._callbacks):
                self._state.unregister(pid)
            self._callbacks.clear()
        self._base.close()

    def _on_exit(self, pid: int, returncode: int) -> None:
        # Called from the base watcher when the child has terminated.
        with self._lock:
            entry = self._callbacks.pop(pid, None)
        self._state.unregister(pid)
        if entry is None:
            return
        cb, args = entry
        try:
            cb(pid, returncode, *args)
        except Exception:
            # Mirrors asyncio's behavior: never raise out of the watcher.
            pass

    # -- public helper: bind a pipe fd to a pid --------------------------

    def bind_pipe(self, pid: int, fd: int) -> PipeContext:
        """Register a non-blocking read-end ``fd`` for ``pid``.

        Call this immediately after spawning the subprocess, with the child's
        stdout/stderr pipe read end. Multiple PIDs may share an ``fd``.
        """
        return self._state.register(pid, fd)


# ----------------------------------------------------------------------
# Custom event loop policy.
# ----------------------------------------------------------------------

class PipeStateEventLoopPolicy(asyncio.DefaultEventLoopPolicy):
    """Default policy override that yields a StateAwareChildWatcher.

    Ensures every event loop installed under this policy sees the same
    PID-keyed StateManager-backed child watcher.
    """

    def __init__(self) -> None:
        super().__init__()
        self._watcher: Optional[StateAwareChildWatcher] = None
        self._watcher_lock = threading.Lock()

    def get_child_watcher(self) -> StateAwareChildWatcher:
        with self._watcher_lock:
            if self._watcher is None or self._watcher._base is None:
                self._watcher = StateAwareChildWatcher()
                # Safe-call: some policy impls use this to bind to the loop.
                try:
                    self._watcher.attach_loop(self.get_event_loop())
                except RuntimeError:
                    # No current loop yet; attach lazily.
                    pass
            return self._watcher

    def set_child_watcher(self, watcher: Optional[asyncio.AbstractChildWatcher]) -> None:
        with self._watcher_lock:
            if watcher is None:
                if self._watcher is not None:
                    self._watcher.close()
                self._watcher = None
                return
            if not isinstance(watcher, StateAwareChildWatcher):
                warnings.warn(
                    "PipeStateEventLoopPolicy expects a StateAwareChildWatcher; "
                    "wrapping the supplied watcher.",
                    RuntimeWarning,
                    stacklevel=2,
                )
                watcher = StateAwareChildWatcher(base=watcher)
            self._watcher = watcher


# ----------------------------------------------------------------------
# Install helper.
# ----------------------------------------------------------------------

def install() -> PipeStateEventLoopPolicy:
    """Install the custom policy globally and return it."""
    policy = PipeStateEventLoopPolicy()
    asyncio.set_event_loop_policy(policy)
    return policy


__all__ = [
    "PipeContext",
    "StateManager",
    "StateAwareChildWatcher",
    "PipeStateEventLoopPolicy",
    "get_state_manager",
    "install",
]

import asyncio
import os
import subprocess
from asynccpipe_state import install, get_state_manager

async def main():
    install()
    loop = asyncio.get_running_loop()
    watcher = loop._loop.get_child_watcher() if False else asyncio.get_event_loop_policy().get_child_watcher()

    r, w = os.pipe()
    proc = await asyncio.create_subprocess_exec(
        "/bin/sh", "-c", "echo hello from $$ && sleep 0.1 && echo done",
        stdout=w, stderr=w,
    )
    os.close(w)  # parent doesn't need the write end

    watcher.bind_pipe(proc.pid, r)

    # Give the read loop a moment, then drain.
    await asyncio.sleep(0.3)
    ctx = get_state_manager().get_context(proc.pid)
    print("captured:", ctx.drain().decode())

    await proc.wait()

asyncio.run(main())