clevermicro/amq-adapter-python
14
0
Fork 0
Code Issues 16 Pull Requests 4 Actions Packages Projects Releases Wiki Activity

Add support for file streaming upload, update to async pika lib

Browse Source
This commit is contained in:
Stanislav Hejny
2025-03-26 13:54:30 +00:00
parent e8f4abdfaf
commit a2ddce1a2e
22 changed files with 1369 additions and 395 deletions
Download Patch File Download Diff File Expand all files Collapse all files
amqp/rabbitmq/rabbit_mq_consumer.py
+37 -95
View File
@@ -1,7 +1,7 @@
import logging
import pika
from pika.exchange_type import ExchangeType
from aio_pika.abc import AbstractRobustQueue, AbstractRobustExchange
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
@@ -16,124 +16,66 @@ class RabbitMQConsumer(RabbitMQProducer):
def __init__(self, configuration):
super().__init__(configuration, RouterConsumer(configuration))
self.reply_to_exchange: AbstractRobustExchange | None = None # when sending response to RPC request
self.reply_to_queue: AbstractRobustQueue | None= None # when receiving response to RPC request
async def init(self, loop) -> AbstractRobustExchange:
await self.setup_amq_channel(loop)
await self.router.init_routing_paths_consumer(self.router.channel)
cm_reply_to_queue = self.router.get_reply_to_queue_name()
self.reply_to_queue = await self.channel.declare_queue(
name=cm_reply_to_queue, durable=True, exclusive=False, auto_delete=False
)
self.reply_to_exchange = await self.channel.declare_exchange(
AMQ_REPLY_TO_EXCHANGE, type=ExchangeType.direct
)
await self.reply_to_queue.bind(exchange=self.reply_to_exchange, routing_key=cm_reply_to_queue)
return self.reply_to_exchange
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:
except Exception as e:
logging.error(" [E] RabbitMQConsumer.PreDestroy cleanup - connection close error: %s", e)
def register_inbound_routes(self, callback):
async def register_inbound_routes(self, route_mapping: str, callback):
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:
try:
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.channel.basic_consume(
queue=queue_name,
on_message_callback=callback,
auto_ack=False
queue: AbstractRobustQueue = await self.channel.declare_queue(name=queue_name,
durable=True,
exclusive=False,
auto_delete=False,
arguments=queue_args)
exchange_name: str = route.exchange if route.exchange else AMQ_RPC_EXCHANGE
exchange: AbstractRobustExchange = await self.channel.declare_exchange(
name=exchange_name, type=ExchangeType.topic
)
self.router.register_route(
await queue.bind(exchange, route.key)
await queue.consume(callback, no_ack=False)
await self.router.register_route(
AMQRoute(
route.component_name, exchange, queue_name,
route.component_name, exchange_name, 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)
except pika.connection.exceptions.AMQPError as err:
route, exchange, queue_name)
except Exception as err:
logging.error(" [E] Callback registration INCOMPLETE, %s", err)
else:
logging.warning(" [W] Channel is null, cannot register routes.")
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
async def register_rpc_handler(self, rpc_consumer):
await self.reply_to_queue.consume(callback=rpc_consumer,
no_ack=False
)
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
logging.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)
"""
amqp/rabbitmq/rabbit_mq_producer.py
+43 -85
View File
@@ -1,10 +1,5 @@
import logging
import time
import pika
from pika.adapters.blocking_connection import BlockingChannel, BlockingConnection
from pika.channel import Channel
from pika.connection import Connection
import aio_pika
from amqp.adapter.amq_message_factory import AMQMessageFactory
from amqp.config.amq_configuration import AMQConfiguration
@@ -14,38 +9,22 @@ 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: BlockingChannel | None = None
self.connection: BlockingConnection | None = None
self.request_for_window_duration = None
self.channel: aio_pika.RobustChannel | None = None
self.connection: aio_pika.RobustConnection | None = None
self.rpc_exchange: aio_pika.RobustExchange | None = None
self.reply_exchange: aio_pika.RobustExchange | None = 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("******** RabbitMQProducer.init(): %s",
self.amq_configuration.dispatch.amq_host)
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
async def init(self, loop):
await self.setup_amq_channel(loop)
def cleanup(self):
if self.connection:
@@ -54,30 +33,27 @@ class RabbitMQProducer:
except Exception as e:
logging.error("RabbitMQConsumer.PreDestroy cleanup - connection close error: %s", e)
def setup_amq_channel(self):
async def setup_amq_channel(self, loop):
"""
Set up the RabbitMQ connection
:return: nothing, it applies changes to instance variables
"""
credentials: pika.credentials.PlainCredentials = pika.credentials.PlainCredentials(
username=self.amq_configuration.dispatch.rabbit_mq_user,
password=self.amq_configuration.dispatch.rabbit_mq_password
)
factory = pika.ConnectionParameters(
self.connection = await aio_pika.connect_robust(
host=self.amq_configuration.dispatch.amq_host,
credentials=credentials
port=self.amq_configuration.dispatch.amq_port,
login=self.amq_configuration.dispatch.rabbit_mq_user,
password=self.amq_configuration.dispatch.rabbit_mq_password,
loop=loop
)
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)
logging.info("RabbitMQProducer.setupAMQ, got connection, open=%s", self.connection.is_closed)
self.channel = await self.connection.channel()
await self.channel.set_qos(prefetch_count=self.PREFETCH_COUNT)
await self.router.init_routing_paths(self.channel)
self.router.set_route_added_notifier(
lambda route: logging.info(" [DBG] ---------- RouteAdded, setting up metrics: amqp.adapter.exchange.%s",
route.exchange))
def send_message(self, message: AMQMessage):
async 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.
@@ -87,27 +63,31 @@ class RabbitMQProducer:
"""
try:
route = self.router.find_route_by_message(message)
if route:
if route.timeout < 0:
if route and self.rpc_exchange:
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)
pika_message: aio_pika.message.Message = aio_pika.message.Message(
body=AMQMessageFactory.serialize(message),
content_encoding='utf-8',
delivery_mode=2
)
await self.rpc_exchange.publish(
message=pika_message,
routing_key=self.router.get_routing_key(message), # context path is a routing key
)
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,
pika_message: aio_pika.message.Message = aio_pika.message.Message(
body=AMQMessageFactory.serialize(message),
content_encoding='utf-8',
delivery_mode=2,
correlation_id=str(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)
await self.rpc_exchange.publish(
message=pika_message,
routing_key=self.router.get_routing_key(message), # context path is a routing key
)
logging.info(" [x] basicPublish (RPC call), messageId: %s, to: %s",
message.id.as_string(), route.as_string())
@@ -117,30 +97,8 @@ class RabbitMQProducer:
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()
async def request_routes(self):
await self.router.request_routes_from_adapters()
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")
def ack_nack_callback(method):
if method == pika.spec.Basic.Ack:
pass
if method == pika.spec.Basic.Nack:
pass