Class Queue

fastdeploy/inter_communicator/fmq.py:196–268

class Queue(BaseComponent):
    def __init__(self, context, name: str, role: str = "producer"):
        endpoint = EndpointManager.get_endpoint(name)
        super().__init__(context, endpoint)

        self.name = name
        self.role = Role(role)
        self.copy = endpoint.copy
        self.socket_conf = EndpointManager.config.socket_config
        self._msg_id = 0

        full_ep = f"{endpoint.protocol}://{endpoint.address}"

        self.socket = self.context.socket(zmq.PUSH if self.role == Role.PRODUCER else zmq.PULL)
        self.socket_conf.apply(self.socket, self.role == Role.PRODUCER)

        if self.role == Role.PRODUCER:
            self.socket.connect(full_ep)
        else:
            self.socket.bind(full_ep)

        fmq_logger.info(f"Queue {name}({role}) initialized on {full_ep}")

    async def put(self, data: Any, shm_threshold: int = 1024 * 1024):
        """
        Send data to the queue.

        Args:
            data: The data to send. Can be any serializable object or bytes.
            shm_threshold: Size threshold in bytes. If the data is of type bytes and its size is
                greater than or equal to this threshold, shared memory will be used to send the message.
                Default is 1MB (1024 * 1024 bytes).

        Raises:
            PermissionError: If called by a non-producer role.
        """
        if self.role != Role.PRODUCER:
            raise PermissionError("Only producers can send messages.")

        desc = None
        payload = data

        if isinstance(data, bytes) and len(data) >= shm_threshold:
            desc = Descriptor.create(data)
            payload = None

        msg = Message(msg_id=self._msg_id, payload=payload, descriptor=desc)
        raw = msg.serialize()

        async with self.lock:
            await self.socket.send(raw, copy=self.copy)
            self._msg_id += 1

    async def get(self, timeout: int = None) -> Optional[Message]:
        # Receive data from queue
        if self.role != Role.CONSUMER:
            raise PermissionError("Only consumers can get messages.")