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

import logging
import pika
from pika.exchange_type import ExchangeType

from amqp.model.model import AMQRoute
from amqp.rabbitmq.rabbit_mq_producer import RabbitMQProducer
from amqp.router.router_base import AMQ_REPLY_TO_EXCHANGE, DLQ_EXCHANGE, AMQ_RPC_EXCHANGE
from amqp.router.router_consumer import RouterConsumer
from amqp.router.utils import sanitize_as_rabbitmq_name


class RabbitMQConsumer(RabbitMQProducer):
    PREFETCH_COUNT = 1
    DLQ_ARGS = {"x-dead-letter-exchange": DLQ_EXCHANGE}

    def __init__(self, configuration):
        super().__init__(configuration, RouterConsumer(configuration))

    def cleanup(self):
        logging.info(" [*] RabbitMQConsumer.cleanup()")
        self.router.cleanup()  # de-register route(s)
        if self.router.is_open(self.channel):
            try:
                self.channel.close()
            except (pika.connection.exceptions.AMQPError, pika.connection.exceptions.ChannelClosedByBroker) as e:
                logging.error(" [E] RabbitMQConsumer.PreDestroy cleanup - connection close error: %s", e)

    def register_routes(self):
        if self.router.is_open(self.channel):
            queue_args = self.DLQ_ARGS if self.amq_configuration.dispatch.use_dlq else None
            route_mapping = self.amq_configuration.amq_adapter.route_mapping
            logging.info(" [*] RabbitMQConsumer register routes, cfg mapping: [%s]", route_mapping)
            requested_routes = self.router.unwrap_route_list(route_mapping)
            for route in requested_routes:
                queue_name = route.queue() if route.queue() else self.router.get_consume_queue_name()
                logging.info(" [*] RabbitMQConsumer register route: [%s]", route)
                self.channel.queue_declare(queue_name, durable=True, exclusive=False,
                                           auto_delete=False, arguments=queue_args)
                exchange = route.exchange() if route.exchange() else AMQ_RPC_EXCHANGE
                self.channel.exchange_declare(exchange, exchange_type=pika.exchange_type.ExchangeType.topic)
                self.channel.queue_bind(queue_name, exchange, route.key())
                self.router.register_route(
                    AMQRoute(
                        route.component_name(), exchange, queue_name,
                        sanitize_as_rabbitmq_name(route.key()),
                        route.timeout(), route.valid_until()
                    )
                )
                logging.info(" [*] AMQ connection initialized for route: %s, listens on: %s/%s",
                             route, exchange, queue_name)
        else:
            logging.warning(" [W] Channel is null, cannot register routes.")

    def register_deliver_callback(self, callback_provider):
        try:
            for route in self.router.get_own_routes():
                self.channel.basic_consume(
                    queue=route.queue(),
                    on_message_callback=callback_provider(self.channel),
                    auto_ack=False
                )
                logging.info(" [*] Consumer Waiting for messages on: %s", route.queue())
        except pika.connection.exceptions.AMQPError as err:
            logging.error(" [E] Callback registration INCOMPLETE, %s", err)

    def register_rpc_handler(self, rpc_consumer):
        cm_reply_to_queue = self.router.get_reply_to_queue_name()
        self.channel.queue_declare(cm_reply_to_queue, durable=True, exclusive=False, auto_delete=False)
        self.channel.exchange_declare(
            AMQ_REPLY_TO_EXCHANGE, exchange_type=pika.exchange_type.ExchangeType.direct
        )
        self.channel.queue_bind(cm_reply_to_queue, AMQ_REPLY_TO_EXCHANGE, cm_reply_to_queue)
        self.channel.basic_consume(
            queue=cm_reply_to_queue,
            on_message_callback=rpc_consumer,
            auto_ack=False
        )

    """
    def limit_rate_inside_time_window(self, message: AMQMessage) -> int:
        next_batch_size = BackpressureSubscriber.NO_OP
        now = time.time_ns() // 1000000
        queue_size = self.getQueueSize(message)

        # 1. Backpressure? if unprocessed count exceeds threshold, wait for some time
        if queue_size > THRESHOLD:
            # TODO - signal queue backpressure explicitly to upstream/metrics
            # wait: estimated-per-message-processing-time * SQRT(over-the-limit)
            self.wait_for_milliseconds(
                (WINDOW / self.request_for_window_duration) * (queue_size - THRESHOLD) ** 0.5)
            queue_size = self.getQueueSize(message)  # refresh queue size after wait

        # 2. batch complete? determine next batch size and wait till window end. reset window.
        if self.fulfilled >= self.request_for_window_duration:
            next_batch_size = self.calculateNextBatchSize(queue_size)
            self.request_for_window_duration = next_batch_size
            remaining_time = (self.window_started_at + WINDOW) - now
            if remaining_time > 0:
                self.wait_for_milliseconds(remaining_time)
                queue_size = self.getQueueSize(message)  # refresh queue size after wait
            self.window_started_at = time.time_ns() // 1000000
            self.fulfilled = 0

        # 3. Window elapsed? Reset window & recalculate batch size
        if now >= self.window_started_at + WINDOW:
            if next_batch_size == BackpressureSubscriber.NO_OP:
                next_batch_size = self.calculateNextBatchSize(queue_size)
                self.request_for_window_duration = next_batch_size
            self.window_started_at = time.time_ns() // 1000000
            self.fulfilled = 0

        # 4. send the message from the current batch
        self.send_message(message)
        self.fulfilled += 1

        # 5. report and setup state for the next call
        end = time.time_ns() // 1000000
        self.on_next_timestamp = end
        return next_batch_size

    def calculateNextBatchSize(self, queue_size: int) -> int:
        next_batch_size = self.request_for_window_duration
        if queue_size <= THRESHOLD:
            # processing proceeds well, can increase the rate
            next_batch_size += self.subscriber.increaseRate(self.request_for_window_duration)
            logging.debug(f"INCREASE rate to {next_batch_size}. qSize={queue_size}")
        else:
            decrease = self.subscriber.backoff(queue_size - THRESHOLD)
            if self.request_for_window_duration > decrease:
                next_batch_size -= decrease
            logging.debug(f"DECREASE rate to {next_batch_size}. qSize={queue_size}")
        return next_batch_size
    """
    """
    def getQueueSize(self, message: AMQMessage) -> int:
        try:
            route_queue_name = self.router.findRoute(message)
            queue_name = route_queue_name.queue if route_queue_name else self.TASK_DISPATCH_QUEUE_NAME
            logger.info(f"Get Queue size for {queue_name}")
            return max(len(self.awaitingACK), self.channel.message_count(queue_name))
        except Exception:
            return len(self.awaitingACK)
    """