Source code for edri.abstract.worker.worker

from logging import getLogger
from multiprocessing.connection import Connection
from pathlib import Path
from time import monotonic
from traceback import format_exc
from typing import Any, Generic, TypeVar
from abc import abstractmethod, ABC

from edri.dataclass.event import Event
from edri.dataclass.worker import WorkerExit
from edri.events.edri.manager import WorkerQuit, WorkerStarted

__all__ = ["Worker", "WorkerCancelled", "T"]


T = TypeVar("T", bound=Event)




class WorkerCancelled(Exception):
    """Internal cooperative-cancellation signal raised inside a worker."""





class Worker(Generic[T], ABC):
    """
    A generic worker class designed to communicate with a manager through events.

    Attributes:
        _name (str): The name of the worker, used for logging purposes.
        _manager_pipe (Connection): The communication pipe to the manager.
        event (EventType): The initial event associated with the worker.
        logger (Logger): Logger instance for logging messages.

    Methods:
        __init__: Initializes a new Worker instance.
        do: Abstract method defining the main logic of the worker; to be implemented by subclasses.
        event_send: Sends an event to the manager.
        event_receive: Receives an event from the manager, blocking if no event is available.
        run: Entry point for the worker's execution, wrapping the `do` method with error handling and cleanup.
    """
    def __init__(self, pipe: Connection, event: T | None, name: str, *args: Any, **kwargs: Any) -> None:
        """
        Initializes a new Worker instance.

        Parameters:
            pipe (Connection): The communication pipe to the manager.
            event (EventType): The initial event associated with this worker.
            name (str): A name for the worker, used for logging.
            *args: Extra positional arguments for a concrete subclass's own
                __init__ override; ignored here. Accepted so WorkerThread/
                WorkerProcess subclasses that don't override __init__ can still
                be started via WorkerController.start(event, cls, *args,
                **kwargs) without this passthrough call raising TypeError.
            **kwargs: Extra keyword arguments, same rationale as *args.
        """
        self._manager_pipe = pipe
        self._name = name
        self.event = event
        self.logger = getLogger(self._name)
        self._cancel_event: Any = None
        # Bound by the controller for thread runtimes only; see
        # _bind_router_queue. Stays None in a process worker, which runs in a
        # different address space and must use event_send() instead.
        self._router_queue: Any = None

    def _bind_lifecycle(self, cancel_event: Any) -> None:
        """Bind controller-owned primitives without changing subclass constructors."""
        self._cancel_event = cancel_event

    def _bind_router_queue(self, router_queue: Any) -> None:
        """Give a *thread* runtime direct access to its manager's router queue.

        Publishing through this queue is not the same thing as
        :meth:`event_send`, which writes to the worker-manager pipe and is
        answered by the manager. It exists for side-channel traffic a worker
        produces for the rest of the system rather than for its own manager --
        the Pydantic AI observability action stream is the first such producer.

        Bound after construction and before ``start()``, so it never changes a
        subclass's constructor contract. It is deliberately *not* bound for
        process runtimes: the manager's queue object does not cross a process
        boundary that way, and a worker must be able to see that it has no
        binding rather than discover it through a broken send.
        """
        self._router_queue = router_queue

    def _request_cancel(self) -> None:
        if self._cancel_event is not None:
            self._cancel_event.set()

    @property
    def cancel_requested(self) -> bool:
        return self._cancel_event is not None and self._cancel_event.is_set()



    def wait_cancelled(self, timeout: float | None = None) -> bool:
        if self._cancel_event is None:
            return False
        return bool(self._cancel_event.wait(timeout))




    def raise_if_cancelled(self) -> None:
        if self.cancel_requested:
            raise WorkerCancelled




    def event_send(self, event: Event) -> None:
        """
        Sends an event to the manager.

        Parameters:
            event (Event): The event to be sent.
        """
        event._timing.stamp(self.__class__.__name__, "Sent")
        self._manager_pipe.send(event)
        self.logger.debug("Event was sent: %s", event)




    def event_receive(self) -> Event:
        """
        Receives an event from the manager. This method blocks until an event is available.

        Returns:
            Event: The received event.
        """
        event: Event = self._manager_pipe.recv()
        event._timing.stamp(self.__class__.__name__, "Received")
        return event


    @staticmethod
    def _fix_environment() -> None:
        import site, sys, os

        exe = Path(sys.executable)
        venv_root = exe.parent.parent


        version = f"python{sys.version_info.major}.{sys.version_info.minor}"
        site_packages = venv_root / "lib" / version / "site-packages"

        # prepend new site-packages:
        prev = set(sys.path)
        site.addsitedir(str(site_packages))
        sys.path[:] = [p for p in sys.path if p not in prev] + list(prev)

        os.environ["VIRTUAL_ENV"] = str(venv_root.resolve())
        sys.real_prefix = sys.prefix  # type: ignore[attr-defined]
        sys.prefix = str(venv_root)



    def setup(self) -> None:
        """
        Hook run once before :meth:`do`, inside the same error-handling scope.

        The default is a no-op. Override it for one-time preparation that must
        happen in the worker's own thread/process (building helpers, opening
        resources). A failure here is treated exactly like a failure in
        :meth:`do`: it is logged and the worker still sends ``WorkerQuit`` so the
        manager can reap it.
        """
        pass




    @abstractmethod
    def do(self) -> None:
        """
        Abstract method that should be implemented by subclasses to define the worker's main logic.
        """
        pass




    def run(self) -> None:
        """
        The main entry point for the worker's execution. It runs :meth:`setup`
        then :meth:`do`, handling any exceptions and ensuring cleanup and
        communication with the manager upon termination.
        """
        self._fix_environment()
        if self.event:
            self.event._timing.stamp(self.__class__.__name__, "Started")
        quit_event = WorkerQuit()
        try:
            self.event_send(WorkerStarted(started=monotonic()))
            self.raise_if_cancelled()
            self.setup()
            self.do()
        except WorkerCancelled:
            quit_event.exit = WorkerExit.CANCELLED
        except BaseException as e:
            self.logger.error("Worker %s was closed unexpectedly", self._name, exc_info=e)
            quit_event.exit = WorkerExit.RAISED
            quit_event.exception_type = type(e).__qualname__
            quit_event.exception_message = str(e)
            quit_event.traceback = format_exc()
        finally:
            if quit_event.exit is WorkerExit.RETURNED and self.cancel_requested:
                quit_event.exit = WorkerExit.CANCELLED
            self.event_send(quit_event)
            self._manager_pipe.close()