Source code for edri.switch.receiver

from logging import getLogger
from queue import Queue
from threading import Thread
from uuid import UUID

from edri.protocol.delivery import decode_delivery
from edri.protocol.errors import ProtocolError
from edri.protocol.frame import FrameType
from edri.switch.connection import Connection
from edri.switch.forwarder import ForwardMessage, ForwardType




class Receiver(Thread):
    def __init__(self, forwarder_queue: Queue[ForwardMessage], router_id: UUID,
                 connection: Connection, connections: dict[UUID, Connection]) -> None:
        super().__init__(daemon=True, name=f"switch-receiver-{router_id}")
        self.forwarder_queue = forwarder_queue
        self.connection = connection
        self.connections = connections
        self.router_id = router_id
        self.logger = getLogger(f"{__name__}.{router_id}")



    def run(self) -> None:
        channel = self.connection.channel
        if channel is None:
            return
        try:
            while not self.connection.stop.is_set():
                frame = channel.receive()
                # An evicted session must not inject more events or ACKs.
                if self.connections.get(self.router_id) is not self.connection:
                    return
                if frame.type == FrameType.ACK:
                    if frame.payload:
                        raise ValueError("ACK frames must have an empty payload")
                    delivery_id = UUID(bytes=frame.target)
                    if not self.connection.mailbox.ack(delivery_id):
                        self.logger.debug("Stale/unknown ACK: %s", delivery_id)
                elif frame.type == FrameType.DEMANDS:
                    self.connection.demands = frame.payload
                    self.forwarder_queue.put(ForwardMessage(self.router_id, ForwardType.OTHERS, frame.payload))
                elif frame.type == FrameType.LAST_EVENTS:
                    self.forwarder_queue.put(ForwardMessage(self.router_id, ForwardType.OTHERS, frame.payload))
                elif frame.type == FrameType.EVENT:
                    delivery = decode_delivery(frame.payload)
                    self.forwarder_queue.put(ForwardMessage(
                        UUID(bytes=frame.target), ForwardType.SPECIFIC, delivery.payload,
                        self.router_id, delivery.id, self.connection,
                    ))
                else:
                    self.logger.error("Unexpected frame type: %s", frame.type.name)
                    return
        except (ProtocolError, ValueError) as e:
            self.logger.error("Invalid connection frame: %s", e)
        except (ConnectionError, OSError):
            self.logger.info("Router disconnected: %s", self.router_id)
        finally:
            self.connection.disconnect()