Source code for edri.switch.forwarder

from dataclasses import dataclass
from enum import IntEnum
from logging import getLogger
from queue import Queue
from threading import Thread
from time import monotonic
from uuid import UUID

from edri.protocol.frame import ZERO_TARGET, FrameType
from edri.switch.connection import Connection




class ForwardType(IntEnum):
    SPECIFIC = 1
    OTHERS = 2





@dataclass(frozen=True)
class ForwardMessage:
    router_id: UUID
    kind: ForwardType
    payload: bytes
    source: UUID | None = None
    delivery_id: UUID | None = None
    source_connection: Connection | None = None





class Forwarder(Thread):
    def __init__(self, forwarder_queue: Queue[ForwardMessage], connections: dict[UUID, Connection]) -> None:
        super().__init__(daemon=True)
        self.forwarder_queue = forwarder_queue
        self.connections = connections
        self.logger = getLogger(__name__)
        # Bounded by the Router cache horizon: retries with the same transport
        # identifier must not enqueue a second copy after a lost source ACK.
        self._seen: dict[tuple[UUID, UUID], float] = {}
        self._last_cleanup = monotonic()

    def _accepted_before(self, source: UUID, delivery_id: UUID) -> bool:
        from edri.config.setting import CACHE_TIMEOUT
        now = monotonic()
        if now - self._last_cleanup > CACHE_TIMEOUT:
            self._seen = {k: t for k, t in self._seen.items() if now - t < CACHE_TIMEOUT * 2}
            self._last_cleanup = now
        key = (source, delivery_id)
        if key in self._seen:
            return True
        self._seen[key] = now
        return False



    def run(self) -> None:
        while True:
            item = self.forwarder_queue.get()
            if item.source is not None and item.delivery_id is not None:
                key = (item.source, item.delivery_id)
                if key in self._seen:
                    self._ack_source(item)
                    continue
            if item.kind == ForwardType.SPECIFIC:
                connection = self.connections.get(item.router_id)
                if connection is None:
                    self.logger.error("Unknown router id: %s; source delivery is unacknowledged", item.router_id)
                    continue
                connection.mailbox.put(item.payload)
            else:
                for router_id, connection in list(self.connections.items()):
                    if router_id != item.router_id:
                        connection.mailbox.put(item.payload)
            if item.source is not None and item.delivery_id is not None:
                self._accepted_before(item.source, item.delivery_id)
                self._ack_source(item)


    def _ack_source(self, item: ForwardMessage) -> None:
        assert item.source is not None and item.delivery_id is not None
        connection = item.source_connection
        if connection is None or self.connections.get(item.source) is not connection or connection.channel is None:
            return
        try:
            connection.channel.send(FrameType.ACK, item.delivery_id.bytes, b"")
        except Exception:
            self.logger.debug("Source ACK failed for %s; sender will retry", item.delivery_id, exc_info=True)