amq-adapter-python/amqp/rabbitmq/rabbit_mq_producer.py

import logging
import time

import pika

from amqp.adapter.amq_message_factory import AMQMessageFactory
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import AMQMessage, AMQRoute
from amqp.router.router_producer import RouterProducer


class RabbitMQProducer:
    PREFETCH_COUNT = 1
    THRESHOLD = 1000
    WINDOW = 5000  # 5 seconds

    def __init__(self, amq_configuration: AMQConfiguration, router: RouterProducer = None):
        self.channel = None
        self.connection = None
        self.request_for_window_duration = None
        self.amq_configuration = amq_configuration
        logging.info("*******************************************")
        logging.info("******** RabbitMQProducer.init() ********** %s - %s - %s",
                     self.amq_configuration.dispatch.amq_host,
                     self.amq_configuration.dispatch.use_confirms,
                     self.WINDOW)
        logging.info("*******************************************")
        self.router = router or RouterProducer(amq_configuration)

        # this.meterRegistry = new PrometheusMeterRegistry(PrometheusConfig.DEFAULT);

        self.setup_amq_channel()
        self.router.init_routing_paths(self.channel)

        #  self.backpressure_per_exchange = defaultdict(lambda: Gauge(
        #    f"amqp.adapter.exchange.{exchange}", registry=CollectorRegistry()))

        self.router.set_route_added_notifier(
            lambda route: logging.info(" [DBG] ---------- RouteAdded, setting up metrics: amqp.adapter.exchange.%s",
                                       route.exchange))

        self.request_for_window_duration = self.THRESHOLD // 2
        self.window_started_at = time.time()
        self.on_next_timestamp = self.window_started_at
        self.fulfilled = 0

    def cleanup(self):
        if self.connection:
            try:
                self.connection.close()
            except Exception as e:
                logging.error("RabbitMQConsumer.PreDestroy cleanup - connection close error: %s", e)

    def setup_amq_channel(self):
        """
        Set up the RabbitMQ connection
        :return: nothing, it applies changes to instance variables
        """
        factory = pika.ConnectionParameters(
            host=self.amq_configuration.dispatch.amq_host,
            username=self.amq_configuration.dispatch.rabbit_mq_user,
            password=self.amq_configuration.dispatch.rabbit_mq_password
        )
        self.connection = pika.BlockingConnection(factory)
        logging.info("RabbitMQProducer.setupAMQ, got connection, open=%s", self.connection.is_open)
        self.connection.add_on_connection_blocked_callback(self.blocked_listener)
        self.connection.add_on_connection_unblocked_callback(self.unblocked_listener)
        self.channel = self.connection.channel()
        self.channel.basic_qos(prefetch_count=self.PREFETCH_COUNT)
        if self.amq_configuration.dispatch.use_confirms():
            self.setup_local_ack()
        self.channel.add_on_return_callback(self.return_listener)

    def send_message(self, message: AMQMessage):
        """
        The API routine that sends a datagram message (AMQMessage) to the destination. The Destination
        is determined from the AMQMessage metadata section, using `domain` and `path` as key.
        :param message: message to send
        :return: nothing (void) as submitting message to RabbitMQ exchange does not return anything
         meaningful. There will be asynch ACK later that will confirm the message was sent to RMQ.
        """
        try:
            route = self.router.find_route_by_message(message)
            if route:
                if route.timeout < 0:
                    # No response expected
                    self.channel.basic_publish(
                        route.exchange,
                        self.router.get_routing_key(message),  # context path is a routing key
                        properties=pika.BasicProperties(delivery_mode=2),
                        body=AMQMessageFactory.serialize(message)
                    )
                    logging.info(f" [x] basicPublish (no-reply), messageId: {message.id}, to: {route}")
                else:
                    # This is RPC call, there should be response
                    props = pika.BasicProperties(
                        correlation_id=message.id,
                        reply_to=self.router.get_reply_to_queue_name()
                    )
                    self.channel.basic_publish(
                        route.exchange,
                        self.router.get_routing_key(message),  # context path is a routing key
                        properties=props,
                        body=AMQMessageFactory.serialize(message)
                    )
                    logging.info(" [x] basicPublish (RPC call), messageId: %s, to: %s",
                                 message.id.as_string(), route.as_string())
            else:
                logging.warning("Message not sent, no route for {}/{}", message.domain, message.path)

        except Exception as err:
            logging.error("Send message failed: %s", err)

    def request_routes(self):
        self.router.request_routes_from_adapters()

    def setup_local_ack(self):
        self.channel.confirm_delivery()

    def find_route(self, domain: str, path: str) -> AMQRoute:
        return self.router.find_route(domain, path)

    def return_listener(self, return_message):
        logging.error(
            f" [E] ****** (NOT DELIVERED) ID={return_message.properties.message_id}, "
            f"Exchg=[{return_message.exchange}], Key=[{return_message.routing_key}], "
            f"replyTxt={return_message.reply_text}, length={len(return_message.body) if return_message.body else 0}"
        )

    def blocked_listener(self, reason: str):
        logging.error(f" [E] CHANNEL BLOCKED, reason={reason}")

    def unblocked_listener(self):
        logging.info(" [*] CHANNEL UNBLOCKED")