Source code for edri.abstract.manager.manager_priority_base

from abc import abstractmethod
from dataclasses import dataclass, field
from datetime import datetime
from time import sleep
from logging import Logger
from multiprocessing import Queue
from multiprocessing.connection import wait, Connection, Pipe
from queue import PriorityQueue, Empty
from threading import Thread, Event as EventThreading
from typing import Any, Type, Optional, cast

from edri.abstract import ManagerBase
from edri.dataclass.event import Event
from edri.abstract.manager.worker import Worker

__all__ = ["ManagerPriorityBase"]


@dataclass(order=True)
class QueueItem:
    priority: int
    item: Event = field(compare=False)




class ManagerPriorityBase(ManagerBase):
    """
    An extension of ManagerBase that introduces priority handling for events. This class
    manages events using a priority queue, ensuring that higher-priority events are processed
    before lower-priority ones.

    Attributes:
        event_queue (PriorityQueue[QueueItem]): A priority queue for managing events based on their priority.
        event_worker (Thread): A thread dedicated to handling command execution based on event priority.
        waking_up (tuple[Connection, Connection]): A pair of connections used to wake up the manager when a new worker starts.

    Methods:
        __init__: Initializes the ManagerPriorityBase with an optional router queue and logger.
        _prepare: Prepares the manager by initializing the event queue, event worker thread, and waking up pipe.
        additional_pipes: Property to include the waking_up pipe in the set of additional pipes the manager listens to.
        command_handler: Thread target function that processes events from the priority queue.
        get_priority: Abstract method that must be implemented to define the priority of events.
        run_resolver: Main event loop, enhanced to handle events based on priority.
        start_worker: Overrides the start_worker method to wake up the manager when a new worker is started.
    """

    # ``router_queue`` is annotated Any: the multiprocessing.Queue factory is
    # not subscriptable at runtime, and PEP 649 lazy annotation evaluation
    # (Python 3.14) evaluates __init__ signature annotations in
    # inspect.signature(); see manager_base.py for details.
    def __init__(self, router_queue: Optional[Any] = None, logger: Optional[Logger] = None,
                 from_time: Optional[datetime] = None, max_workers: int = -1) -> None:
        """
        Initializes the ManagerPriorityBase with an optional router queue and logger, setting up the foundation
        for event priority management.

        Parameters:
            router_queue (Optional[Queue[Event]]): The queue for communicating with the router.
            logger (Optional[Logger]): Logger instance for logging events.
            from_time (Optional[datetime]): Start time for filtering events. Default is None.
        """
        super().__init__(router_queue, logger, from_time, max_workers)
        self.event_queue: "PriorityQueue[QueueItem]"
        self.event_worker: Thread
        self.event_worker_stop: EventThreading
        self.waking_up: tuple[Connection, Connection]

    def __del__(self) -> None:
        if hasattr(self, "event_worker_stop") and self.event_worker_stop:
            self.event_worker_stop.set()



    @abstractmethod
    def get_priority(self, event_type: Type[Event]) -> int:
        """
        Abstract method that determines the priority of an event. Must be implemented by subclasses
        to define how event priorities are assigned.

        Parameters:
            event_type (Type[Event]): The class of the event for which the priority needs to be determined.

        Returns:
            int: The priority of the event.
        """
        pass


    def _prepare(self) -> None:
        """
        Prepares the manager by initializing the event queue, starting the event worker thread,
        and setting up the waking_up pipe used to signal new work.
        """
        super()._prepare()
        self.event_queue = PriorityQueue()
        self.event_worker_stop = EventThreading()
        self.event_worker = Thread(target=self.command_handler)
        self.event_worker.start()
        self.waking_up = Pipe(duplex=False)

    @property
    def additional_pipes(self) -> set[Connection]:
        """
        Includes the waking_up pipe in the set of additional pipes the manager listens to.

        Returns:
            set[Connection]: A set of additional pipes, including the waking_up pipe.
        """
        return {self.waking_up[0]}

    def _worker_runtime_started(self, worker_id: str, worker: Worker) -> None:
        if hasattr(self, "waking_up"):
            self.waking_up[1].send(None)

    def _worker_deadline_changed(self) -> None:
        if hasattr(self, "waking_up"):
            self.waking_up[1].send(None)



    def run_resolver(self) -> None:
        """
        The main event loop for the manager, enhanced to handle events based on priority. It listens on all pipes,
        including the waking_up pipe, to efficiently process and prioritize events.
        """
        self._resolver_running = True
        while self._resolver_running:
            pipes = self.get_pipes()
            if not pipes:
                self.workers._process_timeouts()
                sleep(self.workers._next_timeout(self.resolver_idle_delay) or 0.0)
                continue
            try:
                # multiprocessing.connection.wait() is typed to return a union list that
                # includes socket/int; only Connection pipes are waited on and resolved,
                # so the result is narrowed with an erased cast (no runtime change).
                active_pipes = cast(
                    list[Connection],
                    wait(pipes, timeout=self.workers._next_timeout(self.resolver_wait_timeout)),
                )
                worker_pipes = self.workers.pipes()
                ready_workers = [pipe for pipe in active_pipes if pipe in worker_pipes]
                for active_pipe in ready_workers:
                    self.resolve_non_router_pipe(active_pipe)
                self.workers._process_timeouts()
                for active_pipe in active_pipes:
                    if active_pipe in worker_pipes:
                        continue
                    if active_pipe == self.router_pipe:
                        event: Event = self.router_pipe.recv()
                        self.logger.debug("Received event from router: %s", event)
                        priority = self.get_priority(event.__class__)
                        self.event_queue.put(QueueItem(priority, event))
                    elif active_pipe == self.waking_up[0]:
                        self.waking_up[0].recv()
                        self.logger.debug("New pipes loaded")
                        continue
                    else:
                        self.resolve_non_router_pipe(active_pipe)
                    self.workers._process_timeouts()
            except KeyboardInterrupt:
                break
        self._resolver_running = False




    def command_handler(self) -> None:
        """
         The thread target function that processes events from the priority queue, ensuring that events are handled
         in order of their priority.
         """
        while not self.event_worker_stop.is_set():
            try:
                event = self.event_queue.get(timeout=1).item
            except Empty:
                continue
            self.resolve(event)